This is an automated email from the ASF dual-hosted git repository.
viirya 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 2a261e2c39a9 [SPARK-57423][PYTHON][TEST] Share base mixin across
MapInBatch Arrow/Pandas siblings
2a261e2c39a9 is described below
commit 2a261e2c39a9960db8f22199f791295f98ac6d77
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Sun Jun 14 21:04:45 2026 -0700
[SPARK-57423][PYTHON][TEST] Share base mixin across MapInBatch Arrow/Pandas
siblings
### What changes were proposed in this pull request?
Follow-up to SPARK-57137/57138/57414. Apply the sibling base-mixin pattern
to the last remaining Arrow/Pandas pair in
`python/benchmarks/bench_eval_type.py`:
- `_MapPandasIterBenchMixin` now subclasses `_MapArrowIterBenchMixin`,
dropping its duplicated `_build_scenario` (identical body) and
`_write_scenario` (identical except the eval type constant and the `ret_type`
handling).
- The Arrow base gains the `_eval_type` class attribute and parameterizes
`_write_scenario` on `self._eval_type`. The Pandas `_udfs` values are
normalized from `(func, arg_offsets)` to the Arrow base's `(func, ret_type,
arg_offsets)` shape so the inherited writer works unchanged.
- Drive-by: convert the `_build_scenario` of the newly added
`_TransformWithStatePandasBenchMixin` (SPARK-57020, developed in parallel with
SPARK-57414) from `staticmethod` with a hardcoded class-name lookup to
`classmethod` with `cls._scenario_configs`, matching the now-uniform convention.
Net diff: +16 / -40 lines.
### Why are the changes needed?
The two MapInBatch mixins were the last undeclared sibling pair:
near-identical `_build_scenario` and `_write_scenario` bodies that had to be
kept in lock-step manually. After this change, every Arrow/Pandas sibling pair
in the file shares a base mixin, every `_build_scenario` resolves configs via
`cls`, and no redundant protocol-writing copies remain.
### Does this PR introduce _any_ user-facing change?
No. Test-only change in the benchmark module.
### How was this patch tested?
- Verified the generated worker input is byte-identical to the pre-refactor
output for all 21 `*TimeBench` classes across every (scenario, udf) cell, with
the cloudpickled UDF command excluded and the TWS stub state-server port pinned
(the port is freshly allocated per process; the UDF pickle embeds
`co_firstlineno` and a random cloudpickle class-tracker id -- all
environment/location metadata unaffected by this change).
- Compared the pickle opcode streams of all 63 pickled UDF/UDTF callables
between old and new: only the metadata above differs; code, constants, and
names are unchanged.
- Ran `setup` + `time_worker` end-to-end for the affected `*TimeBench`
classes (MapArrowIter, MapPandasIter, TransformWithStatePandas) across every
UDF, and `peakmem_worker` for MapPandasIter and TransformWithStatePandas.
### Was this patch authored or co-authored using generative AI tooling?
Yes. Generated-by: Claude Code (claude-opus-4-8)
Closes #56483 from viirya/SPARK-57423-map-iter-dedup.
Authored-by: Liang-Chi Hsieh <[email protected]>
Signed-off-by: Liang-Chi Hsieh <[email protected]>
(cherry picked from commit 0802a5a69ea5ad28ada6410c6a534311c2a7dc1f)
Signed-off-by: Liang-Chi Hsieh <[email protected]>
---
python/benchmarks/bench_eval_type.py | 56 +++++++++++-------------------------
1 file changed, 16 insertions(+), 40 deletions(-)
diff --git a/python/benchmarks/bench_eval_type.py
b/python/benchmarks/bench_eval_type.py
index 8e86550689c7..9bf005880378 100644
--- a/python/benchmarks/bench_eval_type.py
+++ b/python/benchmarks/bench_eval_type.py
@@ -1336,6 +1336,9 @@ class _MapArrowIterBenchMixin:
batch_size=batch_size,
)
+ _eval_type = PythonEvalType.SQL_MAP_ARROW_ITER_UDF
+ # Each UDF entry: (func, ret_type, arg_offsets). ret_type=None means
+ # "use the input row schema".
_udfs = {
"identity_udf": (_identity_batch_iter, None, [0]),
"sort_udf": (_sort_batch_iter, None, [0]),
@@ -1348,12 +1351,12 @@ class _MapArrowIterBenchMixin:
batches, schema = self._build_scenario(scenario)
udf_func, ret_type, arg_offsets = self._udfs[udf_name]
if ret_type is None:
- # mapInArrow UDFs return an Iterator[pa.RecordBatch] with the same
+ # map-in-batch UDFs return an iterator of frames with the same
# schema as the input row (the inner struct, since make_batches
# wraps the row schema in a single struct column for the wire).
ret_type = schema.fields[0].dataType
MockProtocolWriter.write_worker_input(
- PythonEvalType.SQL_MAP_ARROW_ITER_UDF,
+ self._eval_type,
lambda b: MockProtocolWriter.write_udf_payload(udf_func, ret_type,
arg_offsets, b),
lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),
buf,
@@ -1372,10 +1375,12 @@ class
MapArrowIterUDFPeakmemBench(_MapArrowIterBenchMixin, _PeakmemBenchBase):
# UDF receives ``Iterator[pandas.DataFrame]``, returns
``Iterator[pandas.DataFrame]``.
-class _MapPandasIterBenchMixin:
+class _MapPandasIterBenchMixin(_MapArrowIterBenchMixin):
"""Provides ``_write_scenario`` for SQL_MAP_PANDAS_ITER_UDF.
- Wraps input batches in a struct column to match the JVM-side wire format
+ Inherits ``_build_scenario`` and ``_write_scenario`` from the Arrow
+ sibling; only the eval type, the UDFs, and the per-scenario row counts
+ differ. The struct-column wrapping matches the JVM-side wire format
(``MapInBatchEvaluatorFactory`` wraps each row in another row, and the
Pandas serializer turns that struct back into a ``pandas.DataFrame``
when ``df_for_struct=True``).
@@ -1406,42 +1411,15 @@ class _MapPandasIterBenchMixin:
"mixed_types": ("mixed", 200_000, 10, 5_000),
}
- @classmethod
- def _build_scenario(cls, name):
- """Build a single scenario by name."""
- np.random.seed(42)
- type_key, num_rows, num_cols, batch_size = cls._scenario_configs[name]
- pool = MockDataFactory.NAMED_TYPE_POOLS[type_key]
- struct_type = MockDataFactory.make_struct_type(
- num_fields=num_cols,
- base_types=pool,
- )
- return MockDataFactory.make_batches(
- num_rows=num_rows,
- num_cols=1,
- spark_type_pool=[struct_type],
- batch_size=batch_size,
- )
-
+ _eval_type = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
_udfs = {
- "identity_udf": (_identity_pdf_iter, [0]),
- "sort_udf": (_sort_pdf_iter, [0]),
- "filter_udf": (_filter_pdf_iter, [0]),
+ "identity_udf": (_identity_pdf_iter, None, [0]),
+ "sort_udf": (_sort_pdf_iter, None, [0]),
+ "filter_udf": (_filter_pdf_iter, None, [0]),
}
params = [list(_scenario_configs), list(_udfs)]
param_names = ["scenario", "udf"]
- def _write_scenario(self, scenario, udf_name, buf):
- batches, schema = self._build_scenario(scenario)
- udf_func, arg_offsets = self._udfs[udf_name]
- ret_type = schema.fields[0].dataType
- MockProtocolWriter.write_worker_input(
- PythonEvalType.SQL_MAP_PANDAS_ITER_UDF,
- lambda b: MockProtocolWriter.write_udf_payload(udf_func, ret_type,
arg_offsets, b),
- lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),
- buf,
- )
-
class MapPandasIterUDFTimeBench(_MapPandasIterBenchMixin, _TimeBenchBase):
pass
@@ -1849,17 +1827,15 @@ class _TransformWithStatePandasBenchMixin:
"nested_struct": (200, 5_000, 4, _NESTED_POOL),
}
- @staticmethod
- def _build_scenario(name):
+ @classmethod
+ def _build_scenario(cls, name):
"""Build a single TWS Pandas scenario.
Returns ``(batches, schema)`` where ``batches`` is a plain list of
Arrow
RecordBatches with rows pre-sorted by the leading int32 key column.
"""
np.random.seed(42)
- num_groups, rows_per_group, num_value_cols, value_pool = (
- _TransformWithStatePandasBenchMixin._scenario_configs[name]
- )
+ num_groups, rows_per_group, num_value_cols, value_pool =
cls._scenario_configs[name]
total_rows = num_groups * rows_per_group
key_array = pa.array(
np.repeat(np.arange(num_groups, dtype=np.int32), rows_per_group),
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
