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