You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by yi...@apache.org on 2026/09/29 00:30:06 UTC
(spark) branch branch-4.x updated: [SPARK-59624][PYTHON][FOLLOWUP] Drop zero-row batches in concat_batches and accept an optional schema
This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new e49d8959907c [SPARK-59624][PYTHON][FOLLOWUP] Drop zero-row batches in concat_batches and accept an optional schema
e49d8959907c is described below
commit e49d8959907c76122029d0bd55352a632ac38dd6
Author: Spenser Sun <hs...@gmail.com>
AuthorDate: Tue Sep 29 00:29:41 2026 +0000
[SPARK-59624][PYTHON][FOLLOWUP] Drop zero-row batches in concat_batches and accept an optional schema
### What changes were proposed in this pull request?
`ArrowBatchTransformer.concat_batches` now drops zero-row batches and takes an optional `schema`. If no rows remain, it returns an empty batch of `schema`, or raises a clear `PySparkValueError` when no schema is given instead of failing a bare assertion; batches with mismatched schemas still fail inside PyArrow as before. The element-wise iterator path passes its output schema, and the grouped and window aggregate handlers now pass each input group straight in.
### Why are the changes needed?
A batch with no columns cannot be converted or concatenated with real batches, so callers that know their output schema can now always get a typed result. This follows review feedback on #58990.
### Does this PR introduce _any_ user-facing change?
No, `concat_batches` is internal.
### How was this patch tested?
Updated `test_concat_batches` and added `test_concat_batches_drops_zero_row_batches`, covering calls with and without a schema, passing on PyArrow 18, 24 and 25; they fail without the change.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 5.5)
Closes #59107 from Spenserrrr/arrow-concat-batches.
Authored-by: Spenser Sun <hs...@gmail.com>
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
(cherry picked from commit 36c9287e21bbad11dc418804c70293dc599dbbe6)
Signed-off-by: Yicong-Huang <17...@users.noreply.github.com>
---
python/pyspark/eval_handlers/_arrow.py | 10 ++-------
python/pyspark/sql/conversion.py | 19 +++++++++++-----
python/pyspark/sql/tests/test_conversion.py | 34 ++++++++++++++++++++++++-----
python/pyspark/worker.py | 4 +++-
4 files changed, 47 insertions(+), 20 deletions(-)
diff --git a/python/pyspark/eval_handlers/_arrow.py b/python/pyspark/eval_handlers/_arrow.py
index 98ba58d1b061..ec2e919b5147 100644
--- a/python/pyspark/eval_handlers/_arrow.py
+++ b/python/pyspark/eval_handlers/_arrow.py
@@ -274,10 +274,7 @@ class ArrowGroupedAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
import pyarrow as pa
for group in data:
- batch_list = list(group)
- if not batch_list:
- continue
- concatenated = ArrowBatchTransformer.concat_batches(batch_list)
+ concatenated = ArrowBatchTransformer.concat_batches(group)
results = [
udf_func(
*[concatenated.column(o) for o in args_offsets],
@@ -353,10 +350,7 @@ class ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
import pyarrow as pa
for group in data:
- batch_list = list(group)
- if not batch_list:
- continue
- concatenated = ArrowBatchTransformer.concat_batches(batch_list)
+ concatenated = ArrowBatchTransformer.concat_batches(group)
num_rows = concatenated.num_rows
result_arrays = []
diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py
index 7d99e574f06b..b0c819a69d7f 100644
--- a/python/pyspark/sql/conversion.py
+++ b/python/pyspark/sql/conversion.py
@@ -182,17 +182,24 @@ class ArrowBatchTransformer:
)
@classmethod
- def concat_batches(cls, batches: Iterable["pa.RecordBatch"]) -> "pa.RecordBatch":
+ def concat_batches(
+ cls, batches: Iterable["pa.RecordBatch"], schema: Optional["pa.Schema"] = None
+ ) -> "pa.RecordBatch":
"""Concatenate same-schema RecordBatches by row.
- A single batch is returned unchanged. PyArrow before 19.0.0 has no ``concat_batches``;
- the fallback concatenates the equivalent StructArrays and converts the result back to a
- RecordBatch.
+ ``batches`` may be any iterable of RecordBatches and is consumed once, and zero-row
+ batches are dropped. If no rows remain, an empty batch of ``schema`` is returned, or an
+ error is raised when ``schema`` is not given. A single remaining batch is returned
+ unchanged. PyArrow before 19.0.0 has no ``concat_batches``; the fallback concatenates
+ the equivalent StructArrays and converts the result back to a RecordBatch.
"""
import pyarrow as pa
- batches = tuple(batches)
- assert batches
+ batches = [batch for batch in batches if batch.num_rows]
+ if not batches:
+ if schema is None:
+ raise PySparkValueError("concat_batches needs a schema when the input has no rows")
+ return pa.RecordBatch.from_pylist([], schema=schema)
if len(batches) == 1:
return batches[0]
if hasattr(pa, "concat_batches"):
diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py
index ba2fb4060a79..d0b3fd0898c3 100644
--- a/python/pyspark/sql/tests/test_conversion.py
+++ b/python/pyspark/sql/tests/test_conversion.py
@@ -118,13 +118,37 @@ class ArrowBatchTransformerTests(unittest.TestCase):
def test_concat_batches(self):
import pyarrow as pa
+ schema = pa.schema([("x", pa.int64())])
batches = [
- pa.RecordBatch.from_arrays([pa.array([1, 2])], ["x"]),
- pa.RecordBatch.from_arrays([pa.array([3])], ["x"]),
+ pa.RecordBatch.from_pylist([{"x": 1}, {"x": 2}], schema=schema),
+ pa.RecordBatch.from_pylist([{"x": 3}], schema=schema),
]
- result = ArrowBatchTransformer.concat_batches(iter(batches))
- self.assertEqual(result.column(0).to_pylist(), [1, 2, 3])
- self.assertIs(ArrowBatchTransformer.concat_batches(iter(batches[:1])), batches[0])
+ for kwargs in ({}, {"schema": schema}):
+ result = ArrowBatchTransformer.concat_batches(iter(batches), **kwargs)
+ self.assertEqual(result.column(0).to_pylist(), [1, 2, 3])
+ result = ArrowBatchTransformer.concat_batches(iter(batches[:1]), **kwargs)
+ self.assertIs(result, batches[0])
+ other = pa.RecordBatch.from_pylist([{"y": 1}])
+ with self.assertRaises(pa.ArrowInvalid):
+ ArrowBatchTransformer.concat_batches(iter([batches[0], other]), schema)
+
+ def test_concat_batches_drops_zero_row_batches(self):
+ import pyarrow as pa
+
+ schema = pa.schema([("x", pa.int64())])
+ batch = pa.RecordBatch.from_pylist([{"x": 1}], schema=schema)
+ empty = pa.RecordBatch.from_pylist([], schema=schema)
+ untyped_empty = pa.RecordBatch.from_pylist([])
+ for kwargs in ({}, {"schema": schema}):
+ result = ArrowBatchTransformer.concat_batches(
+ iter([empty, batch, untyped_empty]), **kwargs
+ )
+ self.assertIs(result, batch)
+ for batches in ([], [empty, untyped_empty]):
+ result = ArrowBatchTransformer.concat_batches(iter(batches), schema)
+ self.assertEqual((result.num_rows, result.schema), (0, schema))
+ with self.assertRaisesRegex(PySparkValueError, "needs a schema"):
+ ArrowBatchTransformer.concat_batches(iter(batches))
def test_wrap_struct_basic(self):
"""Test wrapping columns into a struct."""
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 934796d9cddb..a957a5d673d4 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -3182,7 +3182,9 @@ def read_udfs(pickleSer, udf_info_list, eval_type, runner_conf, eval_conf):
if needed == 0:
flat_batch = pa.RecordBatch.from_pylist([], schema=empty_schema)
else:
- combined = ArrowBatchTransformer.concat_batches(pending_chunks)
+ combined = ArrowBatchTransformer.concat_batches(
+ pending_chunks, empty_schema
+ )
flat_batch = combined.slice(0, needed)
remainder = combined.slice(needed)
pending_chunks = [remainder] if remainder.num_rows else []
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org