sarutak commented on code in PR #59961: URL: https://github.com/apache/spark/pull/59961#discussion_r4429239839
########## 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: The class doc of `NativePartitionReader` says a reader can be closed by `DataSourceRDD` on the task thread while another thread is still reading it, e.g. the stdin writer thread of `PipedRDD`. SPARK-60705 covers the case where that thread is inside `next()`. But once `next()` has returned, `reading` is false, so a close on the task thread runs `closeResources()` right away and releases the imported buffers while the writer thread may still be reading the remaining rows of the current batch (e.g. `df.rdd.pipe("head -1")`, where the task completes as soon as `head` exits). Before this PR, those reads went through Arrow Java, so after `root.close()` they hit the empty buffers (or the reference count check) and threw. With this PR, the vectors keep the raw addresses captured in `refresh()` and read freed memory, so the writer thread can read garbage or crash the executor JVM if the memory was unmapped. The Javadoc above ("the buffers of a batch stay allocated until the reader imports the next batch or closes") also does not cover this case. Off-heap `OffHeapColumnVector` has a similar hazard, but that is opt-in, while this would be the default path for native data sources. Could we keep the buffers alive while a batch is handed out to another thread? For example, when `close()` comes from a thread other than the one that last called `next()`, it could only mark the reader closed and leave `closeResources()` to that thread's next `next()` call, as it already does for a close during `next()`. If you think the current behavior is acceptable, could you at least document it in the Javadoc and the PR description? ########## 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: Since the fast path now returns long-backed decimals, `roundToByte/Short/Int` can give different results from `ArrowColumnVector`, which keeps BigDecimal-backed decimals (SPARK-59443) and decides those conversions on `toDouble`. For example, DECIMAL(18,9) `123456789.999999999`: `ArrowColumnVector` gives `roundToInt` = 123456790 and NativeColumnVector gives 123456789; DECIMAL(18,15) `127.999999999999999`: `roundToByte` fails with CAST_OVERFLOW through `ArrowColumnVector` and returns 127 here. The new results match Spark's own column vectors, so this looks like the right behavior to me, but `checkDecimal` compares only `toJavaBigDecimal`, precision, scale, `toLong` and the UnsafeRow, so it neither shows nor pins this difference. Could you add the `roundTo*` results to the check (as the expected values of Spark's own vectors), and mention the difference in the description, where it says the decimals are tested "with the checks of ArrowColumnVectorSuite"? ########## 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: Each batch here ends up with at most one data buffer (I checked `getDataBuffers.size` for the three batches: 0, 1, 1), so the buffer index is never non-zero. Replacing `buffers[Platform.getInt(null, view + 8)]` with `buffers[0]` in `StringViews.getUTF8String` still passes the suite. Could you add a batch whose long values span more than one data buffer, and assert `getDataBuffers.size > 1` so the case stays covered? -- 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]
