dongjoon-hyun commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4470948911
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -720,7 +761,11 @@ class NativePartitionReader(
array.release()
throw e
}
+ // The import releases the buffers of the previous batch: the batch
has no rows until the
+ // new one passes the check, and the column vectors read its buffers
before the check.
+ batch.setNumRows(0)
Review Comment:
**[5] No test fails if either part of the fix for item 2 is reverted.**
Thank you for the change in 5d9a78b. But the test of the null rejection
(`NativeDataSourceSuite.scala:461`) fails on the first batch, since
`native_range` makes `name` null at id 0, and it only checks the exception from
`collect()`. Before the first batch, the batch already has 0 rows and the
column vectors have not been read, so removing `batch.setNumRows(0)` here, or
moving `columns.foreach(_.refresh())` back after `checkNulls()`, still passes
every test. `NativeColumnVectorSuite` calls `refresh()` directly, without a
reader. Could a reader-level test read a batch that passes the check, get the
failure from a later `next()`, and then check that `reader.get().numRows()` is
0 and that the columns read the rejected batch? For example, with the schema
`id BIGINT, name STRING NOT NULL` and `end` 12, `partitions` 2 and `batch_size`
2, the second partition reads ids 6-7 and 8-9, and its third `next()` fails on
the
null at id 10, while row 1 of that batch holds `n11`.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVectorSuite.scala:
##########
@@ -0,0 +1,395 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi
+
+import java.math.{BigDecimal => JavaBigDecimal}
+import java.nio.charset.StandardCharsets
+
+import scala.util.{Failure, Success, Try, Using}
+
+import org.apache.arrow.memory.{BufferAllocator, RootAllocator}
+import org.apache.arrow.vector._
+import org.apache.arrow.vector.complex.StructVector
+import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType}
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.expressions.UnsafeRow
+import org.apache.spark.sql.catalyst.expressions.codegen.UnsafeRowWriter
+import org.apache.spark.sql.execution.vectorized.{OnHeapColumnVector,
ReusedArrowColumnVector}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{Decimal, DecimalType}
+import org.apache.spark.sql.vectorized.ArrowColumnVector
+
+/**
+ * Checks that a [[NativeColumnVector]] reads the nulls and values that an
[[ArrowColumnVector]]
+ * of the same Arrow vector reads.
+ */
+class NativeColumnVectorSuite extends SparkFunSuite {
+
+ // The rows of the batches, more than 8 so that the validity bitmaps have
several bytes, and the
+ // rows that are null.
+ private val rows = 0 until 20
+ private val nulls = Set(1, 8, 9, 15)
+
+ private def withAllocator(f: BufferAllocator => Unit): Unit = {
+ Using.resource(new RootAllocator())(f)
+ }
+
+ /** The column vector of `vector`, refreshed as the reader refreshes it
after an import. */
+ private def native(vector: ValueVector, nullable: Boolean = true):
NativeColumnVector = {
+ val column = NativeColumnVector.create(vector, nullable)
+ column.refresh()
+ column
+ }
+
+ /** Checks that `column` reads the nulls and the values of `vector` as
ArrowColumnVector does. */
+ private def checkSame(vector: ValueVector, column: ArrowColumnVector)(
+ value: (ArrowColumnVector, Int) => Any): Unit = {
+ val expected = new ArrowColumnVector(vector)
+ (0 until vector.getValueCount).foreach { row =>
+ assert(column.isNullAt(row) === expected.isNullAt(row), s"row $row")
+ if (!expected.isNullAt(row)) {
+ assert(value(column, row) === value(expected, row), s"row $row")
+ }
+ }
+ }
+
+ /** Sets the rows of a vector: nulls, and `set` for the other rows. */
+ private def fill[V <: BaseFixedWidthVector](vector: V)(set: (V, Int) =>
Unit): V = {
+ vector.allocateNew(rows.size)
+ rows.foreach { row => if (nulls(row)) vector.setNull(row) else set(vector,
row) }
+ vector.setValueCount(rows.size)
+ vector
+ }
+
+ /** Checks a vector whose values the column vector reads at the addresses of
its buffers. */
+ private def checkDirect(vector: ValueVector)(value: (ArrowColumnVector, Int)
=> Any): Unit = {
+ Using.resource(vector) { vector =>
+ val column = native(vector)
+ assert(column.getClass !== classOf[NativeColumnVector],
vector.getClass.getName)
+ checkSame(vector, column)(value)
+ }
+ }
+
+ test("booleans, integers, floats, dates and timestamps") {
+ withAllocator { allocator =>
+ val bytes = Seq(Byte.MinValue, -1: Byte, 0: Byte, 1: Byte, Byte.MaxValue)
+ val shorts = Seq(Short.MinValue, -1: Short, 0: Short, 1: Short,
Short.MaxValue)
+ val ints = Seq(Int.MinValue, -1, 0, 1, Int.MaxValue)
+ val longs = Seq(Long.MinValue, -1L, 0L, 1L, Long.MaxValue)
+ val floats = Seq(0.0f, -0.0f, Float.NaN, Float.NegativeInfinity,
Float.MinValue, 1.5f)
+ val doubles = Seq(0.0, -0.0, Double.NaN, Double.PositiveInfinity,
Double.MaxValue, -2.5)
+
+ checkDirect(fill(new BitVector("b", allocator))((v, row) => v.set(row,
row % 3 % 2))) {
+ _.getBoolean(_)
+ }
+ checkDirect(fill(new TinyIntVector("i8", allocator)) { (v, row) =>
+ v.set(row, bytes(row % bytes.size))
+ })(_.getByte(_))
+ checkDirect(fill(new SmallIntVector("i16", allocator)) { (v, row) =>
+ v.set(row, shorts(row % shorts.size))
+ })(_.getShort(_))
+ checkDirect(fill(new IntVector("i32", allocator))((v, row) => v.set(row,
ints(row % 5)))) {
+ _.getInt(_)
+ }
+ checkDirect(fill(new DateDayVector("day", allocator))((v, row) =>
v.set(row, row - 9))) {
+ _.getInt(_)
+ }
+ checkDirect(fill(new BigIntVector("i64", allocator)) { (v, row) =>
+ v.set(row, longs(row % longs.size))
+ })(_.getLong(_))
+ checkDirect(fill(new TimeStampMicroTZVector("ts", allocator, "UTC")) {
(v, row) =>
+ v.set(row, longs(row % longs.size))
+ })(_.getLong(_))
+ checkDirect(fill(new TimeStampMicroVector("ntz", allocator)) { (v, row)
=>
+ v.set(row, longs(row % longs.size))
+ })(_.getLong(_))
+ // The bits, so that NaN equals NaN.
+ checkDirect(fill(new Float4Vector("f32", allocator)) { (v, row) =>
+ v.set(row, floats(row % floats.size))
+ })((column, row) =>
java.lang.Float.floatToRawIntBits(column.getFloat(row)))
+ checkDirect(fill(new Float8Vector("f64", allocator)) { (v, row) =>
+ v.set(row, doubles(row % doubles.size))
+ })((column, row) =>
java.lang.Double.doubleToRawLongBits(column.getDouble(row)))
+ }
+ }
+
+ test("strings and string views, in several batches") {
+ // Values that a view holds, of up to 12 bytes, and longer ones that a
data buffer holds, with
+ // characters of 2 bytes in UTF-8.
+ val (acute, grave) = (0xe9.toChar.toString, 0xe8.toChar.toString)
+ val short = Seq(Some(""), Some("a"), None, Some("twelve bytes"),
Some(acute + grave), None)
+ val long = Seq(Some("thirteen byte"), None, Some(acute * 40), Some("a long
value " * 20))
+ // Values that do not fit in one data buffer of a view vector, which holds
them in several.
+ val spread = Seq(Some("a" * 100000), None, Some(acute * 50000), Some("c" *
100000))
+ withAllocator { allocator =>
+ def check(vector: FieldVector)(set: (Int, Array[Byte]) => Unit, setNull:
Int => Unit)
+ : Unit = {
+ Using.resource(vector) { vector =>
+ val column = NativeColumnVector.create(vector, true)
+ assert(column.getClass !== classOf[NativeColumnVector],
vector.getClass.getName)
+ // Batches without and with data buffers, with values of both kinds,
and with values in
+ // several data buffers.
+ Seq(short, long, short ++ long, spread).foreach { values =>
+ vector.clear()
+ vector.allocateNew()
+ values.zipWithIndex.foreach {
+ case (Some(value), row) => set(row,
value.getBytes(StandardCharsets.UTF_8))
+ case (None, row) => setNull(row)
+ }
+ vector.setValueCount(values.size)
+ vector match {
+ case views: ViewVarCharVector if values eq spread =>
+ assert(views.getDataBuffers.size > 1)
+ case _ =>
+ }
+ column.refresh()
+ checkSame(vector, column)(_.getUTF8String(_))
+ assert(column.getUTF8String(values.indexOf(None)) === null)
+ }
+ }
+ }
+ val strings = new VarCharVector("s", allocator)
+ check(strings)(strings.setSafe(_, _), strings.setNull)
+ val views = new ViewVarCharVector("v", allocator)
+ check(views)(views.setSafe(_, _), views.setNull)
+ }
+ }
+
+ private def withDecimalVector(precision: Int, scale: Int)(f: DecimalVector
=> Unit): Unit = {
+ withAllocator { allocator =>
+ Using.resource(new DecimalVector("decimal", allocator, precision,
scale))(f)
+ }
+ }
+
+ /**
+ * Checks that `column` returns the decimal of a row that `expected`, an
[[ArrowColumnVector]] of
+ * the same vector, returns, as Spark uses it, or fails as it does. Its
integral conversions are
+ * those of the decimal of Spark's own column vectors, which hold the
decimals that a long holds,
+ * of a scale from 0 to their precision, as longs, while ArrowColumnVector
holds BigDecimals,
+ * which Decimal rounds to bytes, shorts and ints through doubles.
+ */
+ private def checkDecimal(
+ expected: ArrowColumnVector,
+ column: ArrowColumnVector,
+ row: Int,
+ precision: Int,
+ scale: Int): Unit = {
+ def result(decimal: Decimal): (JavaBigDecimal, Int, Int, Long, UnsafeRow)
= {
+ val writer = new UnsafeRowWriter(1)
+ writer.write(0, decimal, precision, scale)
+ (decimal.toJavaBigDecimal, decimal.precision, decimal.scale,
decimal.toLong, writer.getRow)
+ }
+ def rounded(decimal: Decimal): Seq[Either[String, Long]] = {
+ Seq[Decimal => Long](_.roundToByte(), _.roundToShort(), _.roundToInt(),
_.roundToLong())
+ .map(round => Try(round(decimal)).toEither.left.map(e =>
s"${e.getClass}: ${e.getMessage}"))
+ }
+ def own(decimal: Decimal): Decimal = {
+ if (0 <= scale && scale <= precision && precision <=
Decimal.MAX_LONG_DIGITS) {
+ Decimal.createUnsafe(decimal.toUnscaledLong, precision, scale)
+ } else {
+ decimal
+ }
+ }
+ Try(result(expected.getDecimal(row, precision, scale))) match {
+ case Success(value) =>
+ val decimal = column.getDecimal(row, precision, scale)
+ assert(result(decimal) === value)
+ assert(rounded(decimal) === rounded(own(expected.getDecimal(row,
precision, scale))))
+ case Failure(error) =>
+ val actual = intercept[Exception](result(column.getDecimal(row,
precision, scale)))
+ assert(actual.getClass === error.getClass)
+ assert(actual.getMessage === error.getMessage)
+ }
+ }
+
+ test("decimals in and out of their precision, at other precisions and
scales") {
+ for (precision <- 1 to Decimal.MAX_LONG_DIGITS; scale <- Set(0, precision
/ 2, precision)) {
+ withDecimalVector(precision, scale) { vector =>
+ val limit = BigInt(10).pow(precision)
+ // The bounds of the precision, values out of it, which the byte
setter permits, and
+ // values that a long does not hold.
+ val values = Seq(BigInt(0), BigInt(1), BigInt(-1), limit - 1, -(limit
- 1), limit, -limit,
+ BigInt(Long.MaxValue), BigInt(Long.MinValue), (BigInt(1) << 64) + 7)
+ vector.allocateNew(values.size + 1)
+ values.zipWithIndex.foreach { case (value, row) =>
+ vector.setBigEndian(row, value.toByteArray)
+ }
+ vector.setNull(values.size)
+ vector.setValueCount(values.size + 1)
+ val expected = new ArrowColumnVector(vector)
+ val column = native(vector)
+ assert(column.getClass !== classOf[NativeColumnVector])
+ val requested = Seq((precision, scale), (Decimal.MAX_LONG_DIGITS,
scale), (38, scale),
+ (precision + 1, scale), (precision, scale + 1))
+ for (row <- values.indices; (p, s) <- requested) {
+ checkDecimal(expected, column, row, p, s)
+ }
+ assert(column.getDecimal(values.size, precision, scale) === null)
+ }
+ }
+ }
+
+ test("decimals of a negative scale") {
+ val conf = new SQLConf
+ val key = SQLConf.LEGACY_ALLOW_NEGATIVE_SCALE_OF_DECIMAL_ENABLED.key
+ SQLConf.withExistingConf(conf) {
+ conf.setConfString(key, "true")
+ withDecimalVector(5, -2) { vector =>
+ vector.allocateNew(1)
+ vector.set(0, new JavaBigDecimal("1.23E+4"))
+ vector.setValueCount(1)
+ val expected = new ArrowColumnVector(vector)
+ val column = native(vector)
+ Seq((5, -2), (8, 0), (5, -3)).foreach { case (precision, scale) =>
+ checkDecimal(expected, column, 0, precision, scale)
+ }
+ assert(column.getDecimal(0, 5, -2).toLong === 12300L)
+ conf.setConfString(key, "false")
+ checkDecimal(expected, column, 0, 5, -2)
+ }
+ }
+ }
+
+ test("decimals that a long holds round to integers as those of Spark's own
column vectors") {
+ // The decimals of Spark's own column vectors and of the native one, of
OnHeapColumnVector and
+ // NativeColumnVector, hold longs, which Decimal rounds exactly.
+ def decimals(value: String, scale: Int)(f: (Decimal, Decimal) => Unit):
Unit = {
+ val precision = Decimal.MAX_LONG_DIGITS
+ withDecimalVector(precision, scale) { vector =>
+ vector.allocateNew(1)
+ vector.set(0, new JavaBigDecimal(value))
+ vector.setValueCount(1)
+ Using.resource(new OnHeapColumnVector(1, DecimalType(precision,
scale))) { own =>
+ own.putDecimal(0, Decimal(value), precision)
+ f(own.getDecimal(0, precision, scale), native(vector).getDecimal(0,
precision, scale))
+ }
+ }
+ }
+ decimals("123456789.999999999", 9) { (own, column) =>
+ assert(column.roundToInt() === own.roundToInt())
+ }
+ decimals("127.999999999999999", 15) { (own, column) =>
+ assert(column.roundToByte() === own.roundToByte())
+ }
+ }
+
+ test("the decimals of the fields of a struct still hold BigDecimals") {
+ // ArrowColumnVector reads the fields of a struct, so that their decimals,
unlike those of a
+ // top-level column, hold BigDecimals, which Decimal rounds to ints
through doubles.
+ withAllocator { allocator =>
+ Using.resource(StructVector.empty("struct", allocator)) { struct =>
+ val field = FieldType.nullable(new ArrowType.Decimal(18, 9, 128))
+ val child = struct.addOrGet("d", field, classOf[DecimalVector])
+ struct.allocateNew()
+ child.setSafe(0, new JavaBigDecimal("123456789.999999999"))
+ struct.setIndexDefined(0)
+ child.setValueCount(1)
+ struct.setValueCount(1)
+ val top = native(child).getDecimal(0, 18, 9)
+ val nested = native(struct).getChild(0).getDecimal(0, 18, 9)
+ assert(top.toDebugString === "Decimal(compact, 123456789999999999, 18,
9)")
+ assert(top.roundToInt() === 123456789)
+ assert(nested.toDebugString === "Decimal(expanded,
123456789.999999999, 18, 9)")
+ assert(nested.roundToInt() === 123456790)
+ }
+ }
+ }
+
+ test("reads the buffers of the batch that the reader imported last") {
+ withAllocator { allocator =>
+ Using.resource(new BigIntVector("id", allocator)) { vector =>
+ val column = NativeColumnVector.create(vector, true)
+ // Each batch has other buffers: with nulls, without nulls, empty, and
larger.
+ val batches = Seq(
+ Seq(Some(1L), None, Some(3L)),
+ Seq(Some(4L), Some(5L)),
+ Seq.empty[Option[Long]],
+ (0L until 1000L).map(id => Option.when(id % 7 != 0)(id)))
+ batches.foreach { values =>
+ vector.clear()
+ vector.allocateNew(values.size)
+ values.zipWithIndex.foreach {
+ case (Some(value), row) => vector.set(row, value)
+ case (None, row) => vector.setNull(row)
+ }
+ vector.setValueCount(values.size)
+ column.refresh()
+ checkSame(vector, column)(_.getLong(_))
+ }
+ }
+ }
+ }
+
+ test("a column that the schema declares not nullable") {
+ withAllocator { allocator =>
+ Using.resource(new IntVector("k", allocator)) { vector =>
+ vector.allocateNew(rows.size)
+ rows.foreach(row => vector.set(row, row))
+ vector.setValueCount(rows.size)
+ val column = native(vector, nullable = false)
+ assert(column.getClass !== classOf[NativeColumnVector])
+ checkSame(vector, column)(_.getInt(_))
+ // The reader rejects a batch with a null in such a column before it
refreshes the column
Review Comment:
**[6] Nit: this comment still describes the order before 5d9a78b.**
`readNext` now refreshes the column vectors right after the import and calls
`checkNulls()` after that (`NativeDataSource.scala:767-769`), so the reader no
longer rejects the batch before it refreshes the column vector. The column
vector does not count the nulls because `Direct.refresh()` skips the count for
a column that is not nullable (`NativeColumnVector.java:130`). The
`columnVector` scaladoc (`NativeDataSource.scala:695-696`), which says that
`checkNulls` runs "on each batch before the column vector reads it", can be
read the same way. Could both say that `refresh()` comes first, and that the
batch has no rows until `checkNulls()` passes?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.datasources.v2.ffi;
+
+import java.nio.ByteOrder;
+import java.util.List;
+
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.vector.BaseVariableWidthViewVector;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.BitVector;
+import org.apache.arrow.vector.DateDayVector;
+import org.apache.arrow.vector.DecimalVector;
+import org.apache.arrow.vector.Float4Vector;
+import org.apache.arrow.vector.Float8Vector;
+import org.apache.arrow.vector.IntVector;
+import org.apache.arrow.vector.SmallIntVector;
+import org.apache.arrow.vector.TimeStampMicroTZVector;
+import org.apache.arrow.vector.TimeStampMicroVector;
+import org.apache.arrow.vector.TinyIntVector;
+import org.apache.arrow.vector.ValueVector;
+import org.apache.arrow.vector.VarCharVector;
+import org.apache.arrow.vector.ViewVarCharVector;
+
+import org.apache.spark.sql.execution.vectorized.ReusedArrowColumnVector;
+import org.apache.spark.sql.types.Decimal;
+import org.apache.spark.sql.vectorized.ArrowColumnVector;
+import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.types.UTF8String;
+
+/**
+ * The column vector of a top-level column of a {@code NativePartitionReader},
which imports every
+ * batch into the same Arrow vectors.
+ *
+ * <p>After each import, the reader calls {@link #refresh()}, and the vector
then reads booleans,
+ * integers, floats, dates, timestamps, decimals, and the strings of the Arrow
types Utf8 and
+ * Utf8View at the addresses of the buffers of the batch. {@link
ArrowColumnVector} reads each
+ * value through Arrow Java, which by default checks on each access the bounds
and the reference
+ * count of the buffer, and whether the value is null.
+ * Spark does not need these checks here: it reads only the rows of the batch,
whose buffers Arrow
+ * Java sizes by the lengths of the arrays, which the reader checks before it
imports them, and
+ * the buffers of a batch stay allocated until the thread that read it calls
next again or closes
+ * the reader: a close on another thread releases them at once only if that
thread has ended, and
+ * otherwise leaves the release to that thread, or to a thread that waits for
it to end. Spark
+ * reads a value of a nullable column only after {@link #isNullAt(int)}, and
the reader checks
+ * that the columns that the schema declares not nullable have no null, so
that the vector of such
+ * a column does not count the nulls of each batch.
+ *
+ * <p>The vector returns the decimals that fit their precision in a long, of a
scale from 0 to
+ * their precision, as the unscaled longs that Spark's own column vectors
hold, which Spark adds as
+ * longs, rather than as BigDecimals. The other decimals, large strings,
binary values, structs,
+ * lists and maps, and the fields of structs are read through {@link
ArrowColumnVector}.
+ */
+class NativeColumnVector extends ReusedArrowColumnVector {
+
+ /**
+ * The column vector of a top-level Arrow vector of a reader, of a column
that the schema to
+ * read declares {@code nullable} or not.
+ */
+ static NativeColumnVector create(ValueVector vector, boolean nullable) {
+ if (vector instanceof BitVector) {
+ return new Booleans(vector, nullable);
+ } else if (vector instanceof TinyIntVector) {
+ return new Bytes(vector, nullable);
+ } else if (vector instanceof SmallIntVector) {
+ return new Shorts(vector, nullable);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ return new Ints(vector, nullable);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ return new Longs(vector, nullable);
+ } else if (vector instanceof Float4Vector) {
+ return new Floats(vector, nullable);
+ } else if (vector instanceof Float8Vector) {
+ return new Doubles(vector, nullable);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ return new Decimals(decimalVector, nullable);
+ } else if (vector instanceof VarCharVector) {
+ return new Strings(vector, nullable);
+ } else if (vector instanceof ViewVarCharVector viewVector) {
+ return new StringViews(viewVector, nullable);
+ } else {
+ return new NativeColumnVector(vector);
+ }
+ }
+
+ NativeColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ /** Reads the addresses of the buffers of the batch that the reader imported
last. */
+ void refresh() {
+ }
+
+ /** A vector that reads its values and their validity at the addresses of
its buffers. */
+ private abstract static class Direct extends NativeColumnVector {
+ final ValueVector vector;
+ // Whether the schema to read declares the column nullable. The reader
checks that the batches
+ // of a column that it does not declare nullable have no null.
+ private final boolean nullable;
+ // The address of the validity bitmap of the batch, or 0 if the batch has
no null.
+ private long validity;
+ // The address of the data buffer of the batch.
+ long data;
+
+ Direct(ValueVector vector, boolean nullable) {
+ super(vector);
+ this.vector = vector;
+ this.nullable = nullable;
+ }
+
+ @Override
+ void refresh() {
+ // Arrow Java counts the nulls in the validity bitmap on each call.
+ boolean hasNulls = nullable && vector.getNullCount() > 0;
Review Comment:
**[7] Minor: `refresh()` counts the nulls of each nullable column again,
right after the import counted them.** `Data.importIntoVectorSchemaRoot`
unloads the imported struct with `new StructVectorUnloader(structVector)`,
which includes the null counts, so it calls `getNullCount()` on every vector
(arrow-java v19.0.0, `Data.java:458` and `StructVectorUnloader.java:42`). As
L129 says, `getNullCount()` scans the bitmap on each call, so each batch now
scans the validity bitmap of every nullable column twice before Spark reads a
value. `checkArrays` already reads the `null_count` that the library exported
for each array (`NativeDataSource.scala:833`). Could `refresh()` take that
count, and scan the bitmap only for a `null_count` of -1? This trusts an
exported count of 0, which the C data interface requires to be exact.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -496,6 +502,8 @@ class NativePartitionReader(
// The task that reads the partition, which the library may use until the
stream is released.
private var task: NativeLibrary.Task = _
private var root: VectorSchemaRoot = _
+ // The column vectors of the batch, which read the buffers of each batch
that `readNext` imports.
+ private var columns: Array[NativeColumnVector] = _
Review Comment:
**[8] Nit: the local `columns` of `checkArrays` now shadows this field.**
`checkArrays` declares `val columns = root.getSchema.getFields.size` on L890,
an `Int`, so inside it `columns` no longer means these column vectors. L890 is
outside this diff, so I left this here. Could the local become `numColumns`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
