#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""

Microbenchmarks for PySpark UDF eval types.

Each benchmark class simulates the full worker pipeline for one eval type

by constructing the complete binary protocol that ``worker.py``'s

``main(infile, outfile)`` expects.

"""

import io

import json

import os

import socket

import struct

import sys

import tempfile

import threading

from typing import Any, Callable, Iterator

import numpy as np

import pyarrow as pa

from pyspark.cloudpickle import dumps as cloudpickle_dumps

from pyspark.serializers import (

BatchedSerializer,

CPickleSerializer,

SpecialLengths,

write_int,

write_long,

)

from pyspark.sql.types import (

BinaryType,

BooleanType,

DoubleType,

IntegerType,

LongType,

StringType,

StructField,

StructType,

TimestampNTZType,

)

from pyspark.util import PythonEvalType

from pyspark.worker import main as worker_main

# ---------------------------------------------------------------------------

# Mock helpers: protocol writer, data factory, UDF factory

# ---------------------------------------------------------------------------

# ``GroupedBatches[i]`` is the list of Arrow RecordBatches making up the ``i``-th

# DataFrame of a group: 1 DF for grouped scenarios, 2 (left, right) for cogrouped.

# Analogous to ``GroupedBatch``/``CoGroupedBatch`` in ``pyspark.sql.pandas._typing``

# but materialized as lists here since benchmarks pre-build the full payload.

GroupedBatches = list[list[pa.RecordBatch]]

class MockProtocolWriter:

"""Constructs the binary wire protocol that ``worker.py``'s ``main()`` expects."""

@classmethod

def write_bool(cls, val: bool, buf: io.BytesIO) -> None:

buf.write(struct.pack("!?", val))

@classmethod

def write_data_payload(

cls, batch_iter: Iterator[pa.RecordBatch], buf: io.BufferedIOBase

) -> None:

"""Write a plain Arrow IPC stream from an iterator of Arrow batches."""

first_batch = next(batch_iter)

writer = pa.RecordBatchStreamWriter(buf, first_batch.schema)

writer.write_batch(first_batch)

for batch in batch_iter:

writer.write_batch(batch)

writer.close()

@classmethod

def write_grouped_data_payload(

cls,

groups: list[GroupedBatches],

buf: io.BufferedIOBase,

) -> None:

"""Write grouped Arrow data in the wire protocol expected by _load_group_dataframes.

Shape: ``groups[g][df]`` is the list of RecordBatches for DataFrame ``df``

of group ``g``. Each DataFrame's batches are written as one Arrow IPC

stream, matching the real wire protocol where a group may be split across

multiple batches by ``spark.sql.execution.arrow.maxRecordsPerBatch``.

"""

for group in groups:

write_int(len(group), buf) # num_dfs in this group

for df_batches in group:

cls.write_data_payload(iter(df_batches), buf)

write_int(0, buf) # end of groups

@classmethod

def write_pickled_data_payload(

cls, rows: list[tuple], buf: io.BufferedIOBase, batch_size: int = 100

) -> None:

"""Write pickled row batches followed by ``END_OF_DATA_SECTION``.

Mirrors ``BatchEvalPythonExec.getInputIterator`` on the JVM side, which

pickles rows (as tuples) in batches of ``batch_size``; the worker's

``BatchedSerializer`` unbatches them on read.

"""

BatchedSerializer(CPickleSerializer(), batch_size).dump_stream(iter(rows), buf)

write_int(SpecialLengths.END_OF_DATA_SECTION, buf)

@classmethod

def write_preamble(cls, buf: io.BytesIO) -> None:

"""Write everything ``main()`` reads before ``eval_type``."""

write_int(0, buf) # split_index

cls.write_utf8(f"{sys.version_info[0]}.{sys.version_info[1]}", buf) # python version

cls.write_utf8(

json.dumps(

{

"isBarrier": False,

"connInfo": None,

"secret": None,

"stageId": 0,

"partitionId": 0,

"attemptNumber": 0,

"taskAttemptId": 0,

"cpus": 1,

# Plain decimal string, matching CpuAmount.toDisplayString on

# the JVM side; TaskContextInfo.from_stream requires the key.

"cpuAmount": "1",

"resources": {},

"localProperties": {},

}

),

buf,

)

