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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
