[viirya]

What changes were proposed in this pull request?

This fixes the sliced Arrow output processors of the Python runners (PythonArrowOutput.scala), which are used when the internal config spark.sql.execution.arrow.maxRecordsPerOutputBatch or spark.sql.execution.arrow.maxBytesPerOutputBatch is set, and disables the slicing for the runners that cannot consume sliced batches:

1. SliceRecordsArrowOutputProcessorImpl.getRoot returned the whole reader root whenever the remaining rows of the current batch fit the limit, even after part of the batch had already been output as slices. It now returns the whole root only when the chunk is the whole batch, and otherwise slices the remaining rows. SliceBytesArrowOutputProcessorImpl.getRoot gets a currentRowIdx == 0 guard; it already sliced every chunk after the first, so this does not change its behavior.
2. BaseSliceArrowOutputProcessor.produceBatch and close closed the previous output root, even when that root was the reader root itself. They now close the previous output only when it is a slice. The reader root is reused for every input batch, and ArrowOutputProcessorImpl.close still closes it at the end of the stream.
3. ApplyInPandasWithStatePythonRunner, ArrowPythonUDTFRunner and ColumnarArrowPythonWithNamedArgumentRunner now ignore both configs, because their consumers depend on the batch boundaries: applyInPandasWithState reads the data and state counts of the whole batch from its first row, EvalPythonUDTFExec joins each output batch with one input row, and ColumnarArrowEvalPythonEvaluatorFactory.evalArrowColumnar pairs each output batch with one input batch. The docs of both configs now say that these three cases ignore them.

Memory behavior is unchanged: a slice is still released when the next batch is produced or when the processor is closed.

Why are the changes needed?

With either config set:

1. When a batch returned from Python had more rows than maxRecordsPerOutputBatch, its last chunk was output as the whole batch again. For example, a 10-row batch with a limit of 3 was output as 3, 3, 3, and then all 10 rows. mapInArrow / mapInPandas and the other operators that flatten the output batches returned duplicated rows. Arrow scalar Python UDFs, which join each output row with the next input row, returned wrong values (the extra rows were joined with the following input rows) and then failed the task once the input rows ran out.
2. When one output batch was the whole reader root and the next one was a slice, producing the slice closed the reader root, which is still being read. The next read then failed, e.g. with IndexOutOfBoundsException: index: 0, length: 1 (expected: range(0, 0)) from StructVector.isSet. Both the records and the bytes slicing hit this.
3. Slicing the batches of the three runners above gives wrong results or failures regardless of 1 and 2: applyInPandasWithState fails with an NPE or emits padding rows, an Arrow Python UDTF in a LATERAL join with several input rows joins output rows with the wrong input rows, and the columnar-input path of Arrow Python UDFs fails its row count assertion.

Does this PR introduce any user-facing change?

Yes, only for users who set one of the internal configs above (both are disabled by default):

• mapInArrow, mapInPandas and the other affected operators no longer return duplicated rows, and Arrow scalar Python UDFs no longer return wrong values or fail.
• They no longer fail when a batch that fits the limit is followed by one that needs slicing.
• applyInPandasWithState, Arrow Python UDTFs and the columnar-input path of Arrow Python UDFs no longer slice their output batches, so they no longer fail or return wrong results under these configs.

How was this patch tested?

• Added PythonArrowOutputSuite, which writes an Arrow IPC stream of several batches (long, string and struct columns, with nulls) and drives the processors the way PythonArrowOutput does. It checks that every input row is output exactly once and in order, that no Arrow memory is leaked, that records slicing outputs chunks of exactly the expected sizes, and that bytes slicing respects the byte limit or outputs a single row. It covers records slicing with batch sizes that divide the limit evenly and unevenly, an empty batch, and batches below and above the limit, plus bytes slicing with batches below and above the limit and with a limit smaller than any single row.
• test_different_output_length in test_arrow_map.py now checks the exact output instead of comparing sets, which hid the duplicated rows.
• Added test_udtf_lateral_join_with_output_batch_slicing (test_udtf.py) and test_apply_in_pandas_with_state_output_batch_slicing (test_pandas_grouped_map_with_state.py), which run under each of the two configs and fail without the runner changes.
• Added test_udf_on_cached_input_with_output_batch_slicing (test_arrow_python_udf_cached.py), which runs an Arrow Python UDF on one 10-row Arrow-cached batch through the columnar input path under each of the two configs and fails without the change to ColumnarArrowPythonWithNamedArgumentRunner.
• test_arrow_map, test_udtf, test_pandas_grouped_map_with_state and test_arrow_python_udf_cached pass.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5.5)

This pull request and its description were written by Isaac.

[dongjoon-hyun]

Thank you for the fix, @viirya. I reviewed the latest commit, 9c4563d, and left 13 inline comments. Here is a summary, ordered by priority.

Correctness