cls.write_utf8("/tmp", buf) # spark_files_dir

write_int(0, buf) # num_python_includes

cls.write_bool(False, buf) # needs_broadcast_decryption_server

write_int(0, buf) # num_broadcast_variables

@classmethod

def write_init_message(

cls,

eval_type: int,

write_udf: Callable[[io.BufferedIOBase], None],

target_buffer: io.BytesIO,

runner_conf: dict[str, str] | None = None,

eval_conf: dict[str, str] | None = None,

) -> None:

"""Write the initial message with header, length + its data."""

# Write everything to a seperate buffer so we can

# determine the length of the initial message.

buf = io.BytesIO()

cls.write_preamble(buf)

write_int(eval_type, buf)

if runner_conf:

write_int(len(runner_conf), buf)

for k, v in runner_conf.items():

cls.write_utf8(k, buf)

cls.write_utf8(v, buf)

else:

write_int(0, buf) # RunnerConf (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 the actual data

# header...

write_int(SpecialLengths.START_OF_INIT_MESSAGE, target_buffer)

write_int(buf.getbuffer().nbytes, target_buffer)

# ... + previously buffered data

buf.seek(0)

target_buffer.write(buf.read())

@classmethod

def write_udf_payload(

cls,

udf_func: Callable[..., Any],

return_type: StructType,

arg_offsets: list[int],

buf: io.BytesIO,

) -> None:

"""Write the ``read_single_udf`` portion of the protocol."""

write_int(1, buf) # num_udfs

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(1, buf) # num_chained

command = cloudpickle_dumps((udf_func, return_type))

write_int(len(command), buf)

buf.write(command)

write_long(0, buf) # result_id

@classmethod

def write_utf8(cls, s: str, buf: io.BytesIO) -> None:

"""Write a length-prefixed UTF-8 string (matches ``UTF8Deserializer.loads``)."""

encoded = s.encode("utf-8")

write_int(len(encoded), buf)

buf.write(encoded)

@classmethod

