You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "Yicong-Huang (via GitHub)" <gi...@apache.org> on 2026/08/20 22:45:57 UTC
[PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Yicong-Huang opened a new pull request, #58182:
URL: https://github.com/apache/spark/pull/58182
### What changes were proposed in this pull request?
Add an internal config `spark.sql.execution.pythonUDF.arrow.maxBytesPerOutputBatch` (bytes, default `-1` = no limit) and a new `ArrowBatchTransformer.resize_batches(Iterator[RecordBatch], max_bytes)` helper. When the config is set, the `applyInPandas` (`SQL_GROUPED_MAP_PANDAS_UDF`) worker slices each output Arrow `RecordBatch` that exceeds the cap into `ceil(nbytes / max_bytes)` row-balanced, zero-copy pieces before sending them to the JVM; a batch already within the cap (or empty) passes through unchanged. The config value is shipped to the worker through `ArrowPythonRunner.getPythonRunnerConfMap`.
The config is intentionally general (not `applyInPandas`-specific) so the same output-batch bound can be extended to other Arrow-based Python UDF eval types later. This PR wires it for `applyInPandas`, which produces one Arrow batch per group and is the case most prone to a single oversized output batch.
### Why are the changes needed?
`applyInPandas` produces one output Arrow `RecordBatch` per group, so a large group yields a single oversized batch. That strains the JVM and downstream operators (large allocations, and columnar operators that do not handle very large batches well). Bounding the output batch size on the worker side lets users cap the batch produced for a big group without changing the UDF.
### Does this PR introduce _any_ user-facing change?
No. The config is internal and defaults to `-1` (no limit), so behavior is unchanged unless it is explicitly set.
### How was this patch tested?
New unit tests for `resize_batches` in `test_conversion.py`: a large batch splits into row-balanced pieces preserving row order, and small and empty batches pass through unchanged.
### Was this patch authored or co-authored using generative AI tooling?
No
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #58182:
URL: https://github.com/apache/spark/pull/58182#issuecomment-5671699104
@viirya reframed the doc and PR description per your feedback: best-effort pre-slicing, no 2GB claim, and it now says explicitly it does not prevent conversion-time offset overflow. I also replied on the wiring-test thread. Could you take another look when you have a moment?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3832810786
##########
python/pyspark/sql/conversion.py:
##########
@@ -73,10 +84,72 @@
class ArrowBatchTransformer:
"""
- Pure functions that transform RecordBatch -> RecordBatch.
+ Pure functions to transform:
+ - RecordBatch -> RecordBatch
+ - Iterator[RecordBatch] -> Iterator[RecordBatch]
+
They should have no side effects (no I/O, no writing to streams).
"""
+ @staticmethod
+ def resize_batches(
+ batches: Iterator["pa.RecordBatch"], max_bytes: int
+ ) -> Iterator["pa.RecordBatch"]:
+ """
+ Slice each RecordBatch down toward ``max_bytes``.
+
+ A batch estimated larger than ``max_bytes`` is split into
+ ``ceil(nbytes / max_bytes)`` slices, with the rows divided as evenly as
+ possible across them; a batch already within ``max_bytes`` (or empty) is
+ yielded unchanged. Slicing is zero-copy: each slice is a view over the
+ input batch's buffers, not a copy.
+
+ The examples below use ``max_bytes = 256 MB``.
+
+ Case 1 - larger than max_bytes, so it is split. A 700 MB batch of
+ 1,000,000 rows -> ceil(700 / 256) = 3 slices; the rows do not divide
+ evenly, so they are balanced to 333,333 / 333,333 / 333,334:
+
+ in +-----------------------------------+
+ | 700 MB, 1,000,000 rows |
+ +-----------------------------------+
+ out +-----------+-----------+-----------+
+ | ~233 MB | ~233 MB | ~233 MB |
+ | 333,333 | 333,333 | 333,334 |
+ +-----------+-----------+-----------+
+
+ Case 2 - within max_bytes, so it is passed through unchanged. An 80 MB
+ batch of 500,000 rows -> 1 batch, identical to the input:
+
+ in +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+ out +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+
+ Case 3 - empty (0 rows), so it is passed through unchanged:
+
+ in +-- 0 rows --+
+ +------------+
+ out +-- 0 rows --+
+ +------------+
+ """
+ for batch in batches:
+ num_rows = batch.num_rows
+ if num_rows == 0:
+ yield batch
+ continue
+
+ nbytes = batch.nbytes
+ if nbytes <= max_bytes:
+ yield batch
+ continue
+
+ num_slices = min(math.ceil(nbytes / max_bytes), num_rows)
Review Comment:
Yeah, row-balanced doesn't strictly bound bytes on skewed rows. I'd keep the fast estimate here for now. Measuring each slice is about 100x slower per batch in a microbench. And this is on the output hot path. I softened the doc to call it best-effort, so a skewed batch can still exceed the cap. Byte-accurate slicing can come later. Does that work?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3841046192
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,26 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_MAX_BYTES_PER_OUTPUT_BATCH =
+ buildConf("spark.sql.execution.pythonUDF.arrow.maxBytesPerOutputBatch")
Review Comment:
Wait, how is this differnt from `spark.sql.execution.arrow.maxBytesPerOutputBatch`? We have another conf `spark.sql.execution.arrow.maxBytesPerBatch`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #58182:
URL: https://github.com/apache/spark/pull/58182#issuecomment-5674573004
Thanks @HyukjinKwon @viirya and @zhengruifeng, merging now.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3982881827
##########
python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py:
##########
@@ -1553,6 +1553,15 @@ class ApplyInPandasTests(ApplyInPandasTestsMixin, ReusedSQLTestCase):
pass
+class ApplyInPandasWithWorkerOutputBatchSlicingTests(ApplyInPandasTests):
Review Comment:
Good point. Replaced the suite re-run with a single focused test: one large group with a small `workerOutputBatchMaxBytes` so the output batch is split into several pieces, asserting the result is unchanged. That keeps the split path covered without re-running the whole suite.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3985458616
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,31 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_WORKER_OUTPUT_BATCH_MAX_BYTES =
+ buildConf("spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes")
+ .internal()
+ .doc("Best-effort byte-size target for a single Arrow RecordBatch produced by an " +
+ "Arrow-based Python UDF worker, applied on the worker before the batch is sent to " +
+ "the JVM. applyInPandas hands each group to the UDF as one batch, so a large group " +
+ "can build a batch past Arrow's 2GB limit that then fails to transfer. When set, " +
+ "the worker splits a batch estimated larger than this into ceil(nbytes / value) " +
+ "row-balanced, zero-copy pieces to keep each one under the limit; the estimate " +
Review Comment:
[P2] Could we clarify which 2GB limitation this addresses? The 32-bit offsets in ordinary string/binary arrays limit the data addressable within an individual array. Since splitting happens after `PandasToArrowConversion.convert` constructs the complete RecordBatch, it cannot prevent conversion-time offset overflow. A batch exceeding 2GB across multiple valid columns is a separate case: Arrow IPC uses a 64-bit `bodyLength`, so the format itself has no general 2GB batch limit. Please identify the specific worker-to-JVM write/read or allocation limitation in supported versions, ideally with a reproducer demonstrating the fix. Otherwise, describe this as best-effort pre-slicing to reduce the batch size received by the JVM, explicitly excluding conversion-time offset overflow.
##########
python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py:
##########
@@ -202,6 +202,21 @@ def test_supported_types(self):
assert_frame_equal(expected2, result2)
assert_frame_equal(expected3, result3)
+ def test_output_batch_split_preserves_result(self):
+ # A small worker output-batch cap splits a group's output Arrow batch into several
+ # pieces before it is sent to the JVM. The result must be unchanged by the split.
+ df = self.spark.range(1000).selectExpr("id", "1 as k")
+
+ def add_one(pdf):
+ return pdf.assign(id=pdf.id + 1)
+
+ conf = {"spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes": 128}
+ with self.sql_conf(conf):
+ result = df.groupby("k").applyInPandas(add_one, "id long, k int").sort("id").toPandas()
+
+ expected = df.toPandas().assign(id=lambda p: p.id + 1)
+ assert_frame_equal(expected.reset_index(drop=True), result.reset_index(drop=True))
Review Comment:
[P2] This verifies result correctness with the config enabled, but would also pass if the config were omitted from `ArrowPythonRunner` or the worker's splitting wrapper were removed. The helper unit tests would still pass, leaving that wiring regression undetected. Could we also assert that the worker emits multiple batches, or observe their sizes, either here or in a small worker/serializer-level test? The focused-test approach is appropriate; the missing piece is an assertion that fails when worker-side splitting is disconnected.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #58182:
URL: https://github.com/apache/spark/pull/58182#issuecomment-5674578995
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/eca170f2650abf80d4db42b4d95352a82f268606
- merged into branch-4.x https://github.com/apache/spark/commit/016d9c89b5b28bd83e9b2958bdc08b118e9c8755
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang closed pull request #58182: [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size
URL: https://github.com/apache/spark/pull/58182
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "viirya (via GitHub)" <gi...@apache.org>.
viirya commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r4010675921
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,31 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_WORKER_OUTPUT_BATCH_MAX_BYTES =
+ buildConf("spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes")
+ .internal()
+ .doc("Best-effort byte-size target for a single Arrow RecordBatch produced by an " +
+ "Arrow-based Python UDF worker, applied on the worker before the batch is sent to " +
+ "the JVM. applyInPandas hands each group to the UDF as one batch, so a large group " +
+ "can build a batch past Arrow's 2GB limit that then fails to transfer. When set, " +
+ "the worker splits a batch estimated larger than this into ceil(nbytes / value) " +
+ "row-balanced, zero-copy pieces to keep each one under the limit; the estimate " +
Review Comment:
Thanks, this addresses my concern. The documentation now distinguishes pre-slicing an already constructed batch from preventing conversion-time offset overflow. With the 2GB claim removed, I don't think a reproducer for that particular failure is needed for this PR.
##########
python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py:
##########
@@ -202,6 +202,21 @@ def test_supported_types(self):
assert_frame_equal(expected2, result2)
assert_frame_equal(expected3, result3)
+ def test_output_batch_split_preserves_result(self):
+ # A small worker output-batch cap splits a group's output Arrow batch into several
+ # pieces before it is sent to the JVM. The result must be unchanged by the split.
+ df = self.spark.range(1000).selectExpr("id", "1 as k")
+
+ def add_one(pdf):
+ return pdf.assign(id=pdf.id + 1)
+
+ conf = {"spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes": 128}
+ with self.sql_conf(conf):
+ result = df.groupby("k").applyInPandas(add_one, "id long, k int").sort("id").toPandas()
+
+ expected = df.toPandas().assign(id=lambda p: p.id + 1)
+ assert_frame_equal(expected.reset_index(drop=True), result.reset_index(drop=True))
Review Comment:
That's reasonable. I checked the existing `mapInArrow` output-slicing tests and agree that this follows the same pattern. The current test still won't detect disconnected wiring, but I don't think adding intrusive instrumentation is warranted here. I'm happy to leave this as a non-blocking coverage improvement for a future worker/serializer-level test.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #58182:
URL: https://github.com/apache/spark/pull/58182#issuecomment-5605393707
Small update on the test: I scoped it to just the worker-side split (this PR's config) instead of also setting the existing JVM byte cap. This config is only read worker-side, so the test targets exactly what the PR changes.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3974193439
##########
python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py:
##########
@@ -1553,6 +1553,15 @@ class ApplyInPandasTests(ApplyInPandasTestsMixin, ReusedSQLTestCase):
pass
+class ApplyInPandasWithWorkerOutputBatchSlicingTests(ApplyInPandasTests):
Review Comment:
this will run the whole tests, while our ci resouce is limited, wondering can we reduce the resouce in someway
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #58182:
URL: https://github.com/apache/spark/pull/58182#issuecomment-5363396202
@zhengruifeng PTAL!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #58182:
URL: https://github.com/apache/spark/pull/58182#issuecomment-5417459026
The worker-side split isn't about peak memory (the group is materialized in full either way); it's keeping a single output batch under Arrow's 2GB limit so it can transfer at all (SPARK-32294, ARROW-4890). I renamed the config to `spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes` so it no longer reads as a sibling of `spark.sql.execution.arrow.maxBytesPerOutputBatch`, and the doc now names that one and states the relationship: mine is a best-effort worker-side pre-split; the existing one slices JVM-side with a byte-accurate measure, but only after the batch is read back, so it can't help a batch too large to transfer. They compose, worker split first. I kept the algorithm as-is per your point about not duplicating the JVM slicer, and added a test that runs the suite with both caps small so the composition is exercised.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3826415579
##########
python/pyspark/sql/conversion.py:
##########
@@ -73,10 +84,72 @@
class ArrowBatchTransformer:
"""
- Pure functions that transform RecordBatch -> RecordBatch.
+ Pure functions to transform:
+ - RecordBatch -> RecordBatch
+ - Iterator[RecordBatch] -> Iterator[RecordBatch]
+
They should have no side effects (no I/O, no writing to streams).
"""
+ @staticmethod
+ def resize_batches(
+ batches: Iterator["pa.RecordBatch"], max_bytes: int
+ ) -> Iterator["pa.RecordBatch"]:
+ """
+ Slice each RecordBatch down toward ``max_bytes``.
+
+ A batch estimated larger than ``max_bytes`` is split into
+ ``ceil(nbytes / max_bytes)`` slices, with the rows divided as evenly as
+ possible across them; a batch already within ``max_bytes`` (or empty) is
+ yielded unchanged. Slicing is zero-copy: each slice is a view over the
+ input batch's buffers, not a copy.
+
+ The examples below use ``max_bytes = 256 MB``.
+
+ Case 1 - larger than max_bytes, so it is split. A 700 MB batch of
+ 1,000,000 rows -> ceil(700 / 256) = 3 slices; the rows do not divide
+ evenly, so they are balanced to 333,333 / 333,333 / 333,334:
+
+ in +-----------------------------------+
+ | 700 MB, 1,000,000 rows |
+ +-----------------------------------+
+ out +-----------+-----------+-----------+
+ | ~233 MB | ~233 MB | ~233 MB |
+ | 333,333 | 333,333 | 333,334 |
+ +-----------+-----------+-----------+
+
+ Case 2 - within max_bytes, so it is passed through unchanged. An 80 MB
+ batch of 500,000 rows -> 1 batch, identical to the input:
+
+ in +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+ out +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+
+ Case 3 - empty (0 rows), so it is passed through unchanged:
+
+ in +-- 0 rows --+
+ +------------+
+ out +-- 0 rows --+
+ +------------+
+ """
+ for batch in batches:
+ num_rows = batch.num_rows
+ if num_rows == 0:
+ yield batch
+ continue
+
+ nbytes = batch.nbytes
+ if nbytes <= max_bytes:
+ yield batch
+ continue
+
+ num_slices = min(math.ceil(nbytes / max_bytes), num_rows)
Review Comment:
The `min(..., num_rows)` cap here is the only guard for rows individually larger than the cap - the one-row-batch case the config doc promises. None of the three new tests reach `num_slices == num_rows`, so dropping the cap would still pass CI while emitting zero-length slices. Consider a test with rows each larger than `max_bytes`. Non-blocking.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,22 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_MAX_BYTES_PER_OUTPUT_BATCH =
+ buildConf("spark.sql.execution.pythonUDF.arrow.maxBytesPerOutputBatch")
+ .internal()
+ .doc("Byte-size cap for a single Arrow RecordBatch produced by an Arrow-based Python UDF " +
Review Comment:
This internal config's doc reads as general (any "Arrow-based Python UDF worker"), but only `applyInPandas` (`SQL_GROUPED_MAP_PANDAS_UDF`) is wired to honor it in `worker.py`. A value set for other Arrow-based Python UDFs (mapInArrow, mapInPandas, Arrow scalar UDFs) is shipped to the worker but silently ignored. Since this doc is the durable contract (the PR body isn't), consider noting that only applyInPandas honors it today. Non-blocking.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3826735298
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,22 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_MAX_BYTES_PER_OUTPUT_BATCH =
+ buildConf("spark.sql.execution.pythonUDF.arrow.maxBytesPerOutputBatch")
+ .internal()
+ .doc("Byte-size cap for a single Arrow RecordBatch produced by an Arrow-based Python UDF " +
Review Comment:
Good catch -- updated the doc to note that only `applyInPandas` (`SQL_GROUPED_MAP_PANDAS_UDF`) honors this today; other Arrow-based Python UDFs ship the value but ignore it. Generalizing to more eval types is tracked as a follow-up.
##########
python/pyspark/sql/conversion.py:
##########
@@ -73,10 +84,72 @@
class ArrowBatchTransformer:
"""
- Pure functions that transform RecordBatch -> RecordBatch.
+ Pure functions to transform:
+ - RecordBatch -> RecordBatch
+ - Iterator[RecordBatch] -> Iterator[RecordBatch]
+
They should have no side effects (no I/O, no writing to streams).
"""
+ @staticmethod
+ def resize_batches(
+ batches: Iterator["pa.RecordBatch"], max_bytes: int
+ ) -> Iterator["pa.RecordBatch"]:
+ """
+ Slice each RecordBatch down toward ``max_bytes``.
+
+ A batch estimated larger than ``max_bytes`` is split into
+ ``ceil(nbytes / max_bytes)`` slices, with the rows divided as evenly as
+ possible across them; a batch already within ``max_bytes`` (or empty) is
+ yielded unchanged. Slicing is zero-copy: each slice is a view over the
+ input batch's buffers, not a copy.
+
+ The examples below use ``max_bytes = 256 MB``.
+
+ Case 1 - larger than max_bytes, so it is split. A 700 MB batch of
+ 1,000,000 rows -> ceil(700 / 256) = 3 slices; the rows do not divide
+ evenly, so they are balanced to 333,333 / 333,333 / 333,334:
+
+ in +-----------------------------------+
+ | 700 MB, 1,000,000 rows |
+ +-----------------------------------+
+ out +-----------+-----------+-----------+
+ | ~233 MB | ~233 MB | ~233 MB |
+ | 333,333 | 333,333 | 333,334 |
+ +-----------+-----------+-----------+
+
+ Case 2 - within max_bytes, so it is passed through unchanged. An 80 MB
+ batch of 500,000 rows -> 1 batch, identical to the input:
+
+ in +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+ out +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+
+ Case 3 - empty (0 rows), so it is passed through unchanged:
+
+ in +-- 0 rows --+
+ +------------+
+ out +-- 0 rows --+
+ +------------+
+ """
+ for batch in batches:
+ num_rows = batch.num_rows
+ if num_rows == 0:
+ yield batch
+ continue
+
+ nbytes = batch.nbytes
+ if nbytes <= max_bytes:
+ yield batch
+ continue
+
+ num_slices = min(math.ceil(nbytes / max_bytes), num_rows)
Review Comment:
Added `test_resize_batches_row_larger_than_cap` (3 rows of ~1KB each, `max_bytes=1`) so `num_slices == num_rows` is exercised -- guards against the `min(..., num_rows)` cap regressing into empty slices.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3991155055
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,31 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_WORKER_OUTPUT_BATCH_MAX_BYTES =
+ buildConf("spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes")
+ .internal()
+ .doc("Best-effort byte-size target for a single Arrow RecordBatch produced by an " +
+ "Arrow-based Python UDF worker, applied on the worker before the batch is sent to " +
+ "the JVM. applyInPandas hands each group to the UDF as one batch, so a large group " +
+ "can build a batch past Arrow's 2GB limit that then fails to transfer. When set, " +
+ "the worker splits a batch estimated larger than this into ceil(nbytes / value) " +
+ "row-balanced, zero-copy pieces to keep each one under the limit; the estimate " +
Review Comment:
You're right, the 2GB framing was too strong. The split runs after PandasToArrowConversion builds the batch, so it can't prevent conversion-time offset overflow, and Arrow IPC bodyLength is 64-bit so there's no general 2GB batch limit. I reframed the doc as best-effort pre-slicing that reduces the size of each batch the JVM receives and allocates, and it now says explicitly it does not prevent conversion-time offset overflow. Updated the PR description the same way.
##########
python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py:
##########
@@ -202,6 +202,21 @@ def test_supported_types(self):
assert_frame_equal(expected2, result2)
assert_frame_equal(expected3, result3)
+ def test_output_batch_split_preserves_result(self):
+ # A small worker output-batch cap splits a group's output Arrow batch into several
+ # pieces before it is sent to the JVM. The result must be unchanged by the split.
+ df = self.spark.range(1000).selectExpr("id", "1 as k")
+
+ def add_one(pdf):
+ return pdf.assign(id=pdf.id + 1)
+
+ conf = {"spark.sql.execution.pythonUDF.arrow.workerOutputBatchMaxBytes": 128}
+ with self.sql_conf(conf):
+ result = df.groupby("k").applyInPandas(add_one, "id long, k int").sort("id").toPandas()
+
+ expected = df.toPandas().assign(id=lambda p: p.id + 1)
+ assert_frame_equal(expected.reset_index(drop=True), result.reset_index(drop=True))
Review Comment:
Fair point that the result assertion passes with or without the wiring. The split logic itself is covered directly by the resize_batches unit tests, and this end-to-end test confirms results stay correct under a small cap. Observing the worker's emitted batch count from the DataFrame API isn't exposed cleanly (the collect/exchange re-batches, so it wouldn't reflect the worker split), and the existing mapInArrow output-slicing tests are correctness-only for the same reason, so I'd keep this consistent with that pattern.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3826963636
##########
python/pyspark/sql/conversion.py:
##########
@@ -73,10 +84,72 @@
class ArrowBatchTransformer:
"""
- Pure functions that transform RecordBatch -> RecordBatch.
+ Pure functions to transform:
+ - RecordBatch -> RecordBatch
+ - Iterator[RecordBatch] -> Iterator[RecordBatch]
+
They should have no side effects (no I/O, no writing to streams).
"""
+ @staticmethod
+ def resize_batches(
+ batches: Iterator["pa.RecordBatch"], max_bytes: int
+ ) -> Iterator["pa.RecordBatch"]:
+ """
+ Slice each RecordBatch down toward ``max_bytes``.
+
+ A batch estimated larger than ``max_bytes`` is split into
+ ``ceil(nbytes / max_bytes)`` slices, with the rows divided as evenly as
+ possible across them; a batch already within ``max_bytes`` (or empty) is
+ yielded unchanged. Slicing is zero-copy: each slice is a view over the
+ input batch's buffers, not a copy.
+
+ The examples below use ``max_bytes = 256 MB``.
+
+ Case 1 - larger than max_bytes, so it is split. A 700 MB batch of
+ 1,000,000 rows -> ceil(700 / 256) = 3 slices; the rows do not divide
+ evenly, so they are balanced to 333,333 / 333,333 / 333,334:
+
+ in +-----------------------------------+
+ | 700 MB, 1,000,000 rows |
+ +-----------------------------------+
+ out +-----------+-----------+-----------+
+ | ~233 MB | ~233 MB | ~233 MB |
+ | 333,333 | 333,333 | 333,334 |
+ +-----------+-----------+-----------+
+
+ Case 2 - within max_bytes, so it is passed through unchanged. An 80 MB
+ batch of 500,000 rows -> 1 batch, identical to the input:
+
+ in +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+ out +-- 80 MB, 500,000 rows --+
+ +-------------------------+
+
+ Case 3 - empty (0 rows), so it is passed through unchanged:
+
+ in +-- 0 rows --+
+ +------------+
+ out +-- 0 rows --+
+ +------------+
+ """
+ for batch in batches:
+ num_rows = batch.num_rows
+ if num_rows == 0:
+ yield batch
+ continue
+
+ nbytes = batch.nbytes
+ if nbytes <= max_bytes:
+ yield batch
+ continue
+
+ num_slices = min(math.ceil(nbytes / max_bytes), num_rows)
Review Comment:
To enforce the byte cap, this needs to measure the byte size of each emitted slice, not only derive a row-balanced slice count from the total batch size. With variable-width rows, one slice can contain most of the bytes and still exceed `max_bytes`; the existing JVM-side byte cap keeps shrinking and measuring candidate slices until it is below the cap or down to one row.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58923][PYTHON] Bound applyInPandas output Arrow batch size [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #58182:
URL: https://github.com/apache/spark/pull/58182#discussion_r3841046192
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5371,6 +5371,26 @@ object SQLConf {
.version("4.0.0")
.fallbackConf(BUFFER_SIZE)
+ val PYTHON_UDF_ARROW_MAX_BYTES_PER_OUTPUT_BATCH =
+ buildConf("spark.sql.execution.pythonUDF.arrow.maxBytesPerOutputBatch")
Review Comment:
Wait, how is this differnt from `spark.sql.execution.arrow.maxBytesPerOutputBatch`? We have another conf `spark.sql.execution.arrow.maxBytesPerBatch`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org