1. SliceRecordsArrowOutputProcessorImpl L325: ApplyInPandasWithStatePythonRunner reads the batch-wide data/state counts from row 0 of every chunk, so for a batch of more than M but at most 2M rows (M = maxRecordsPerOutputBatch), the last chunk can now fail with an NPE or emit all-null padding rows, where it used to re-emit the whole batch.
2. SliceRecordsArrowOutputProcessorImpl L322: the Arrow UDTF joins each output batch with one input row (EvalPythonUDTFExec L125-127), so slicing misjoins LATERAL output, and for partitions with several input rows this PR turns the previous IndexOutOfBoundsException into silently wrong results.
3. ColumnarArrowEvalPythonEvaluatorFactory L220: the columnar-input path asserts one result batch per input batch, so any output slicing still fails it; the path is on by default in 4.2.x, and this is pre-existing.
4. SliceBytesArrowOutputProcessorImpl L271: getBatchBytes can overflow for struct-wrapped outputs over 2 GiB, because NonNullableStructVector.getBufferSize returns an int; pre-existing.

Tests

5. test_arrow_map.py L154: test_different_output_length still compares the output as a set, so the end-to-end test passes with the duplicated-rows bug reverted.
6. PythonArrowOutputSuite L156: the bytes tests check neither the per-output byte bound nor the one-row fallback.

Docs and comments

7. PR description: for Arrow scalar UDFs, the old records bug produced wrong values and then a task failure, not duplicated rows (EvalPythonEvaluatorFactory L125).
8. SliceBytesArrowOutputProcessorImpl L271: the conf doc says "equal or less than this value", but both byte checks use <; pre-existing.

Cleanup

9. BaseSliceArrowOutputProcessor L229: the batchRoot != prevRoot guard is now redundant, and it uses != while closePrevSlice uses ne.
10. SliceRecordsArrowOutputProcessorImpl L322: comparing the chunk size with rowCount states the whole-root rule directly, without the currentRowIdx == 0 clause.
11. PythonArrowOutputSuite L57: Value, toRow and fromRow can be one InternalRow generator plus ColumnarBatchRow.copy().
12. PythonArrowOutputSuite L71: writeStream can use ArrowWriter.create(schema, "UTC").
13. PythonArrowOutputSuite L106: the BufferAllocator annotation is the only use of its import.

Items 1 and 2 matter most: this PR changes how those two runners fail under these configs, from duplicated rows to an NPE, and from a task failure to silently wrong LATERAL output. Could we disable output slicing for them, here or in a follow-up, and narrow the user-facing claim in the description? Items 3 and 4 are pre-existing and fine as follow-ups, item 7 only needs the PR description, and items 9-13 are optional.

[dongjoon-hyun]

ApplyInPandasWithStatePythonRunner cannot consume sliced batches, and this line changes how it fails. The runner mixes in PythonArrowOutput without overriding arrowMaxRecordsPerOutputBatch or arrowMaxBytesPerOutputBatch, but its batches only make sense whole:

• The worker writes the batch-wide (dataCount, stateCount) into every row of the count column, and pads the shorter of data and state with null rows up to max(D, S) (python/pyspark/eval_handlers/_stateful.py L322-343).
• deserializeColumnarBatch reads the counts from row 0 of each chunk and applies take(count) to that chunk (ApplyInPandasWithStatePythonRunner.scala:206-207, 231, 240).

With maxRecordsPerOutputBatch=3 and a group that returns 5 rows with one state (one 5-row batch with counts (5, 1)):

• Before this PR, the chunks were [0, 3) and then the whole root: 3 rows were duplicated and the state was applied twice, but nothing failed.
• With this PR, they are [0, 3) and [3, 5). The second chunk takes padding state row 3, so parse(row.getUTF8String(0).toString) at ApplyInPandasWithStatePythonRunner.scala:245 throws an NPE. When 0 < D < S, all-null data rows are emitted instead.

Batches with D == S are fixed by this PR, while batches of more than twice the limit and bytes slicing have failed this way since 4.1.0. Could we override both limits to -1 in ApplyInPandasWithStatePythonRunner, here or in a follow-up, and add an applyInPandasWithState test under the slicing configs?

[viirya]

Done in c0490ad: ApplyInPandasWithStatePythonRunner now overrides both limits to -1. Added test_apply_in_pandas_with_state_output_batch_slicing, which runs under both configs; without the override it fails with the NPE you described.

[dongjoon-hyun]

The Arrow-optimized Python UDTF cannot consume sliced batches either, and with this PR it fails silently. EvalPythonUDTFExec.scala:125-127 pairs each output ColumnarBatch with one queued input row, because the worker sends exactly one batch per eval() call (see the comment at python/pyspark/eval_handlers/_udtf.py L1139). ArrowPythonUDTFRunner does not override the limits, and spark.sql.execution.pythonUDTF.arrow.enabled is true by default since 4.2.0.