def write_worker_input(

cls,

eval_type: int,

write_udf: Callable[[io.BufferedIOBase], None],

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_init_message(

eval_type, write_udf, buf, runner_conf=runner_conf, eval_conf=eval_conf

)

write_data(buf)

write_int(SpecialLengths.END_OF_STREAM, buf)

@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."""

MAX_RECORDS_PER_BATCH = 10_000

TYPE_REGISTRY: dict[str, tuple[Callable, Any]] = {

"int": (lambda r: pa.array(np.random.randint(0, 1000, r, dtype=np.int32)), IntegerType()),

"double": (lambda r: pa.array(np.random.rand(r)), DoubleType()),

"string": (lambda r: pa.array([f"s{j}" for j in range(r)]), StringType()),

"binary": (lambda r: pa.array([f"b{j}".encode() for j in range(r)]), BinaryType()),

"boolean": (lambda r: pa.array(np.random.choice([True, False], r)), BooleanType()),

}

MIXED_TYPES = [

TYPE_REGISTRY["int"],

TYPE_REGISTRY["string"],

TYPE_REGISTRY["binary"],

TYPE_REGISTRY["boolean"],

]

NUMERIC_TYPES = [TYPE_REGISTRY["int"], TYPE_REGISTRY["double"]]

NAMED_TYPE_POOLS: dict[str, list[tuple[Callable, Any]]] = {

"mixed": MIXED_TYPES,

"pure_ints": [TYPE_REGISTRY["int"]],

"pure_floats": [TYPE_REGISTRY["double"]],

"pure_strings": [TYPE_REGISTRY["string"]],

"pure_ts": [

(

lambda r: pa.array(

np.arange(0, r, dtype="datetime64[us]"),

type=pa.timestamp("us", tz=None),

),

TimestampNTZType(),

)

],

}

@classmethod

def make_struct_type(

cls,

*,

num_fields: int,

base_types: list[tuple[Callable, Any]],

) -> tuple[Callable, StructType]:

"""Compose flat ``base_types`` into a single struct type pool entry.

Returns ``(factory_fn, StructType)`` suitable for inclusion in a

``spark_type_pool`` list. The factory produces a ``pa.StructArray`` whose

sub-fields cycle through ``base_types``.

"""

fields = [

StructField(f"col_{i}", base_types[i % len(base_types)][1]) for i in range(num_fields)

]

def factory(r):

arrays = [base_types[i % len(base_types)][0](r) for i in range(num_fields)]

names = [f.name for f in fields]

return pa.StructArray.from_arrays(arrays, names=names)

return (factory, StructType(fields))

@classmethod

def make_batches(

cls,

*,

num_rows: int,

num_cols: int,

batch_size: int = MAX_RECORDS_PER_BATCH,

spark_type_pool: list[tuple[Callable, Any]],

) -> tuple[list[pa.RecordBatch], StructType]:

"""Create RecordBatches with columns cycling through ``spark_type_pool``.

Splits ``num_rows`` into batches of at most ``batch_size`` rows.

Each pool entry is ``(factory_fn, SparkType)``; if an entry produces

a ``pa.StructArray``, it becomes a struct column naturally.

"""

batches = []

for offset in range(0, num_rows, batch_size):

rows = min(batch_size, num_rows - offset)

arrays = [spark_type_pool[i % len(spark_type_pool)][0](rows) for i in range(num_cols)]

names = [f"col_{i}" for i in range(num_cols)]

batches.append(pa.RecordBatch.from_arrays(arrays, names=names))

schema = StructType(

[

StructField(f"col_{i}", spark_type_pool[i % len(spark_type_pool)][1])

for i in range(num_cols)

]

)

return batches, schema

@classmethod

def make_rows(

cls,

*,

num_rows: int,

num_cols: int,

spark_type_pool: list[tuple[Callable, Any]],

) -> tuple[list[tuple], StructType]:

"""Create Python row tuples with columns cycling through ``spark_type_pool``."""

batches, schema = cls.make_batches(

num_rows=num_rows, num_cols=num_cols, spark_type_pool=spark_type_pool

)

rows = [

row for batch in batches for row in zip(*(col.to_pylist() for col in batch.columns))

]

return rows, schema

@classmethod

def make_grouped_batches(

cls,

*,

num_groups: int,

num_rows: int,

num_cols: int,

batch_size: int = MAX_RECORDS_PER_BATCH,

spark_type_pool: list[tuple[Callable, Any]],

) -> tuple[list[GroupedBatches], StructType]:

"""Create groups, each containing a single DataFrame.

Each group is ``[batches]`` where ``batches`` is the list of Arrow

RecordBatches for that group's DataFrame (``num_rows`` rows split into

batches of ``batch_size``).

"""

groups = []

for _ in range(num_groups):

batches, _ = cls.make_batches(

num_rows=num_rows,

num_cols=num_cols,

batch_size=batch_size,

spark_type_pool=spark_type_pool,

)

groups.append([list(batches)])

schema = StructType(

[

StructField(f"col_{i}", spark_type_pool[i % len(spark_type_pool)][1])

for i in range(num_cols)

]

)

return groups, schema

@classmethod

def make_cogrouped_batches(

cls,

*,

num_groups: int,

num_rows: int,

num_cols: int,

batch_size: int = MAX_RECORDS_PER_BATCH,

spark_type_pool: list[tuple[Callable, Any]],

) -> tuple[list[GroupedBatches], StructType]:

"""Create cogroups, each containing two DataFrames (left, right).

Each cogroup is ``[left_batches, right_batches]``. The two DataFrames

share the same schema but have independent data, each with ``num_rows``

rows and ``num_cols`` flat columns.

"""

left_groups, schema = cls.make_grouped_batches(

num_groups=num_groups,

num_rows=num_rows,

num_cols=num_cols,

batch_size=batch_size,

spark_type_pool=spark_type_pool,

)

right_groups, _ = cls.make_grouped_batches(

num_groups=num_groups,

num_rows=num_rows,

num_cols=num_cols,

batch_size=batch_size,

spark_type_pool=spark_type_pool,

)

cogroups = [[left_groups[i][0], right_groups[i][0]] for i in range(num_groups)]

return cogroups, schema

class MockUDFFactory:

"""Constructs UDF command payloads for the worker protocol."""

@classmethod

def make_grouped_arg_offsets(cls, num_key_cols: int, num_value_cols: int) -> list[int]:

"""Build arg_offsets: ``[total_len, num_keys, key_idx..., value_idx...]``."""

key_idxs = list(range(num_key_cols))

value_idxs = list(range(num_key_cols, num_key_cols + num_value_cols))

offsets = [num_key_cols] + key_idxs + value_idxs

return [len(offsets)] + offsets

# ---------------------------------------------------------------------------

# General benchmark base classes

# ---------------------------------------------------------------------------

class _TimeBenchBase:

"""Base for ``time_*`` benchmarks (any eval type).

