viirya commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4471391219
##########
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))
Review Comment:
Done in cafd4cf047c: the restored test is merged into the reusing-reader
test, which runs with columnar input on and off, and its comment says only that
the default reads the scan as rows.
##########
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
Review Comment:
Done in cafd4cf047c: this comment now says that the queue is polled by the
consumer of the output, possibly on another thread, and the class doc no longer
says that the pass-through columns are released on the task thread.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowEvalPythonEvaluatorFactory.scala:
##########
@@ -262,7 +290,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.
val bufferedIter = new WrittenBatchReleasingIterator(inputIter, {
batch =>
batch.rowIterator().asScala.foreach { row =>
queue.add(unsafeProj(row).copy())
Review Comment:
Agreed. It is pre-existing, so I'll fix it separately, together with the
other issues of the row queue path.
##########
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
+ * than the one that requested its batch;</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 that keeps such vectors does not account for their memory, and
cannot spill 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:
Keeping it internal would leave Path 1 unreachable for every source outside
Spark, so enabling the conf could only send Arrow batches to the row queue
path. The opt-in is what lets an Arrow-backed connector that returns new
vectors for every batch use the path that SPARK-56350 added for it.
`ColumnVector` is `@Evolving`, so the method can still change, but I'd like
HyukjinKwon's opinion before it ships. I also changed L103 in cafd4cf047c as
you suggest: "A consumer may keep such vectors without accounting for their
memory or spilling them."
--
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]
