viirya commented on code in PR #59582:
URL: https://github.com/apache/spark/pull/59582#discussion_r4471394046
##########
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") {
Review Comment:
Done in cafd4cf047c, except the last point:
- The fresh-vector and not-retainable tests are folded into the SPARK-60306
grid with a `retainable` field, so the not-retainable case runs in both codegen
modes and checks `maxUnreleasedBatches`.
- The reusing-reader test runs with columnar input on and off, and replaces
the restored default-configuration test.
- The cache test is now the columnar-input case of the test in
`ArrowCachedPythonUDFSuite`, and the serializer fixture is gone from this suite.
I kept `ArrowColumnarPythonUDFPipelinedSuite` as it is, since it comes from
SPARK-60306, and the 200,000-row cache test no longer runs in it.
--
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]