Setup materializes full input bytes in memory so that disk I/O does not

affect latency measurements.

Subclasses provide ``params``, ``param_names``, and ``_write_scenario``.

"""

def setup(self, *args):

buf = io.BytesIO()

self._write_scenario(*args, buf)

self._input = buf.getvalue()

def time_worker(self, *args):

worker_main(io.BytesIO(self._input), io.BytesIO())

class _PeakmemBenchBase:

"""Base for ``peakmem_*`` benchmarks (any eval type).

Benchmark streams input to a temp file and replays from disk so that

setup memory does not inflate peak-memory measurements.

Subclasses provide ``params``, ``param_names``, and ``_write_scenario``.

"""

def setup(self, *args):

fd, path = tempfile.mkstemp(prefix="spark-bench-replay-", suffix=".bin")

with os.fdopen(fd, "wb") as f:

self._write_scenario(*args, f)

self._replay_path = path

def teardown(self, *args):

try:

os.remove(self._replay_path)

except (FileNotFoundError, AttributeError):

pass

def peakmem_worker(self, *args):

with open(self._replay_path, "rb") as infile:

worker_main(infile, io.BytesIO())

# -- SQL_ARROW_BATCHED_UDF --------------------------------------------------

# UDF receives individual Python values per row, returns a single Python value.

class _ArrowBatchedBenchMixin:

"""Provides ``_write_scenario`` for SQL_ARROW_BATCHED_UDF.

Passes the ``input_type`` (StructType JSON) through ``EvalConf``, which the

non-legacy and legacy paths both read via ``eval_conf.input_type``.

"""

# Row-by-row processing is ~100x slower than columnar Arrow UDFs,

# so batch counts are much smaller to keep benchmarks under 60s.

_scenario_configs = {

"sm_batch_few_col": ("mixed", 20_000, 5, 1_000),

"sm_batch_many_col": ("mixed", 5_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_floats": ("pure_floats", 50_000, 10, 5_000),

"pure_strings": ("pure_strings", 50_000, 10, 5_000),

"mixed_types": ("mixed", 50_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]

return MockDataFactory.make_batches(

num_rows=num_rows,

num_cols=num_cols,

spark_type_pool=pool,

batch_size=batch_size,

)

# UDFs for SQL_ARROW_BATCHED_UDF operate on individual Python values.

# arg_offsets=[0] means the UDF receives column 0 value per row.

_udfs = {

"identity_udf": (lambda x: x, None, [0]),

"stringify_udf": (lambda x: str(x), StringType(), [0]),

"nullcheck_udf": (lambda x: x is not None, BooleanType(), [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, ret_type, arg_offsets = self._udfs[udf_name]

if ret_type is None:

ret_type = schema.fields[0].dataType

MockProtocolWriter.write_worker_input(

PythonEvalType.SQL_ARROW_BATCHED_UDF,

lambda b: MockProtocolWriter.write_udf_payload(udf_func, ret_type, arg_offsets, b),

lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),

buf,

eval_conf={"input_type": schema.json()},

)

class ArrowBatchedUDFTimeBench(_ArrowBatchedBenchMixin, _TimeBenchBase):

pass

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),

}

@classmethod

def _build_scenario(cls, 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)

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_ARROW_TABLE_UDF ----------------------------------------------------

# Python UDTF (``@udtf(useArrow=True)``): handler is a class with ``eval(self, *args)``

# that yields output rows. Each input row triggers one ``eval`` call; yielded rows

# are converted to Arrow via ``LocalDataToArrowConversion``.

class _ArrowTableUDFIdentity:

def eval(self, x):

yield (x,)

class _ArrowTableUDFExplode:

def eval(self, x):

for _ in range(3):

yield (x,)

class _ArrowTableUDFFilter:

def eval(self, x):

if x is not None and (hash(x) & 1):

yield (x,)

class _ArrowTableUDFStringify:

def eval(self, x):

yield (str(x),)

class _ArrowTableUDFBenchMixin:

"""Provides ``_write_scenario`` for SQL_ARROW_TABLE_UDF (Python UDTF, useArrow=True).

