viirya commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4461163635
##########
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:
Agreed. This PR does not change the pairing, and it happens with the default
configuration of 4.2 and 4.3 on Path 1, so I'll file it as a bug and fix it
separately.
##########
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:
Agreed. Like item 5, it happens with the default configuration of 4.2 and
4.3, so I'll file it as a bug and fix it separately.
##########
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:
Fixed in cfce388a0f5: the cache is back to master and closes each root when
it returns the next batch, and its vectors are not retainable, so nothing is
pinned until the task ends. The default also stays false.
##########
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:
This PR no longer changes the default, so non-Arrow columnar children keep
reading rows as on master. With columnar input enabled, `supportsColumnar`
still ignores the vector types, as on 4.2 and 4.3; I'll include it with the
separate issues.
##########
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:
Fixed in cfce388a0f5: input batches without rows are released and skipped on
all columnar paths.
##########
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:
Covered by the contract in cfce388a0f5: the doc of `isRetainable` makes the
producer responsible for freeing the vectors the consumer has not released when
the task completes, e.g. when the consumer stops early.
`ArrowBackedDataSourceV2` does this with a task completion listener.
##########
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:
You're right, thanks. Fixed in cfce388a0f5: the test now sets
`ARROW_EXECUTION_MAX_RECORDS_PER_BATCH` to 1000, so the 200,000 rows are cached
in 200 batches, and it asserts the row queue path.
##########
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:
Fixed in cfce388a0f5 by the revert of `ArrowCachedBatchSerializer` to
master, and the comments of the cache test were updated.
--
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]
