dongjoon-hyun commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4473274683
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +231,21 @@ 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 by the consumer of the output, possibly on another thread.
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
-
- val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
- batch
- }
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, { batch
=>
+ val passThruCols = passThruColumns(batch)
+ if (canHold(passThruCols)) {
+ passThruQueue.add((passThruCols, batch.numRows()))
+ false
+ } else {
+ passThruQueue.add((copyColumns(passThruCols, batch.numRows()),
batch.numRows()))
Review Comment:
The path is chosen from the first batch only, so if only the first batch is
retainable (like the `freshBatches` scenario in the test), the pass-through
columns of every later batch are copied into on-heap vectors and queued in
`passThruQueue`. The config doc mentions that they are kept without accounting
for their memory, but nothing bounds them either: the queue holds as many
batches as the writer runs ahead of the consumer (all of them for an iterator
UDF that buffers its input), so wide rows and large batches may OOM the
executor. If the same input took Path 2, `HybridRowQueue` would account for
these rows with the task memory manager and spill them to disk. Since this risk
is new in this PR, could we account for (or bound) the copies, or fall back to
the row queue path instead of copying?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +231,21 @@ 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 by the consumer of the output, possibly on another thread.
val passThruQueue =
- new ArrayDeque[(Array[ColumnVector], Int)]()
-
- val bufferedIter = inputIter.map { batch =>
- val passThruCols = childOutput.indices.map(
- i => batch.column(i)).toArray
- passThruQueue.add((passThruCols, batch.numRows()))
- batch
- }
+ new ConcurrentLinkedQueue[(Array[ColumnVector], Int)]()
+
+ val bufferedIter = new WrittenBatchReleasingIterator(inputIter, { batch
=>
+ val passThruCols = passThruColumns(batch)
+ if (canHold(passThruCols)) {
+ passThruQueue.add((passThruCols, batch.numRows()))
Review Comment:
Path 1 pairs each Python result batch with exactly one input batch
([`passThruQueue.poll()` and `assert(passThruRows ==
numRows)`](https://github.com/apache/spark/blob/cafd4cf047c657783548bc5c39c9d150a07a7e53/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala#L264-L267)).
However, this exec also takes `SQL_SCALAR_PANDAS_ITER_UDF` and
`SQL_SCALAR_ARROW_ITER_UDF`, for which the worker verifies only the total row
count (`verify_iter_result_row_count`), so the boundaries of the result batches
may differ from those of the input batches. For example, with a retainable
Arrow source and columnar input enabled, a valid iterator UDF like the
following fails with `Batch size mismatch` once a partition has more than one
input batch:
```python
@pandas_udf("long")
def f(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
yield pd.concat(list(it)) + 1
```
The row queue and row paths pair the results by row, so they are not
affected. This is pre-existing (since SPARK-56350), but since this PR touches
this function and documents the 1:1 assumption, could we route the iterator
eval types to the row queue path, or at least file a JIRA to track it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -262,7 +299,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.
Review Comment:
Path 2 also uses `ColumnarArrowPythonWithNamedArgumentRunner`, which
[overrides `arrowMaxRecordsPerOutputBatch` / `arrowMaxBytesPerOutputBatch` to
`-1`](https://github.com/apache/spark/blob/cafd4cf047c657783548bc5c39c9d150a07a7e53/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala#L55-L58)
because Path 1 pairs batches. Path 2 pairs the results by row
(`queue.remove()` per output row), so slicing would be safe here, but
`spark.sql.execution.arrow.maxRecordsPerOutputBatch` / `maxBytesPerOutputBatch`
are silently ignored. For example, a user who sets `maxBytesPerOutputBatch` to
bound the JVM memory for large UDF results gets the whole result of each input
batch at once after enabling columnar input. This is pre-existing, but since no
built-in source returns `isRetainable() == true`, Path 2 effectively becomes
the main columnar path after this PR. Could we keep the slicing for Path 2,
e.g., by passing whether to disable it to the runner?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -167,13 +360,14 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
// path. Both release each batch once it has been written to Python.
for {
codegen <- Seq(true, false)
- (input, suffix, arrowIdColumn) <- Seq(
- ("name", "", true), ("name", "", false), ("concat(name, '_')", "_",
true))
+ (input, suffix, arrowIdColumn, retainable) <- Seq(
+ ("name", "", true, true), ("name", "", false, true), ("name", "",
true, false),
Review Comment:
The comment above says only that the batches take the row queue path if the
`id` column is not an `ArrowColumnVector`. Could we update it for the new
`("name", "", true, false)` case, which takes the row queue path because the
vectors are not retainable?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -346,3 +513,9 @@ class ArrowColumnarPythonUDFSuite extends
SharedSparkSession {
}
}
}
+
+/** Runs the tests of [[ArrowColumnarPythonUDFSuite]] with pipelined Python
UDF execution. */
+class ArrowColumnarPythonUDFPipelinedSuite extends ArrowColumnarPythonUDFSuite
{
Review Comment:
`ArrowColumnarPythonUDFPipelinedSuite` reruns the whole suite, including the
tests that don't depend on pipelined execution. In addition,
`takesArrowColumnarPath` / `outputColumnClasses` re-execute the query on every
call, so the SPARK-60306 test runs each 20000-row query up to four times per
combination, two of them only to inspect the output column classes, and this is
repeated in both suites. Could we limit the pipelined suite to the relevant
tests, or determine the path once per query, to reduce the CI time?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/vectorized/ColumnVector.java:
##########
@@ -84,6 +84,34 @@ 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, concurrently with the producer
returning the next
+ * batches or being closed;</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 may keep such vectors without accounting for their memory or
spilling 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:
This adds a public API with a non-trivial contract (release from another
thread concurrently with the producer, and cleanup on task completion) for a
single internal consumer, the `private[python]` evaluator, while no in-tree
producer implements it, as the config doc says. So Path 1 is unreachable from
the built-in sources, and the contract cannot be validated with a real source.
In addition, `ReusedArrowColumnVector` already exists as the opposite "do not
keep" marker (with a no-op `closeIfFreeable`), so we would have two overlapping
lifetime signals. Since a public API is hard to revert, would an internal
marker/trait be enough for now? Alternatively, we could drop Path 1 and always
handle the pass-through columns with the row queue (or an accounted columnar
copy).
##########
python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py:
##########
@@ -134,8 +134,10 @@ def
test_udf_on_cached_numeric_column_alongside_unselected_string(self):
df.unpersist()
def test_udf_on_cached_input_with_output_batch_slicing(self):
- # The columnar path pairs each result batch with one input batch, so
the 10-row result
- # of the single cached batch must not be sliced under the configs
below.
+ # The 10-row result of the single cached batch is sliced under the
configs below, and
Review Comment:
This comment says that the 10-row result is sliced under the configs below
and paired by row on the row queue path. However, the row queue path uses
`ColumnarArrowPythonWithNamedArgumentRunner`, which disables output batch
slicing (the `-1` overrides), so the result still comes back as a single 10-row
batch even with `maxRecordsPerOutputBatch=3`. As a result, this test doesn't
exercise the slicing at all and would keep passing if the row-wise pairing were
broken. Could we fix the comment, or make the row queue path honor the slicing
(see my comment on Path 2)?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +351,17 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
}
+ /** Copies the first `numRows` rows of the pass-through columns into
on-heap vectors. */
+ private def copyColumns(
Review Comment:
If any pass-through column of a batch cannot be held, this copies all the
pass-through columns, including the retainable ones, row by row through
`ColumnarBatchRow` and `RowToColumnConverter`. The allocation and conversion
are also almost the same as in `rowsToColumnarBatches`. For a batch of 100
columns where only one is a reused vector, all 100 columns are copied with
per-row virtual calls for every batch. Would it be cheaper to keep the
retainable columns as they are and copy only the others, or to copy column by
column?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowBackedDataSourceV2.scala:
##########
@@ -250,6 +277,58 @@ private class ReusingArrowPartitionReader(start: Int, end:
Int, batchSize: Int)
}
}
+/** Returns the batches of `first`, then those of `second`. */
+private class SwitchingArrowPartitionReader(
+ first: PartitionReader[ColumnarBatch],
+ second: PartitionReader[ColumnarBatch])
+ extends PartitionReader[ColumnarBatch] {
+ private var current = first
+
+ override def next(): Boolean = {
+ if (current.next()) return true
+ if (current eq second) return false
+ current = second
+ second.next()
+ }
+
+ override def get(): ColumnarBatch = current.get()
+
+ override def close(): Unit = {
+ first.close()
+ second.close()
+ }
+}
+
+/**
+ * Reads the batches of `fresh` and `reused` in step, which cover the same
rows, and returns the
+ * `id` column of the first with the other columns of the second: a batch
whose first column is
+ * retainable while the others are not.
+ */
+private class MixedArrowPartitionReader(
+ fresh: PartitionReader[ColumnarBatch],
+ reused: PartitionReader[ColumnarBatch])
+ extends PartitionReader[ColumnarBatch] {
+
+ override def next(): Boolean = {
+ val hasNext = fresh.next()
+ assert(reused.next() == hasNext)
Review Comment:
nit: `reused.next()` has a side effect, but it is called inside
`assert(...)`, which is elidable in Scala (e.g., with `-Xelide-below`). Then
the reused reader would not advance, and the other columns would be empty or
stale. Could we evaluate it first?
```scala
val reusedHasNext = reused.next()
assert(reusedHasNext == hasNext)
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5708,15 +5708,16 @@ object SQLConf {
val ARROW_PYSPARK_UDF_COLUMNAR_INPUT_ENABLED =
buildConf("spark.sql.execution.arrow.pythonUDF.columnarInput.enabled")
- .doc("When true, Arrow-based Python UDFs (pandas UDFs) can accept " +
- "columnar input directly from upstream operators that produce " +
- "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.")
+ .doc("When true, Arrow-based Python UDFs (pandas UDFs) accept columnar
input " +
+ "from upstream operators that produce ColumnarBatch, and write the UDF
inputs " +
+ "of Arrow-backed batches to Python directly. The columns that the UDFs
pass " +
+ "through skip the conversion to rows and back only if all the columns
of the " +
Review Comment:
The doc says that the pass-through columns skip the conversion to rows and
back if all the columns of the first batch are retainable `ArrowColumnVector`s,
but it also requires the UDF inputs to be plain column references; otherwise,
it takes the row fallback path (Path 3). The migration guide entry has the same
gap. For example, a DSv2 source that implements `isRetainable` would still go
through rows with `udf(concat(name, '_'))`, which can be confusing when
diagnosing performance. Could we mention this condition?
--
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]