With maxRecordsPerOutputBatch=3, a UDTF that yields 5 rows per input row, and SELECT t.id, u.* FROM range(0, 2, 1, 1) t, LATERAL f(t.id) u:

• Before this PR, the first slice of the batch for the second input row closed the reader root, and the next chunk failed with IndexOutOfBoundsException.
• With this PR, the output batches have 3, 2, 3 and 2 rows, so id = 0 gets 3 rows and id = 1 gets 7, two of which belong to id = 0, and nothing fails.

Partitions with a single input row are fixed by this PR. Could we disable output slicing in ArrowPythonUDTFRunner (override both limits to -1) and add a LATERAL UDTF test under the slicing configs?

[viirya]

Done in c0490ad: ArrowPythonUDTFRunner now overrides both limits to -1. Added test_udtf_lateral_join_with_output_batch_slicing, which runs a LATERAL UDTF over 3 input rows under both configs, in both the Arrow and the legacy Arrow UDTF test classes; without the override it gets the misjoined rows.

[dongjoon-hyun]

Minor: this one is pre-existing and not changed here, but the description says Arrow Python UDF output is fixed for these configs, and the columnar-input path still cannot take sliced output. ColumnarArrowEvalPythonEvaluatorFactory.scala:219-220 pairs each result batch with exactly one input batch and asserts that the row counts match. For example, in a session started with spark.sql.cache.serializer=org.apache.spark.sql.execution.columnar.ArrowCachedBatchSerializer, with spark.sql.execution.arrow.pythonUDF.columnarInput.enabled=true (the default in 4.2.x) and maxRecordsPerOutputBatch=3, a scalar Arrow UDF on id of spark.range(10).coalesce(1).cache() gets its 10-row result back as 3, 3, 3 and 1 rows, and the first chunk fails with AssertionError: ... pass-through has 10 rows but UDF result has 3 rows. Could we narrow the user-facing claim here, and disable output slicing for that runner in a follow-up?

[viirya]

Since the fix is the same, I did it here: ColumnarArrowPythonWithNamedArgumentRunner also overrides both limits to -1 in c0490ad, and the description now lists it with the other two runners. I also checked the remaining consumers of PythonArrowOutput (cogroup, transformWithState, window, map, scalar Arrow UDFs, grouped and incremental agg): they consume the output row by row or take the only row of a one-row batch, which slicing leaves as is.

[dongjoon-hyun]

Minor: this is pre-existing, but getBatchBytes can still overflow for the struct-wrapped outputs of mapInArrow, mapInPandas, applyInPandas and UDTFs. Their root holds one top-level struct vector, and in Arrow 19 NonNullableStructVector.getBufferSize sums its children in a long and returns (int) buffer (StructVector then adds the validity size as an int), so the Long accumulator at L305-307 does not help. For example, a mapInArrow UDF that yields one batch of three int64 columns with 2^27 rows (about 3.06 GiB) measures -1,006,632,960 bytes, so with maxBytesPerOutputBatch=256m this check passes and the whole root goes out unsliced. The check at L284 can accept an oversized slice the same way, and ArrowStreamReader does accept record batch bodies over 2 GiB. Could we sum the sizes of the leaf vectors as Long in a follow-up?

[viirya]

Good catch. Filed SPARK-60395 for this, together with the < vs. "equal or less" mismatch.

[dongjoon-hyun]

