gengliangwang commented on code in PR #59961: URL: https://github.com/apache/spark/pull/59961#discussion_r4429506605
########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java: ########## @@ -0,0 +1,314 @@ +/* + * 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 strings 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 reader imports the next batch or closes. 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. + * + * <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, 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. */ + static NativeColumnVector create(ValueVector vector) { + if (vector instanceof BitVector) { + return new Booleans(vector); + } else if (vector instanceof TinyIntVector) { + return new Bytes(vector); + } else if (vector instanceof SmallIntVector) { + return new Shorts(vector); + } else if (vector instanceof IntVector || vector instanceof DateDayVector) { + return new Ints(vector); + } else if (vector instanceof BigIntVector || vector instanceof TimeStampMicroTZVector || + vector instanceof TimeStampMicroVector) { + return new Longs(vector); + } else if (vector instanceof Float4Vector) { + return new Floats(vector); + } else if (vector instanceof Float8Vector) { + return new Doubles(vector); + } else if (vector instanceof DecimalVector decimalVector) { + return new Decimals(decimalVector); + } else if (vector instanceof VarCharVector) { + return new Strings(vector); + } else if (vector instanceof ViewVarCharVector viewVector) { + return new StringViews(viewVector); + } 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; + // 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) { + super(vector); + this.vector = vector; + } + + @Override + void refresh() { + validity = vector.getNullCount() == 0 ? 0 : vector.getValidityBuffer().memoryAddress(); + data = vector.getDataBuffer().memoryAddress(); Review Comment: Done in 5f55102a6ed0: Good catch. I did what you suggested: the reader now remembers the thread that called next() last. A close() on any other thread, even when no next() is running, only marks the reader closed. That thread's next call to next() then does the release (it returns false). The exception is a thread that has already ended, where close() releases right away. Release happens once, behind a `released` flag. I documented this in the class docs of NativePartitionReader and NativeColumnVector, including the tradeoff: a thread that stops mid-batch without calling next() again, and stays alive, keeps the reader's resources. For PipedRDD this doesn't happen, because PrintWriter swallows the broken pipe and the writer goes on to next(). I added the test "close a reader while another thread reads the batch that it returned". In it, the batch is still readable after close() and the stream is released only on the next next(). I'll also note this in the PR description. <sub>Automated reply (Opus)</sub> ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVectorSuite.scala: ########## @@ -0,0 +1,308 @@ +/* + * 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.ReusedArrowColumnVector +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.Decimal +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): NativeColumnVector = { + val column = NativeColumnVector.create(vector) + 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)) + withAllocator { allocator => + def check(vector: FieldVector)(set: (Int, Array[Byte]) => Unit, setNull: Int => Unit) + : Unit = { + Using.resource(vector) { vector => + val column = NativeColumnVector.create(vector) + assert(column.getClass !== classOf[NativeColumnVector], vector.getClass.getName) + // Batches without and with data buffers, and with values of both kinds. + Seq(short, long, short ++ long).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) + 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) Review Comment: Done in 5f55102a6ed0: Added a fourth batch of values around 100KB each, too large to share one data buffer, and it asserts `getDataBuffers.size > 1` for the view vector. The values differ, so reading from `buffers[0]` instead of the right buffer would fail. <sub>Automated reply (Opus)</sub> ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVectorSuite.scala: ########## @@ -0,0 +1,308 @@ +/* + * 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.ReusedArrowColumnVector +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.Decimal +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): NativeColumnVector = { + val column = NativeColumnVector.create(vector) + 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)) + withAllocator { allocator => + def check(vector: FieldVector)(set: (Int, Array[Byte]) => Unit, setNull: Int => Unit) + : Unit = { + Using.resource(vector) { vector => + val column = NativeColumnVector.create(vector) + assert(column.getClass !== classOf[NativeColumnVector], vector.getClass.getName) + // Batches without and with data buffers, and with values of both kinds. + Seq(short, long, short ++ long).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) + 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. + */ + private def checkDecimal( Review Comment: Done in 5f55102a6ed0: `checkDecimal` now also compares roundToByte/Short/Int/Long. The expected values come from the decimal as Spark's own vectors hold it: long-backed for precision <= 18 and scale 0..precision. I also added a test that pins your two examples against OnHeapColumnVector and ArrowColumnVector: 123456789.999999999 gives roundToInt 123456789 (Arrow: 123456790), and 127.999999999999999 gives roundToByte 127 (Arrow: CAST_OVERFLOW). I'll mention the difference in the description. <sub>Automated reply (Opus)</sub> ########## sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java: ########## @@ -0,0 +1,314 @@ +/* + * 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 strings 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 reader imports the next batch or closes. 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. + * + * <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, 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. */ + static NativeColumnVector create(ValueVector vector) { + if (vector instanceof BitVector) { + return new Booleans(vector); + } else if (vector instanceof TinyIntVector) { + return new Bytes(vector); + } else if (vector instanceof SmallIntVector) { + return new Shorts(vector); + } else if (vector instanceof IntVector || vector instanceof DateDayVector) { + return new Ints(vector); + } else if (vector instanceof BigIntVector || vector instanceof TimeStampMicroTZVector || + vector instanceof TimeStampMicroVector) { + return new Longs(vector); + } else if (vector instanceof Float4Vector) { + return new Floats(vector); + } else if (vector instanceof Float8Vector) { + return new Doubles(vector); + } else if (vector instanceof DecimalVector decimalVector) { + return new Decimals(decimalVector); + } else if (vector instanceof VarCharVector) { + return new Strings(vector); + } else if (vector instanceof ViewVarCharVector viewVector) { + return new StringViews(viewVector); + } 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; + // 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) { + super(vector); + this.vector = vector; + } + + @Override + void refresh() { + validity = vector.getNullCount() == 0 ? 0 : vector.getValidityBuffer().memoryAddress(); + data = vector.getDataBuffer().memoryAddress(); + } + + @Override + public final boolean isNullAt(int rowId) { + return validity != 0 && !bit(validity, rowId); + } + + /** The bit of a bitmap, whose bits Arrow orders from the least significant of each byte. */ + static boolean bit(long address, int index) { + return (Platform.getByte(null, address + (index >> 3)) & (1 << (index & 7))) != 0; + } + } + + private static final class Booleans extends Direct { + Booleans(ValueVector vector) { + super(vector); + } + + @Override + public boolean getBoolean(int rowId) { + return bit(data, rowId); + } + } + + private static final class Bytes extends Direct { + Bytes(ValueVector vector) { + super(vector); + } + + @Override + public byte getByte(int rowId) { + return Platform.getByte(null, data + rowId); + } + } + + private static final class Shorts extends Direct { + Shorts(ValueVector vector) { + super(vector); + } + + @Override + public short getShort(int rowId) { + return Platform.getShort(null, data + 2L * rowId); + } + } + + private static final class Ints extends Direct { + Ints(ValueVector vector) { + super(vector); + } + + @Override + public int getInt(int rowId) { + return Platform.getInt(null, data + 4L * rowId); + } + } + + private static final class Longs extends Direct { + Longs(ValueVector vector) { + super(vector); + } + + @Override + public long getLong(int rowId) { + return Platform.getLong(null, data + 8L * rowId); Review Comment: Done in 5f55102a6ed0: Same issue as the thread above, fixed the same way. A close() on a thread other than the one that called next() last now leaves the release to that thread's next next() call, unless that thread has ended. There's a new test that reads the batch on the other thread after close(). <sub>Automated reply (Opus)</sub> -- 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]
