This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 9db0e6edcafe [SPARK-58132][PYTHON][TEST] Reduce memory footprint of
grouped-agg UDF batch-slicing tests
9db0e6edcafe is described below
commit 9db0e6edcafe0e7a44b512e0f4cd9d3466a82707
Author: Gurpreet Singh Multani
<[email protected]>
AuthorDate: Thu Jul 16 07:22:36 2026 +0900
[SPARK-58132][PYTHON][TEST] Reduce memory footprint of grouped-agg UDF
batch-slicing tests
### What changes were proposed in this pull request?
This PR shrinks the input size and batch-size configs used by the
`test_arrow_batch_slicing` test in the Arrow and pandas grouped-aggregate
UDF test suites:
- `python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py`
- `python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py`
Specifically, in both files the identical change is applied (a proportional
10x shrink that preserves the multi-slice behavior the test exercises):
- `range(10000000)` -> `range(1000000)`
- `assert len(v) == 10000000 / 2` -> `assert len(v) == 1000000 / 2`
- batch-size configs `[(1000, 2**31 - 1), (0, 1048576), (1000, 1048576)]`
-> `[(100, 2**31 - 1), (0, 104858), (100, 104858)]`
### Why are the changes needed?
The tests materialize a 10-million-row DataFrame and slice it into batches.
Under constrained driver heap this is unnecessarily memory-hungry and can
flake. A 10x smaller dataset with proportionally smaller
`maxRecordsPerBatch` / `maxBytesPerBatch` still produces multiple batches
per group, so the batch-slicing path is exercised exactly as before while
using far less memory.
### Does this PR introduce _any_ user-facing change?
No. Test-only change.
### How was this patch tested?
Existing tests (`test_arrow_batch_slicing` in both suites) continue to pass
with the reduced sizes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #57263 from gurpreetmultanii/reduce-grouped-agg-udf-test-memory.
Authored-by: Gurpreet Singh Multani
<[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
(cherry picked from commit 653779253713fc5babd8a441d6511a46c0192d2e)
Signed-off-by: Hyukjin Kwon <[email protected]>
---
python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py | 6 +++---
python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py | 6 +++---
2 files changed, 6 insertions(+), 6 deletions(-)
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py
b/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py
index 5ba00f558640..0227524d4341 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_udf_grouped_agg.py
@@ -997,18 +997,18 @@ class GroupedAggArrowUDFTestsMixin:
def test_arrow_batch_slicing(self):
import pyarrow as pa
- df = self.spark.range(10000000).select(
+ df = self.spark.range(1000000).select(
(sf.col("id") % 2).alias("key"), sf.col("id").alias("v")
)
@arrow_udf("long", ArrowUDFType.GROUPED_AGG)
def arrow_max(v):
- assert len(v) == 10000000 / 2, len(v)
+ assert len(v) == 1000000 / 2, len(v)
return pa.compute.max(v)
expected =
(df.groupby("key").agg(sf.max("v").alias("res")).sort("key")).collect()
- for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 1048576), (1000,
1048576)]:
+ for maxRecords, maxBytes in [(100, 2**31 - 1), (0, 104858), (100,
104858)]:
with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes):
with self.sql_conf(
{
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py
b/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py
index 73a44cacd5aa..47decf731a77 100644
--- a/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py
+++ b/python/pyspark/sql/tests/pandas/test_pandas_udf_grouped_agg.py
@@ -828,18 +828,18 @@ class GroupedAggPandasUDFTestsMixin:
self.assertEqual(expected2.collect(), result2.collect())
def test_arrow_batch_slicing(self):
- df = self.spark.range(10000000).select(
+ df = self.spark.range(1000000).select(
(sf.col("id") % 2).alias("key"), sf.col("id").alias("v")
)
@pandas_udf("long", PandasUDFType.GROUPED_AGG)
def pandas_max(v):
- assert len(v) == 10000000 / 2, len(v)
+ assert len(v) == 1000000 / 2, len(v)
return v.max()
expected =
(df.groupby("key").agg(sf.max("v").alias("res")).sort("key")).collect()
- for maxRecords, maxBytes in [(1000, 2**31 - 1), (0, 1048576), (1000,
1048576)]:
+ for maxRecords, maxBytes in [(100, 2**31 - 1), (0, 104858), (100,
104858)]:
with self.subTest(maxRecords=maxRecords, maxBytes=maxBytes):
with self.sql_conf(
{
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
