dongjoon-hyun commented on code in PR #59961:
URL: https://github.com/apache/spark/pull/59961#discussion_r4488238390
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -1053,14 +1109,25 @@ class NativePartitionReader(
if (task != null) {
task.close()
}
- Utils.tryLogNonFatalError(allocator.close())
+ if (freeBatch) {
+ Utils.tryLogNonFatalError(allocator.close())
+ } else {
+ // The action must not refer to the reader, which would then never
become unreachable.
+ val (batchRoot, batchAllocator) = (root, allocator)
+ NativePartitionReader.cleaner.register(this, () =>
Review Comment:
**[9] A reader closed before the end of its stream frees its last batch only
at a GC that the batch itself never triggers, and these batches add up across
tasks.** Thank you for the rework in d56b02d. Following up on the trade-off you
asked about in
[r4483289144](https://github.com/apache/spark/pull/59961#discussion_r4483289144),
and on @sarutak's
[r4445437083](https://github.com/apache/spark/pull/59961#discussion_r4445437083):
the class doc (L656-657) and the PR description bound this to one batch per
reader, but not per executor. Every task whose reader is closed with
`batchReturned` set registers one more action here, and the buffers that the
library allocated for all of them stay until those readers are collected. A
LIMIT, a sort-merge join that stops reading its streamed side, and a failed or
killed task all close the reader this way. In a long task, the reader, its
`ColumnarBatch` and its column vectors have often been promoted, as @sarutak
noted, so they wait for an old-gen
eration collection, which a large heap with a low allocation rate may not run
for a long time. Since the JVM does not count this memory, the RSS of the
executor can grow until YARN or Kubernetes kills the container for exceeding
its memory overhead, while the heap looks fine.
Could you share how much this can hold in practice, e.g. for
`df.limit(n).write` over many partitions of a source with large batches? One
way to make the multi-threaded case deterministic: a `next()` after `close()`
(L891-895) means that its thread is done with the batch that it got, so it
could free the batch there before it throws, e.g. with the `Cleanable` that
`register` returns, just as `close()` frees at once when `next()` was called
again. That leaves the GC only for readers that no thread calls again. Item 14
asks about a more general alternative.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -847,20 +875,35 @@ class NativePartitionReader(
}
override def next(): Boolean = {
- // A reader that was closed while it read has released the batch that it
read.
- startReading() && Utils.tryWithSafeFinally(readNext())(stopReading()) &&
synchronized(!closed)
+ startReading()
+ val hasBatch = Utils.tryWithSafeFinally(readNext())(stopReading())
+ synchronized {
+ // A reader that was closed while it read has released the batch that it
read.
+ batchReturned = hasBatch && !closed
Review Comment:
**[10] A `next()` that is inside `get_next` when another thread closes the
reader still returns false, so that thread takes the close for the end of the
data.** L643-645 says that a `next()` after the close fails "so that such a
thread does not take the close for the end of the data: script transformation,
for example, then destroys its process". But this line returns false for a
`next()` that was already in `get_next` when `close()` ran. After the first
batch, the writer thread of script transformation is the one that calls
`hasNext`. So if the task completes while that thread is inside `next()`, it
ends `processRows()` normally, closes the stdin of the script and waits in
`proc.waitFor()` (`BaseScriptTransformationExec.scala:349-350`). If it is
between two calls instead, it gets the `IllegalStateException` and destroys the
process (`BaseScriptTransformationExec.scala:344`). So the outcome depends on
timing. Could this `next()` also throw once `stopReading()` has released the rea
der, e.g. `if (closed) throw ...` here rather than `hasBatch && !closed`?
Nit, on the same path: the `IllegalStateException` of L893 has no error
condition, unlike the other failures of the reader (`nativeDataSourceError`,
`error(name, ...)`). Now that every thread that outlives its task gets it, and
script transformation logs it at ERROR (`Feed exit cause by:`), could it use an
error condition, or `SparkException.internalError` as `NativeHandle.get` does
(L1435)?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -155,11 +157,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val inputColumnIndices = resolveColumnIndices(allInputs.toSeq)
- // Peek at first batch to check if Arrow-backed.
+ // Peek at first batch to check if Arrow-backed. Path 1 holds the
pass-through columns past
+ // the next batch, which a ReusedArrowColumnVector does not allow: its
producer overwrites or
+ // releases it, so such batches take path 2, which copies them.
val peekIter = new PeekableIterator(inputIter)
val isArrow = peekIter.peek().exists { batch =>
batch.numCols() > 0 &&
- batch.column(0).isInstanceOf[ArrowColumnVector]
+ batch.column(0).isInstanceOf[ArrowColumnVector] &&
+ !(0 until batch.numCols()).exists(
+ i => batch.column(i).isInstanceOf[ReusedArrowColumnVector])
Review Comment:
**[12] No test reaches this new check with a `ReusedArrowColumnVector`, and
it also moves the outputs of Python UDFs to path 2.** With
`spark.sql.execution.arrow.pythonUDF.columnarInput.enabled=true`, this sends to
path 2 every batch with a `ReusedArrowColumnVector`, which includes the result
columns of an upstream Python UDF: `ArrowOutputProcessorImpl` builds them as
`ReusedArrowColumnVector`s (`PythonArrowOutput.scala:186`, `:254`). That looks
right for them too, but the PR description only mentions native data sources.
As far as I can see, no test gets here with such a batch: 'Arrow UDFs on the
columns of a native data source' (`NativeDataSourceSuite.scala:648`) runs with
the default configuration, which keeps the plan row-based, and the sources of
`ArrowColumnarPythonUDFSuite` build plain `ArrowColumnVector`s
(`ArrowBackedDataSourceV2.scala:221`, `:362-366`). Could the native test also
run with the configuration enabled, including the 200000-row case, and could
the PR descript
ion mention the outputs of Python UDFs?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVectorSuite.scala:
##########
@@ -0,0 +1,416 @@
+/*
+ * 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, null)
+ column.refresh(-1)
+ 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, null)
+ 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(-1)
+ 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)
Review Comment:
**[11] This PR changes the results of queries over top-level decimal
columns, but the PR description says that it has no user-facing change.** Thank
you for pinning the nested result, as you planned in
[r4470203152](https://github.com/apache/spark/pull/59961#discussion_r4470203152).
The other half of this test is a change: before this PR, a top-level
DECIMAL(18,9) column read through `ArrowColumnVector` was BigDecimal-backed as
well. So with ANSI on, the default, `CAST(d AS INT)` of `123456789.999999999`
returned 123456790, and now it returns 123456789, as Spark's own sources do.
Likewise, `CAST(d AS TINYINT)` of a DECIMAL(18,15) `127.999999999999999` failed
with `CAST_OVERFLOW`, and now returns 127. Could the user-facing change section
of the PR description describe this, together with the nested decimals that
keep the old results until the follow-up? Could this test, or the class doc of
`NativeColumnVector` (L65-68), also name the JIRA of that follow-up?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -1053,14 +1109,25 @@ class NativePartitionReader(
if (task != null) {
task.close()
}
- Utils.tryLogNonFatalError(allocator.close())
+ if (freeBatch) {
+ Utils.tryLogNonFatalError(allocator.close())
+ } else {
+ // The action must not refer to the reader, which would then never
become unreachable.
+ val (batchRoot, batchAllocator) = (root, allocator)
+ NativePartitionReader.cleaner.register(this, () =>
+ Utils.tryLogNonFatalError(
+
Utils.tryWithSafeFinally(batchRoot.close())(batchAllocator.close())))
+ }
}
}
}
object NativePartitionReader {
import NativeCalls._
+ // Frees the buffers of the batches that may still be read when their
readers are closed.
+ private val cleaner = Cleaner.create()
Review Comment:
**[17] Nit: this adds a second `Cleaner`, with its own thread, next to
`NativeHandle.cleaner` in the same file (L1444).** Each `Cleaner.create()`
starts a daemon thread, and the two use different idioms for the rule that the
action must not refer to its owner: a `State` that implements `Runnable` there
(L1446-1447), and a lambda over locals here (L1115-1116). Could both share one,
e.g. by making `NativeHandle.cleaner` `private[ffi]` or moving it to a shared
object?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSource.scala:
##########
@@ -1014,32 +1062,40 @@ class NativePartitionReader(
"but Spark only reads batches without null rows: mark the values of
the columns as null " +
"instead.")
}
- val columns = root.getSchema.getFields.size
- children(batch, columns, "The batch", s"the schema has $columns columns")
- .zip(root.getSchema.getFields.asScala).foreach { case (column, field)
=>
+ val numColumns = root.getSchema.getFields.size
+ val arrays = children(batch, numColumns, "The batch", s"the schema has
$numColumns columns")
+ arrays.zip(root.getSchema.getFields.asScala).foreach { case (column,
field) =>
if (column.length < batch.length) {
throw error(name, "read", s"The column ${field.getName} has a length
of " +
s"${column.length}, but the batch has a length of ${batch.length}.")
}
checkColumn(column, field, field.getName)
}
+ arrays.map(_.null_count).toArray
}
override def get(): ColumnarBatch = batch
override def close(): Unit = synchronized {
if (!closed) {
closed = true
+ // A thread in next releases the reader once get_next returns, see
stopReading. Otherwise, a
+ // consumer may still read the batch that next returned, unless next was
called since.
if (!reading) {
- closeResources()
+ closeResources(freeBatch = !batchReturned)
Review Comment:
**[14] Design question: could the consumers that outlive their task stop
reading before the reader is closed, rather than the reader guessing whether
its batch is still in use?** The Cleaner exists for threads that read a batch
after the task completed, such as the writer thread of script transformation
and the pipelined Python writer. But the reader cannot tell them from the
common case where nothing reads the batch anymore, e.g. a single-threaded
LIMIT, so every such close now waits for a GC (item 9). The hazard is not
specific to native data sources either: with
`spark.sql.columnVector.offheap.enabled`, the vectorized Parquet reader frees
its `OffHeapColumnVector`s in `close()`
(`VectorizedParquetRecordReader.java:234`), and such a thread can read freed
memory there too. Task completion listeners run in reverse order of
registration, so a listener that such a consumer registers when it starts its
thread runs before the one of `DataSourceRDD` that closes the reader. If that
list
ener destroyed the script, or stopped the writer, and joined the thread, every
columnar source would be safe, and this reader could free its last batch at
once again. Did you consider this? It could also be a follow-up that removes
the Cleaner later; I don't want to block on it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/ffi/NativeDataSourceSuite.scala:
##########
@@ -141,6 +156,31 @@ class NativeDataSourceSuite extends QueryTest with
SharedSparkSession {
}.get
}
+ /** The allocators of the readers of native data sources that are not freed
yet. */
+ private def readerAllocators(): Set[BufferAllocator] = {
+ ArrowUtils.rootAllocator.getChildAllocators.asScala
+ .filter(_.getName.startsWith("native data source ")).toSet
+ }
+
+ /**
+ * Creates the reader of a partition of the batch scan of `df`, and returns
it with the allocator
+ * of its buffers, which the reader closes once it frees them.
+ */
+ private def readerAndAllocator(
+ df: DataFrame,
+ partition: Int = 0): (PartitionReader[ColumnarBatch], BufferAllocator) =
{
+ val batch = scan(df).toBatch
+ val (factory, partitions) = (batch.createReaderFactory(),
batch.planInputPartitions())
+ val before = readerAllocators()
+ val reader = factory.createColumnarReader(partitions(partition))
+ val created = readerAllocators() -- before
Review Comment:
**[13] Nit: `readerAndAllocator` and `afterAll` look at every reader
allocator of the JVM, so they can fail for a reader of another test.** This
line takes the allocators that appeared while L175 ran, across the whole root
allocator. So a reader that a streaming query or a leftover task of an earlier
test creates at that moment fails this test with "The reader has 2 allocators."
And `afterAll` (L132-136) fails the whole suite after a minute if any test
keeps a reference to a reader, or never closes one, which then registers no
Cleaner. These allocators all have the same name, `native data source
native_range`, so the failure does not say which test left one. Could the check
run per test instead, e.g. in `afterEach` or in the tests that close a reader
early, so that a leak fails the test that made it?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,331 @@
+/*
+ * 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(long)}, 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 while the batch may be read: until
next is called again,
+ * and if the reader is closed before then, until the reader is unreachable,
which the vector keeps
+ * it from being. 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. The
vector of a nullable
+ * column takes the null_count that the library exported, and only counts the
nulls when the
+ * library did not.
+ *
+ * <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 reader, which the vector keeps reachable: a reader closed while its
batch may still be
+ // read frees the buffers of the batch only once it is unreachable.
+ private Object reader;
+
+ /**
+ * The column vector of a top-level Arrow vector of {@code reader}, of a
column that the schema to
+ * read declares {@code nullable} or not.
+ */
+ static NativeColumnVector create(ValueVector vector, boolean nullable,
Object reader) {
+ NativeColumnVector column;
+ if (vector instanceof BitVector) {
+ column = new Booleans(vector, nullable);
+ } else if (vector instanceof TinyIntVector) {
+ column = new Bytes(vector, nullable);
+ } else if (vector instanceof SmallIntVector) {
+ column = new Shorts(vector, nullable);
+ } else if (vector instanceof IntVector || vector instanceof DateDayVector)
{
+ column = new Ints(vector, nullable);
+ } else if (vector instanceof BigIntVector || vector instanceof
TimeStampMicroTZVector ||
+ vector instanceof TimeStampMicroVector) {
+ column = new Longs(vector, nullable);
+ } else if (vector instanceof Float4Vector) {
+ column = new Floats(vector, nullable);
+ } else if (vector instanceof Float8Vector) {
+ column = new Doubles(vector, nullable);
+ } else if (vector instanceof DecimalVector decimalVector) {
+ column = new Decimals(decimalVector, nullable);
+ } else if (vector instanceof VarCharVector) {
+ column = new Strings(vector, nullable);
+ } else if (vector instanceof ViewVarCharVector) {
+ column = new StringViews(vector, nullable);
+ } else {
+ column = new NativeColumnVector(vector);
+ }
+ column.reader = reader;
+ return column;
+ }
+
+ NativeColumnVector(ValueVector vector) {
+ super(vector);
+ }
+
+ /**
+ * Reads the addresses of the buffers of the batch that the reader imported
last, whose array of
+ * this column the library exported with a null_count of {@code nullCount}.
+ */
+ void refresh(long nullCount) {
+ }
+
+ /** A vector that reads its values and their validity at the addresses of
its buffers. */
+ private abstract static class Direct extends NativeColumnVector {
+ // 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.nullable = nullable;
+ }
+
+ @Override
+ void refresh(long nullCount) {
+ ValueVector vector = getValueVector();
+ // The C data interface requires a null_count other than -1 to be exact,
so the nulls are
+ // only counted, which Arrow Java does in the validity bitmap on each
call, for a negative
+ // one.
+ boolean hasNulls = nullable && (nullCount < 0 ? vector.getNullCount() >
0 : nullCount > 0);
+ validity = hasNulls ? vector.getValidityBuffer().memoryAddress() : 0;
+ data = vector.getDataBuffer().memoryAddress();
Review Comment:
**[16] Nit: the cached addresses assume that nothing changes the buffers of
these vectors between two imports, which the class doc does not say.**
`ArrowColumnVector` reads the buffer fields of its vector on each access, so it
followed a consumer that reallocated them, while `refresh()` keeps their
addresses until the next import. The vectors are still handed out through
`getValueVector()`, and with `columnarInput.enabled`, path 2 of
`ColumnarArrowEvalPythonEvaluatorFactory` now writes them with
`ColumnarArrowPythonInput.writeArrowDirect`, whose
`VectorSchemaRoot.of(...).setRowCount(batch.numRows())` calls `setValueCount`
on the reader's own vectors. That does not reallocate today, since the row
count is the length of the imported arrays, but it would be easy to break.
Could the class doc state this assumption, next to the one about the buffer
lengths (L55-57)?
##########
sql/core/src/main/java/org/apache/spark/sql/execution/datasources/v2/ffi/NativeColumnVector.java:
##########
@@ -0,0 +1,331 @@
+/*
+ * 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(long)}, 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 while the batch may be read: until
next is called again,
+ * and if the reader is closed before then, until the reader is unreachable,
which the vector keeps
+ * it from being. 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. The
vector of a nullable
+ * column takes the null_count that the library exported, and only counts the
nulls when the
+ * library did not.
+ *
+ * <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 {
Review Comment:
**[15] Design question: could the direct reads live in
`ReusedArrowColumnVector`, rather than in a subclass for native data sources
only?** `ArrowOutputProcessorImpl` also loads every batch into the same vectors
and reads them through `ReusedArrowColumnVector`
(`PythonArrowOutput.scala:175-187`), so the results of Python UDFs keep the
checked reads. A refresh hook on `ReusedArrowColumnVector`, which each producer
calls after it loads a batch, would give both producers these reads, and any
later one, as long as it keeps the buffers alive as this reader does. It would
also keep a single copy of the Arrow layout decoding: `Decimals` (L235-266) and
`StringViews` (L293-330) repeat the word order, the sign check and the view
decoding of `ArrowColumnVector.SmallDecimalAccessor` and `StringViewAccessor`
with raw reads, so a later fix there would not reach this class. A follow-up is
fine; I'd like to hear what you think.
--
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]
