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