viirya commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4471387253
##########
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:
Thanks, fixed in cafd4cf047c. The two checks now share one helper,
`canHold`, which requires the pass-through columns to be `ArrowColumnVector`s
whose `isRetainable` is true, for the first batch and for every later one, so
they agree now; my reply on item 2 overstated it before. The migration guide
now says that the pass-through columns of a batch are kept only if all its
columns are retainable `ArrowColumnVector`s, and that a source can return
`true` if it keeps the promises in the documentation of
`ColumnVector.isRetainable`. The conf doc says the same.
##########
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:
Done in cafd4cf047c: the bullet now says that `closeIfFreeable` "may be
called from another thread than the one that requested its batch, concurrently
with the producer returning the next batches or being closed".
##########
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:
Done in cafd4cf047c: `takesArrowColumnarPath` checks the last column, the
UDF result, and the SPARK-60306 grid also checks that column 0 is not copied on
Path 1. With the first-batch check changed to accept any `ArrowColumnVector`,
five tests now fail: the reusing reader, the empty batches over vectors that
are not retainable, the stacked UDFs, the SPARK-60306 grid and the cache.
##########
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:
Done in cafd4cf047c: the later batches of the switching reader now have a
fresh, retainable `id` column next to the reused columns, so each of them mixes
both, and the test asserts that all 200 batches are released, which counts the
198 copied ones through their `id` column. With `true` changed to `false` the
test fails on the release count, and with `forall` changed to `exists` it fails
on the values of the reused columns.
##########
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:
Added "The Arrow columnar path does not slice the results of Python" in
cafd4cf047c, which runs a UDF over 100-row batches of the test source with
`ARROW_EXECUTION_MAX_RECORDS_PER_OUTPUT_BATCH` and then
`ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH` set to 3. Without the two
overrides it fails with the "Batch size mismatch" assertion. The comment of the
PySpark test now says that the cache takes the row queue path.
--
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]