Minor: this exact check is what the end-to-end test lacks. As the description says, test_different_output_length (python/pyspark/sql/tests/arrow/test_arrow_map.py L154) compares the output as a set, and this PR leaves it as is, so with the getRoot fix reverted, MapInArrowWithOutputArrowBatchSlicingRecordsTests gets 199 rows and still passes. A sorted comparison would not work, because MapInArrowWithArrowBatchSlicingTestsAndReducedBatchSizeTests sends 10 input batches (1000 rows). But there is one partition and the UDF ignores its input, so the order is deterministic in all five classes that run this test. Could we make that check exact too, e.g. self.assertEqual([r.a for r in actual], list(range(100)) * (len(actual) // 100))?

[viirya]

Done in c0490ad with your check.

[dongjoon-hyun]

Nit: since closePrevSlice() now skips the reader root, this guard no longer changes the outcome. getRoot returns either root or a new root.slice(...), so batchRoot and prevRoot are the same object only for a whole root after a whole root, where closePrevSlice() is already a no-op. The guard also uses != while closePrevSlice uses ne; != is a reference check here only because VectorSchemaRoot does not override equals(Object), although it has a content-comparing equals(VectorSchemaRoot) overload. Could we simplify it as follows?

    val batchRoot = getRoot
    closePrevSlice()
    prevRoot = batchRoot

[viirya]

Done in c0490ad.

[dongjoon-hyun]

Nit: computing the chunk size once states the rule directly (return the reader root only when this chunk is the whole batch), and it needs neither the currentRowIdx == 0 clause nor the second comparison inside math.min:

    val numRows = math.min(rowCount - currentRowIdx, arrowMaxRecordsPerOutputBatch)
    val rootSlice = if (numRows == rowCount) root else root.slice(currentRowIdx, numRows)

When currentRowIdx > 0, numRows < rowCount, so this behaves the same for empty, exact-fit, oversized and partly-output batches. The guard at L271 can stay, since it also skips recomputing getBatchBytes(root). Could we use this form?

[viirya]

Done in c0490ad.

[dongjoon-hyun]

Nit: Value, toRow and fromRow describe the schema three more times, and nulls are encoded two ways (null for s and b, Option for st). ColumnarBatchRow.copy() already deep-copies the strings and the nested struct, and GenericInternalRow compares field by field, so one generator could replace all three:

  private def row(i: Long): InternalRow = {
    val s = if (i % 3 == 0) null else UTF8String.fromString(s"str_$i")
    val st = if (i % 5 == 0) null else {
      InternalRow(i.toInt * 2, if (i % 2 == 0) null else UTF8String.fromString(s"b_$i"))
    }
    InternalRow(i, s, st)
  }

Then writer.write(row(next)), output += batch.rowIterator().asScala.map(_.copy()).toSeq and assert(output.flatten === (0L until batchSizes.sum).map(row)) check the same thing, and PythonStreamingSourceRunner.scala:368 uses the same _.copy() idiom. Could we simplify it this way?

[viirya]

Done in c0490ad.

[dongjoon-hyun]

Nit: ArrowWriter.create(schema, "UTC") (ArrowWriter.scala:36-45) builds the same Arrow schema (toArrowSchema(schema, "UTC", true, false)) and the root on ArrowUtils.rootAllocator, as ArrowWriterSuite does. The dedicated "writer" allocator only checks the buffers of the fixture itself, and the "reader" allocator still catches processor leaks. Could we use val writer = ArrowWriter.create(schema, "UTC"), pass writer.root to ArrowStreamWriter, and close writer.root in finally, which also makes the VectorSchemaRoot import unnecessary?

[viirya]

Done in c0490ad.

[dongjoon-hyun]

Nit: newChildAllocator already returns BufferAllocator, so this can be the same one-liner as L72, val allocator = ArrowUtils.rootAllocator.newChildAllocator("reader", 0, Long.MaxValue) (90 characters), and then the BufferAllocator import at L25 is unused. Could we drop the annotation?

[viirya]

Done in c0490ad.

[dongjoon-hyun]

Thank you for addressing items 1-13, @viirya. I reviewed the latest commit, c0490ad, and left 15 more inline comments, continuing the numbering from items 1-13. Here is a summary, ordered by priority.

Correctness (pre-existing, outside this PR's diff)

14. EvalPythonUDTFExec L104 (posted on ArrowPythonUDTFRunner L54): in pipelined mode, count is updated by the writer and task threads without synchronization, so a lost update can silently join UDTF output with the wrong input row.
15. ColumnarArrowEvalPythonEvaluatorFactory L220 (posted on ColumnarArrowPythonWithNamedArgumentRunner L56): scalar iterator UDFs reach evalArrowColumnar and can change the batch boundaries, which fails its row count assertion without any slicing.
16. ColumnarArrowEvalPythonEvaluatorFactory L196 (posted on ColumnarArrowPythonWithNamedArgumentRunner L56): in pipelined mode, passThruQueue is a plain ArrayDeque shared by the writer and task threads.

Tests

17. ColumnarArrowPythonWithNamedArgumentRunner L57: this override is the only one of the three without a test that fails without it.
18. test_arrow_map.py L155: the exact check I suggested in item 5 passes on an empty result.
19. PythonArrowOutputSuite L125: the records tests bound the output sizes but do not pin them.

Docs and comments

20. ApplyInPandasWithStatePythonRunner L77: the docs of both configs do not say that these three runners now ignore them.
21. ColumnarArrowPythonWithNamedArgumentRunner L55: only evalArrowColumnar pairs each output batch with one input batch.
22. test_pandas_grouped_map_with_state.py L198: the worker packs all groups of a partition into one batch, and both keys share a partition here.
23. PythonArrowOutputSuite L133: this comment and the one at L102 overstate what the case and the leak check cover.

Performance and cleanup

24. SliceBytesArrowOutputProcessorImpl L271: the shrink search restarts from all remaining rows for every output batch, which copies O(R^2/c) offsets; pre-existing.
25. test_udtf.py L3808: the rest of the file compares rows with assertDataFrameEqual.
26. test_pandas_grouped_map_with_state.py L201: the loop over pdf_iter is not needed.
27. SliceBytesArrowOutputProcessorImpl L271: if (sliceBatchCount == 0) is always true; pre-existing.
28. SliceRecordsArrowOutputProcessorImpl L321: the index could advance by numRows.

The processor fixes and the three overrides look correct to me, and your analysis of the other consumers holds. Item 17 matters most for this PR, and items 18 and 20-23 are small test, comment and doc fixes. Items 14-16 are pre-existing and outside this PR, so separate JIRAs would fit; items 14 and 16 affect the pipelined mode in the 4.3.0 RCs. Items 24 and 27 could go with SPARK-60395, and items 19, 25, 26 and 28 are optional.

[dongjoon-hyun]

Minor: this is pre-existing and outside this PR's diff, but it breaks the pairing that this comment relies on. In pipelined mode (spark.python.udf.pipelined.enabled, new in 4.3.0, default false), BasePythonRunner.compute pulls the input iterator on a writer thread (PythonRunner.scala:464, :543, :1322), so count += 1 in projectedRowIter (EvalPythonUDTFExec.scala:104) runs on that thread, while if (count > 0) { queue.remove(); count -= 1 } (EvalPythonUDTFExec.scala:125-127) runs on the task thread. count is a captured local var, i.e. a plain scala.runtime.LongRef, and only the queue was made safe for this split (lockFree = pipelined, EvalPythonUDTFExec.scala:63-67). If the task reads c, the writer stores c + 1 and the task then stores c - 1, the increment is lost: when the task reaches the last queued input row, count is already 0, so that row's output is joined with the previous input row without any error. A lost decrement instead makes the first terminate() batch call queue.remove() on an exhausted queue. This is on master, branch-4.x and branch-4.3, not on branch-4.2, and the new LATERAL test runs in non-pipelined mode. Could we file a separate JIRA to make count an AtomicLong, or to split it into a writer-only @volatile counter and a task-only one?

[viirya]

Good catch. SPARK-56642 made the queue lock-free for the writer/task split but left count shared. Filed SPARK-60414 for it, together with the passThruQueue race (item 16).

[dongjoon-hyun]

Minor: this is pre-existing and outside this PR's diff, but scalar iterator UDFs break the premise of this comment even without any slicing. ExtractPythonUDFs sends SQL_SCALAR_PANDAS_ITER_UDF and SQL_SCALAR_ARROW_ITER_UDF to ArrowEvalPython, and ArrowEvalPythonExec.supportsColumnar (ArrowEvalPythonExec.scala:104-105) has no eval-type check, so they reach evalArrowColumnar. The worker checks only their total row count, and the API allows prefetching, so their result batches need not line up with the input batches. For example, with the Arrow cache serializer, spark.sql.execution.arrow.pythonUDF.columnarInput.enabled=true (the default on branch-4.3) and spark.sql.execution.arrow.maxRecordsPerBatch=5, the following returns the two 5-row cached batches as one 10-row result, and ColumnarArrowEvalPythonEvaluatorFactory.scala:220-222 fails the task with Batch size mismatch: pass-through has 5 rows but UDF result has 10 rows.

df = spark.range(10).coalesce(1).cache()
df.count()

@pandas_udf("long")
def f(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
    yield pd.concat(list(it)) + 1

df.select(f("id")).collect()

Could we file a follow-up to send the two iterator eval types to the row-queue path, with a test over two cached batches?

[viirya]

Agreed, the columnar path should not take the iterator eval types. Filed SPARK-60415 to send them to the row-queue path, with a test over two cached batches.

[dongjoon-hyun]

Minor: also pre-existing and outside this PR's diff, on the same Arrow columnar path. In pipelined mode, passThruQueue (ColumnarArrowEvalPythonEvaluatorFactory.scala:195-196) is a plain java.util.ArrayDeque: the writer thread adds to it while it pulls the input (ColumnarArrowEvalPythonEvaluatorFactory.scala:198-203, via PythonRunner.scala:1322), and the task thread polls it (ColumnarArrowEvalPythonEvaluatorFactory.scala:219). The row-queue paths in the same file use HybridRowQueue(..., lockFree = pipelined) for exactly this split (ColumnarArrowEvalPythonEvaluatorFactory.scala:243-250). The socket orders each add before the poll for its own result, but nothing orders a later add and its grow() against that poll: if pollFirst reads head before grow() moves the elements and writes head + 1 after it, the next poll() returns null and ColumnarArrowEvalPythonEvaluatorFactory.scala:219 throws scala.MatchError: null; in the other order, one batch comes out twice and the next result gets the wrong pass-through columns. It needs more than 16 batches in flight, so it is rare, and it is on master, branch-4.x and branch-4.3. Could we use a ConcurrentLinkedQueue there in a follow-up?

[viirya]

Agreed. Filed it in SPARK-60414 together with the count race in EvalPythonUDTFExec (item 14), since both come from the pipelined writer thread.

[dongjoon-hyun]

This override is the only one of the three without a test that fails without it. Reverting L57-58 keeps every test green, since nothing runs the columnar path under either config, and the testing section of the description names only the UDTF and applyInPandasWithState tests. The repro of my item 3 still applies: with the Arrow cache serializer, columnar input enabled and maxRecordsPerOutputBatch=3, a scalar Arrow UDF on id of spark.range(0, 10, 1, 1).cache() takes evalArrowColumnar and fails at ColumnarArrowEvalPythonEvaluatorFactory.scala:220-222 with pass-through has 10 rows but UDF result has 3 rows (1 rows with maxBytesPerOutputBatch=3). ArrowPythonUDFCachedInputTests (test_arrow_python_udf_cached.py) already sets the static serializer, enables columnar input and has _assert_columnar_arrow_eval, but its 2-3-row inputs form 1-row batches, so the test needs one multi-row cached batch. Could we add a test like the following there and name it in the description? It also matters for backports, since columnar input is on by default on branch-4.3 and branch-4.2.

    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()

[viirya]

Added test_udf_on_cached_input_with_output_batch_slicing as you suggested in eced9b0. Without this override it fails with pass-through has 10 rows but UDF result has 3 rows (1 rows with the bytes config). The description now names it.

[dongjoon-hyun]

Nit: correction to my item 5. The check I suggested passes when actual is empty, because list(range(100)) * (0 // 100) is [], while the old set check failed in that case, so the test no longer catches a lost output. Every class here sends at least one input batch, so could we use list(range(100)) * max(1, len(actual) // 100)?

[viirya]

Done in eced9b0.

[dongjoon-hyun]

Minor: this is pre-existing (SPARK-51931), but the search below restarts from all remaining rows for every output batch (L276), and each probe that starts at a non-zero row allocates and refills its offset buffers (BaseVariableWidthVector.splitAndTransferOffsetBuffer in Arrow 19). So splitting one batch of R rows into chunks of c rows copies O(R^2/c) offsets instead of O(R). By my estimate, a 1M-row struct<long, string> batch (about 32 MiB) with maxBytesPerOutputBatch=1m takes about 824 probes for 33 outputs and copies about 567 MiB of offsets, versus about 4 MiB for a single pass. Also, the first probe, slice(0, rowCount), always fails because it measures the same as the root that L268 just checked, and the one-row fallback rebuilds the 1-row slice that the loop just closed. Could we start each search from the last row count that fit, e.g. in SPARK-60395, which touches this loop anyway?

[viirya]

Agreed, added it to SPARK-60395.

[dongjoon-hyun]

Nit: the rest of this file compares rows with assertDataFrameEqual (imported at L70 and used by the LATERAL test at L420-423), which ignores the row order, keeps duplicates and shows the differing rows on failure. Could we use it here too?

                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)

[viirya]

Done in eced9b0.

[dongjoon-hyun]

Nit: this loop is not needed. The worker skips the input chunks that the function does not read (_stateful.py:370-380) and builds the output batch only from the yielded DataFrame and the state, and test_apply_in_pandas_with_state_basic_more_data (L147-166) also calls state.update((1,)) and yields without touching pdf_iter. Could we drop these two lines?

[viirya]

Dropped in eced9b0.

[dongjoon-hyun]

Nit: also pre-existing, in the same branch. The loop below ends only through return (L284) or with sliceBatchCount == 0, because (n * 0.9).toInt is in [0, n) for n >= 1, so if (sliceBatchCount == 0) at L292-294 is always true, and the initial root of var rootSlice (L274) is never read. Could we replace L292-295 with rootSlice = root.slice(currentRowIdx, 1), e.g. together with SPARK-60395?

[viirya]

Added it to SPARK-60395 as well.

[dongjoon-hyun]

Nit (optional): now that numRows is computed above, could we advance by it, currentRowIdx = currentRowIdx + numRows, so that the chunk choice and the index update use the same value? It behaves the same for every batch, since numRows == rowCount == root.getRowCount in the whole-root case and a slice has numRows rows.

[viirya]

Done in eced9b0.

[dongjoon-hyun]

Thank you for addressing items 14-28, @viirya. I reviewed the latest commit, eced9b0, and left 15 more inline comments, continuing the numbering from items 14-28. Here is a summary, ordered by priority.

Correctness (pre-existing, outside this PR's diff)

29. PythonArrowOutput L49 (posted on SQLConf L5477): .toInt wraps a negative maxBytesPerOutputBatch below Int.MinValue, e.g. -3g, to a positive limit, so a value that the doc calls no limit turns byte slicing on.

API and design

30. PythonArrowOutput L47-49 (posted on ArrowPythonUDTFRunner L56): nothing at the two defaults says when a runner has to override them, and none of the three runners that need it had the opt-out before this PR.

Tests

31. SliceRecordsArrowOutputProcessorImpl L319: no test runs an Arrow scalar UDF on the row path under either config, although the description lists it as a user-facing fix.

Docs and comments

32. SliceBytesArrowOutputProcessorImpl L268: the PR description says this guard changes nothing because the old code "already sliced every chunk after the first", but the old code returned the closed root itself as the next chunk.
33. SQLConf L5456: in a branch-4.1 backport, the sentence from my item 20 would name a columnar-input path that branch-4.1 does not have.
34. BaseSliceArrowOutputProcessor L245: "input batch" here is a batch read from Python, while the runner comments use "input" for what is sent to Python.
35. ApplyInPandasWithStatePythonRunner L78: the PR title names only the processor fixes, not that three runners now ignore both configs.
36. PythonArrowOutputSuite L114: the description's "batch sizes that divide the limit evenly and unevenly" has the relation reversed.

Cleanup

37. BaseSliceArrowOutputProcessor L249: prevVectors is kept only to be closed here, and prevRoot.close() already releases the same vectors.
38. test_arrow_python_udf_cached.py L110: the new test sits between _assert_udf_on_cached_strings and its only two callers.
39. BaseSliceArrowOutputProcessor L242: getVectors builds new wrappers for every whole-root output, which could now reuse the base vectors; pre-existing.
40. BaseSliceArrowOutputProcessor L217 (posted on L226): rowCount > 0 && never changes the result; pre-existing.
41. PythonArrowOutputSuite L145: with L146, this assertion cannot fail on its own for the current cases.
42. test_pandas_grouped_map_with_state.py L219: _test_apply_in_pandas_with_state_basic leaves this_query running after each call; pre-existing.
43. PythonArrowOutputSuite L100: readAll closes the processor once, while production closes it twice.

Items 14-28 are all addressed, and the new tests ran and passed on eced9b0, including the Connect parity classes. Item 29 is the only behavior issue, and it is pre-existing, so it could go with SPARK-60395 like items 24 and 27. Items 30 and 31 are a short comment and an optional end-to-end test, items 32, 35 and 36 only need the PR description or title, item 33 only matters for a branch-4.1 backport, and items 34 and 37-43 are optional.

[dongjoon-hyun]

Minor: this is pre-existing (SPARK-51931) and outside this PR's diff, but the checkValue below (L5478) bounds only the upper side, while PythonArrowOutput.scala:49 narrows the value with .toInt, so some negative values below Int.MinValue, which the doc above calls no limit, wrap to a positive limit. byteFromString (ConfigBuilder.scala:82-90) accepts a leading -, so -3g parses to -3221225472 and passes the check, and its low 32 bits are 1073741824, so PythonArrowOutput.scala:130 builds a SliceBytesArrowOutputProcessorImpl with a 1 GiB limit. -4294967295 becomes a 1-byte limit, which outputs every row as its own batch, and if maxRecordsPerOutputBatch is also set, PythonArrowOutput.scala:51-54 throws for the runners that keep the trait defaults, although the bytes limit is off. The same .toInt is on branch-4.1, branch-4.2 and branch-4.x. Could we clamp the value before narrowing, e.g. math.max(SQLConf.get.arrowMaxBytesPerOutputBatch, 0L).toInt at PythonArrowOutput.scala:49 (or bound the lower side here), together with SPARK-60395?

[dongjoon-hyun]

Nit: the two defaults at PythonArrowOutput.scala:47-49 do not say when a runner has to override them. Once either internal config is set, every runner that mixes in the trait slices its output (PythonArrowOutput.scala:127-135), so these overrides are opt-outs, and none of the three runners that need one had it before this PR. ApplyInPandasWithStatePythonRunner (SPARK-40434, 3.4.0) and this runner (SPARK-43964, 3.5.0) predate the configs (SPARK-51769 and SPARK-51931, 4.1.0), and ColumnarArrowPythonWithNamedArgumentRunner (SPARK-56350, 4.2.0) was added after them, although its consumer asserts one result batch per input batch (ColumnarArrowEvalPythonEvaluatorFactory.scala:219-222). Could we add a short comment on the two defs, so that the next runner author sees the contract there? For example:

  // 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.

[dongjoon-hyun]

Nit (optional): the first user-facing bullet of the description says that Arrow scalar Python UDFs no longer return wrong values or fail, but no test runs one on the row path under either config. Only test_arrow_map.py (L416, L424) and the three new tests set them, and test_udf_on_cached_input_with_output_batch_slicing takes evalArrowColumnar, which now ignores them. PythonArrowOutputSuite already pins this getRoot, so this would only add end-to-end coverage of the row-by-row join at EvalPythonEvaluatorFactory.scala:124-125. Could we add something like the following to ArrowPythonUDFTestsMixin (test_arrow_python_udf.py), so that it also runs in the legacy, non-legacy and Connect classes? Without this fix it fails as in my item 7, with NoSuchElementException from HybridQueue.scala:165.

    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)])

A maxBytesPerOutputBatch=3 case like in the other three tests would pass without this PR, since it cuts every 82-byte result batch down to one-row slices and never outputs the whole root.

[dongjoon-hyun]

Nit (PR description): item 1 of the description says that this guard does not change the behavior because the bytes getRoot "already sliced every chunk after the first", but before this PR it did not once produceBatch had closed the reader root. The closed vectors measure 0 bytes, so the old getBatchBytes(root) < arrowMaxBytesPerOutputBatch check returned the closed root itself as the next chunk, with currentRowIdx += root.getRowCount, since VectorSchemaRoot.close() keeps the row count; that is how the bytes slicing hit item 2 under "Why are the changes needed?". With the second fix, every chunk after the first is a slice anyway, so the guard is a no-op now. Could we say "SliceBytesArrowOutputProcessorImpl.getRoot gets a currentRowIdx == 0 guard; since 2 keeps the reader root open, every chunk after the first is a slice anyway, so this guard does not change the behavior."?

[dongjoon-hyun]

Nit (backport): the sentence I suggested in item 20 is accurate on master, branch-4.x and branch-4.3, but branch-4.1 has no columnar-input path: neither spark.sql.execution.arrow.pythonUDF.columnarInput.enabled nor ColumnarArrowPythonWithNamedArgumentRunner, which SPARK-56350 added in 4.2.0. A branch-4.1 backport would drop the ColumnarArrowPythonRunner.scala and test_arrow_python_udf_cached.py hunks, while both doc strings there are the same as at the merge base, so this hunk would apply cleanly and name an exception that branch-4.1 does not have. Could a branch-4.1 backport say "applyInPandasWithState and Arrow Python UDTFs ignore this config." in both docs? Also, branch-4.2 has neither test_arrow_python_udf_cached.py nor ArrowCachedBatchSerializer, so the columnar override would go there without a test.

[dongjoon-hyun]

Nit: also pre-existing, but this method builds a new set of ArrowColumnVector wrappers, including those of struct children, for every output, also when getRoot returns the reader root itself, although ArrowOutputProcessorImpl already holds wrappers over that root (vectors, L174-176). Since closePrevSlice no longer closes the root, they could be reused, as the non-slicing processor does for every batch (L193). The gain is small, a few wrapper objects and ArrowUtils.fromArrowField calls per whole-root output. Could we reuse them, e.g. as follows, here or together with #59581, which changes this exact method?

  protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = {
    if (root eq this.root) {
      super.getVectors(root)
    } else {
      root.getFieldVectors.asScala.map { vector =>
        new ArrowColumnVector(vector)
      }.toArray[ColumnVector]
    }
  }

[dongjoon-hyun]

Nit: pre-existing and outside this PR's diff (posted here since loadBatch at L216-224 has no diff line), but rowCount > 0 && at L217 never changes the result. currentRowIdx is -1 only while rowCount is -1 (L178, L212) and otherwise stays within [0, rowCount], so currentRowIdx < rowCount alone already implies rowCount >= 1. It gives the same answer for the initial state, an empty batch, a partly or fully output batch, and after the end of the stream. Could we drop the clause, e.g. with SPARK-60395?

[dongjoon-hyun]

Nit (optional): with L146, this assertion cannot fail on its own for any case in the table. Each case has a multi-row input batch of at least maxBytes, e.g. the 57-row and 33-row batches of Seq(1, 0, 57, 2, 33) measure 1188 and 693 bytes even without their string data, so an unsliced batch already fails L146. It would only matter for a future case whose batches all fit. Could we drop it, or say in the comment above that it guards the choice of cases?

[dongjoon-hyun]

Nit: pre-existing and not changed here, but _test_apply_in_pandas_with_state_basic (L62-99) neither stops its query nor removes its input directory, so each call leaves this_query running until the next call stops it (L79-80). This test now does that twice, and in the Connect parity class each live query also keeps a server-side foreachBatch worker. The cost is small, since at most one such query is alive at a time. Could we add a try/finally with q.stop() and shutil.rmtree(input_path, ignore_errors=True) to the helper, here or in a follow-up, so that it covers all its callers?

[dongjoon-hyun]

Nit (optional): readAll closes the processor once, while production closes it twice, in read() at the end of the stream (PythonArrowOutput.scala:114-116) and again in the task completion listener (PythonArrowOutput.scala:77-85), the second time after the allocator is closed. Since closePrevSlice does not reset prevRoot, the second close() closes the last slice, the base vectors and the root again. That is safe in Arrow 19, because released buffers become the allocator's empty buffer, whose release is a no-op, and BaseAllocator.close() returns early once closed, so this is not a bug, but the suite does not pin it. Could we call processor.close() once more after allocator.close() at L103, so that a later change to close() that touches the allocator would fail here?