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