[SPARK-60367][SQL] Fix sliced Arrow output processors and disable slicing for runners that need whole batches

python/pyspark/sql/tests/arrow/test_arrow_map.py

@@ -151,7 +151,8 @@ def func(iterator):
 
         df = self.spark.range(10)
         actual = df.repartition(1).mapInArrow(func, "a long").collect()
-        self.assertEqual(set((r.a for r in actual)), set(range(100)))
+        # Each input batch produces 0..99 once, in order, since there is one partition.
+        self.assertEqual([r.a for r in actual], list(range(100)) * max(1, len(actual) // 100))
 
     def test_other_than_recordbatch_iter(self):
         with self.quiet():

python/pyspark/sql/tests/arrow/test_arrow_python_udf.py

@@ -385,6 +385,14 @@ def create_struct_with_interval(interval_val, name_val):
         self.assertIsNotNone(result["interval_field"])
         self.assertEqual(result["name"], "test_name")
 
+    def test_udf_with_output_batch_slicing(self):
+        # Each result row is joined with the next queued input row, so the sliced 10-row result
+        # batch must be output once and in order.
+        df = self.spark.range(0, 10, 1, 1)
+        with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"}):
+            result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id")))
+            assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)])
+
 
 @unittest.skipIf(
     not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message

python/pyspark/sql/tests/arrow/test_arrow_python_udf_cached.py

@@ -133,6 +133,22 @@ def test_udf_on_cached_numeric_column_alongside_unselected_string(self):
         finally:
             df.unpersist()
 
+    def test_udf_on_cached_input_with_output_batch_slicing(self):
+        # The columnar path pairs each result batch with one input batch, so the 10-row result
+        # of the single cached batch must not be sliced under the configs below.
+        df = self.spark.range(0, 10, 1, 1).cache()
+        try:
+            for conf in [
+                {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"},
+                {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"},
+            ]:
+                with self.subTest(conf=conf), self.sql_conf(conf):
+                    result = df.select(col("id"), udf(lambda x: x + 1, "long")(col("id")))
+                    self._assert_columnar_arrow_eval(result, expected_nodes=1)
+                    assertDataFrameEqual(result, [Row(i, i + 1) for i in range(10)])
+        finally:
+            df.unpersist()
+
     def test_udf_on_cached_map_column(self):
         # Canonical map vectors from the cache pass the gate (their entry-struct children carry
         # the canonical "key"/"value" names in order) and ride the columnar path verbatim. The

python/pyspark/sql/tests/pandas/test_pandas_grouped_map_with_state.py

@@ -93,10 +93,14 @@ def prepare_test_resource():
             .start()
         )
 
-        self.assertEqual(q.name, "this_query")
-        self.assertTrue(q.isActive)
-        q.processAllAvailable()
-        self.assertTrue(q.exception() is None)
+        try:
+            self.assertEqual(q.name, "this_query")
+            self.assertTrue(q.isActive)
+            q.processAllAvailable()
+            self.assertTrue(q.exception() is None)
+        finally:
+            q.stop()
+            shutil.rmtree(input_path, ignore_errors=True)
 
     def test_apply_in_pandas_with_state_basic(self):
         def func(key, pdf_iter, state):
@@ -194,6 +198,30 @@ def check_results(batch_df, _):
 
         self._test_apply_in_pandas_with_state_basic(func, check_results)
 
+    def test_apply_in_pandas_with_state_output_batch_slicing(self):
+        # The worker packs the data rows (5 per group) and state rows (1 per group) of the groups
+        # in a partition into one output batch. Its data and state counts cover the whole batch,
+        # so it must not be sliced by the configs below.
+        def func(key, pdf_iter, state):
+            state.update((1,))
+            yield pd.DataFrame(
+                {"key": [f"{key[0]}_{i}" for i in range(5)], "countAsString": list("01234")}
+            )
+
+        expected = sorted(
+            Row(key=f"{k}_{i}", countAsString=str(i)) for k in ["hello", "this"] for i in range(5)
+        )
+
+        def check_results(batch_df, _):
+            assert sorted(batch_df.collect()) == expected
+
+        for conf in [
+            {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"},
+            {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"},
+        ]:
+            with self.subTest(conf=conf), self.sql_conf(conf):
+                self._test_apply_in_pandas_with_state_basic(func, check_results)
+
     def test_apply_in_pandas_with_state_python_worker_random_failure(self):
         input_path = tempfile.mkdtemp()
         output_path = tempfile.mkdtemp()

python/pyspark/sql/tests/test_udtf.py

@@ -3786,6 +3786,27 @@ def eval(self):
                 ):
                     udtf(TestUDTF, returnType=ret_type)().collect()
 
+    def test_udtf_lateral_join_with_output_batch_slicing(self):
+        # Each output batch is joined with one input row, so it must not be sliced by the
+        # configs below.
+        @udtf(returnType="x: int")
+        class TestUDTF:
+            def eval(self, n: int):
+                for i in range(5):
+                    yield (n * 10 + i,)
+
+        self.spark.udtf.register("test_udtf_slicing", TestUDTF)
+        expected = [Row(id=n, x=n * 10 + i) for n in range(3) for i in range(5)]
+        for conf in [
+            {"spark.sql.execution.arrow.maxRecordsPerOutputBatch": "3"},
+            {"spark.sql.execution.arrow.maxBytesPerOutputBatch": "3"},
+        ]:
+            with self.subTest(conf=conf), self.sql_conf(conf):
+                df = self.spark.sql(
+                    "SELECT t.id, u.x FROM range(0, 3, 1, 1) t, LATERAL test_udtf_slicing(t.id) u"
+                )
+                assertDataFrameEqual(df, expected)
+
 
 class LegacyUDTFArrowTests(LegacyUDTFArrowTestsMixin, ReusedSQLTestCase):
     @classmethod

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala

@@ -5452,7 +5452,9 @@ object SQLConf {
         "there is no limit. Note that the complete ArrowRecordBatch is actually created but " +
         "the number of records is limited when sending it to the downstream operator. This is " +
         "used to avoid large batches being sent to the downstream operator including " +
-        "the columnar-based operator implemented by third-party libraries.")
+        "the columnar-based operator implemented by third-party libraries. " +
+        "applyInPandasWithState, Arrow Python UDTFs and Arrow Python UDFs that take " +
+        "columnar input directly ignore this config.")
       .version("4.1.0")
       .internal()
       .intConf
@@ -5468,7 +5470,8 @@ object SQLConf {
         "the columnar-based operator implemented by third-party libraries. Spark will try to " +
         "create batches with the size equal or less than this value. Normally it should not " +
         "happen, but if in extreme case that even one record is still larger than this value, " +
-        "Spark will create a batch with one record.")
+        "Spark will create a batch with one record. applyInPandasWithState, Arrow Python " +
+        "UDTFs and Arrow Python UDFs that take columnar input directly ignore this config.")
       .version("4.1.0")
       .internal()
       .bytesConf(ByteUnit.BYTE)

sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala

@@ -51,6 +51,11 @@ class ArrowPythonUDTFRunner(
   with BasicPythonArrowOutput {
   ArrowUtils.failDuplicatedFieldNames(schema)
 
+  // `EvalPythonUDTFExec` joins each output batch with one input row, so the batches returned
+  // from Python must not be sliced.
+  override protected def arrowMaxRecordsPerOutputBatch: Int = -1
+  override protected def arrowMaxBytesPerOutputBatch: Int = -1
+
   override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf
 
   override protected def evalConf: Map[String, String] = {

sql/core/src/main/scala/org/apache/spark/sql/execution/python/ColumnarArrowPythonRunner.scala

@@ -52,6 +52,11 @@ private[python] class ColumnarArrowPythonWithNamedArgumentRunner(
   with ColumnarArrowPythonInput
   with BasicPythonArrowOutput {
 
+  // `ColumnarArrowEvalPythonEvaluatorFactory.evalArrowColumnar` pairs each output batch with
+  // one input batch, so the batches returned from Python must not be sliced.
+  override protected def arrowMaxRecordsPerOutputBatch: Int = -1
+  override protected def arrowMaxBytesPerOutputBatch: Int = -1
+
   override protected def runnerConf: Map[String, String] = super.runnerConf ++ pythonRunnerConf
 
   // The input schema (and per-UDF nesting depth) travel in evalConf, exactly like

sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala

@@ -44,6 +44,9 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] { self: BasePythonRunner[
 
   protected def deserializeColumnarBatch(batch: ColumnarBatch, schema: StructType): OUT
 
+  // The limits for slicing the batches returned from Python. A runner whose consumer cannot take
+  // a batch split into several, e.g. one that pairs each batch with one input row, overrides both
+  // to -1, as `ArrowPythonUDTFRunner` does.
   protected def arrowMaxRecordsPerOutputBatch: Int = SQLConf.get.arrowMaxRecordsPerOutputBatch
 
   protected def arrowMaxBytesPerOutputBatch: Int = SQLConf.get.arrowMaxBytesPerOutputBatch.toInt
@@ -210,8 +213,7 @@ abstract class BaseSliceArrowOutputProcessor(
   extends ArrowOutputProcessorImpl(reader, pythonMetrics) {
 
   protected var currentRowIdx = -1
-  protected var prevRoot: VectorSchemaRoot = null
-  protected var prevVectors: Array[ColumnVector] = _
+  protected var prevRoot: VectorSchemaRoot = root
 
   override def loadBatch(): Boolean = {
     if (rowCount > 0 && currentRowIdx < rowCount) {
@@ -225,34 +227,34 @@ abstract class BaseSliceArrowOutputProcessor(
 
   override def produceBatch(): ColumnarBatch = {
     val batchRoot = getRoot
+    closePrevSlice()
+    prevRoot = batchRoot
 
-    if (batchRoot != prevRoot) {
-      if (prevRoot != null) {
-        prevVectors.foreach(_.close())
-        prevRoot.close()
-      }
-      prevRoot = batchRoot
-    }
-
-    val vectors = getVectors(batchRoot)
-    prevVectors = vectors
-
-    val batch = new ColumnarBatch(vectors)
+    val batch = new ColumnarBatch(getVectors(batchRoot))
     batch.setNumRows(batchRoot.getRowCount)
     batch
   }
 
   protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = {
-    root.getFieldVectors.asScala.map { vector =>
-      new ArrowColumnVector(vector)
-    }.toArray[ColumnVector]
+    if (root eq this.root) {
+      super.getVectors(root)
+    } else {
+      root.getFieldVectors.asScala.map { vector =>
+        new ArrowColumnVector(vector)
+      }.toArray[ColumnVector]
+    }
   }
 
-  override def close(): Unit = {
-    if (prevRoot != null) {
-      prevVectors.foreach(_.close())
+  // Releases the previous output batch if it is a slice. The reader loads every batch from Python
+  // into the same root, so it must not be closed here; `ArrowOutputProcessorImpl.close` closes it.
+  private def closePrevSlice(): Unit = {
+    if (prevRoot ne root) {
       prevRoot.close()
     }
+  }
+
+  override def close(): Unit = {
+    closePrevSlice()
     super.close()
   }
 }
@@ -264,7 +266,8 @@ class SliceBytesArrowOutputProcessorImpl(
   extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
 
   protected override def getRoot: VectorSchemaRoot = {
-    if (getBatchBytes(root) < arrowMaxBytesPerOutputBatch) {
+    // Output the whole root only if no row of this batch has been output yet.
+    if (currentRowIdx == 0 && getBatchBytes(root) < arrowMaxBytesPerOutputBatch) {
       currentRowIdx = currentRowIdx + root.getRowCount
       root
     } else {
@@ -313,14 +316,11 @@ class SliceRecordsArrowOutputProcessorImpl(
   extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
 
   protected override def getRoot: VectorSchemaRoot = {
-    val remainingRows = rowCount - currentRowIdx
-    val rootSlice = if (remainingRows > arrowMaxRecordsPerOutputBatch) {
-      root.slice(currentRowIdx, arrowMaxRecordsPerOutputBatch)
-    } else {
-      root
-    }
+    val numRows = math.min(rowCount - currentRowIdx, arrowMaxRecordsPerOutputBatch)
+    // Output the whole root only if this chunk is the whole batch.
+    val rootSlice = if (numRows == rowCount) root else root.slice(currentRowIdx, numRows)
 
-    currentRowIdx = currentRowIdx + rootSlice.getRowCount
+    currentRowIdx = currentRowIdx + numRows
 
     rootSlice
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/python/streaming/ApplyInPandasWithStatePythonRunner.scala

@@ -72,6 +72,11 @@ class ApplyInPandasWithStatePythonRunner(
   with PythonArrowOutput[OutType] {
   ArrowUtils.failDuplicatedFieldNames(inputSchema)
 
+  // The data and state counts of each batch returned from Python cover the whole batch and are
+  // read from its first row, so the batch must not be sliced.
+  override protected def arrowMaxRecordsPerOutputBatch: Int = -1
+  override protected def arrowMaxBytesPerOutputBatch: Int = -1
+
   // Installed at worker launch. This runner is constructed per task, so each launch reads the
   // configuration that task carries. A running streaming query carries a snapshot:
   // `StreamExecution` runs its batches on `sparkSession.cloneSession()`, whose conf is copied at