Writes the extra ``input_type`` (StructType JSON) into ``EvalConf`` that the

non-legacy path requires, and uses the UDTF wire protocol (no num_udfs/result_id).

"""

# Per-input-row ``LocalDataToArrowConversion.convert`` call makes this path

# ~15-20x slower than SQL_ARROW_BATCHED_UDF, so row counts are scaled down

# accordingly to keep each measurement under ASV's per-sample budget.

_scenario_configs = {

"sm_batch_few_col": ("mixed", 2_000, 5, 500),

"sm_batch_many_col": ("mixed", 500, 50, 500),

"lg_batch_few_col": ("mixed", 5_000, 5, 2_500),

"lg_batch_many_col": ("mixed", 2_000, 50, 2_000),

"pure_ints": ("pure_ints", 5_000, 10, 2_500),

"pure_strings": ("pure_strings", 5_000, 10, 2_500),

}

@classmethod

def _build_scenario(cls, name):

np.random.seed(42)

type_key, num_rows, num_cols, batch_size = cls._scenario_configs[name]

pool = MockDataFactory.NAMED_TYPE_POOLS[type_key]

return MockDataFactory.make_batches(

num_rows=num_rows,

num_cols=num_cols,

spark_type_pool=pool,

batch_size=batch_size,

)

# Each entry: (handler_class, return_type_or_None, arg_offsets).

# ``None`` return_type means "use input column 0's type".

_udtfs = {

"identity_udtf": (_ArrowTableUDFIdentity, None, [0]),

"explode_udtf": (_ArrowTableUDFExplode, None, [0]),

"filter_udtf": (_ArrowTableUDFFilter, None, [0]),

"stringify_udtf": (_ArrowTableUDFStringify, StringType(), [0]),

}

# "arrow": non-legacy pure Arrow path; "legacy_pandas": legacy pandas

# conversion path (flag on), which goes through pandas Series/DataFrame.

_conversions = ["arrow", "legacy_pandas"]

params = [list(_scenario_configs), list(_udtfs), list(_conversions)]

param_names = ["scenario", "udtf", "conversion"]

def _write_scenario(self, scenario, udtf_name, conversion, buf):

batches, schema = self._build_scenario(scenario)

handler, ret_type, arg_offsets = self._udtfs[udtf_name]

if ret_type is None:

ret_type = schema.fields[0].dataType

return_type = StructType([StructField("c0", ret_type)])

runner_conf = None

if conversion == "legacy_pandas":

runner_conf = {

"spark.sql.legacy.execution.pythonUDTF.pandas.conversion.enabled": "true"

}

MockProtocolWriter.write_worker_input(

PythonEvalType.SQL_ARROW_TABLE_UDF,

lambda b: MockProtocolWriter.write_arrow_udtf_payload(

handler, return_type, arg_offsets, b

),

lambda b: MockProtocolWriter.write_data_payload(iter(batches), b),

buf,

runner_conf=runner_conf,

eval_conf={"input_type": schema.json()},

)

class ArrowTableUDFTimeBench(_ArrowTableUDFBenchMixin, _TimeBenchBase):

pass

class ArrowTableUDFPeakmemBench(_ArrowTableUDFBenchMixin, _PeakmemBenchBase):

pass

# -- SQL_TABLE_UDF ----------------------------------------------------------

# Python UDTF (``@udtf``, useArrow=False): input rows arrive as pickled tuples and

# each row triggers one ``eval(self, *args)`` call; yielded rows are converted

# with ``toInternal`` and pickled back one input row at a time.

class _TableUDFIdentity:

def eval(self, *args):

yield args

class _TableUDFExplode:

def eval(self, *args):

for _ in range(3):

yield args

class _TableUDFCount:

"""Emits nothing per row; yields a single-row count from ``terminate``."""

def __init__(self):

self._count = 0

def eval(self, *args):

self._count += 1

def terminate(self):

yield (self._count,)

class _TableUDFStringify:

def eval(self, *args):

yield tuple(str(a) for a in args)

class _TableUDFBenchMixin:

"""Provides ``_write_scenario`` for SQL_TABLE_UDF (pickled Python UDTF).

