dongjoon-hyun commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4437566761
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,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.exists(_.isInstanceOf[ReusedColumnVector])) {
Review Comment:
**2. [Medium] The check of later batches is weaker than the check of the
first one, so an unmarked `OrcColumnVector` is held without a copy**
The first batch must be all unmarked `ArrowColumnVector`s (L165-L171), but a
later batch is copied only if a column has the marker, so an unmarked vector of
any other class is held. `OrcColumnVector` has no marker and the next batch
overwrites its data in place (`OrcColumnarBatchReader.java` L244,
`OrcAtomicColumnVector.java` L117), so "A vector without it may be held by the
consumer past the next batch" (`ReusedColumnVector.java` L28-L29) is false for
it.
One task can read both, with the Arrow cache serializer:
```sql
-- c: a materialized df.repartition(8, "k").cache()
-- t: an ORC table CLUSTERED BY (k) INTO 8 BUCKETS, with more than 4096 rows
per file
SELECT k, count(*), max(u)
FROM (SELECT k, v, nondeterministic_udf(v) u
FROM (SELECT k, v FROM c UNION ALL SELECT k, v FROM t))
GROUP BY k
```
`UnionExec` concatenates co-partitioned children in one task
(`basicPhysicalOperators.scala` L1641, `PartitionerAwareUnionRDD.scala`
L88-L89). The first batch comes from the cache, so Path 1 is taken and the ORC
batches reach L229. The writer reads the next batch before the results of the
current one (`PythonRunner.scala` L1589), so this does not depend on timing.
`k` and `v` of an ORC batch are read from a later batch and paired with the UDF
results of the earlier one, and `GROUP BY k` uses the wrong keys. There is no
error, since the assert at L249 compares row counts taken at enqueue time. This
comes from reading the code, I did not run it.
The trigger is narrow. Parquet is safe because its vectors are marked, and
the UDF must be nondeterministic because a deterministic project is pushed
below the union (`Optimizer.scala` L1499). Marking the cache vectors removes
this route too. The two checks should still be one predicate, since the fix is
small: copy a batch unless all its columns can be held. Marking
`OrcColumnVector` alone would leave third-party vectors.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,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.exists(_.isInstanceOf[ReusedColumnVector])) {
+ passThruQueue.add((copyColumns(batch), batch.numRows()))
+ true
+ } else {
+ passThruQueue.add((passThruCols, batch.numRows()))
Review Comment:
**3. [Medium] Holding on Path 1 has no memory accounting, no spill, and no
bound tied to the size of the pass-through columns**
You noted that the held roots stay in memory until their results come back.
The part that I think needs a decision is how many. The writer stops only when
the socket stops taking bytes (`PythonRunner.scala` L1589 and L1601-L1602), and
the socket only sees the UDF inputs (`ColumnarArrowPythonInput.scala` L115),
while this queue has no limit and holds all of `childOutput`.
Take `SELECT big_string, slow_udf(small_int)` over the Arrow cache. An INT
input of 10,000 rows is about 41 KB per batch. Assuming a few MB of socket
buffers, which I did not measure, that is dozens to over a hundred batches in
flight. At 1 KB of pass-through per row, that is around 1 GB per task. An
iterator UDF that consumes its input before yielding, which the docs allow,
holds the whole partition.
This is off-heap memory in a `Long.MaxValue` child allocator (L1277-L1280 of
`ArrowCachedBatchSerializer.scala`) that `TaskMemoryManager` does not know
about and that cannot spill, while the row path and the row queue path use
`HybridRowQueue`. So a query that spills on master can fail with a direct
memory OOM. I did not reproduce that.
For the Arrow cache this is new, as v4.3.0 closes the previous root in
`next()`, and marking the vectors keeps that. Sources with fresh vectors
remain. Could the conf doc or the description at least say so, since
`columnarInput.enabled=false` is the way out? If you prefer to bound it, one
way is to account for the held bytes with a `MemoryConsumer` and copy to the
row queue when memory is denied.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +335,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)
Review Comment:
**5. [High] `RowToColumnConverter` has no UDT or NullType case, so the row
queue path fails with `UNSUPPORTED_DATATYPE`**
Both uses, this one and `rowsToColumnarBatches` (L362, over `outputSchema`),
reach `case unknown => throw ExecutionErrors.unsupportedDataTypeError(unknown)`
(`Columnar.scala` L313), and the converter is built eagerly.
`ArrowEvalPythonExec.supportsColumnar` has no type check, so with the new
default this fails queries that succeed on master:
```python
# features is a VectorUDT column, which Parquet reads in batches.
spark.read.parquet(p).select("id", "features", some_udf("id"))
```
Parquet supports batch reads of a UDT (`ParquetUtils.scala` L233-L234), a
plain `udf()` is an Arrow UDF by default (`SQLConf.scala` L5908), and with no
data filter there is no `FilterExec` above the scan, so this takes Path 2. The
same failure comes from:
- An Arrow UDF returning a UDT (`fallbackOnUDT` is false by default,
`SQLConf.scala` L5945) or `void`, over a vectorized scan or the default cache
of a primitive schema.
- A Parquet void column (`SQLConf.scala` L2158).
- The Arrow cache with a UDT column and a UDF over a complex expression
(Path 3).
- The Arrow cache with a UDT column and simple column references on an empty
partition, where `peek()` is `None`. A non-empty partition takes Path 1, so
this depends on the data distribution. With the cache vectors marked, it
becomes every partition.
`test_udf_with_udt` uses `createDataFrame`, which is row-based, so no test
takes this path. This comes from reading the code, I did not run it. Cases in
`RowToColumnConverter` for NullType, and for a UDT through `udt.sqlType`, look
like the direct fix.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +335,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
Review Comment:
**6. [Medium] With `maxRecordsPerBatch` of zero or less,
`rowsToColumnarBatches` never ends or throws**
`rowsToColumnarBatches` uses the raw `conf.arrowMaxRecordsPerBatch`
(`ArrowEvalPythonExec.scala` L145) as its capacity and loop bound, while the
conf says "If set to zero or negative there is no limit" (`SQLConf.scala`
L5583) and the row path normalizes it (`PythonArrowInput.scala` L140-L143). For
an Arrow UDF directly over a vectorized Parquet or ORC scan or the default
cache of a primitive schema:
- `0`: `while (count < batchSize && rowIter.hasNext)` (L374) consumes
nothing while `hasNext` stays true, so it returns empty batches forever. From
reading the code, the loop never reaches the `InterruptibleIterator`
(`PythonRunner.scala` L806), since `flatMap` answers from its current iterator,
so a task kill does not stop it either.
- A negative value: `NegativeArraySizeException` from `allocateColumns`
(L364), even on an empty partition.
- A large value: that capacity is allocated up front for every output column.
Copying the normalization of `PythonArrowInput` would make
`allocateColumns(Int.MaxValue, ...)` fail, so the initial capacity and the row
limit need to be separate. `RowToColumnarExec` uses `conf.columnBatchSize` for
this (`Columnar.scala` L530).
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -155,11 +158,17 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val inputColumnIndices = resolveColumnIndices(allInputs.toSeq)
- // Peek at first batch to check if Arrow-backed.
+ // Peek at the first batch to check whether its columns are Arrow vectors
+ // that may be held until the result of the batch comes back.
val peekIter = new PeekableIterator(inputIter)
val isArrow = peekIter.peek().exists { batch =>
- batch.numCols() > 0 &&
- batch.column(0).isInstanceOf[ArrowColumnVector]
+ batch.numCols() > 0 && (0 until batch.numCols()).forall { i =>
+ batch.column(i) match {
+ case _: ReusedColumnVector => false
+ case _: ArrowColumnVector => true
Review Comment:
**4. [High] A source that reuses its vectors and does not know this internal
marker gets wrong results with the default configuration**
Without the marker, any `ArrowColumnVector` is held past the next batch. But
the marker is under `org.apache.spark.sql.execution`, which the public docs
exclude (`SparkBuild.scala` L1712), and the public contracts say the opposite:
"This method should return same value until `next` is called"
(`PartitionReader.java` L49), "ColumnVector is expected to be reused during the
entire data loading process" (`ColumnVector.java` L49), and
`ColumnarDataSource.scala` L91-L92. `docs/sql-migration-guide.md` L34 tells
such a source only to override `closeIfFreeable`.
A reader that follows that guide, refilling one `VectorSchemaRoot` per batch
with a subclass whose `closeIfFreeable` is a no-op, takes Path 1 for `select
id, name, udf(name)`. The writer pulls the next batch before the result of the
current one is read (`PythonRunner.scala` L1589), so the queue holds the same
vectors many times, and the pass-through values of every batch but the last one
are those of a later batch. The job succeeds, and the assert at L249 cannot see
it, since it compares row counts taken at enqueue time. This comes from reading
the code, I did not run it.
Other readers that follow the public contract break too:
- A plain `new ArrowColumnVector(v)` over a refilled root: a full scan fails
with `IndexOutOfBoundsException`, since `ColumnarToRowExec` closes the shared
vector, while `LIMIT`, `show` or `take` can return mispaired rows with no error.
- Fresh vectors whose allocator is closed in `PartitionReader.close()`:
`DataSourceRDD.scala` L275-L279 closes the reader on exhaustion while batches
are still held, so it fails with "Memory was leaked".
This is what #59551 guarded against. This PR changes the test reader to
`ReusedArrowColumnVector` (`ArrowBackedDataSourceV2.scala` L233), and removes
the test of the plain reader and both migration guide entries, so nothing
covers this shape any more. "this PR makes it correct for sources that reuse or
release their vectors" in the description holds only for sources that carry the
marker.
On the options in [your question to
HyukjinKwon](https://github.com/apache/spark/pull/59582#issuecomment-6561237694):
options 2 and 3 fix only the no-op subclass. The plain vector and the reader
that frees on close have no override, and `ReleaseCountingArrowColumnVector`
(`ArrowBackedDataSourceV2.scala` L344-L350) overrides it and still releases, so
it would be misclassified. To be safe by default, I would flip the polarity: a
public way for a vector to opt in to being held, with every other vector not
held. At the least, the migration guide needs to tell these sources what to do,
by a means they can use.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -314,6 +335,15 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
rowsToColumnarBatches(rowIter, context)
Review Comment:
**7. [Medium] The output buffer of `rowsToColumnarBatches` has no byte
limit, so large values hit the 2 GB limit of a vector**
Rows are buffered in on-heap vectors up to `maxRecordsPerBatch` (10,000 by
default, L374-L375) across result batches, including all the pass-through
columns, with no byte limit.
For a Parquet table with a 250 KB binary column and `withColumn("pred",
some_udf("image"))`, the input is fine, at most 4096 rows or about 1 GB per
batch. The output overflows near row 8590 with `Cannot reserve additional
contiguous bytes ... (integer overflow)` (`WritableColumnVector.java` L143 and
L163), and the advice of that message, a smaller reader batch, does not help.
The threshold is an average of about 215 KB per value, and a small heap fails
earlier, since the vector grows by doubling. This comes from reading the code,
I did not run it.
On master this query works, because the row path streams rows and keeps the
pass-through columns in `HybridRowQueue`. Item 11 avoids this for non-Arrow
children, and a byte limit per output batch would fix it for the rest.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,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
=>
Review Comment:
**12. [Low] A zero-row input batch calls the UDF with empty input**
The columnar input paths write whatever batch they get
(`ColumnarArrowPythonInput.scala` L65-L73, `PythonArrowInput.scala` L153-L161),
and the worker does not skip it (`_pandas.py` L93), so a UDF such as
`np.stack(s.values)` or `s.iloc[0]` fails with `ValueError` or `IndexError` on
a zero-row batch. The row path writes a batch only when it has rows
(`PythonArrowInput.scala` L197 and L203).
I am not sure how common the trigger is. Built-in Parquet, ORC and the cache
of row input do not produce such batches, but
`ArrowCachedBatchSerializerSuite.scala` L2581-L2582 calls one "legal input from
a columnar source", and `NativeDataSource.scala` L725 can return one. If this
matters, such a batch could be skipped and released before it reaches the queue
and the writer.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5676,14 +5676,14 @@ 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,
unless the " +
+ "upstream operator may reuse or release them before then, as it marks
with " +
+ "ReusedColumnVector: such columns are read as rows or copied.")
.version("4.2.0")
.withBindingPolicy(ConfigBindingPolicy.SESSION)
.booleanConf
- .createWithDefault(false)
+ .createWithDefault(true)
Review Comment:
**11. [Medium] The new default also moves non-Arrow columnar children to the
row queue path, with no benefit**
`ArrowEvalPythonExec.supportsColumnar` (`ArrowEvalPythonExec.scala`
L125-L126) only asks whether the child can output batches, so this default also
applies to vectorized Parquet and ORC scans and to the default cache. With
default confs, no built-in source returns unmarked Arrow vectors, since the
Arrow cache is opt-in, so they can only take paths 2 and 3.
Compared with the row path, Path 2 saves the `MutableProjection` of the UDF
inputs. It adds `RowToColumnConverter` over all output columns
(`vectors.toArray` per row at `ColumnarArrowEvalPythonEvaluatorFactory.scala`
L375, two copies per string at `Columnar.scala` L373-L374) and a
`ColumnarToRowExec` above. The description and `bench_arrow_columnar_udf.py`
cover only Arrow-backed sources, so I do not know the size of this cost. Items
5, 6, 7, 9 and 10 apply to these sources as well, and they get nothing in
return.
Could this default apply only to Arrow children? For example, in
`supportsColumnar`:
```scala
child.supportsColumnar && conf.arrowPySparkUDFColumnarInputEnabled &&
child.vectorTypes.forall(_.forall(_ == classOf[ArrowColumnVector].getName))
```
This returns V1 Parquet and the default cache without AQE to the row path,
and keeps the Arrow cache. It is not enough on its own, as ORC V1, `UnionExec`,
`BatchScanExec` and `QueryStageExec` report `None` and pass. `QueryStageExec`
would need to delegate `vectorTypes`, and ORC would need its own. If that is
too much for this PR, I would move the change of the default to its own PR.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala:
##########
@@ -1348,19 +1350,24 @@ private class ArrowCachedBatchToColumnarBatchIterator(
deserializeToRoot(batchIter.next().asInstanceOf[ArrowCachedBatch])
}
- previousRoot = root
-
// When projected on load, the root already holds only the selected
columns in output order,
// so wrap its vectors directly. Otherwise it holds all cache columns and
must be selected.
- val selectedColumns = if (canProjectOnLoad) {
- root.getFieldVectors.asScala.map(v => new
ArrowColumnVector(v)).toArray[ColumnVector]
+ val selectedVectors = if (canProjectOnLoad) {
+ root.getFieldVectors.asScala.toArray
+ } else {
+ columnIndices.map(root.getFieldVectors.get(_))
+ }
+ val numRows = root.getRowCount
+ val batch = if (selectedVectors.isEmpty) {
+ // No column can release the root, and it holds no buffer the batch
needs.
+ root.close()
+ new ColumnarBatch(Array.empty[ColumnVector], numRows)
} else {
- val allColumns = root.getFieldVectors.asScala.map { vector =>
- new ArrowColumnVector(vector)
- }.toArray[ColumnVector]
- columnIndices.map(allColumns(_))
+ unreleasedRoots.synchronized(unreleasedRoots.add(root))
Review Comment:
**1. [High] The row fallback of the Arrow cache never releases a batch, so
every root stays open until the task ends**
`next()` no longer closes the previous root and leaves that to the consumer,
but the fallback in `convertCachedBatchToInternalRow` (L214-L225) never calls
`closeIfFreeable`:
```scala
batchIter.flatMap { batch =>
val numRows = batch.numRows()
new Iterator[InternalRow] {
...
val row = batch.getRow(rowIdx)
rowIdx += 1
toUnsafe(row)
```
`InMemoryTableScanExec.supportsRowBased` is true
(`InMemoryTableScanExec.scala` L90), so `Columnar.scala` L581 adds no
`ColumnarToRowExec` and a row-based parent reads through `doExecute`, with AQE
on or off (`QueryStageExec.scala` L116-L117). Selecting an array, struct, map,
calendar interval, variant, void, UDT, geo or nanosecond timestamp column takes
this branch (L191-L205). For example, with the Arrow serializer:
```scala
val cached = spark.range(0, n, 1, 1).selectExpr("id", "array(id, id + 1)
a").cache()
cached.selectExpr("a[0]").count()
```
Before this PR at most one root was open, or two with prefetch. Now all the
roots of the partition stay in the child allocator until the listener at L1330
runs, outside the memory manager and with no spill, so a large partition can
hit the direct memory limit or a container kill. I did not reproduce the OOM.
The results are correct. This has nothing to do with Python UDFs or
`columnarInput.enabled`, so that conf cannot avoid it. Everything is closed at
the end of the task, so no leak check or existing test sees it
(`ArrowCachedBatchSerializerSuite.scala` L185 only uses `checkAnswer`).
Marking the vectors, as I suggest in the review body, removes this. If you
keep the release in the consumer, this needs the same shape as
`ColumnarEvaluatorFactory.scala` L54-L55, and a test that bounds the open
roots. `toUnsafe` copies the values, so it is safe:
```scala
CompletionIterator[InternalRow, Iterator[InternalRow]](
batch.rowIterator().asScala.map(toUnsafe), batch.closeIfFreeable())
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -43,25 +44,25 @@ import org.apache.spark.util.Utils
* Three execution paths based on input characteristics:
*
* 1. '''Arrow columnar path''' (UDF inputs are simple column refs AND
- * columns are [[ArrowColumnVector]]): Arrow FieldVectors are extracted
- * directly and serialized to IPC. Pass-through columns are kept as
- * [[ColumnVector]] references, which is only correct if the batches stay
- * valid until their results come back from Python: a columnar reader may
- * reuse or release the vectors of a batch once it is asked for the next
- * one. So `spark.sql.execution.arrow.pythonUDF.columnarInput.enabled`
- * is disabled by default. Output is produced by columnar combining:
- * passThruCols ++ resultCols -> ColumnarBatch, so the consumer of the
- * output releases the pass-through columns with closeIfFreeable on the
- * task thread. With `spark.python.udf.pipelined.enabled`, the Python
+ * columns are [[ArrowColumnVector]]s that are not [[ReusedColumnVector]]s):
+ * Arrow FieldVectors are extracted directly and serialized to IPC.
+ * Pass-through columns are kept as [[ColumnVector]] references until the
+ * result of their batch comes back from Python, which the producer allows
+ * by not marking them as reused. Output is produced by columnar
+ * combining: passThruCols ++ resultCols -> ColumnarBatch, so the consumer
+ * of the output releases the pass-through columns with closeIfFreeable on
+ * the task thread. With `spark.python.udf.pipelined.enabled`, the Python
* writer reads the next input batches on its own thread meanwhile, so
* the upstream operator must allow its vectors to be released while it
* returns the next batches.
*
* 2. '''Non-Arrow columnar path''' (UDF inputs are simple column refs
- * BUT columns are NOT [[ArrowColumnVector]]): Non-Arrow columnar
- * readers (Parquet vectorized, InMemoryTableScan) reuse ColumnVector
- * objects across batches. Uses [[HybridRowQueue]] for pass-through
- * buffering and converts joined rows back to [[ColumnarBatch]].
+ * BUT columns are NOT such [[ArrowColumnVector]]s): e.g. Parquet
+ * vectorized readers reuse [[WritableColumnVector]]s across batches, and
Review Comment:
**10. [Medium] The default cache now pins every decompressed batch until the
task ends**
`InMemoryRelation.scala` L215 allocates new vectors per cached batch and
L228 registers `_ => columnarBatch.close()`, which keeps each batch reachable
until the task completes. `closeIfFreeable` is a no-op for them
(`WritableColumnVector.java` L128-L130), so the release in
`WrittenBatchReleasingIterator` does nothing here.
On master, a row-based parent reads this relation through
`GenerateColumnAccessor`, so these vectors do not exist. With the new default,
an Arrow UDF directly over the default cache of a primitive schema reads it as
batches:
```python
spark.range(n).selectExpr("id", "id * 2 v").cache().select("id",
some_udf("v"))
```
For 100 million rows in a partition and two long columns, that is about 1.8
GB by calculation, held on the heap for the whole task, outside the memory
manager, or off-heap with `spark.sql.columnVector.offheap.enabled`. These
vectors are marked, so the query takes Path 2 and gets no benefit from the
columnar input. See item 11.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/ArrowColumnarPythonUDFSuite.scala:
##########
@@ -130,6 +155,173 @@ 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 marks its vectors as reused, since
it refills them for
+ // the next batch and closes them once it is closed, as soon as it is
exhausted, so 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)
+ }
+ }
+
+ test("SPARK-60325: Reused vectors after fresh ones: the pass-through columns
are copied") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true") {
+ val pandasUDF = TestScalarPandasUDF(
+ name = "arrow_switch_udf", returnType = Some(StringType))
+ registerTestUDF(pandasUDF, spark)
+
+ // The first batches have fresh vectors, so the Arrow columnar path is
taken. It holds the
+ // pass-through columns of those as they are, and copies those of the
later batches, whose
+ // vectors the reader refills for the next batch.
+ val numRows = 20000
+ val freshBatches = 2
+ val result = readArrowSource(
+ numRows = numRows, reuseVectors = true, freshBatches = freshBatches,
batchSize = 100)
+ .selectExpr("id", "name", "value", "data", "arrow_switch_udf(name) as
udf_name")
+ val arrowExec =
singleArrowEvalPythonExec(result.queryExecution.executedPlan)
+ val copied = arrowExec.executeColumnar().mapPartitionsInternal { batches
=>
+ batches.map(_.column(0).isInstanceOf[OnHeapColumnVector])
+ }.collect().toSeq
+ val numBatches = numRows / 100
+ assert(copied == Seq.fill(freshBatches)(false) ++ Seq.fill(numBatches -
freshBatches)(true))
+
+ 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.getString(4) == s"row_$id", s"UDF result mismatch for id
$id")
+ }
+ // The consumer releases the fresh batches, which were held.
+ assert(ArrowBackedDataSourceV2.returnedBatches.get == freshBatches)
+ assert(ArrowBackedDataSourceV2.releasedBatches.get == freshBatches)
+ }
+ }
+
+ test("SPARK-60325: Arrow cache: the pass-through columns match the UDF
results") {
+ assume(shouldTestPandasUDFs)
+ withSQLConf(
+ SQLConf.ARROW_PYSPARK_EXECUTION_ENABLED.key -> "true",
+ SQLConf.COLUMN_BATCH_SIZE.key -> "1000") {
Review Comment:
**14. [Low] The Arrow serializer does not read `COLUMN_BATCH_SIZE`, so this
test does not have 200 batches**
`ArrowCachedBatchSerializer` cuts batches by `conf.arrowMaxRecordsPerBatch`
(`ArrowCachedBatchSerializer.scala` L93) and does not read `columnBatchSize`.
So the cache here has 20 batches of 10,000 rows, not the 200 batches of 1,000
rows in the description and in [your
comment](https://github.com/apache/spark/pull/59582#issuecomment-6561237904).
Could this use `ARROW_EXECUTION_MAX_RECORDS_PER_BATCH`, like
`ArrowCachedPythonUDFSuite.scala` L58?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -155,11 +158,17 @@ private[python] class
ColumnarArrowEvalPythonEvaluatorFactory(
val inputColumnIndices = resolveColumnIndices(allInputs.toSeq)
- // Peek at first batch to check if Arrow-backed.
+ // Peek at the first batch to check whether its columns are Arrow vectors
+ // that may be held until the result of the batch comes back.
val peekIter = new PeekableIterator(inputIter)
val isArrow = peekIter.peek().exists { batch =>
Review Comment:
**16. [Nit] `isArrow` and the comments on paths 2 and 3 no longer match**
`isArrow` is now false for a `ReusedArrowColumnVector`, which is an Arrow
vector, and "non-Arrow" at L179, L257 and L286, and in the heading at L59, have
the same problem, since a stacked `ArrowEvalPythonExec` sends its reused Arrow
vectors to Path 2. Could this be something like `canHoldInput`, with those
comments saying that the vectors cannot be held?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/ArrowCachedBatchSerializer.scala:
##########
@@ -1257,6 +1257,12 @@ private class ColumnarBatchToArrowCachedBatchIterator(
/**
* Iterator that converts ArrowCachedBatch to ColumnarBatch.
+ *
+ * Every batch is loaded into its own root, and its columns are plain
[[ArrowColumnVector]]s, which
+ * a consumer may hold past the next batch: the consumer releases a batch with
closeIfFreeable once
Review Comment:
**15. [Low] Docs and test comments still describe the old lifetime of the
roots**
If you keep the release in the consumer, these need an update. Marking the
vectors keeps them right.
- `docs/sql-arrow-cache-format.md` L229 says the roots are "released as soon
as the encode/decode completes", and L235 that the off-heap buffers are
"proportional to a single batch", which is the reason it gives for not growing
`spark.executor.memoryOverhead`.
- `ArrowCachedPythonUDFSuite.scala` L28-L29 says the scan "closes the root
of the previous batch when it returns the next one".
- `ArrowCachedBatchSerializerSuite.scala` L2613-L2614 says the previous root
"must only be closed on the consumer thread (in next())", and L2639 speaks of
holding exactly one batch.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,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)]()
Review Comment:
**8. [Medium] Path 1 pairs output and input batches one to one, which
iterator UDFs do not promise**
L248-L251 take one queue entry per result batch and assert equal row counts,
but `SQL_SCALAR_PANDAS_ITER_UDF` and `SQL_SCALAR_ARROW_ITER_UDF` only verify
the total (`_pandas.py` L214-L215, `_arrow.py` L769-L772), and the docs allow
it: "it can prefetch the data from the input iterator as long as the lengths
are the same" (`arrow_pandas.rst` L150-L151).
Over the Arrow cache or a source with fresh Arrow vectors, with simple
column references, this fails with `AssertionError: assertion failed: Batch
size mismatch`:
```python
@pandas_udf("long", PandasUDFType.SCALAR_ITER)
def f(it):
yield pd.concat(list(it)) + 1
```
So does the in-tree `predict_batch_udf(fn, rt, batch_size=100)("a")` when an
input batch has more than `batch_size` rows (`functions.py` L807). If all sizes
matched and the UDF yields one more empty result, `poll()` returns null, which
gives a `MatchError`. The row path and the row queue path pair by row and
succeed.
The existing tests, such as
`test_scalar_iter_pandas_udf_with_single_output_batch`, use `spark.range`,
which is row-based, and the suites here have no iterator UDF. This comes from
reading the code, I did not run it. Excluding the iterator eval types from Path
1 looks like the smallest fix.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,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.exists(_.isInstanceOf[ReusedColumnVector])) {
+ passThruQueue.add((copyColumns(batch), batch.numRows()))
+ true
+ } else {
+ passThruQueue.add((passThruCols, batch.numRows()))
+ false
Review Comment:
**13. [Low] Held batches are not released on early termination, although
`ReusedColumnVector` says the consumer releases them**
The new Javadoc says of an unmarked vector that "the consumer releases it
with `closeIfFreeable` once done with it" (`ReusedColumnVector.java` L29-L30),
but nothing drains `passThruQueue` when a task stops early on `LIMIT`,
`show()`, a Python error or a kill.
The row path leaves its current batch unreleased in that case too, so a
source needs its own cleanup anyway, as the test reader has
(`ArrowBackedDataSourceV2.scala` L327-L336). The difference is how many. A
reader whose listener closes its last batch and its child allocator passes on
the row path, while on Path 1 the rest of the queue makes `allocator.close()`
throw "Memory was leaked by query", which fails a task that should succeed.
Also, since only the columns are kept here, a `ColumnarBatch` subclass that
overrides `closeIfFreeable` never gets that call.
I found no trigger in the tree, because the Arrow cache has its own
listener, and I have not seen such a third-party connector. A sentence about
early termination in the Javadoc, or a listener that drains the queue, would
cover it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -200,15 +214,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.exists(_.isInstanceOf[ReusedColumnVector])) {
+ passThruQueue.add((copyColumns(batch), batch.numRows()))
+ true
+ } else {
+ passThruQueue.add((passThruCols, batch.numRows()))
+ false
+ }
+ })
val pyRunner = new ColumnarArrowPythonWithNamedArgumentRunner(
Review Comment:
**9. [Medium] The columnar input paths ignore `maxRecordsPerBatch` and
`maxBytesPerBatch`**
This runner sends each upstream batch as one Arrow batch
(`ColumnarArrowPythonInput.scala` L65-L73 and L141-L144, and
`PythonArrowInput.scala` L157-L159 for Path 3), with no split on this side or
in the worker. The row path limits both (`PythonArrowInput.scala` L248-L251).
With the new default, a user who sets
`spark.sql.execution.arrow.maxRecordsPerBatch=32` for a GPU pandas UDF over a
Parquet scan gets 4096 rows per call and no error, while `arrow_pandas.rst`
L414-L420 documents this conf as the way to bound each batch and avoid OOM. A
batch-dependent UDF such as `s.rank()` also changes its values with no error,
even with default confs, since a batch goes from 10,000 rows to 4096. The bytes
limit of 64 MB does not apply either.
The Arrow cache cuts its batches with the conf at cache time, so the limit
holds there unless the conf changes. Path 1 needs the one to one pairing, but
paths 2 and 3 pair by row, so they could size their batches like
`BatchedPythonArrowInput`. Item 11 avoids this for non-Arrow children.
--
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]
