You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ru...@apache.org on 2026/05/06 10:03:29 UTC
(spark) branch branch-4.x updated: [SPARK-56716][PYTHON][TESTS][4.X] Add ASV microbenchmark for SQL_ARROW_UDTF
This is an automated email from the ASF dual-hosted git repository.
zhengruifeng 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 e6a1bd0c69ae [SPARK-56716][PYTHON][TESTS][4.X] Add ASV microbenchmark for SQL_ARROW_UDTF
e6a1bd0c69ae is described below
commit e6a1bd0c69aeda06bcd9d26625e9c55c4881f86d
Author: Yicong Huang <17...@users.noreply.github.com>
AuthorDate: Wed May 6 18:03:17 2026 +0800
[SPARK-56716][PYTHON][TESTS][4.X] Add ASV microbenchmark for SQL_ARROW_UDTF
### What changes were proposed in this pull request?
Backport of #55691 to `branch-4.x`.
This PR adds ASV microbenchmarks for `SQL_ARROW_UDTF` (PyArrow-native Python UDTFs created via `arrow_udtf`) to `python/benchmarks/bench_eval_type.py`.
The new `_ArrowUDTFBenchMixin` produces `ArrowUDTFTimeBench` and `ArrowUDTFPeakmemBench`, parametrized by scenario (batch size, column count, type pool) and three handler variants:
- `identity_udtf` - yields the input batch as a `pa.Table`
- `filter_udtf` - keeps rows whose first column is non-null (vectorized)
- `count_udtf` - aggregates each batch into a single-row count table
To support this, two helpers are added to `MockProtocolWriter`:
- `write_worker_input` gains an optional `eval_conf` parameter (UDTF needs `table_arg_offsets` in EvalConf)
- `write_arrow_udtf_payload` mirrors `PythonUDTFRunner.writeUDTF` on the JVM side: argument offsets, partition-child indexes, optional pickled `AnalyzeResult`, the cloudpickled UDTF class, the result schema, and the UDTF name
The wire batch carries one struct column `_0` whose fields are the table's schema; `table_arg_offsets=[0]` tells `ArrowStreamArrowUDTFSerializer` to flatten that struct into a `pa.RecordBatch` for the UDTF's `eval(batch)` method.
### Why are the changes needed?
This is part of [SPARK-55724](https://issues.apache.org/jira/browse/SPARK-55724) (Micro-benchmark PySpark Eval Types). Establishing a stable baseline for `SQL_ARROW_UDTF` is a prerequisite for the upcoming serializer refactor so we can detect any regression objectively.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Clean cherry-pick from master (commit 257c99dc750). New ASV microbenchmarks; same as the master PR.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #55704 from Yicong-Huang/SPARK-56716-4.x.
Authored-by: Yicong Huang <17...@users.noreply.github.com>
Signed-off-by: Ruifeng Zheng <ru...@foxmail.com>
---
python/benchmarks/bench_eval_type.py | 139 ++++++++++++++++++++++++++++++++++-
1 file changed, 138 insertions(+), 1 deletion(-)
diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py
index 2cc467e1d1b0..5b648f8e3aaf 100644
--- a/python/benchmarks/bench_eval_type.py
+++ b/python/benchmarks/bench_eval_type.py
@@ -160,6 +160,7 @@ class MockProtocolWriter:
write_data: Callable[[io.BufferedIOBase], None],
buf: io.BufferedIOBase,
runner_conf: dict[str, str] | None = None,
+ eval_conf: dict[str, str] | None = None,
) -> None:
"""Write the full worker binary stream: preamble + command + data + end."""
cls.write_preamble(buf)
@@ -171,11 +172,43 @@ class MockProtocolWriter:
cls.write_utf8(v, buf)
else:
write_int(0, buf) # RunnerConf (0 key-value pairs)
- write_int(0, buf) # EvalConf (0 key-value pairs)
+ if eval_conf:
+ write_int(len(eval_conf), buf)
+ for k, v in eval_conf.items():
+ cls.write_utf8(k, buf)
+ cls.write_utf8(v, buf)
+ else:
+ write_int(0, buf) # EvalConf (0 key-value pairs)
write_udf(buf)
write_data(buf)
write_int(-4, buf) # SpecialLengths.END_OF_STREAM
+ @classmethod
+ def write_arrow_udtf_payload(
+ cls,
+ udtf_cls: type,
+ return_type: StructType,
+ arg_offsets: list[int],
+ buf: io.BytesIO,
+ ) -> None:
+ """Write the ``read_udtf`` portion of the protocol for SQL_ARROW_UDTF.
+
+ Mirrors ``PythonUDTFRunner.writeUDTF`` on the JVM side: argument offsets,
+ partition-child indexes, optional pickled AnalyzeResult, the cloudpickled
+ UDTF class, the result schema, and the UDTF name.
+ """
+ write_int(len(arg_offsets), buf) # num_arg
+ for offset in arg_offsets:
+ write_int(offset, buf)
+ cls.write_bool(False, buf) # is_kwarg
+ write_int(0, buf) # num_partition_child_indexes (no PARTITION BY)
+ cls.write_bool(False, buf) # has_pickled_analyze_result
+ command = cloudpickle_dumps(udtf_cls)
+ write_int(len(command), buf)
+ buf.write(command)
+ cls.write_utf8(return_type.json(), buf) # result schema
+ cls.write_utf8(udtf_cls.__name__, buf) # udtf_name
+
class MockDataFactory:
"""Creates mock Arrow batches and group structures for benchmarks."""
@@ -471,6 +504,110 @@ class ArrowBatchedUDFPeakmemBench(_ArrowBatchedBenchMixin, _PeakmemBenchBase):
pass
+# -- SQL_ARROW_UDTF -------------------------------------------------------------
+# UDTF receives a single ``pa.RecordBatch`` (the TABLE argument flattened from a
+# struct column at offset 0) per worker batch and yields ``pa.Table`` /
+# ``pa.RecordBatch``.
+
+
+class _ArrowUDTFIdentity:
+ """Identity Arrow UDTF: yields the input batch as a ``pa.Table``."""
+
+ def eval(self, batch):
+ import pyarrow as pa
+
+ yield pa.table(batch)
+
+
+class _ArrowUDTFFilter:
+ """Filter Arrow UDTF: keeps rows whose first column is non-null."""
+
+ def eval(self, batch):
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ table = pa.table(batch)
+ mask = pc.is_valid(table.column(0))
+ yield table.filter(mask)
+
+
+class _ArrowUDTFCount:
+ """Aggregating Arrow UDTF: yields a single-row count of the input batch."""
+
+ def eval(self, batch):
+ import pyarrow as pa
+
+ yield pa.table({"count": pa.array([batch.num_rows], type=pa.int32())})
+
+
+_ARROW_UDTF_COUNT_RETURN = StructType([StructField("count", IntegerType())])
+
+
+class _ArrowUDTFBenchMixin:
+ """Provides ``_write_scenario`` for SQL_ARROW_UDTF.
+
+ The wire batch carries one struct column ``_0`` whose fields are the table's
+ schema. ``table_arg_offsets=[0]`` (passed via EvalConf) tells the serializer
+ to flatten that struct into a ``pa.RecordBatch`` which is forwarded to the
+ UDTF's ``eval(batch)`` method.
+ """
+
+ _scenario_configs = {
+ "sm_batch_few_col": ("mixed", 5_000, 5, 1_000),
+ "sm_batch_many_col": ("mixed", 2_000, 50, 1_000),
+ "lg_batch_few_col": ("mixed", 50_000, 5, 10_000),
+ "lg_batch_many_col": ("mixed", 20_000, 50, 10_000),
+ "pure_ints": ("pure_ints", 50_000, 10, 5_000),
+ "pure_strings": ("pure_strings", 50_000, 10, 5_000),
+ }
+
+ @staticmethod
+ def _build_scenario(name):
+ np.random.seed(42)
+ type_key, num_rows, num_cols, batch_size = _ArrowUDTFBenchMixin._scenario_configs[name]
+ pool = MockDataFactory.NAMED_TYPE_POOLS[type_key]
+ struct_type = MockDataFactory.make_struct_type(num_fields=num_cols, base_types=pool)
+ batches, schema = MockDataFactory.make_batches(
+ num_rows=num_rows,
+ num_cols=1,
+ spark_type_pool=[struct_type],
+ batch_size=batch_size,
+ )
+ # Inner table schema (what the serializer hands to the UDTF after
+ # flattening the struct at offset 0).
+ inner_schema = schema.fields[0].dataType
+ return batches, inner_schema
+
+ _udtfs = {
+ "identity_udtf": (_ArrowUDTFIdentity, lambda inner_schema: inner_schema),
+ "filter_udtf": (_ArrowUDTFFilter, lambda inner_schema: inner_schema),
+ "count_udtf": (_ArrowUDTFCount, lambda inner_schema: _ARROW_UDTF_COUNT_RETURN),
+ }
+ params = [list(_scenario_configs), list(_udtfs)]
+ param_names = ["scenario", "udtf"]
+
+ def _write_scenario(self, scenario, udtf_name, buf):
+ batches, inner_schema = self._build_scenario(scenario)
+ udtf_cls, return_type_fn = self._udtfs[udtf_name]
+ return_type = return_type_fn(inner_schema)
+
+ MockProtocolWriter.write_worker_input(
+ PythonEvalType.SQL_ARROW_UDTF,
+ lambda b: MockProtocolWriter.write_arrow_udtf_payload(udtf_cls, return_type, [0], b),
+ lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),
+ buf,
+ eval_conf={"table_arg_offsets": "0"},
+ )
+
+
+class ArrowUDTFTimeBench(_ArrowUDTFBenchMixin, _TimeBenchBase):
+ pass
+
+
+class ArrowUDTFPeakmemBench(_ArrowUDTFBenchMixin, _PeakmemBenchBase):
+ pass
+
+
# -- SQL_COGROUPED_MAP_ARROW_UDF ------------------------------------------------
# UDF receives two ``pa.Table`` (left, right) per co-group, returns ``pa.Table``.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org