Every input column is passed to ``eval`` as a positional argument, so the

``*_args`` scenarios scale the per-row argument count.

"""

_scenario_configs = {

"few_args_ints": ("pure_ints", 50_000, 2),

"many_args_ints": ("pure_ints", 20_000, 20),

"few_args_strings": ("pure_strings", 50_000, 2),

"many_args_strings": ("pure_strings", 20_000, 20),

"mixed_types": ("mixed", 20_000, 5),

}

@classmethod

def _build_scenario(cls, name):

np.random.seed(42)

type_key, num_rows, num_cols = cls._scenario_configs[name]

pool = MockDataFactory.NAMED_TYPE_POOLS[type_key]

return MockDataFactory.make_rows(num_rows=num_rows, num_cols=num_cols, spark_type_pool=pool)

_udtfs = {

"identity_udtf": (_TableUDFIdentity, lambda schema: schema),

"explode_udtf": (_TableUDFExplode, lambda schema: schema),

"count_udtf": (_TableUDFCount, lambda schema: _ARROW_UDTF_COUNT_RETURN),

"stringify_udtf": (

_TableUDFStringify,

lambda schema: StructType([StructField(f.name, StringType()) for f in schema]),

),

}

params = [list(_scenario_configs), list(_udtfs)]

param_names = ["scenario", "udtf"]

def _write_scenario(self, scenario, udtf_name, buf):

rows, schema = self._build_scenario(scenario)

udtf_cls, return_type_fn = self._udtfs[udtf_name]

return_type = return_type_fn(schema)

arg_offsets = list(range(len(schema)))

MockProtocolWriter.write_worker_input(

PythonEvalType.SQL_TABLE_UDF,

lambda b: MockProtocolWriter.write_arrow_udtf_payload(

udtf_cls, return_type, arg_offsets, b

),

lambda b: MockProtocolWriter.write_pickled_data_payload(rows, b),

buf,

)

class TableUDFTimeBench(_TableUDFBenchMixin, _TimeBenchBase):

pass

class TableUDFPeakmemBench(_TableUDFBenchMixin, _PeakmemBenchBase):

pass

# -- SQL_COGROUPED_MAP_ARROW_UDF ------------------------------------------------

# UDF receives two ``pa.Table`` (left, right) per co-group, returns ``pa.Table``.

class _CogroupedMapArrowBenchMixin:

"""Provides _write_scenario for SQL_COGROUPED_MAP_ARROW_UDF."""

def _cogrouped_map_arrow_identity(left, right):

"""Identity cogroup UDF: returns left table as-is."""

return left

def _cogrouped_map_arrow_concat(left, right):

"""Concat cogroup UDF: vertically concatenates left and right tables."""

import pyarrow as pa

return pa.concat_tables([left, right])

def _cogrouped_map_arrow_left_semi(left, right):

"""Left-semi cogroup UDF: filters left rows whose key exists in right."""

key_col = left.column_names[0]

return left.join(right.select([key_col]), keys=key_col, join_type="left semi")

_scenario_configs = {

"few_groups_sm": (50, 5_000, 1, 4),

"few_groups_lg": (50, 50_000, 1, 4),

"many_groups_sm": (2_000, 500, 1, 4),

"many_groups_lg": (500, 10_000, 1, 4),

"wide_cols": (200, 5_000, 1, 20),

"multi_key": (200, 5_000, 3, 5),

}

@classmethod

def _build_scenario(cls, name):

"""Build a cogroup scenario: two DataFrames with the same grouping structure.

Unlike grouped map (which wraps columns in a struct), cogroup batches

have flat columns: [key_col_0, ..., key_col_k, val_col_0, ..., val_col_v].

