dongjoon-hyun commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4462405213
##########
sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnVector.java:
##########
@@ -84,6 +84,33 @@ public void closeIfFreeable() {
close();
}
+ /**
+ * Returns whether a consumer may keep this column vector after it asks the
producer of its batch
+ * for the next batch, until the consumer releases it with {@link
#closeIfFreeable()}. By default,
+ * a column vector may be reused or released by its producer once the next
batch is requested, so
+ * a consumer has to copy what it needs past that point.
+ * <p>
+ * A producer that returns new vectors for every batch can override this to
return true, which
+ * promises that:
+ * <ul>
+ * <li>the producer neither overwrites nor frees this vector when the next
batch is requested,
+ * or when the producer is closed, e.g. a {@code PartitionReader} once it
is exhausted;</li>
+ * <li>{@link #closeIfFreeable()} frees this vector, and may be called
from another thread
Review Comment:
**18. [Medium] The contract does not say that `closeIfFreeable` can run
while the producer returns the next batches**
With `spark.python.udf.pipelined.enabled`, the writer thread pulls the
producer through `ColumnarArrowPythonInput.scala` L65-L66, which calls its
`next()` and `get()`, and `close()` once it is exhausted (`DataSourceRDD.scala`
L273-L279). Meanwhile, the task thread releases a held batch downstream
(`Columnar.scala` L229-L231). Nothing orders the two (`PythonRunner.scala`
L1620: "without any inter-thread queues or locks"). So a retainable vector is
released on another thread while the producer returns later batches or is
closed, which "may be called from another thread than the one that requested
its batch" does not say.
The third bullet makes this matter. To free what the consumer has not
released when the task completes, a producer has to track its outstanding
vectors. A connector that adds them to a plain `ArrayList` or `HashSet` in
`next()` and removes them in `closeIfFreeable` loses entries when the two race,
and after an early stop `allocator.close()` reports "Memory was leaked". The
test source needs `allBatches.synchronized` in both places for exactly this
(`ArrowBackedDataSourceV2.scala` L345-L348 and L416-L421), and the class doc of
the evaluator states the requirement
(`ColumnarArrowEvalPythonEvaluatorFactory.scala` L56-L59), but the public
Javadoc does not. Nothing in the tree is affected, since only the test source
opts in and Arrow allocators are thread-safe.
Could this bullet say something like "... than the one that requested its
batch, concurrently with the producer returning the next batches or being
closed"?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +221,22 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
outputTypes: Seq[DataType],
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
+ // In pipelined mode the batches are added in the writer thread and
+ // polled in the task thread.
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
- val bufferedIter = inputIter.map { batch =>
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, { batch
=>
val passThruCols = childOutput.indices.map(
i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
- batch
- }
+ if (passThruCols.forall(_.isRetainable)) {
Review Comment:
**20. [Low] Nothing pins the later-batch check or the release of a copied
batch**
Only "Reused vectors after fresh ones" reaches the copy branch. The batches
it copies come from `ReusingArrowPartitionReader`, whose
`ReusedArrowColumnVector.closeIfFreeable` is a no-op and is not counted
(`ArrowBackedDataSourceV2.scala` L238-L239), and it asserts the counters of the
2 fresh batches only (L287-L288 of the suite). Also, every batch in the suite
has pass-through columns that are either all retainable or all not. So both of
these pass every test:
- `true` at L237 changed to `false`. A copied batch is then never released.
With a retainable first batch, later batches that are fresh but not retainable
and rely on the consumer's `closeIfFreeable`, such as a plain
`ArrowColumnVector` per batch, stay allocated until the task ends.
- `forall` on this line changed to `exists`, or to a check of column 0 only.
A later batch that mixes retainable columns with a reused one is then held
without a copy, and the reused column shows the values of a later batch next to
the UDF results of this one, with no error, which is [the failure of item
2](https://github.com/apache/spark/pull/59582#discussion_r4437566761).
Could the switching test count the releases of the copied batches too, for
example with a second reader that is an `ArrowBackedPartitionReader` with
`retainable = false`, or by counting the calls to the no-op `closeIfFreeable`
of the reusing reader, which keeps the in-place overwrite? One later batch that
mixes a retainable column with a reused one would pin this line.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -155,19 +162,29 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val inputColumnIndices = resolveColumnIndices(allInputs.toSeq)
- // Peek at first batch to check if Arrow-backed.
- val peekIter = new PeekableIterator(inputIter)
- val isArrow = peekIter.peek().exists { batch =>
- batch.numCols() > 0 &&
- batch.column(0).isInstanceOf[ArrowColumnVector]
+ val nonEmptyIter = inputIter.filter { batch =>
Review Comment:
**22. [Low] The skip and release of zero-row batches has no test**
This fixes [item
12](https://github.com/apache/spark/pull/59582#discussion_r4437566812), but no
test feeds a batch without rows. Both test readers return `false` before they
would build one (`if (current >= end) return false` comes before `val count =
math.min(batchSize, end - current)`, `ArrowBackedDataSourceV2.scala` L242-L243
and L393-L395), and the cached data is built from rows. So replacing this
filter with `val nonEmptyIter = inputIter`, or removing only the
`closeIfFreeable` at L166, passes every test:
- The first brings item 12 back for a source that returns an empty batch,
such as `NativeDataSource` (`batch.setNumRows(root.getRowCount)` at L725 of
`NativeDataSource.scala`). The worker calls the UDF for every batch
(`_pandas.py` L93 and L110), so `s.iloc[0]` raises `IndexError`.
- The second keeps a retainable empty batch until the task ends.
Could the test source get an option that returns empty batches, first, in
the middle and last, with checks that they are released and that no empty batch
reaches the output? The test UDF, `x.apply(...)`, does not fail on empty input,
so the results alone would not show it.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -130,6 +158,203 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60325: Arrow-backed source that reuses its vectors: Arrow
columnar path not taken") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_reuse_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // The Arrow columnar path would keep the pass-through columns of a
batch until its result
+ // comes back from Python. The reader refills its vectors for the next
batch and closes them
+ // once it is closed, as soon as it is exhausted, so they are not
retainable, and the
+ // batches take the row queue path. With two batches, the reader is
closed before their
+ // results come back. With many, it refills its vectors with the next
batches before then.
+ // The test UDF casts its input to string: on a string column, the cast
is removed, so the
+ // inputs are column references.
+ Seq(200, 20000).foreach { numRows =>
+ val result = readArrowSource(numRows = numRows, reuseVectors = true,
batchSize = 100)
+ .selectExpr("id", "name", "value", "data", "arrow_reuse_udf(name) as
udf_name")
+ val arrowExec =
singleArrowEvalPythonExec(result.queryExecution.executedPlan)
+ assert(!takesArrowColumnarPath(arrowExec))
+
+ val rows = result.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ rows.foreach { row =>
+ val id = row.getInt(0)
+ assert(row.getString(1) == s"row_$id", s"name mismatch for id $id")
+ assert(row.getString(3) == s"data_$id", s"data mismatch for id $id")
+ assert(row.getString(4) == s"row_$id", s"UDF result mismatch for id
$id")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60325: Arrow-backed source with fresh vectors: held without a
copy and released") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_fresh_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ val numRows = 2000
+ val numBatches = 20
+ val result = readArrowSource(numRows = numRows, batchSize = numRows /
numBatches)
+ .selectExpr("id", "name", "value", "data", "arrow_fresh_udf(name) as
udf_name")
+ // The pass-through columns are output as the reader returned them.
+ assert(takesArrowColumnarPath(
+ singleArrowEvalPythonExec(result.queryExecution.executedPlan)))
+
+ ArrowBackedDataSourceV2.resetBatchCounters()
+ val rows = result.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ rows.foreach { row =>
+ val id = row.getInt(0)
+ assert(row.getString(1) == s"row_$id", s"name mismatch for id $id")
+ assert(row.getDouble(2) == id * 0.1, s"value mismatch for id $id")
+ assert(row.getString(3).length == 100, s"data mismatch for id $id")
+ assert(row.getString(4) == s"row_$id", s"UDF result mismatch for id
$id")
+ }
+ // The consumer of the output releases each input batch once its rows
are consumed, before
+ // the task completes. How many are held at a time depends on how far
the writer runs
+ // ahead of the results, which may be all of them for a small input.
+ assert(ArrowBackedDataSourceV2.returnedBatches.get == numBatches)
+ assert(ArrowBackedDataSourceV2.releasedBatches.get == numBatches)
Review Comment:
**23. [Low] The limit check does not run the UDF, so no test stops early
over held batches**
The third promise of `isRetainable` has the producer free what the consumer
has not released, "e.g. when the consumer stops early", but no test stops Path
1 early. The only limit check, `result.limit(10).selectExpr("count(*)")` in the
SPARK-60306 test
([L432-L434](https://github.com/apache/spark/blob/cfce388a0f5a4b90090d8dd00b1195d72152743c/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala#L432-L434)),
references no column of `result`, and the test UDF is deterministic
(`udfDeterministic = true`, `IntegratedUDFTestUtils.scala` L1392). So
`ColumnPruning` drops `udf_out` (`Optimizer.scala` L1544 for the aggregate,
L1618 through the limits and L1520 for the project) in the operator
optimizations of `super.defaultBatches` (`SparkOptimizer.scala` L67), before
`ExtractPythonUDFs` runs (L101), and the plan has no `ArrowEvalPythonExec`. Its
comment, "A limit stops consuming the output batches early, and the generated
code of ColumnarToRowExec
closes the current one", then does not apply to the Python UDF.
That check comes from SPARK-60306, but the contract of this PR relies on it
now. I found this by reading the optimizer rules and did not run it, so
printing that plan would confirm it. Could the check use
`result.limit(10).collect()`, or an aggregate over `udf_out`, and could this
test check the counters after such an early stop?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -48,23 +65,34 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
numRows: Int = 100,
numPartitions: Int = 1,
reuseVectors: Boolean = false,
+ freshBatches: Int = 0,
+ retainable: Boolean = true,
batchSize: Int = ArrowBackedDataSourceV2.BATCH_SIZE,
arrowIdColumn: Boolean = true) = {
spark.read
.format(arrowSource)
.option("numRows", numRows.toString)
.option("numPartitions", numPartitions.toString)
.option("reuseVectors", reuseVectors.toString)
+ .option("freshBatches", freshBatches.toString)
+ .option("retainable", retainable.toString)
.option("batchSize", batchSize.toString)
.option("arrowIdColumn", arrowIdColumn.toString)
.load()
}
+ private def singleArrowEvalPythonExec(plan: SparkPlan): ArrowEvalPythonExec
= {
+ val execs = collectNodes[ArrowEvalPythonExec](plan)
+ assert(execs.length == 1, s"Expected one ArrowEvalPythonExec in
plan:\n$plan")
+ execs.head
+ }
+
/**
- * Whether `exec` takes the Arrow columnar path, which outputs its input
columns as they are,
+ * Whether `exec` takes the Arrow columnar path, which outputs the input
columns as they are,
* while the other paths output on-heap vectors.
*/
private def takesArrowColumnarPath(exec: ArrowEvalPythonExec): Boolean = {
+ assert(exec.supportsColumnar)
val columnClasses = exec.executeColumnar().mapPartitionsInternal { batches
=>
batches.map(_.column(0).getClass.getName)
Review Comment:
**19. [Medium] `takesArrowColumnarPath` cannot tell a copying Path 1 from
the row queue path**
The helper decides by the class of column 0, but Path 1 now also puts an
`OnHeapColumnVector` there when it copies a batch
(`passThruQueue.add((copyColumns(batch), batch.numRows()))` at L236 of the
evaluator). So `!takesArrowColumnarPath(...)` only shows that no batch was
held, not that the row queue path was taken, and no test pins the first-batch
check.
For example, change L176 of the evaluator to `case v: ArrowColumnVector =>
true`. The reusing reader (L179), the reader that is not retainable (L240, and
its counts at L251-L252, since a copied batch is released once it has been
written), the Arrow cache (L310) and the stacked outer exec (L345) then all
take Path 1 and copy every batch. Column 0 is on-heap, and the values and
counts stay the same, so both suites and `test_arrow_python_udf_cached.py`
pass. In production, that change would send the Arrow cache, stacked Python
output and `ReusedArrowColumnVector` sources to Path 1. Every batch would then
be copied into `passThruQueue`, which is neither accounted for nor spillable,
instead of the `HybridRowQueue` of the row queue path. The description's
"Making both `isRetainable` checks always true fails five of the tests above"
holds for the two checks together, but not for this one alone. I did not run
this mutation; it comes from tracing each test.
Could the helper decide by the last column instead, the UDF result? It is a
`ReusedArrowColumnVector` from the Python output on Path 1, copied or not, and
an `OnHeapColumnVector` on paths 2 and 3. The fresh-vector test could keep a
check of column 0, so that "without a copy" stays pinned too.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +342,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
}
+ /** Copies the pass-through columns of `batch` into on-heap vectors. */
+ private def copyColumns(batch: ColumnarBatch): Array[ColumnVector] = {
+ val vectors = OnHeapColumnVector.allocateColumns(
+ math.max(batch.numRows(), 1),
childSchema).toArray[WritableColumnVector]
+ val converter = new RowToColumnConverter(childSchema)
+ batch.rowIterator().asScala.foreach(converter.convert(_, vectors))
Review Comment:
**25. [Low] `copyColumns` copies the rows of `rowIterator` over the whole
width of the batch**
The writer sends `batch.numRows()` rows of the selected columns
(`ColumnarArrowPythonInput.scala` L114-L120, and `selectColumns` at L160-L164
for the row-by-row case), and L236 records `batch.numRows()`. This line,
however, copies whatever `batch.rowIterator()` returns, over every column of
the batch:
- `RowToColumnConverter.convert` loops over `row.numFields`, the width of
the batch (`Columnar.scala` L254-L255), while `converters` and `vectors` have
the width of `childSchema`. A later batch with one more column than
`childOutput` fails with `ArrayIndexOutOfBoundsException` here only; the held
branch, paths 2 and 3, and `ColumnarToRowExec` ignore the extra column.
- A data source may extend `ColumnarBatch` (`ColumnarBatch.java` L27), and
`numRows()` includes filtered rows (L96). If the `rowIterator` of a later batch
skips rows 0 and 2 of 4, the copy holds rows 1 and 3 and two empty slots, which
are paired with the UDF results of rows 0 to 3, with no error.
Both need a third-party source that opts in, so this is minor. Copying `new
ColumnarBatch(passThruCols, batch.numRows())`, as `selectColumns` does in the
writer, avoids both. In the same method, the converter could be built once next
to `childSchema`, and `math.max(batch.numRows(), 1)` at L348 can no longer take
its second branch, since `nonEmptyIter` drops the empty batches.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -130,6 +158,203 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60325: Arrow-backed source that reuses its vectors: Arrow
columnar path not taken") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_reuse_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // The Arrow columnar path would keep the pass-through columns of a
batch until its result
+ // comes back from Python. The reader refills its vectors for the next
batch and closes them
+ // once it is closed, as soon as it is exhausted, so they are not
retainable, and the
+ // batches take the row queue path. With two batches, the reader is
closed before their
+ // results come back. With many, it refills its vectors with the next
batches before then.
+ // The test UDF casts its input to string: on a string column, the cast
is removed, so the
+ // inputs are column references.
+ Seq(200, 20000).foreach { numRows =>
+ val result = readArrowSource(numRows = numRows, reuseVectors = true,
batchSize = 100)
+ .selectExpr("id", "name", "value", "data", "arrow_reuse_udf(name) as
udf_name")
+ val arrowExec =
singleArrowEvalPythonExec(result.queryExecution.executedPlan)
+ assert(!takesArrowColumnarPath(arrowExec))
Review Comment:
**28. [Nit] The restored default-config test gives the reason for the
default that this PR removes**
The test "Arrow-backed source that reuses its vectors: correct with the
default configuration", which this PR restores as on master, still says that
reused vectors are why columnar input is disabled by default, and that its
batches "would take the Arrow columnar path if columnar input were enabled"
([L463-L466](https://github.com/apache/spark/blob/cfce388a0f5a4b90090d8dd00b1195d72152743c/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala#L463-L466)
and
[L475-L479](https://github.com/apache/spark/blob/cfce388a0f5a4b90090d8dd00b1195d72152743c/sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala#L475-L479)).
This line now asserts the opposite for the same reader and UDF with columnar
input enabled, and [your
comment](https://github.com/apache/spark/pull/59582#issuecomment-6616652366)
gives another reason for the default: nothing in the tree opts in, and items 5
to 11 remain. Could those comments s
ay only that the default reads the scan as rows?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5676,10 +5676,10 @@ object SQLConf {
"Arrow-backed ColumnarBatch (e.g., DataSource V2 connectors), " +
"bypassing the ColumnarToRow and ArrowWriter conversion. " +
Review Comment:
**27. [Low] The conf doc still describes the feature as an optimization that
bypasses ColumnarToRow**
With the opt-in, L5676-L5678 ("Arrow-backed ColumnarBatch (e.g., DataSource
V2 connectors), bypassing the ColumnarToRow and ArrowWriter conversion. This
optimization reduces data transfer overhead ...") describe what happens only
for vectors that override a method new in 4.4.0, which no built-in operator and
no existing connector does. For every other Arrow child, the Arrow cache
included, enabling the conf sends each batch through the row queue path. That
path projects and copies each row into `HybridRowQueue`, joins it with the
result and appends it to on-heap vectors, and the `ColumnarToRowExec` above
converts it back to rows. By my count from the code, which I did not measure,
that is about 6 copies of a pass-through string against 3 or 4 with the conf
off. The UDF inputs are still written as Arrow directly, but the doc does not
say that the conf can cost more than leaving it disabled.
Could it say that the conf helps only when the upstream vectors return
`true` from `isRetainable`, which no built-in operator does, and that other
batches are converted to rows and back? The cost itself is [item
11](https://github.com/apache/spark/pull/59582#discussion_r4437566808), so this
is only about the doc.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnVector.java:
##########
@@ -84,6 +84,33 @@ public void closeIfFreeable() {
close();
}
+ /**
+ * Returns whether a consumer may keep this column vector after it asks the
producer of its batch
+ * for the next batch, until the consumer releases it with {@link
#closeIfFreeable()}. By default,
+ * a column vector may be reused or released by its producer once the next
batch is requested, so
+ * a consumer has to copy what it needs past that point.
+ * <p>
+ * A producer that returns new vectors for every batch can override this to
return true, which
+ * promises that:
+ * <ul>
+ * <li>the producer neither overwrites nor frees this vector when the next
batch is requested,
+ * or when the producer is closed, e.g. a {@code PartitionReader} once it
is exhausted;</li>
+ * <li>{@link #closeIfFreeable()} frees this vector, and may be called
from another thread
+ * than the one that requested its batch;</li>
+ * <li>the producer frees the vectors that the consumer has not released
when the task
+ * completes, e.g. when the consumer stops early.</li>
+ * </ul>
+ * A consumer that keeps such vectors does not account for their memory, and
cannot spill them.
+ * <p>
+ * This applies to the top-level columns of a batch: the children of a
struct, array or map
+ * column share the lifetime of that column.
+ *
+ * @since 4.4.0
+ */
+ public boolean isRetainable() {
Review Comment:
**31. [Question] Should `isRetainable` be public before a producer uses it?**
I asked for a public way to opt in in [item
4](https://github.com/apache/spark/pull/59582#discussion_r4437566771), when the
default was being turned on. Now the default stays `false` and nothing in the
tree returns `true` here. Enabling columnar input sends the Arrow cache,
`NativeDataSource` and stacked UDFs to the row queue path, as if every simple
column reference went to `evalWithRowQueue`, and Path 1 and `copyColumns` run
only for the test source. The three promises then become a public contract that
no shipped code exercises yet. Would keeping the opt-in internal until its
first producer exists be an option, or is a connector waiting for it? You asked
HyukjinKwon about the API, so I am fine either way. Also, L103 states a
limitation of this consumer as a fact about every consumer; "may keep such
vectors without accounting for their memory or spilling them" would stay true
if a consumer accounts for them later.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5676,10 +5676,10 @@ object SQLConf {
"Arrow-backed ColumnarBatch (e.g., DataSource V2 connectors), " +
"bypassing the ColumnarToRow and ArrowWriter conversion. " +
"This optimization reduces data transfer overhead between " +
- "the JVM and Python worker processes. It keeps the columns that the
UDFs pass " +
- "through until the results of their batch come back from Python, so it
requires " +
- "an upstream operator whose batches stay valid until then, which a
columnar " +
- "reader does not promise.")
+ "the JVM and Python worker processes. The columns that the UDFs pass
through " +
+ "are kept until the results of their batch come back from Python only
if the " +
+ "upstream operator allows it with ColumnVector.isRetainable, and their
memory is " +
+ "not accounted for until then. Other columns are read as rows or
copied.")
Review Comment:
**26. [Low] The docs do not say that the copies on Path 1 are held without
accounting too**
This says that the held columns are "not accounted for until then" and that
"Other columns are read as rows or copied", and `ColumnVector.java` L103 says
the same of the retainable vectors. But on Path 1, a later batch that cannot be
held is copied into on-heap vectors (L236 and L346-L348 of the evaluator). The
copies stay in `passThruQueue` the same way until their results come back,
outside `TaskMemoryManager` and not spillable. Had the same batches come first
in the partition, they would have gone to the `HybridRowQueue` of the row queue
path (L282-L287), which spills.
For example, take a reader that returns fresh retainable vectors for its
first batches and then refills one root, as `SwitchingArrowPartitionReader`
does (198 of the 200 batches are copied), under `select id, wide_string,
udf(small_int)`. The writer runs ahead by the buffered bytes of the UDF input
only ([item
3](https://github.com/apache/spark/pull/59582#discussion_r4437566767)), so
dozens of copies can be held. Each is about 2 to 4 times the string bytes,
since every reallocation of the string data also allocates a null byte per data
byte (`OnHeapColumnVector.java` L673-L675) and the capacity doubles. The answer
to item 3, that "a producer opts in knowing that", does not cover these copies,
since their producer did not opt in. A phrase here and in the Javadoc would
cover it, for example "Other columns are read as rows or, in a partition whose
first batch was kept, copied on the heap and kept the same way".
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -262,7 +290,7 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
childOutput, childOutput)
val resultIter = if (inputColumnIndices.isDefined) {
- // Path 2: simple column refs, non-Arrow vectors.
+ // Path 2: simple column refs, columns that cannot be held.
val bufferedIter = new WrittenBatchReleasingIterator(inputIter, {
batch =>
batch.rowIterator().asScala.foreach { row =>
queue.add(unsafeProj(row).copy())
Review Comment:
**30. [Low] Pre-existing: `.copy()` before `HybridRowQueue.add` is not
needed**
This is not from this PR, but more inputs take this path now.
`HybridRowQueue.add` copies the row into its page (`RowQueue.scala` L89) or
writes it to its file (L143) and keeps no reference to it, so `.copy()` here
and at L317 allocates an `UnsafeRow` and a `byte[]` per row for nothing. The
row path passes the row as is (`EvalPythonEvaluatorFactory.scala` L135). Also,
`vectors.toArray` at L382 allocates an array per output row. With columnar
input enabled, every Arrow child in the tree takes this path now, so dropping
the `.copy()`, and building that array once as `copyColumns` does, would help.
It does not need to block this PR.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -130,6 +158,203 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60325: Arrow-backed source that reuses its vectors: Arrow
columnar path not taken") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_reuse_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // The Arrow columnar path would keep the pass-through columns of a
batch until its result
+ // comes back from Python. The reader refills its vectors for the next
batch and closes them
+ // once it is closed, as soon as it is exhausted, so they are not
retainable, and the
+ // batches take the row queue path. With two batches, the reader is
closed before their
+ // results come back. With many, it refills its vectors with the next
batches before then.
+ // The test UDF casts its input to string: on a string column, the cast
is removed, so the
+ // inputs are column references.
+ Seq(200, 20000).foreach { numRows =>
+ val result = readArrowSource(numRows = numRows, reuseVectors = true,
batchSize = 100)
+ .selectExpr("id", "name", "value", "data", "arrow_reuse_udf(name) as
udf_name")
+ val arrowExec =
singleArrowEvalPythonExec(result.queryExecution.executedPlan)
+ assert(!takesArrowColumnarPath(arrowExec))
+
+ val rows = result.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ rows.foreach { row =>
+ val id = row.getInt(0)
+ assert(row.getString(1) == s"row_$id", s"name mismatch for id $id")
+ assert(row.getString(3) == s"data_$id", s"data mismatch for id $id")
+ assert(row.getString(4) == s"row_$id", s"UDF result mismatch for id
$id")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60325: Arrow-backed source with fresh vectors: held without a
copy and released") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
Review Comment:
**21. [Low] No test guards the output slicing overrides of the columnar
runner any more**
Path 1 pairs each result batch with one input batch, so
`ColumnarArrowPythonWithNamedArgumentRunner` turns off output slicing with
`arrowMaxRecordsPerOutputBatch = -1` and `arrowMaxBytesPerOutputBatch = -1`
([`ColumnarArrowPythonRunner.scala`
L55-L58](https://github.com/apache/spark/blob/cfce388a0f5a4b90090d8dd00b1195d72152743c/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala#L55-L58)).
The test of SPARK-60367 that guards them,
`test_udf_on_cached_input_with_output_batch_slicing`, runs over the Arrow
cache. Its vectors are not retainable now, so it takes the row queue path,
which pairs by row, and passes with or without the overrides;
`_assert_columnar_arrow_eval` checks only `supportsColumnar()`, which is true
on both paths. Both confs default to -1 (`SQLConf.scala` L5605 and L5627), and
no other test sets them on Path 1.
So removing L57-L58 passes every test. After that, a source whose vectors
are retainable fails a scalar UDF over a 100-row batch under
`spark.sql.execution.arrow.maxRecordsPerOutputBatch=3` with "Batch size
mismatch: pass-through has 100 rows but UDF result has 3 rows" (L256-L258 of
the evaluator).
Could this test also run once with
`ARROW_EXECUTION_MAX_RECORDS_PER_OUTPUT_BATCH` set to 3, and once with
`ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH` set to 3, around
`result.collect()`? They cannot be set together (`PythonArrowOutput.scala`
L56-L59). The comment of the PySpark test, "The columnar path pairs each result
batch with one input batch"
([L137-L138](https://github.com/apache/spark/blob/cfce388a0f5a4b90090d8dd00b1195d72152743c/python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py#L137-L138)),
is no longer true for the cache either.
##########
docs/sql-migration-guide.md:
##########
@@ -31,7 +31,7 @@ license: |
- Since Spark 4.4, Data Source V2 file scans correctly evaluate deterministic
Python UDF filter predicates that reference only partition columns. Previously,
these predicates could be dropped, returning rows that did not satisfy the
filter (SPARK-60486). Data derived from affected query results may need to be
recomputed.
- Since Spark 4.4, quoted Variant path keys decode `\\`, `\/`, `\b`, `\f`,
`\n`, `\r`, `\t`, `\uXXXX` (including valid surrogate pairs), and escapes of
the enclosing quote (`\'` or `\"`). Previously, backslashes were literal, so
existing paths can match different keys or become invalid. For example, use
`$['a\\']` instead of `$['a\']`. Unknown escapes and unpaired surrogate escapes
stay literal. Dot-notation keys are unchanged. To keep the old behavior, set
`spark.sql.legacy.variantPathBackslashAsLiteral` to `true` before creating the
query.
-- Since Spark 4.4, `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`
is `false` by default: Arrow-based Python UDFs read rows from an Arrow-backed
upstream operator, such as a Data Source V2 scan or the Arrow cache, as before
Spark 4.2. With columnar input, they kept the columns that they pass through
until the results of their batch came back from Python, and failed or returned
wrong results when the operator reused or released the vectors of the batch in
the meantime. To enable columnar input again, set
`spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` to `true`.
+- Since Spark 4.4, `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`
is `false` by default: Arrow-based Python UDFs read rows from an Arrow-backed
upstream operator, such as a Data Source V2 scan or the Arrow cache, as before
Spark 4.2. With columnar input, they kept the columns that they pass through
until the results of their batch came back from Python, and failed or returned
wrong results when the operator reused or released the vectors of the batch in
the meantime. To enable columnar input again, set
`spark.sql.execution.arrow.pythonUDF.columnarInput.enabled` to `true`. They
then keep only the columns whose `ColumnVector.isRetainable` returns `true`,
and read the others as rows or copy them: a Data Source V2 source that returns
new Arrow vectors for every batch can override `isRetainable` to return `true`.
Review Comment:
**17. [Medium] The migration guide describes per-column holding and only one
of the promises of `isRetainable`**
This sentence is the only place that tells a source author how to opt in, so
I think it is worth making it exact before `isRetainable` ships. Two parts of
it do not match the code or the Javadoc:
- The decision is per batch, not per column, while the guide says that the
UDFs "keep only the columns whose `ColumnVector.isRetainable` returns `true`,
and read the others as rows or copy them". Path 1 is taken only if every column
of the first non-empty batch is an `ArrowColumnVector` that is retainable
(`ColumnarArrowEvalPythonEvaluatorFactory.scala` L173-L180), and a later batch
is held only if all its pass-through columns are retainable, otherwise all of
them are copied (L232-L237). So a source that marks its Arrow data columns but
also returns one `ConstantColumnVector` or another non-Arrow column gets none
of its columns held.
- "returns new Arrow vectors for every batch" is needed but not enough. The
Javadoc (`ColumnVector.java` L96-L101) also requires that the producer does not
free the vectors when it is closed, that it frees the unreleased ones when the
task completes, and that `closeIfFreeable` may run on another thread.
For example, a reader that allocates each batch from its own child allocator
and closes the allocator in `PartitionReader.close()` meets the condition in
the guide. `DataSourceRDD` closes the reader as soon as it is exhausted
(`DataSourceRDD.scala` L273-L279), while Path 1 still holds the batches that
the writer ran ahead with. So `allocator.close()` fails the task with "Memory
was leaked", or, if `close()` frees the vectors, the held pass-through columns
are read after they are freed. A reader that closes the previous root in
`next()`, as the Arrow cache does, also returns new vectors for every batch.
Could the guide say that a batch is held only if all its columns are
retainable `ArrowColumnVector`s, and point to the Javadoc of `isRetainable` for
what returning `true` promises? The conf doc at `SQLConf.scala` L5679-L5682 has
the same per-column wording. Relatedly, the first-batch check requires an
`ArrowColumnVector` while the later-batch check does not, so [your reply on
item 2](https://github.com/apache/spark/pull/59582#discussion_r4461161671) that
"both checks now use the same predicate" is not quite the case. One helper used
in both places would keep the two checks and the docs in line.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -130,6 +158,203 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
+ test("SPARK-60325: Arrow-backed source that reuses its vectors: Arrow
columnar path not taken") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_reuse_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // The Arrow columnar path would keep the pass-through columns of a
batch until its result
+ // comes back from Python. The reader refills its vectors for the next
batch and closes them
+ // once it is closed, as soon as it is exhausted, so they are not
retainable, and the
+ // batches take the row queue path. With two batches, the reader is
closed before their
+ // results come back. With many, it refills its vectors with the next
batches before then.
+ // The test UDF casts its input to string: on a string column, the cast
is removed, so the
+ // inputs are column references.
+ Seq(200, 20000).foreach { numRows =>
+ val result = readArrowSource(numRows = numRows, reuseVectors = true,
batchSize = 100)
+ .selectExpr("id", "name", "value", "data", "arrow_reuse_udf(name) as
udf_name")
+ val arrowExec =
singleArrowEvalPythonExec(result.queryExecution.executedPlan)
+ assert(!takesArrowColumnarPath(arrowExec))
+
+ val rows = result.collect()
+ assert(rows.map(_.getInt(0)).toSeq == (0 until numRows))
+ rows.foreach { row =>
+ val id = row.getInt(0)
+ assert(row.getString(1) == s"row_$id", s"name mismatch for id $id")
+ assert(row.getString(3) == s"data_$id", s"data mismatch for id $id")
+ assert(row.getString(4) == s"row_$id", s"UDF result mismatch for id
$id")
+ }
+ }
+ }
+ }
+
+ test("SPARK-60325: Arrow-backed source with fresh vectors: held without a
copy and released") {
Review Comment:
**24. [Nit] Some of the new tests and the cache fixture repeat existing
ones**
- This test repeats the `("name", "", true)` case of the SPARK-60306 grid
(L395-L396), which checks the same path, values and counters for the same
source and batches, in both codegen modes, with the plan shape and the limit.
Adding `retainable` to that tuple, with `arrowColumnarPath = suffix.isEmpty &&
arrowIdColumn && retainable`, would replace this test. It would also give the
next test, the only one with Arrow vectors that are not retainable, both
codegen modes and the check of `maxUnreleasedBatches`. Its counters alone
(L251-L252) would also match if every batch were released late, so they do not
show the release once written that its comment describes (L233-L234).
- The new reusing-reader test (L161-L191) copies the body of the
default-config test (L461-L496), and one test over both values of the conf
would do.
- The cache test and the static serializer fixture (L43-L59 and L292-L321)
copy `ArrowCachedPythonUDFSuite`, which caches the same data with the same
batch size and differs only in the columnar input conf, a session conf. Moving
the test there as its columnar-input case would keep the Arrow serializer out
of the other 15 tests here.
- `ArrowColumnarPythonUDFPipelinedSuite` reruns all 16 tests, including ones
that only build a plan (L110, L522 and L534) or run no Python (L358), and the
200,000-row cache test. Toggling the mode in the tests that need it, as
`PythonUDFSuite.scala` L391-L408 does, would keep the coverage.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +221,22 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
outputTypes: Seq[DataType],
columnIndices: Array[Int]): Iterator[ColumnarBatch] = {
+ // In pipelined mode the batches are added in the writer thread and
Review Comment:
**29. [Nit] The queue is not always polled on the task thread**
This comment, and "releases the pass-through columns with closeIfFreeable on
the task thread" at L52-L54, do not hold for stacked execs in pipelined mode,
as in the new stacked test under `ArrowColumnarPythonUDFPipelinedSuite`. The
outer exec takes the row queue path, and its pipelined writer pulls the
`resultIter.map` of the inner one. So the inner `poll()` at L255 runs on the
outer writer thread, except for the first batch, which the outer `peek()` pulls
on the task thread. The outer `releaseWritten()` (L408) then releases the
source vectors held by the inner one on that writer thread too. The
`ConcurrentLinkedQueue` and the contract allow this, so it is only the comment.
Something like "polled by the consumer of the output, possibly on another
thread" would cover it.
--
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]
