viirya commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4471389309
##########
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:
Added in cafd4cf047c: the test source has an `emptyBatches` option, which
returns batches without rows first, in the middle and last, and "Input batches
without rows are released and skipped" checks, on both paths, that no output
batch is empty and that all 23 batches are released. Replacing the filter with
`inputIter` or removing the `closeIfFreeable` fails 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:
You're right, thanks. In cafd4cf047c the check is
`result.limit(10).collect()` over all the columns, with an assertion that its
plan has an `ArrowEvalPythonExec` and a check of the first 10 UDF results. On
Path 1, the source frees the batches still held when the task completes, and
closing its allocator fails the task if any of them is left. I did not add an
assertion on the counters after the stop, since how many batches are released
before it depends on how far the writer runs ahead.
##########
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:
Done in cafd4cf047c: `copyColumns` copies `new ColumnarBatch(passThruCols,
batch.numRows())`, the converter is built once with `childSchema`, and the
`math.max` is gone.
##########
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:
Done in cafd4cf047c: the conf doc says that on-heap copies of the columns of
a later batch that cannot all be kept are kept without accounting too, and so
does the class doc of the evaluator. I left the Javadoc of `isRetainable` out
of it, since the copies are made by this consumer from vectors that do not opt
in, while the Javadoc is about what a retainable vector promises.
##########
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:
Done in cafd4cf047c: the conf doc now says that the pass-through columns
skip the conversion to rows and back only if all the columns of the first batch
are retainable `ArrowColumnVector`s, which no built-in operator returns, and
that the batches are otherwise converted to rows and back, which can cost more
than leaving the conf disabled.
--
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]