"""

np.random.seed(42)

num_groups, rows_per_group, num_key_cols, num_value_cols = cls._scenario_configs[name]

n_cols = num_key_cols + num_value_cols

type_pool = MockDataFactory.MIXED_TYPES[:n_cols]

while len(type_pool) < n_cols:

type_pool = type_pool + MockDataFactory.MIXED_TYPES[: n_cols - len(type_pool)]

cogroups, schema = MockDataFactory.make_cogrouped_batches(

num_groups=num_groups,

num_rows=rows_per_group,

num_cols=n_cols,

spark_type_pool=type_pool,

batch_size=rows_per_group,

)

return_type = StructType(schema.fields[num_key_cols:])

return (cogroups, return_type, num_key_cols, num_value_cols)

_eval_type = PythonEvalType.SQL_COGROUPED_MAP_ARROW_UDF

# Each UDF entry: (func, n_args). n_args=2 -> func(left, right);

# n_args=3 -> func(key, left, right). The Arrow path has no 3-arg variant,

# but the tuple shape is shared with the Pandas sibling so ``_write_scenario``

# can be inherited unchanged.

_udfs = {

"identity_udf": (_cogrouped_map_arrow_identity, 2),

"concat_udf": (_cogrouped_map_arrow_concat, 2),

"left_semi_udf": (_cogrouped_map_arrow_left_semi, 2),

}

params = [list(_scenario_configs), list(_udfs)]

param_names = ["scenario", "udf"]

def _write_scenario(self, scenario, udf_name, buf):

groups, schema, num_key_cols, num_value_cols = self._build_scenario(scenario)

udf_func, _ = self._udfs[udf_name]

left_offsets = MockUDFFactory.make_grouped_arg_offsets(num_key_cols, num_value_cols)

right_offsets = MockUDFFactory.make_grouped_arg_offsets(num_key_cols, num_value_cols)

arg_offsets = left_offsets + right_offsets

MockProtocolWriter.write_worker_input(

self._eval_type,

lambda b: MockProtocolWriter.write_udf_payload(udf_func, schema, arg_offsets, b),

lambda b: MockProtocolWriter.write_grouped_data_payload(groups, buf=b),

buf,

)

class CogroupedMapArrowUDFTimeBench(_CogroupedMapArrowBenchMixin, _TimeBenchBase):

pass

class CogroupedMapArrowUDFPeakmemBench(_CogroupedMapArrowBenchMixin, _PeakmemBenchBase):

pass

# -- SQL_COGROUPED_MAP_PANDAS_UDF ----------------------------------------------

# UDF receives two ``pandas.DataFrame`` (left, right) per co-group, returns

# ``pandas.DataFrame``. Optional 3-arg variant ``(key, left, right)``.

class _CogroupedMapPandasBenchMixin(_CogroupedMapArrowBenchMixin):

"""Provides _write_scenario for SQL_COGROUPED_MAP_PANDAS_UDF.

Inherits ``_build_scenario`` and ``_write_scenario`` from the Arrow

sibling; only the eval type, the UDFs, and the per-scenario row counts

differ. Adds a 3-arg ``key_identity_udf`` variant the Arrow path lacks

(``_write_scenario`` ignores the ``n_args`` slot, so the extra entry is

handled by the inherited writer).

"""

def _cogrouped_map_pandas_identity(left, right):

"""Identity cogroup UDF: returns left DataFrame as-is."""

return left

def _cogrouped_map_pandas_concat(left, right):

"""Concat cogroup UDF: vertically concatenates left and right DataFrames."""

import pandas as pd

return pd.concat([left, right], ignore_index=True)

def _cogrouped_map_pandas_left_semi(left, right):

"""Left-semi cogroup UDF: filters left rows whose key exists in right."""

key_col = left.columns[0]

return left[left[key_col].isin(right[key_col])]

def _cogrouped_map_pandas_key_identity(key, left, right):

"""3-arg cogroup UDF that consumes ``key`` and returns left DataFrame."""

return left

# Scaled down vs SQL_COGROUPED_MAP_ARROW_UDF: pandas conversion adds

# per-group Arrow<->Pandas overhead on both left and right sides.

_scenario_configs = {

"few_groups_sm": (50, 1_000, 1, 4),