This is an automated email from the ASF dual-hosted git repository.
Yicong-Huang 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 ed89e8578f4e [SPARK-59727][PYTHON] Migrate Arrow incremental aggregate
and window UDF eval types to the EvalTypeHandler pipeline
ed89e8578f4e is described below
commit ed89e8578f4e34879d8a925cf1f4b9c8db866a0d
Author: Yicong Huang <[email protected]>
AuthorDate: Thu Oct 1 00:30:06 2026 +0000
[SPARK-59727][PYTHON] Migrate Arrow incremental aggregate and window UDF
eval types to the EvalTypeHandler pipeline
### What changes were proposed in this pull request?
This PR migrates the three Arrow-native incremental aggregate/window UDF
eval types off the `if/elif` chain in `worker.py`'s `read_udfs` onto the
`EvalTypeHandler` pipeline (SPARK-59425, extended in SPARK-59622 /
SPARK-59704): `SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF` (255),
`SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF` (256) and
`SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF` (257).
- `ArrowGroupedAggIncrementalPartialUDFHandler` (`BatchEvalTypeHandler`,
plain `ArrowStreamSerializer`): map-side hash-combine via `reduce`, with the
`maxRecordsPerBatch` buffer cap.
- `ArrowGroupedAggIncrementalFinalUDFHandler` (`GroupedEvalTypeHandler`):
`merge` + `finish` per group, skipping null partial buffers.
- `ArrowWindowAggIncrementalUDFHandler` (`GroupedEvalTypeHandler`):
unbounded/bounded dispatch, extending the running buffer when consecutive
frames only grow on the right.
`_hashable_grouping_key` moves to `pyspark.eval_handlers.utils` as
`hashable_grouping_key`. The corresponding branches and serializer-selection
entries are removed from `worker.py`. This is a pure code move with no behavior
change.
### Why are the changes needed?
Completes moving the Arrow-native eval types onto one registry-dispatched
pipeline. Part of SPARK-59415.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests pass: `pyspark.sql.tests.arrow.test_arrow_python_aggregator`
and `pyspark.sql.tests.connect.arrow.test_parity_arrow_python_aggregator`.
Added handler unit tests in
`pyspark.eval_handlers.tests.test_arrow_eval_type_handlers` (buffer-cap flush
boundary, NaN/complex key canonicalization, null buffer skipping, running vs
sliding frame folding).
Worker microbenchmark via the ASV benchmarks added in #59005 (`./python/asv
run --python=same -a repeat=8`, `time_worker` over the mock wire protocol, ms),
before this migration (`if/elif` in `read_udfs`) vs after (`EvalTypeHandler`):
| `time_worker` (ms) | Before | After |
|---|---|---|
| partial, few_keys | 223±3 | 220±2 |
| partial, many_keys_cap_flush | 284±2 | 279±2 |
| final, few_groups | 154±0.6 | 156±0.4 |
| final, many_groups | 212±2 | 211±2 |
| window, unbounded | 46.9±0.3 | 46.6±0.3 |
| window, running | 101±0.2 | 101±0.4 |
| window, sliding | 333±1 | 331±1 |
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.
Closes #58990 from Yicong-Huang/SPARK-59727-arrow-incremental-agg.
Authored-by: Yicong Huang <[email protected]>
Signed-off-by: Yicong-Huang <[email protected]>
(cherry picked from commit 2b1d78be7d8023fd68d35226d04b915abcbdb742)
Signed-off-by: Yicong-Huang <[email protected]>
---
python/pyspark/eval_handlers/_arrow.py | 277 ++++++++++++++++++-
.../tests/test_arrow_eval_type_handlers.py | 212 +++++++++++++-
python/pyspark/eval_handlers/utils.py | 38 +++
python/pyspark/worker.py | 303 +--------------------
4 files changed, 526 insertions(+), 304 deletions(-)
diff --git a/python/pyspark/eval_handlers/_arrow.py
b/python/pyspark/eval_handlers/_arrow.py
index ec2e919b5147..2f2ccd93a55f 100644
--- a/python/pyspark/eval_handlers/_arrow.py
+++ b/python/pyspark/eval_handlers/_arrow.py
@@ -37,7 +37,7 @@ from pyspark.eval_handlers._base import (
CoGroupedEvalTypeHandler,
GroupedEvalTypeHandler,
)
-from pyspark.eval_handlers.utils import extract_key_value_indexes
+from pyspark.eval_handlers.utils import extract_key_value_indexes,
hashable_grouping_key
from pyspark.eval_handlers.verification import (
verify_iter_result_row_count,
verify_iterator_exhausted,
@@ -389,6 +389,281 @@ class
ArrowWindowAggUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
yield ArrowBatchTransformer.enforce_schema(batch,
self._return_schema)
+class
ArrowGroupedAggIncrementalPartialUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF: map-side PARTIAL stage
of an
+ incremental ``Aggregator``. Hash-combine input rows into a per-group
buffer via
+ the aggregator's ``reduce`` and emit one row per key: the grouping key
columns
+ followed by one buffer struct column per aggregator.
+
+ Ordinary (multi-group) batches are streamed in; the worker keeps one
running
+ buffer per distinct grouping key -- never whole groups of rows. Because
the FINAL
+ stage re-groups these authoritatively after the shuffle, the grouping here
only
+ needs to be a best-effort combine: any keys it fails to collapse (e.g.
NaN, which
+ compares unequal to itself) are merged downstream."""
+
+ eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf:
EvalConf
+ ) -> None:
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ # The leading `num_grouping_keys` input columns are the grouping keys
(see the operator);
+ # `grouping_key_schema` carries their names/types so the emitted key
columns round-trip.
+ grouping_key_schema = eval_conf.grouping_key_schema
+ self._num_grouping_keys = (
+ len(grouping_key_schema.fields) if grouping_key_schema is not None
else 0
+ )
+
+ self._buffer_col_names = ["_%d" % i for i in range(len(udfs))]
+ self._buffer_arrow_types = [
+ to_arrow_type(
+ agg.bufferSchema,
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
+ )
+ for agg, _, _, _ in udfs
+ ]
+ # Buffer field names are invariant across groups; compute them once
per aggregator.
+ self._field_names_by_udf = [
+ [f.name for f in agg.bufferSchema.fields] for agg, _, _, _ in udfs
+ ]
+
+ # The aggregator's `reduce` receives a single positional tuple, so any
named arguments at
+ # the call site are appended after the positional ones, in call order
(kwargs_offsets
+ # preserves that order). This mirrors how a Python call `f(*args,
**kwargs)` would order
+ # them into one value tuple.
+ self._input_offsets_by_udf = [
+ list(args_offsets) + list(kwargs_offsets.values())
+ for _, args_offsets, kwargs_offsets, _ in udfs
+ ]
+ # Input columns actually consumed per batch: the leading grouping keys
plus every
+ # aggregator input, deduplicated. A UDF input may reuse a grouping-key
column (the operator
+ # dedups its projection), so converting by distinct offset avoids
repeated `to_pylist()`.
+ self._needed_offsets = sorted(
+ set(range(self._num_grouping_keys))
+ | {o for offsets in self._input_offsets_by_udf for o in offsets}
+ )
+
+ # Cap the map-side buffer so a high-cardinality partition -- exactly
where partial
+ # aggregation degenerates -- cannot grow the per-key dict without
bound and OOM the worker.
+ # When the distinct-key count reaches the cap we flush the whole map
as one batch and start
+ # fresh; end-of-partition buffers are emitted in equally bounded
chunks. This is safe
+ # because the FINAL stage re-groups the emitted partial buffers
authoritatively after the
+ # shuffle and merges any duplicate keys that early flushes produce. A
non-positive
+ # maxRecordsPerBatch means "unbounded" (mirroring the reader side).
+ max_records = runner_conf.arrow_max_records_per_batch
+ self._cap = max_records if max_records > 0 else None
+
+ def run(self, split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ import pyarrow as pa
+
+ num_grouping_keys = self._num_grouping_keys
+ cap = self._cap
+
+ # hashable key -> (representative key value tuple, list of
per-aggregator buffers)
+ groups: dict[Any, tuple] = {}
+ key_field_types: list | None = None
+
+ def make_batch(entries: list) -> pa.RecordBatch:
+ arrays = []
+ names = []
+ for j in range(num_grouping_keys):
+ arrays.append(
+ pa.array(
+ [e[0][j] for e in entries],
+ type=key_field_types[j], # type: ignore[index]
+ )
+ )
+ names.append("k_%d" % j)
+ for i, field_names in enumerate(self._field_names_by_udf):
+ structs = [
+ {name: e[1][i][t] for t, name in enumerate(field_names)}
for e in entries
+ ]
+ arrays.append(pa.array(structs,
type=self._buffer_arrow_types[i]))
+ names.append(self._buffer_col_names[i])
+ return pa.RecordBatch.from_arrays(arrays, names)
+
+ for batch in data:
+ if key_field_types is None:
+ key_field_types = [batch.schema.field(j).type for j in
range(num_grouping_keys)]
+ pylist_by_offset = {o: batch.column(o).to_pylist() for o in
self._needed_offsets}
+ key_cols = [pylist_by_offset[j] for j in range(num_grouping_keys)]
+ udf_cols = [
+ [pylist_by_offset[o] for o in offsets] for offsets in
self._input_offsets_by_udf
+ ]
+ for r in range(batch.num_rows):
+ key_values = tuple(key_cols[j][r] for j in
range(num_grouping_keys))
+ hashable_key = hashable_grouping_key(key_values)
+ entry = groups.get(hashable_key)
+ if entry is None:
+ buffers = [agg.zero() for agg, _, _, _ in self._udfs]
+ groups[hashable_key] = (key_values, buffers)
+ else:
+ buffers = entry[1]
+ for i, (agg, _, _, _) in enumerate(self._udfs):
+ cols_i = udf_cols[i]
+ buffers[i] = agg.reduce(buffers[i], tuple(c[r] for c in
cols_i))
+ if cap is not None and len(groups) >= cap:
+ yield make_batch(list(groups.values()))
+ groups = {}
+
+ if not groups:
+ # Empty partition (or fully flushed above): emit nothing more. The
FINAL stage
+ # supplies the global-aggregation identity row when there is no
input at all.
+ return
+
+ entries = list(groups.values())
+ if cap is None:
+ yield make_batch(entries)
+ else:
+ for start in range(0, len(entries), cap):
+ yield make_batch(entries[start : start + cap])
+
+
+class
ArrowGroupedAggIncrementalFinalUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF: post-shuffle FINAL stage
of an
+ incremental ``Aggregator``. Merge each group's partial buffers via the
aggregator's
+ ``merge`` and produce the output via ``finish``, streaming the buffers one
batch at
+ a time.
+
+ Every group the JVM sends yields exactly one output row; null
partial-buffer rows
+ are skipped, so an empty global aggregation (a single all-null buffer row
injected
+ by the operator) still produces ``finish(zero)``."""
+
+ eval_type = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf:
EvalConf
+ ) -> None:
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ self._col_names = ["_%d" % i for i in range(len(udfs))]
+ self._return_schema = to_arrow_schema(
+ StructType([StructField(n, rt) for n, (_, _, _, rt) in
zip(self._col_names, udfs)]),
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
+ )
+ # Buffer field names are invariant across groups and batches; compute
once per aggregator.
+ self._field_names_by_udf = [
+ [f.name for f in agg.bufferSchema.fields] for agg, _, _, _ in udfs
+ ]
+
+ def run(self, split_index: int, data: Iterator[GroupedBatch]) ->
Iterator[pa.RecordBatch]:
+ import pyarrow as pa
+
+ for group in data:
+ merged: list = [None] * len(self._udfs)
+ for batch in group:
+ for i, (agg, args_offsets, _, _) in enumerate(self._udfs):
+ field_names = self._field_names_by_udf[i]
+ m = merged[i]
+ for row in batch.column(args_offsets[0]).to_pylist():
+ if row is None:
+ continue
+ partial = tuple(row[name] for name in field_names)
+ m = partial if m is None else agg.merge(m, partial)
+ merged[i] = m
+ results = []
+ for i, (agg, _, _, _) in enumerate(self._udfs):
+ m = merged[i] if merged[i] is not None else agg.zero()
+ results.append(agg.finish(m))
+ # Type each output array explicitly (mirroring the PARTIAL stage)
so a non-trivial
+ # outputType or an all-None column does not depend on Arrow type
inference.
+ result_arrays = [
+ pa.array([r], type=self._return_schema.field(i).type) for i, r
in enumerate(results)
+ ]
+ batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names)
+ yield ArrowBatchTransformer.enforce_schema(batch,
self._return_schema)
+
+
+class
ArrowWindowAggIncrementalUDFHandler(GroupedEvalTypeHandler["pa.RecordBatch"]):
+ """SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF: window aggregation with an
incremental
+ ``Aggregator``. The operator sends each frame -- the whole partition for an
+ unbounded frame, or per-row ``[begin, end)`` slices for a bounded one --
and each
+ frame's rows are folded with ``reduce`` (from a fresh ``zero``) and
finished with
+ ``finish``, one output value per input row. A window has no shuffle, so the
+ intermediate buffer never leaves the worker; ``merge`` is not used here."""
+
+ eval_type = PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF
+
+ def __init__(
+ self, udfs: list[tuple[Any, ...]], runner_conf: RunnerConf, eval_conf:
EvalConf
+ ) -> None:
+ require_minimum_pyarrow_version()
+ super().__init__(udfs, runner_conf, eval_conf)
+ self._window_bound_types = runner_conf.window_bound_types
+ self._col_names = ["_%d" % i for i in range(len(udfs))]
+ self._return_schema = to_arrow_schema(
+ StructType([StructField(n, rt) for n, (_, _, _, rt) in
zip(self._col_names, udfs)]),
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
+ )
+
+ def run(self, split_index: int, data: Iterator[GroupedBatch]) ->
Iterator[pa.RecordBatch]:
+ import pyarrow as pa
+
+ def fold(agg: Any, buffer: Any, value_cols: list, start: int, end:
int) -> Any:
+ # Fold rows ``[start, end)`` (each a tuple across ``value_cols``,
matching the
+ # call-site argument order) into ``buffer`` via the aggregator's
``reduce``.
+ for r in range(start, end):
+ buffer = agg.reduce(buffer, tuple(c[r] for c in value_cols))
+ return buffer
+
+ for group in data:
+ concatenated = ArrowBatchTransformer.concat_batches(group)
+ num_rows = concatenated.num_rows
+
+ result_arrays = []
+ for udf_index, (agg, args_offsets, kwargs_offsets, _) in
enumerate(self._udfs):
+ bound_type = self._window_bound_types[udf_index]
+ result_type = self._return_schema.field(udf_index).type
+ if bound_type == "unbounded":
+ # One frame spanning the whole partition: compute once,
repeat per row.
+ value_cols = [concatenated.column(o).to_pylist() for o in
args_offsets] + [
+ concatenated.column(v).to_pylist() for v in
kwargs_offsets.values()
+ ]
+ result = agg.finish(fold(agg, agg.zero(), value_cols, 0,
num_rows))
+ result_arrays.append(pa.array([result] * num_rows,
type=result_type))
+ elif bound_type == "bounded":
+ # Per-row frame ``[begin, end)``. Materialize the
aggregator's input columns
+ # once; frames index into them by row.
+ begin_col = concatenated.column(args_offsets[0])
+ end_col = concatenated.column(args_offsets[1])
+ data_offsets = list(args_offsets[2:]) +
list(kwargs_offsets.values())
+ value_cols = [concatenated.column(o).to_pylist() for o in
data_offsets]
+ # When consecutive frames share the same lower bound and
only grow on the
+ # right (e.g. rowsBetween(unboundedPreceding,
currentRow)), extend the
+ # running buffer by the newly-included rows instead of
refolding from
+ # ``zero`` -- O(n) overall rather than O(n^2). Otherwise
-- the lower bound
+ # advanced (a row left the window, which ``reduce`` cannot
subtract) or the
+ # frame shrank -- refold the frame from ``zero``.
+ results = []
+ have_running = False
+ running: Any = None
+ prev_begin = -1
+ prev_end = 0
+ for i in range(num_rows):
+ begin = begin_col[i].as_py()
+ end = end_col[i].as_py()
+ if have_running and begin == prev_begin and end >=
prev_end:
+ running = fold(agg, running, value_cols, prev_end,
end)
+ else:
+ running = fold(agg, agg.zero(), value_cols, begin,
end)
+ have_running = True
+ prev_begin, prev_end = begin, end
+ results.append(agg.finish(running))
+ result_arrays.append(pa.array(results, type=result_type))
+ else:
+ raise PySparkRuntimeError(
+ errorClass="INVALID_WINDOW_BOUND_TYPE",
+ messageParameters={"window_bound_type": bound_type},
+ )
+
+ batch = pa.RecordBatch.from_arrays(result_arrays, self._col_names)
+ yield ArrowBatchTransformer.enforce_schema(batch,
self._return_schema)
+
+
class ArrowMapUDFHandler(BatchEvalTypeHandler["pa.RecordBatch"]):
"""SQL_MAP_ARROW_ITER_UDF (mapInArrow): the single UDF receives the input
RecordBatch stream and yields a RecordBatch stream, exchanged as flattened
diff --git
a/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py
b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py
index fb84b1020f8d..5c47ae609f41 100644
--- a/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py
+++ b/python/pyspark/eval_handlers/tests/test_arrow_eval_type_handlers.py
@@ -22,20 +22,28 @@ same wire format the serializers produce, so every test
constructs its input the
same way: ``run(0, <input>)`` and assert on the output batches.
"""
+import os
import unittest
+from unittest.mock import patch
from pyspark.errors import PySparkRuntimeError
from pyspark.eval_handlers._base import get_eval_type_handler
+from pyspark.eval_handlers.utils import hashable_grouping_key
from pyspark.sql.pandas.serializers import ArrowStreamCoGroupSerializer,
ArrowStreamSerializer
-from pyspark.sql.types import LongType, StructField, StructType
+from pyspark.sql.types import DoubleType, LongType, StructField, StructType
from pyspark.testing.utils import have_pyarrow, pyarrow_requirement_message
from pyspark.util import PythonEvalType
+with patch.dict(os.environ, {"SPARK_PYTHON_RUNTIME": "PYTHON_WORKER"}):
+ from pyspark.worker_util import EvalConf, RunnerConf
+
if have_pyarrow:
import pyarrow as pa
from pyspark.eval_handlers._arrow import (
ArrowCoGroupedMapUDFHandler,
+ ArrowGroupedAggIncrementalFinalUDFHandler,
+ ArrowGroupedAggIncrementalPartialUDFHandler,
ArrowGroupedAggIterUDFHandler,
ArrowGroupedAggUDFHandler,
ArrowGroupedMapIterUDFHandler,
@@ -43,6 +51,7 @@ if have_pyarrow:
ArrowMapUDFHandler,
ArrowScalarIterUDFHandler,
ArrowScalarUDFHandler,
+ ArrowWindowAggIncrementalUDFHandler,
ArrowWindowAggUDFHandler,
)
from pyspark.sql.conversion import ArrowBatchTransformer
@@ -59,6 +68,33 @@ class _RunnerConf:
window_bound_types = ["unbounded"]
+class _SumAggregator:
+ """Incremental ``Aggregator`` stub: sums non-null inputs, with a ``(sum,
count)`` buffer.
+
+ ``reduce_calls`` counts ``reduce`` invocations so tests can assert how a
window frame was
+ folded (extended vs refolded from ``zero``).
+ """
+
+ bufferSchema = StructType([StructField("s", LongType()), StructField("c",
LongType())])
+
+ def __init__(self):
+ self.reduce_calls = 0
+
+ def zero(self):
+ return (0, 0)
+
+ def reduce(self, buffer, values):
+ self.reduce_calls += 1
+ (v,) = values
+ return buffer if v is None else (buffer[0] + v, buffer[1] + 1)
+
+ def merge(self, left, right):
+ return (left[0] + right[0], left[1] + right[1])
+
+ def finish(self, buffer):
+ return buffer[0]
+
+
def _batch(**columns):
"""A RecordBatch of int64 columns, one per ``name=values`` kwarg."""
return pa.RecordBatch.from_arrays(
@@ -136,6 +172,18 @@ class
ArrowEvalTypeHandlerRegistrationTests(unittest.TestCase):
(PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
ArrowGroupedAggUDFHandler),
(PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
ArrowGroupedAggIterUDFHandler),
(PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
ArrowWindowAggUDFHandler),
+ (
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
+ ArrowGroupedAggIncrementalPartialUDFHandler,
+ ),
+ (
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ ArrowGroupedAggIncrementalFinalUDFHandler,
+ ),
+ (
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
+ ArrowWindowAggIncrementalUDFHandler,
+ ),
):
self.assertIs(get_eval_type_handler(eval_type), handler_cls)
@@ -322,6 +370,168 @@ class ArrowWindowAggUDFHandlerTests(unittest.TestCase):
self.assertEqual(out[0].column("_0").to_pylist(), [10, 30])
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowGroupedAggIncrementalPartialUDFHandlerTests(unittest.TestCase):
+ # Input columns are ``[k, v]``: the leading grouping key, then the
aggregator input. Output
+ # is one row per key: ``k_0`` followed by the ``_0`` buffer struct.
+ _KEY_SCHEMA = StructType([StructField("k", LongType())])
+
+ def _handler(self, cap=10000):
+ return ArrowGroupedAggIncrementalPartialUDFHandler(
+ udfs=[(_SumAggregator(), [1], {}, LongType())],
+
runner_conf=RunnerConf({"spark.sql.execution.arrow.maxRecordsPerBatch":
str(cap)}),
+ eval_conf=EvalConf({"grouping_key_schema":
self._KEY_SCHEMA.json()}),
+ )
+
+ @staticmethod
+ def _rows(out):
+ return [
+ (k, buf["s"], buf["c"])
+ for b in out
+ for k, buf in zip(b.column("k_0").to_pylist(),
b.column("_0").to_pylist())
+ ]
+
+ def test_combines_rows_per_key_across_batches(self):
+ out = list(
+ self._handler().run(0, iter([_batch(k=[1, 2, 1], v=[1, 2, 3]),
_batch(k=[2], v=[4])]))
+ )
+ self.assertEqual(len(out), 1)
+ self.assertEqual(sorted(self._rows(out)), [(1, 4, 2), (2, 6, 2)])
+
+ def test_empty_partition_emits_nothing(self):
+ self.assertEqual(list(self._handler().run(0, iter([]))), [])
+
+ def test_flushes_when_key_count_reaches_cap(self):
+ # cap=2: the check runs after each input batch. The first batch brings
the map to exactly
+ # 2 keys, so it flushes (>= cap); key 1 then reappears in a fresh map,
and the FINAL stage
+ # merges the duplicate downstream. A mid-stream flush is not chunked,
so the second batch
+ # flushes all 3 of its keys at once.
+ out = list(
+ self._handler(cap=2).run(
+ 0, iter([_batch(k=[1, 2], v=[1, 2]), _batch(k=[1, 3, 4],
v=[10, 3, 4])])
+ )
+ )
+ self.assertEqual([b.num_rows for b in out], [2, 3])
+ self.assertEqual(sorted(self._rows(out[:1])), [(1, 1, 1), (2, 2, 1)])
+ self.assertEqual(sorted(self._rows(out[1:])), [(1, 10, 1), (3, 3, 1),
(4, 4, 1)])
+
+ def test_below_cap_is_not_flushed(self):
+ out = list(
+ self._handler(cap=3).run(0, iter([_batch(k=[1, 2], v=[1, 2]),
_batch(k=[1], v=[5])]))
+ )
+ self.assertEqual([b.num_rows for b in out], [2])
+ self.assertEqual(sorted(self._rows(out)), [(1, 6, 2), (2, 2, 1)])
+
+ def test_non_positive_cap_is_unbounded(self):
+ out = list(self._handler(cap=0).run(0, iter([_batch(k=[1, 2, 3], v=[1,
2, 3])])))
+ self.assertEqual([b.num_rows for b in out], [3])
+
+ def test_nan_keys_are_combined(self):
+ batch = pa.RecordBatch.from_arrays(
+ [pa.array([float("nan"), float("nan")]), pa.array([1, 2],
type=pa.int64())],
+ ["k", "v"],
+ )
+ handler = ArrowGroupedAggIncrementalPartialUDFHandler(
+ udfs=[(_SumAggregator(), [1], {}, LongType())],
+ runner_conf=RunnerConf({}),
+ eval_conf=EvalConf(
+ {"grouping_key_schema": StructType([StructField("k",
DoubleType())]).json()}
+ ),
+ )
+ out = list(handler.run(0, iter([batch])))
+ self.assertEqual(out[0].num_rows, 1)
+ self.assertEqual(out[0].column("_0").to_pylist(), [{"s": 3, "c": 2}])
+
+
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowGroupedAggIncrementalFinalUDFHandlerTests(unittest.TestCase):
+ # Each group's rows carry one partial-buffer struct column; the handler
merges them and
+ # emits one ``finish`` value per group.
+ _BUFFER_TYPE = pa.struct([("s", pa.int64()), ("c", pa.int64())]) if
have_pyarrow else None
+
+ def _buffers(self, buffers):
+ return pa.RecordBatch.from_arrays([pa.array(buffers,
type=self._BUFFER_TYPE)], ["buf"])
+
+ def _handler(self):
+ return ArrowGroupedAggIncrementalFinalUDFHandler(
+ udfs=[(_SumAggregator(), [0], {}, LongType())],
+ runner_conf=RunnerConf({}),
+ eval_conf=None,
+ )
+
+ def test_merges_partial_buffers_per_group(self):
+ groups = iter(
+ [
+ iter([self._buffers([{"s": 1, "c": 1}]), self._buffers([{"s":
2, "c": 1}])]),
+ iter([self._buffers([{"s": 10, "c": 2}])]),
+ ]
+ )
+ out = list(self._handler().run(0, groups))
+ self.assertEqual([b.column("_0").to_pylist() for b in out], [[3],
[10]])
+
+ def test_null_buffers_are_skipped(self):
+ # An all-null buffer group (the operator's empty global aggregation)
finishes ``zero``.
+ out = list(self._handler().run(0, _one_group(self._buffers([None]))))
+ self.assertEqual(out[0].column("_0").to_pylist(), [0])
+ out = list(self._handler().run(0, _one_group(self._buffers([None,
{"s": 5, "c": 1}]))))
+ self.assertEqual(out[0].column("_0").to_pylist(), [5])
+
+
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowWindowAggIncrementalUDFHandlerTests(unittest.TestCase):
+ def _handler(self, agg, bound_type, args_offsets):
+ return ArrowWindowAggIncrementalUDFHandler(
+ udfs=[(agg, args_offsets, {}, LongType())],
+ runner_conf=RunnerConf({"window_bound_types": bound_type}),
+ eval_conf=None,
+ )
+
+ def test_unbounded_frame_repeats_one_value(self):
+ handler = self._handler(_SumAggregator(), "unbounded", [0])
+ out = list(handler.run(0, _one_group(_batch(v=[1, 2]), _batch(v=[3]))))
+ self.assertEqual(out[0].column("_0").to_pylist(), [6, 6, 6])
+
+ def test_growing_frame_extends_running_buffer(self):
+ # rowsBetween(unboundedPreceding, currentRow): same lower bound,
growing upper bound,
+ # so each row folds only its newly-included row.
+ agg = _SumAggregator()
+ handler = self._handler(agg, "bounded", [0, 1, 2])
+ batch = _batch(begin=[0, 0, 0, 0], end=[1, 2, 3, 4], v=[1, 2, 3, 4])
+ out = list(handler.run(0, _one_group(batch)))
+ self.assertEqual(out[0].column("_0").to_pylist(), [1, 3, 6, 10])
+ self.assertEqual(agg.reduce_calls, 4)
+
+ def test_sliding_frame_refolds_from_zero(self):
+ # rowsBetween(-1, currentRow): the lower bound advances, so each frame
refolds.
+ agg = _SumAggregator()
+ handler = self._handler(agg, "bounded", [0, 1, 2])
+ batch = _batch(begin=[0, 0, 1, 2], end=[1, 2, 3, 4], v=[1, 2, 3, 4])
+ out = list(handler.run(0, _one_group(batch)))
+ self.assertEqual(out[0].column("_0").to_pylist(), [1, 3, 5, 7])
+ self.assertEqual(agg.reduce_calls, 1 + 1 + 2 + 2)
+
+ def test_invalid_bound_type_raises(self):
+ handler = self._handler(_SumAggregator(), "bogus", [0])
+ with self.assertRaises(PySparkRuntimeError):
+ list(handler.run(0, _one_group(_batch(v=[1]))))
+
+
+class HashableGroupingKeyTests(unittest.TestCase):
+ def test_nested_values_hash_by_value(self):
+ self.assertEqual(
+ hashable_grouping_key(([1, [2]], {"a": 1})),
hashable_grouping_key(([1, [2]], {"a": 1}))
+ )
+
+ def test_nan_values_collapse(self):
+ self.assertEqual(
+ hashable_grouping_key((float("nan"), [float("nan")])),
+ hashable_grouping_key((float("nan"), [float("nan")])),
+ )
+
+ def test_unhashable_value_gets_a_unique_key(self):
+ self.assertNotEqual(hashable_grouping_key(({1},)),
hashable_grouping_key(({1},)))
+
+
@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
class CoGroupedBatchTests(unittest.TestCase):
def test_deserialized_co_group_is_a_pair_of_lists(self):
diff --git a/python/pyspark/eval_handlers/utils.py
b/python/pyspark/eval_handlers/utils.py
index 48a757020337..66f20872c5a9 100644
--- a/python/pyspark/eval_handlers/utils.py
+++ b/python/pyspark/eval_handlers/utils.py
@@ -21,6 +21,8 @@ A leaf module (no pyspark imports), so it is safe to import
from both the worker
and the handlers, on the driver and the executor alike.
"""
+from typing import Any
+
def extract_key_value_indexes(grouped_arg_offsets: list) -> list:
"""
@@ -50,3 +52,39 @@ def extract_key_value_indexes(grouped_arg_offsets: list) ->
list:
parsed.append([offset_keys, offset_values])
idx += offsets_len
return parsed
+
+
+# Sentinel standing in for NaN grouping-key values in the map-side
incremental-aggregate combine.
+# ``float('nan') != float('nan')``, so distinct NaN objects would never
collide in a dict; mapping
+# them to one sentinel lets the map-side combine collapse them (correctness
does not depend on it,
+# as the FINAL stage re-groups authoritatively).
+_NAN_GROUPING_KEY = object()
+
+
+def hashable_grouping_key(key_values: tuple[Any, ...]) -> Any:
+ """
+ Canonicalize a grouping-key value tuple (extracted from Arrow via
``to_pylist``) into a
+ hashable form usable as a ``dict`` key for the map-side PARTIAL combine of
incremental Python
+ aggregators.
+
+ Lists (from ``array`` columns) become tuples and dicts (from ``struct``
columns) become tuples
+ of ``(name, value)`` pairs, so nested complex keys hash by value; NaN
floats map to a sentinel.
+ This is a best-effort combine only: an exotic unhashable value falls back
to a unique object so
+ the row forms its own group, and the FINAL stage merges any keys left
uncollapsed here.
+ """
+
+ def canon(v: Any) -> Any:
+ if isinstance(v, float) and v != v:
+ return _NAN_GROUPING_KEY
+ if isinstance(v, list):
+ return tuple(canon(x) for x in v)
+ if isinstance(v, dict):
+ return tuple((k, canon(x)) for k, x in v.items())
+ return v
+
+ try:
+ canonical = tuple(canon(v) for v in key_values)
+ hash(canonical)
+ return canonical
+ except TypeError:
+ return object()
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 7345f45e1ceb..e92fa5d7afc0 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -88,7 +88,7 @@ from pyspark.sql.pandas.serializers import (
ArrowStreamGroupSerializer,
ArrowStreamSerializer,
)
-from pyspark.sql.pandas.types import to_arrow_schema, to_arrow_type
+from pyspark.sql.pandas.types import to_arrow_type
from pyspark.sql.streaming.stateful_processor_api_client import
StatefulProcessorApiClient
from pyspark.sql.streaming.stateful_processor_util import
TransformWithStateInPandasFuncMode
from pyspark.sql.types import (
@@ -152,42 +152,6 @@ def chain(f, g):
return lambda *a: g(f(*a))
-# Sentinel standing in for NaN grouping-key values in the map-side
incremental-aggregate combine.
-# ``float('nan') != float('nan')``, so distinct NaN objects would never
collide in a dict; mapping
-# them to one sentinel lets the map-side combine collapse them (correctness
does not depend on it,
-# as the FINAL stage re-groups authoritatively).
-_NAN_GROUPING_KEY = object()
-
-
-def _hashable_grouping_key(key_values: Tuple[Any, ...]) -> Any:
- """
- Canonicalize a grouping-key value tuple (extracted from Arrow via
``to_pylist``) into a
- hashable form usable as a ``dict`` key for the map-side PARTIAL combine of
incremental Python
- aggregators.
-
- Lists (from ``array`` columns) become tuples and dicts (from ``struct``
columns) become tuples
- of ``(name, value)`` pairs, so nested complex keys hash by value; NaN
floats map to a sentinel.
- This is a best-effort combine only: an exotic unhashable value falls back
to a unique object so
- the row forms its own group, and the FINAL stage merges any keys left
uncollapsed here.
- """
-
- def canon(v: Any) -> Any:
- if isinstance(v, float) and v != v:
- return _NAN_GROUPING_KEY
- if isinstance(v, list):
- return tuple(canon(x) for x in v)
- if isinstance(v, dict):
- return tuple((k, canon(x)) for k, x in v.items())
- return v
-
- try:
- canonical = tuple(canon(v) for v in key_values)
- hash(canonical)
- return canonical
- except TypeError:
- return object()
-
-
def _verify_column_schema(
actual_names: list, expected_names: list, *, assign_cols_by_name: bool
) -> None:
@@ -1867,22 +1831,14 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PANDAS_INIT_STATE_UDF,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_UDF,
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF,
- PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
- PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
- PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
):
# NOTE: if timezone is set here, that implies respectSessionTimeZone
is True
if eval_type in (
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF,
- # The map-side PARTIAL stage streams ordinary (multi-group)
batches and hash-combines
- # inside the worker, so it uses the plain (non-grouped) stream
serializer below. Only
- # the post-shuffle FINAL stage receives one Arrow stream per group.
- PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
- PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
):
ser = ArrowStreamGroupSerializer(write_start_stream=True)
elif eval_type == PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF:
@@ -1900,170 +1856,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
num_udfs = len(udfs)
- if eval_type ==
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF:
- import pyarrow as pa
-
- # Map-side PARTIAL stage: hash-combine input rows into a per-group
buffer via the
- # aggregator's `reduce`. Ordinary (multi-group) batches are streamed
in; the worker keeps
- # one running buffer per distinct grouping key -- never whole groups
of rows -- and, at end
- # of partition, emits one row per key: the grouping key columns
followed by one buffer
- # struct column per aggregator. Because the FINAL stage re-groups
these authoritatively
- # after the shuffle, the worker's grouping only needs to be a
best-effort combine: any keys
- # it fails to collapse (e.g. NaN, which compares unequal to itself)
are merged downstream.
- #
- # The leading `num_grouping_keys` input columns are the grouping keys
(see the operator);
- # `grouping_key_schema` carries their names/types so the emitted key
columns round-trip.
- grouping_key_schema = eval_conf.grouping_key_schema
- num_grouping_keys = (
- len(grouping_key_schema.fields) if grouping_key_schema is not None
else 0
- )
-
- buffer_col_names = ["_%d" % i for i in range(len(udfs))]
- buffer_arrow_types = [
- to_arrow_type(
- agg.bufferSchema,
- timezone="UTC",
- prefers_large_types=runner_conf.use_large_var_types,
- )
- for agg, _, _, _ in udfs
- ]
- # Buffer field names are invariant across groups; compute them once
per aggregator.
- field_names_by_udf = [[f.name for f in agg.bufferSchema.fields] for
agg, _, _, _ in udfs]
-
- # The aggregator's `reduce` receives a single positional tuple, so any
named arguments at
- # the call site are appended after the positional ones, in call order
(kwargs_offsets
- # preserves that order). This mirrors how a Python call `f(*args,
**kwargs)` would order
- # them into one value tuple.
- input_offsets_by_udf = [
- list(args_offsets) + list(kwargs_offsets.values())
- for _, args_offsets, kwargs_offsets, _ in udfs
- ]
- # Input columns actually consumed per batch: the leading grouping keys
plus every
- # aggregator input, deduplicated. A UDF input may reuse a grouping-key
column (the operator
- # dedups its projection), so converting by distinct offset avoids
repeated `to_pylist()`.
- needed_offsets = sorted(
- set(range(num_grouping_keys)) | {o for offsets in
input_offsets_by_udf for o in offsets}
- )
-
- # Cap the map-side buffer so a high-cardinality partition -- exactly
where partial
- # aggregation degenerates -- cannot grow the per-key dict without
bound and OOM the worker.
- # When the distinct-key count reaches the cap we flush the whole map
as one batch and start
- # fresh; end-of-partition buffers are emitted in equally bounded
chunks. This is safe
- # because the FINAL stage re-groups the emitted partial buffers
authoritatively after the
- # shuffle and merges any duplicate keys that early flushes produce. A
non-positive
- # maxRecordsPerBatch means "unbounded" (mirroring the reader side).
- max_records = runner_conf.arrow_max_records_per_batch
- cap = max_records if max_records > 0 else None
-
- def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
- # hashable key -> (representative key value tuple, list of
per-aggregator buffers)
- groups: "dict[Any, tuple]" = {}
- key_field_types: Optional[list] = None
-
- def make_batch(entries: list) -> pa.RecordBatch:
- arrays = []
- names = []
- for j in range(num_grouping_keys):
- arrays.append(
- pa.array(
- [e[0][j] for e in entries],
- type=key_field_types[j], # type: ignore
- )
- )
- names.append("k_%d" % j)
- for i, (agg, _, _, _) in enumerate(udfs):
- field_names = field_names_by_udf[i]
- structs = [
- {name: e[1][i][t] for t, name in
enumerate(field_names)} for e in entries
- ]
- arrays.append(pa.array(structs,
type=buffer_arrow_types[i]))
- names.append(buffer_col_names[i])
- return pa.RecordBatch.from_arrays(arrays, names)
-
- for batch in data:
- if key_field_types is None:
- key_field_types = [batch.schema.field(j).type for j in
range(num_grouping_keys)]
- pylist_by_offset = {o: batch.column(o).to_pylist() for o in
needed_offsets}
- key_cols = [pylist_by_offset[j] for j in
range(num_grouping_keys)]
- udf_cols = [
- [pylist_by_offset[o] for o in input_offsets_by_udf[i]] for
i in range(len(udfs))
- ]
- for r in range(batch.num_rows):
- key_values = tuple(key_cols[j][r] for j in
range(num_grouping_keys))
- hashable_key = _hashable_grouping_key(key_values)
- entry = groups.get(hashable_key)
- if entry is None:
- buffers = [agg.zero() for agg, _, _, _ in udfs]
- groups[hashable_key] = (key_values, buffers)
- else:
- buffers = entry[1]
- for i, (agg, _, _, _) in enumerate(udfs):
- cols_i = udf_cols[i]
- buffers[i] = agg.reduce(buffers[i], tuple(c[r] for c
in cols_i))
- if cap is not None and len(groups) >= cap:
- yield make_batch(list(groups.values()))
- groups = {}
-
- if not groups:
- # Empty partition (or fully flushed above): emit nothing more.
The FINAL stage
- # supplies the global-aggregation identity row when there is
no input at all.
- return
-
- entries = list(groups.values())
- if cap is None:
- yield make_batch(entries)
- else:
- for start in range(0, len(entries), cap):
- yield make_batch(entries[start : start + cap])
-
- return func, ser
-
- if eval_type == PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF:
- import pyarrow as pa
-
- # Post-shuffle FINAL stage: merge each group's partial buffers via the
aggregator's `merge`
- # and produce the output via `finish`. Buffers are streamed and merged
one batch at a time.
- # Every group the JVM sends yields exactly one output row; null
partial-buffer rows are
- # skipped, so an empty global aggregation (a single all-null buffer
row injected by the
- # operator) still produces `finish(zero)`.
- col_names = ["_%d" % i for i in range(len(udfs))]
- return_schema = to_arrow_schema(
- StructType([StructField(name, rt) for name, (_, _, _, rt) in
zip(col_names, udfs)]),
- timezone="UTC",
- prefers_large_types=runner_conf.use_large_var_types,
- )
- # Buffer field names are invariant across groups and batches; compute
once per aggregator.
- field_names_by_udf = [[f.name for f in agg.bufferSchema.fields] for
agg, _, _, _ in udfs]
-
- def grouped_func(
- split_index: int, data: Iterator["GroupedBatch"]
- ) -> Iterator[pa.RecordBatch]:
- for group in data:
- merged: list = [None] * len(udfs)
- for batch in group:
- for i, (agg, args_offsets, _, _) in enumerate(udfs):
- field_names = field_names_by_udf[i]
- m = merged[i]
- for row in batch.column(args_offsets[0]).to_pylist():
- if row is None:
- continue
- partial = tuple(row[name] for name in field_names)
- m = partial if m is None else agg.merge(m, partial)
- merged[i] = m
- results = []
- for i, (agg, _, _, _) in enumerate(udfs):
- m = merged[i] if merged[i] is not None else agg.zero()
- results.append(agg.finish(m))
- # Type each output array explicitly (mirroring the PARTIAL
stage) so a non-trivial
- # outputType or an all-None column does not depend on Arrow
type inference.
- result_arrays = [
- pa.array([r], type=return_schema.field(i).type) for i, r
in enumerate(results)
- ]
- batch = pa.RecordBatch.from_arrays(result_arrays, col_names)
- yield ArrowBatchTransformer.enforce_schema(batch,
return_schema)
-
- return grouped_func, ser
-
if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF:
import pandas as pd
import pyarrow as pa
@@ -2152,99 +1944,6 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
return grouped_func, ser
- if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF:
- import pyarrow as pa
-
- # Window aggregation with an incremental ``Aggregator``. The operator
sends each frame --
- # the whole partition for an unbounded frame, or per-row ``[begin,
end)`` slices for a
- # bounded one -- and the worker folds the frame's rows with ``reduce``
(from a fresh
- # ``zero``) and produces the value with ``finish``, one output value
per input row. A window
- # has no shuffle, so the intermediate buffer never leaves the worker
(unlike the two-stage
- # groupBy path); ``merge`` is not used here.
- window_bound_types = runner_conf.window_bound_types
-
- col_names = ["_%d" % i for i in range(len(udfs))]
- return_schema = to_arrow_schema(
- StructType([StructField(name, rt) for name, (_, _, _, rt) in
zip(col_names, udfs)]),
- timezone="UTC",
- prefers_large_types=runner_conf.use_large_var_types,
- )
-
- def fold(agg: Any, buffer: Any, value_cols: list, start: int, end:
int) -> Any:
- # Fold rows ``[start, end)`` (each a tuple across ``value_cols``,
matching the call-site
- # argument order) into ``buffer`` via the aggregator's ``reduce``.
- for r in range(start, end):
- buffer = agg.reduce(buffer, tuple(c[r] for c in value_cols))
- return buffer
-
- def grouped_func(
- split_index: int, data: Iterator["GroupedBatch"]
- ) -> Iterator[pa.RecordBatch]:
- for group in data:
- batch_list = list(group)
- if not batch_list:
- continue
- if hasattr(pa, "concat_batches"):
- concatenated = pa.concat_batches(batch_list)
- else:
- # pyarrow.concat_batches not supported before 19.0.0
- # remove this once we drop support for old versions
- concatenated = pa.RecordBatch.from_struct_array(
- pa.concat_arrays([b.to_struct_array() for b in
batch_list])
- )
- num_rows = concatenated.num_rows
-
- result_arrays = []
- for udf_index, (agg, args_offsets, kwargs_offsets, _) in
enumerate(udfs):
- bound_type = window_bound_types[udf_index]
- result_type = return_schema.field(udf_index).type
- if bound_type == "unbounded":
- # One frame spanning the whole partition: compute
once, repeat per row.
- value_cols = [concatenated.column(o).to_pylist() for o
in args_offsets] + [
- concatenated.column(v).to_pylist() for v in
kwargs_offsets.values()
- ]
- result = agg.finish(fold(agg, agg.zero(), value_cols,
0, num_rows))
- result_arrays.append(pa.array([result] * num_rows,
type=result_type))
- elif bound_type == "bounded":
- # Per-row frame ``[begin, end)``. Materialize the
aggregator's input columns
- # once; frames index into them by row.
- begin_col = concatenated.column(args_offsets[0])
- end_col = concatenated.column(args_offsets[1])
- data_offsets = list(args_offsets[2:]) +
list(kwargs_offsets.values())
- value_cols = [concatenated.column(o).to_pylist() for o
in data_offsets]
- # When consecutive frames share the same lower bound
and only grow on the
- # right (e.g. rowsBetween(unboundedPreceding,
currentRow)), extend the
- # running buffer by the newly-included rows instead of
refolding from
- # ``zero`` -- O(n) overall rather than O(n^2).
Otherwise -- the lower bound
- # advanced (a row left the window, which ``reduce``
cannot subtract) or the
- # frame shrank -- refold the frame from ``zero``.
- results = []
- have_running = False
- running: Any = None
- prev_begin = -1
- prev_end = 0
- for i in range(num_rows):
- begin = begin_col[i].as_py()
- end = end_col[i].as_py()
- if have_running and begin == prev_begin and end >=
prev_end:
- running = fold(agg, running, value_cols,
prev_end, end)
- else:
- running = fold(agg, agg.zero(), value_cols,
begin, end)
- have_running = True
- prev_begin, prev_end = begin, end
- results.append(agg.finish(running))
- result_arrays.append(pa.array(results,
type=result_type))
- else:
- raise PySparkRuntimeError(
- errorClass="INVALID_WINDOW_BOUND_TYPE",
- messageParameters={"window_bound_type":
bound_type},
- )
-
- batch = pa.RecordBatch.from_arrays(result_arrays, col_names)
- yield ArrowBatchTransformer.enforce_schema(batch,
return_schema)
-
- return grouped_func, ser
-
if eval_type == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF:
import pandas as pd
import pyarrow as pa
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
