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: [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]
