This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new b18b7d316721 [SPARK-58736][SQL][PYTHON][CONNECT] Support incremental
Python aggregators via Arrow
b18b7d316721 is described below
commit b18b7d316721f9d7f479fc08e8d8f4383b51111a
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Wed Aug 19 19:22:26 2026 +0900
[SPARK-58736][SQL][PYTHON][CONNECT] Support incremental Python aggregators
via Arrow
### What changes were proposed in this pull request?
Adds a Python analog of the Scala typed
`org.apache.spark.sql.expressions.Aggregator[IN, BUF, OUT]`
with **true incremental (partial) aggregation** — i.e. map-side combine,
not whole-group
materialization.
Users subclass a new `Aggregator` base class (`zero` / `reduce` / `merge` /
`finish` +
`bufferSchema`) and wrap it with `udaf(...)` for use in
`groupBy().agg(...)`:
```python
from pyspark.sql.aggregator import Aggregator, udaf
from pyspark.sql.types import StructType, StructField, DoubleType, LongType
class Mean(Aggregator):
property
def bufferSchema(self):
return StructType([StructField("sum", DoubleType()),
StructField("count", LongType())])
property
def outputType(self):
return DoubleType()
def zero(self): return (0.0, 0)
def reduce(self, buf, v): return (buf[0] + v[0], buf[1] + 1)
def merge(self, a, b): return (a[0] + b[0], a[1] + b[1])
def finish(self, buf): return buf[0] / buf[1] if buf[1] else None
df.groupBy("k").agg(udaf(Mean())(df.v))
```
Unlike grouped-agg pandas/arrow UDFs (`PythonUDAF` +
`ArrowAggregatePythonExec`), which collect the
whole group and call Python once, this is planned as a **two-stage
aggregation**:
- a map-side **PARTIAL** stage folds each group's input rows into a
per-group buffer via `reduce`;
- the buffers are shuffled by the grouping key (as an Arrow struct column);
- a **FINAL** stage merges the partial buffers via `merge` and produces the
output via `finish`.
Because `merge` is associative/commutative, the result is independent of
partition count.
Class hierarchy / trace:
- `PythonEvalType`: new `SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF`
(255) and `..._FINAL_UDF`
(256), added on both the Python (`pyspark.util`) and JVM
(`api.python.PythonEvalType`) sides.
- Catalyst: new `PythonAggregate` expression (an
`UnevaluableAggregateFunc`, like `PythonUDAF`)
carrying the intermediate `bufferSchema`.
- Planning: `SparkStrategies.Aggregation` routes an all-`PythonAggregate`
aggregate to
`PythonIncrementalAggregateExec.plan(...)`, which builds
`PythonIncrementalAggregatePartialExec` -> (Exchange, inserted by
`EnsureRequirements`) ->
`PythonIncrementalAggregateFinalExec`. Both operators reuse
`ArrowPythonWithNamedArgumentRunner`
+ `GroupedPythonArrowInput`.
- Worker (`worker.py`): a PARTIAL handler that folds input batches into a
buffer via `reduce`, and
a FINAL handler that merges partial-buffer rows via `merge` then `finish`.
- The buffer schema is threaded to the JVM via a new nullable `bufferType`
on
`UserDefinedPythonFunction` (an auxiliary constructor preserves the
existing Py4J arity).
**Spark Connect**: also supported. A new optional `buffer_type` field on
the `PythonUDF` proto
message carries the buffer schema to the server; the Connect client
(`connect/udf.py`,
`connect/expressions.py`) serializes it and `udaf` dispatches on
`is_remote()`; the server
`SparkConnectPlanner` threads `buffer_type` into
`UserDefinedPythonFunction` and builds
`PythonAggregate`, after which execution reuses the same operators/worker
code as classic.
**SQL registration**: `spark.udf.register("my_agg", udaf(agg))` works in
both classic and Connect,
so the aggregator is usable from SQL text (`SELECT my_agg(v) FROM t GROUP
BY k`) — the counterpart
of Scala's `spark.udf.register(name, functions.udaf(agg))`.
Out of scope (planned follow-ups): `DISTINCT`, mixing with SQL aggregate
functions in
one `Aggregate`, window/streaming, real disk spill (currently the map side
bounds memory by
per-partition grouping; associativity makes early partial emission safe),
and a typed-columnar vs.
pickled buffer performance variant.
### Why are the changes needed?
PySpark has no incremental user-defined aggregator: every
custom-aggregation path (grouped-agg
`pandas_udf`/`arrow_udf`, `applyInPandas`) materializes the whole group and
invokes Python once,
with no map-side combine or partial/merge across the shuffle. This adds the
missing
`Aggregator`-style abstraction with genuine partial aggregation, matching
the Scala typed
`Aggregator`.
### Does this PR introduce _any_ user-facing change?
Yes — a new public API: `pyspark.sql.aggregator.Aggregator` and
`udaf(...)`, usable in
`groupBy().agg(...)` and registrable via `spark.udf.register(...)` for use
in SQL. No existing
behavior changes.
### How was this patch tested?
- Compilation verified: `sql/compile` (catalyst + core + sql) builds
cleanly with the new
expression, operators, and planner routing.
- Added `python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py`
(`ArrowPythonAggregatorTests`): checks the incremental aggregator matches
built-in `avg`/`sum`,
a no-group case, a custom buffer, and that results are independent of
partition count (exercising
partial + merge), plus `test_sql_registration` (register via
`spark.udf.register`, invoke from
SQL text). A Connect parity suite (`ArrowPythonAggregatorParityTests`)
runs the same mixin
under `ReusedConnectTestCase`.
- Compilation verified for both classic and Connect: `sql/compile` and
`connect/compile` build
cleanly (including proto regeneration). Full test execution runs in this
PR's CI.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #57952 from HyukjinKwon/SPARK-python-arrow-incremental-aggregator.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
(cherry picked from commit 9c5094adda0008bd4d683c387b613da9afab46ef)
Signed-off-by: Hyukjin Kwon <[email protected]>
---
.../src/main/resources/error/error-conditions.json | 12 +
.../org/apache/spark/api/python/PythonRunner.scala | 12 +
dev/sparktestsupport/modules.py | 2 +
python/pyspark/sql/aggregator.py | 223 ++++++++++
python/pyspark/sql/connect/client/core.py | 3 +
python/pyspark/sql/connect/expressions.py | 5 +
.../pyspark/sql/connect/proto/expressions_pb2.py | 48 +--
.../pyspark/sql/connect/proto/expressions_pb2.pyi | 24 +-
python/pyspark/sql/connect/udf.py | 30 +-
python/pyspark/sql/pandas/_typing/__init__.pyi | 2 +
.../tests/arrow/test_arrow_python_aggregator.py | 456 +++++++++++++++++++++
.../arrow/test_parity_arrow_python_aggregator.py | 29 ++
.../sql/tests/pandas/test_pandas_grouped_map.py | 3 +-
python/pyspark/sql/udf.py | 45 +-
python/pyspark/util.py | 9 +
python/pyspark/worker.py | 217 ++++++++++
.../spark/sql/catalyst/analysis/Analyzer.scala | 3 +-
.../sql/catalyst/analysis/FunctionResolution.scala | 8 +-
.../sql/catalyst/analysis/WindowResolution.scala | 10 +
.../spark/sql/catalyst/expressions/PythonUDF.scala | 49 +++
.../spark/sql/catalyst/plans/NormalizePlan.scala | 2 +
.../spark/sql/errors/QueryCompilationErrors.scala | 17 +
.../main/protobuf/spark/connect/expressions.proto | 3 +
.../sql/connect/planner/SparkConnectPlanner.scala | 13 +-
.../spark/sql/execution/SparkStrategies.scala | 46 ++-
.../python/PythonIncrementalAggregateExec.scala | 393 ++++++++++++++++++
.../python/UserDefinedPythonFunction.scala | 27 +-
27 files changed, 1644 insertions(+), 47 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index 1d29a257e312..1053b7c5a8d1 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -4897,6 +4897,18 @@
],
"sqlState" : "42602"
},
+ "INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA" : {
+ "message" : [
+ "The incremental Python aggregate function <functionName> requires a
struct (StructType) intermediate buffer schema, but got <bufferType>."
+ ],
+ "sqlState" : "42000"
+ },
+ "INVALID_PYTHON_UDF_PLACEMENT" : {
+ "message" : [
+ "The Python user-defined aggregate function(s) <functionList> cannot be
invoked together with other kinds of aggregate function in the same
aggregation. Multiple such functions may be used together, but not mixed with
SQL or differently-typed aggregate functions."
+ ],
+ "sqlState" : "0A000"
+ },
"INVALID_QUERY_MIXED_QUERY_PARAMETERS" : {
"message" : [
"Parameterized query must either use positional, or named parameters,
but not both."
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index d7800b1147e1..015067766c9c 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -87,6 +87,14 @@ private[spark] object PythonEvalType {
val SQL_WINDOW_AGG_ARROW_UDF = 253
val SQL_GROUPED_AGG_ARROW_ITER_UDF = 254
+ // Incremental (partial + final) Arrow aggregator. Unlike the whole-group
grouped-agg UDFs
+ // above, these support true partial aggregation: the PARTIAL eval type
folds input rows into a
+ // per-group buffer (via the aggregator's `reduce`) on the map side, and the
FINAL eval type
+ // merges partial buffers across the shuffle (via `merge`) and produces the
output (via `finish`).
+ // See PythonIncrementalAggregateExec and the Python `Aggregator` API.
+ val SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF = 255
+ val SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF = 256
+
val SQL_TABLE_UDF = 300
val SQL_ARROW_TABLE_UDF = 301
val SQL_ARROW_UDTF = 302
@@ -130,6 +138,10 @@ private[spark] object PythonEvalType {
case SQL_GROUPED_AGG_ARROW_UDF => "SQL_GROUPED_AGG_ARROW_UDF"
case SQL_WINDOW_AGG_ARROW_UDF => "SQL_WINDOW_AGG_ARROW_UDF"
case SQL_GROUPED_AGG_ARROW_ITER_UDF => "SQL_GROUPED_AGG_ARROW_ITER_UDF"
+ case SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF =>
+ "SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF"
+ case SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF =>
+ "SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF"
}
// The eval types produced by ExtractPythonUDFFromLambda: a scalar UDF
lifted out of a
diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py
index e70a5983d931..d5ef8e7b54ab 100644
--- a/dev/sparktestsupport/modules.py
+++ b/dev/sparktestsupport/modules.py
@@ -611,6 +611,7 @@ pyspark_sql = Module(
"pyspark.sql.tests.arrow.test_arrow_cogrouped_map",
"pyspark.sql.tests.arrow.test_arrow_cogrouped_map_misc",
"pyspark.sql.tests.arrow.test_arrow_grouped_map",
+ "pyspark.sql.tests.arrow.test_arrow_python_aggregator",
"pyspark.sql.tests.arrow.test_arrow_python_udf",
"pyspark.sql.tests.arrow.test_arrow_python_udf_cached",
"pyspark.sql.tests.arrow.test_arrow_udf",
@@ -1256,6 +1257,7 @@ pyspark_connect = Module(
"pyspark.sql.tests.connect.arrow.test_parity_arrow_grouped_map",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_cogrouped_map",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_cogrouped_map_misc",
+ "pyspark.sql.tests.connect.arrow.test_parity_arrow_python_aggregator",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_python_udf",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_udf",
"pyspark.sql.tests.connect.arrow.test_parity_arrow_udf_scalar",
diff --git a/python/pyspark/sql/aggregator.py b/python/pyspark/sql/aggregator.py
new file mode 100644
index 000000000000..77110b8f88a3
--- /dev/null
+++ b/python/pyspark/sql/aggregator.py
@@ -0,0 +1,223 @@
+#
+# 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.
+#
+"""
+Incremental user-defined aggregators for PySpark, the Python analog of Scala's
+``org.apache.spark.sql.expressions.Aggregator``.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Tuple
+
+from pyspark.errors import PySparkNotImplementedError, PySparkTypeError,
PySparkValueError
+from pyspark.sql.types import DataType, StructType
+from pyspark.util import PythonEvalType
+
+__all__ = ["Aggregator", "udaf"]
+
+
+class Aggregator(ABC):
+ """
+ Base class for a user-defined *incremental* aggregator, the Python analog
of Scala's
+ :class:`org.apache.spark.sql.expressions.Aggregator`.
+
+ Unlike a grouped-aggregate ``pandas_udf`` (which materializes the whole
group and is invoked
+ once), an :class:`Aggregator` is executed as a genuine two-stage
aggregation with map-side
+ combine: :meth:`reduce` folds input rows into a per-group *buffer* on the
map side, the buffers
+ are shuffled by the grouping key, :meth:`merge` combines the partial
buffers of each group, and
+ :meth:`finish` produces the final output value.
+
+ The buffer is represented as a Python :class:`tuple` whose elements
correspond, in order, to the
+ fields of :attr:`bufferSchema`. An input row is likewise a tuple of the
argument values passed
+ to the aggregator call. :meth:`merge` must be associative and commutative
(the framework may
+ combine partial buffers in any order), and :meth:`zero` must be its
identity element -- see
+ :meth:`zero` for the identity law that makes the result independent of the
partition count.
+
+ .. versionadded:: 4.4.0
+
+ Examples
+ --------
+ A mean aggregator::
+
+ from pyspark.sql.aggregator import Aggregator, udaf
+ from pyspark.sql.types import StructType, StructField, DoubleType,
LongType
+
+ class Mean(Aggregator):
+ @property
+ def bufferSchema(self):
+ return StructType([
+ StructField("sum", DoubleType()),
+ StructField("count", LongType()),
+ ])
+
+ @property
+ def outputType(self):
+ return DoubleType()
+
+ def zero(self):
+ return (0.0, 0)
+
+ def reduce(self, buffer, value):
+ (v,) = value
+ if v is None: # ignore null inputs, like SQL aggregates do
+ return buffer
+ return (buffer[0] + v, buffer[1] + 1)
+
+ def merge(self, b1, b2):
+ return (b1[0] + b2[0], b1[1] + b2[1])
+
+ def finish(self, buffer):
+ return buffer[0] / buffer[1] if buffer[1] else None
+
+ mean = udaf(Mean())
+ df.groupBy("k").agg(mean(df.v)).show()
+ """
+
+ @property
+ @abstractmethod
+ def bufferSchema(self) -> StructType:
+ """The schema of the intermediate buffer that crosses the shuffle."""
+ ...
+
+ @property
+ @abstractmethod
+ def outputType(self) -> DataType:
+ """The data type of the aggregator's output value."""
+ ...
+
+ @abstractmethod
+ def zero(self) -> Tuple[Any, ...]:
+ """The initial (identity) buffer value, as a tuple matching
:attr:`bufferSchema`.
+
+ This must be the identity element for :meth:`merge`::
+
+ merge(buffer, zero()) == buffer
+ merge(zero(), buffer) == buffer
+
+ A fresh ``zero()`` seeds every partition -- and every early-flushed
chunk of the map-side
+ combine -- so associativity and commutativity of :meth:`merge` alone
do not guarantee a
+ partition-independent result; the identity law above is what makes the
aggregate value
+ independent of how the input is split across partitions and batches.
+ """
+ ...
+
+ @abstractmethod
+ def reduce(self, buffer: Tuple[Any, ...], value: Tuple[Any, ...]) ->
Tuple[Any, ...]:
+ """Fold a single input row ``value`` into ``buffer`` and return the
updated buffer."""
+ ...
+
+ @abstractmethod
+ def merge(self, buffer1: Tuple[Any, ...], buffer2: Tuple[Any, ...]) ->
Tuple[Any, ...]:
+ """Merge two partial buffers into one. Must be associative and
commutative."""
+ ...
+
+ @abstractmethod
+ def finish(self, buffer: Tuple[Any, ...]) -> Any:
+ """Produce the output value from the final merged buffer."""
+ ...
+
+ # The aggregator instance is shipped to the worker as the UDF "function";
making it callable
+ # lets it satisfy ``UserDefinedFunction``'s ``callable`` check. It is
never actually invoked as
+ # a function -- the worker calls
:meth:`zero`/:meth:`reduce`/:meth:`merge`/:meth:`finish`.
+ def __call__(self, *args: Any, **kwargs: Any) -> Any:
+ raise PySparkNotImplementedError(
+ errorClass="NOT_IMPLEMENTED",
+ messageParameters={"feature": "calling an Aggregator directly;
wrap it with udaf(...)"},
+ )
+
+
+def udaf(agg: "Aggregator") -> Any:
+ """
+ Turn an :class:`Aggregator` instance into a callable usable in
``groupBy().agg(...)``, the
+ Python counterpart of Scala's ``functions.udaf``.
+
+ The aggregator is executed with true incremental (partial) aggregation and
transfers its
+ intermediate buffer as Arrow; PyArrow is therefore required.
+
+ .. versionadded:: 4.4.0
+
+ Parameters
+ ----------
+ agg : :class:`Aggregator`
+ The aggregator instance.
+
+ Returns
+ -------
+ function
+ A callable that, applied to input columns, produces an aggregate
:class:`Column`.
+
+ Raises
+ ------
+ :class:`PySparkImportError`
+ If a supported version of PyArrow is not installed.
+ :class:`PySparkTypeError`
+ If ``agg`` is not an :class:`Aggregator`, or its ``bufferSchema`` is
not a
+ :class:`StructType`.
+ """
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+ from pyspark.sql.utils import is_remote
+
+ require_minimum_pyarrow_version()
+
+ if is_remote():
+ from pyspark.sql.connect.udf import UserDefinedFunction
+ else:
+ # The classic UserDefinedFunction is a distinct class from the Connect
one above;
+ # both provide the same interface used below, so silence mypy's
reassignment check.
+ from pyspark.sql.udf import UserDefinedFunction # type:
ignore[assignment]
+
+ if not isinstance(agg, Aggregator):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": "agg",
+ "expected_type": "Aggregator",
+ "arg_type": type(agg).__name__,
+ },
+ )
+ if not isinstance(agg.bufferSchema, StructType):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": "bufferSchema",
+ "expected_type": "StructType",
+ "arg_type": type(agg.bufferSchema).__name__,
+ },
+ )
+ # The buffer crosses the shuffle as an Arrow struct whose children are
matched by name, and the
+ # worker keys the buffer tuple back by field name. Duplicate names would
silently collapse
+ # fields on the map side and then fail with an opaque Arrow error
post-shuffle, so reject them
+ # up front where the aggregator is created.
+ field_names = [field.name for field in agg.bufferSchema.fields]
+ if len(field_names) != len(set(field_names)):
+ duplicates = sorted({name for name in field_names if
field_names.count(name) > 1})
+ raise PySparkValueError(
+ errorClass="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT",
+ messageParameters={"field_names": ", ".join(duplicates)},
+ )
+
+ # ``bufferSchema`` is a first-class ``UserDefinedFunction`` field
(threaded to the JVM in
+ # ``_create_judf`` so ``PythonAggregate`` can plan the two-stage
aggregation), so it survives
+ # ``_wrapped()`` and ``spark.udf.register`` without being re-attached.
+ udf_obj = UserDefinedFunction(
+ agg,
+ returnType=agg.outputType,
+ name=agg.__class__.__name__,
+ evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ deterministic=True,
+ bufferSchema=agg.bufferSchema,
+ )
+ return udf_obj._wrapped()
diff --git a/python/pyspark/sql/connect/client/core.py
b/python/pyspark/sql/connect/client/core.py
index 0ff6f0ae90f9..462ad8fe2670 100644
--- a/python/pyspark/sql/connect/client/core.py
+++ b/python/pyspark/sql/connect/client/core.py
@@ -1081,6 +1081,7 @@ class SparkConnectClient(object):
name: Optional[str] = None,
eval_type: int = PythonEvalType.SQL_BATCHED_UDF,
deterministic: bool = True,
+ buffer_type: Optional["DataType"] = None,
) -> str:
"""
Create a temporary UDF in the session catalog on the other side. We
generate a
@@ -1096,6 +1097,8 @@ class SparkConnectClient(object):
eval_type=eval_type,
func=function,
python_ver="%d.%d" % sys.version_info[:2],
+ # Set for the incremental aggregator (see pyspark.sql.aggregator).
+ buffer_type=buffer_type,
)
# construct a CommonInlineUserDefinedFunction
diff --git a/python/pyspark/sql/connect/expressions.py
b/python/pyspark/sql/connect/expressions.py
index 57270398118f..bbe0dfa39478 100644
--- a/python/pyspark/sql/connect/expressions.py
+++ b/python/pyspark/sql/connect/expressions.py
@@ -741,6 +741,7 @@ class PythonUDF:
eval_type: int,
func: Callable[..., Any],
python_ver: str,
+ buffer_type: Optional[DataType] = None,
) -> None:
self._output_type: DataType = (
UnparsedDataType(output_type) if isinstance(output_type, str) else
output_type
@@ -748,6 +749,8 @@ class PythonUDF:
self._eval_type = eval_type
self._func = func
self._python_ver = python_ver
+ # Intermediate buffer schema for an incremental Python aggregator;
None otherwise.
+ self._buffer_type = buffer_type
def to_plan(self, session: "SparkConnectClient") -> proto.PythonUDF:
if isinstance(self._output_type, UnparsedDataType):
@@ -763,6 +766,8 @@ class PythonUDF:
expr.eval_type = self._eval_type
expr.command = CloudPickleSerializer().dumps((self._func, output_type))
expr.python_ver = self._python_ver
+ if self._buffer_type is not None:
+
expr.buffer_type.CopyFrom(pyspark_types_to_proto_types(self._buffer_type))
return expr
def __repr__(self) -> str:
diff --git a/python/pyspark/sql/connect/proto/expressions_pb2.py
b/python/pyspark/sql/connect/proto/expressions_pb2.py
index aa51c393c043..5bb6335cfe3b 100644
--- a/python/pyspark/sql/connect/proto/expressions_pb2.py
+++ b/python/pyspark/sql/connect/proto/expressions_pb2.py
@@ -41,7 +41,7 @@ from pyspark.sql.connect.proto import common_pb2 as
spark_dot_connect_dot_common
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
-
b'\n\x1fspark/connect/expressions.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x19spark/connect/types.proto\x1a\x1aspark/connect/common.proto"\x90<\n\nExpression\x12\x37\n\x06\x63ommon\x18\x12
\x01(\x0b\x32\x1f.spark.connect.ExpressionCommonR\x06\x63ommon\x12=\n\x07literal\x18\x01
\x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x07literal\x12\x62\n\x14unresolved_attribute\x18\x02
\x01(\x0b\x32-.spark.connect.Expression.UnresolvedAttributeH\x00R\x13unresolvedAtt
[...]
+
b'\n\x1fspark/connect/expressions.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x19spark/connect/types.proto\x1a\x1aspark/connect/common.proto"\x90<\n\nExpression\x12\x37\n\x06\x63ommon\x18\x12
\x01(\x0b\x32\x1f.spark.connect.ExpressionCommonR\x06\x63ommon\x12=\n\x07literal\x18\x01
\x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x07literal\x12\x62\n\x14unresolved_attribute\x18\x02
\x01(\x0b\x32-.spark.connect.Expression.UnresolvedAttributeH\x00R\x13unresolvedAtt
[...]
)
_globals = globals()
@@ -135,27 +135,27 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals["_COMMONINLINEUSERDEFINEDFUNCTION"]._serialized_start = 7899
_globals["_COMMONINLINEUSERDEFINEDFUNCTION"]._serialized_end = 8296
_globals["_PYTHONUDF"]._serialized_start = 8299
- _globals["_PYTHONUDF"]._serialized_end = 8503
- _globals["_SCALARSCALAUDF"]._serialized_start = 8506
- _globals["_SCALARSCALAUDF"]._serialized_end = 8720
- _globals["_JAVAUDF"]._serialized_start = 8723
- _globals["_JAVAUDF"]._serialized_end = 8872
- _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_start = 8874
- _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_end = 8973
- _globals["_CALLFUNCTION"]._serialized_start = 8975
- _globals["_CALLFUNCTION"]._serialized_end = 9083
- _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_start = 9085
- _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_end = 9177
- _globals["_MERGEACTION"]._serialized_start = 9180
- _globals["_MERGEACTION"]._serialized_end = 9692
- _globals["_MERGEACTION_ASSIGNMENT"]._serialized_start = 9402
- _globals["_MERGEACTION_ASSIGNMENT"]._serialized_end = 9508
- _globals["_MERGEACTION_ACTIONTYPE"]._serialized_start = 9511
- _globals["_MERGEACTION_ACTIONTYPE"]._serialized_end = 9678
- _globals["_SUBQUERYEXPRESSION"]._serialized_start = 9695
- _globals["_SUBQUERYEXPRESSION"]._serialized_end = 10404
- _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_start = 10001
- _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_end = 10235
- _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_start = 10238
- _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_end = 10382
+ _globals["_PYTHONUDF"]._serialized_end = 8582
+ _globals["_SCALARSCALAUDF"]._serialized_start = 8585
+ _globals["_SCALARSCALAUDF"]._serialized_end = 8799
+ _globals["_JAVAUDF"]._serialized_start = 8802
+ _globals["_JAVAUDF"]._serialized_end = 8951
+ _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_start = 8953
+ _globals["_TYPEDAGGREGATEEXPRESSION"]._serialized_end = 9052
+ _globals["_CALLFUNCTION"]._serialized_start = 9054
+ _globals["_CALLFUNCTION"]._serialized_end = 9162
+ _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_start = 9164
+ _globals["_NAMEDARGUMENTEXPRESSION"]._serialized_end = 9256
+ _globals["_MERGEACTION"]._serialized_start = 9259
+ _globals["_MERGEACTION"]._serialized_end = 9771
+ _globals["_MERGEACTION_ASSIGNMENT"]._serialized_start = 9481
+ _globals["_MERGEACTION_ASSIGNMENT"]._serialized_end = 9587
+ _globals["_MERGEACTION_ACTIONTYPE"]._serialized_start = 9590
+ _globals["_MERGEACTION_ACTIONTYPE"]._serialized_end = 9757
+ _globals["_SUBQUERYEXPRESSION"]._serialized_start = 9774
+ _globals["_SUBQUERYEXPRESSION"]._serialized_end = 10483
+ _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_start = 10080
+ _globals["_SUBQUERYEXPRESSION_TABLEARGOPTIONS"]._serialized_end = 10314
+ _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_start = 10317
+ _globals["_SUBQUERYEXPRESSION_SUBQUERYTYPE"]._serialized_end = 10461
# @@protoc_insertion_point(module_scope)
diff --git a/python/pyspark/sql/connect/proto/expressions_pb2.pyi
b/python/pyspark/sql/connect/proto/expressions_pb2.pyi
index c613ade2f43f..22b65357c234 100644
--- a/python/pyspark/sql/connect/proto/expressions_pb2.pyi
+++ b/python/pyspark/sql/connect/proto/expressions_pb2.pyi
@@ -1826,6 +1826,7 @@ class PythonUDF(google.protobuf.message.Message):
COMMAND_FIELD_NUMBER: builtins.int
PYTHON_VER_FIELD_NUMBER: builtins.int
ADDITIONAL_INCLUDES_FIELD_NUMBER: builtins.int
+ BUFFER_TYPE_FIELD_NUMBER: builtins.int
@property
def output_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
"""(Required) Output type of the Python UDF"""
@@ -1840,6 +1841,11 @@ class PythonUDF(google.protobuf.message.Message):
self,
) ->
google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]:
"""(Optional) Additional includes for the Python UDF."""
+ @property
+ def buffer_type(self) -> pyspark.sql.connect.proto.types_pb2.DataType:
+ """(Optional) Intermediate buffer schema for an incremental Python
aggregator
+ (see PythonAggregate). Set only for the incremental aggregator eval
types.
+ """
def __init__(
self,
*,
@@ -1848,15 +1854,28 @@ class PythonUDF(google.protobuf.message.Message):
command: builtins.bytes = ...,
python_ver: builtins.str = ...,
additional_includes: collections.abc.Iterable[builtins.str] | None =
...,
+ buffer_type: pyspark.sql.connect.proto.types_pb2.DataType | None = ...,
) -> None: ...
def HasField(
- self, field_name: typing_extensions.Literal["output_type",
b"output_type"]
+ self,
+ field_name: typing_extensions.Literal[
+ "_buffer_type",
+ b"_buffer_type",
+ "buffer_type",
+ b"buffer_type",
+ "output_type",
+ b"output_type",
+ ],
) -> builtins.bool: ...
def ClearField(
self,
field_name: typing_extensions.Literal[
+ "_buffer_type",
+ b"_buffer_type",
"additional_includes",
b"additional_includes",
+ "buffer_type",
+ b"buffer_type",
"command",
b"command",
"eval_type",
@@ -1867,6 +1886,9 @@ class PythonUDF(google.protobuf.message.Message):
b"python_ver",
],
) -> None: ...
+ def WhichOneof(
+ self, oneof_group: typing_extensions.Literal["_buffer_type",
b"_buffer_type"]
+ ) -> typing_extensions.Literal["buffer_type"] | None: ...
global___PythonUDF = PythonUDF
diff --git a/python/pyspark/sql/connect/udf.py
b/python/pyspark/sql/connect/udf.py
index c9848634eb30..7acf255e2f06 100644
--- a/python/pyspark/sql/connect/udf.py
+++ b/python/pyspark/sql/connect/udf.py
@@ -32,7 +32,7 @@ from pyspark.sql.connect.expressions import (
PythonUDF,
)
from pyspark.sql.connect.column import Column
-from pyspark.sql.types import DataType, StringType, _parse_datatype_string
+from pyspark.sql.types import DataType, StringType, StructType,
_parse_datatype_string
from pyspark.sql.udf import (
UDFRegistration as PySparkUDFRegistration,
UserDefinedFunction as PySparkUserDefinedFunction,
@@ -113,10 +113,16 @@ def _create_udf(
evalType: int,
name: Optional[str] = None,
deterministic: bool = True,
+ bufferSchema: Optional[StructType] = None,
) -> "UserDefinedFunctionLike":
# Set the name of the UserDefinedFunction object to be the name of
function f
udf_obj = UserDefinedFunction(
- f, returnType=returnType, name=name, evalType=evalType,
deterministic=deterministic
+ f,
+ returnType=returnType,
+ name=name,
+ evalType=evalType,
+ deterministic=deterministic,
+ bufferSchema=bufferSchema,
)
return udf_obj._wrapped()
@@ -139,6 +145,7 @@ class UserDefinedFunction:
name: Optional[str] = None,
evalType: int = PythonEvalType.SQL_BATCHED_UDF,
deterministic: bool = True,
+ bufferSchema: Optional[StructType] = None,
):
if not callable(func):
raise PySparkTypeError(
@@ -178,6 +185,10 @@ class UserDefinedFunction:
)
self.evalType = evalType
self.deterministic = deterministic
+ # Intermediate aggregation buffer schema, set only for an incremental
Python aggregator
+ # (see :class:`pyspark.sql.aggregator.Aggregator`); ``None``
otherwise. A first-class field
+ # so it survives ``_wrapped()``, ``asNondeterministic()`` and
``spark.udf.register``.
+ self.bufferSchema = bufferSchema
@property
def returnType(self) -> DataType:
@@ -210,6 +221,8 @@ class UserDefinedFunction:
eval_type=self.evalType,
func=self.func,
python_ver="%d.%d" % sys.version_info[:2],
+ # Set for incremental Python aggregators (see
pyspark.sql.aggregator).
+ buffer_type=self.bufferSchema,
)
return CommonInlineUserDefinedFunction(
function_name=self._name,
@@ -253,6 +266,7 @@ class UserDefinedFunction:
wrapper.returnType = self.returnType # type: ignore[attr-defined]
wrapper.evalType = self.evalType # type: ignore[attr-defined]
wrapper.deterministic = self.deterministic # type:
ignore[attr-defined]
+ wrapper.bufferSchema = self.bufferSchema # type: ignore[attr-defined]
wrapper.asNondeterministic = functools.wraps( # type:
ignore[attr-defined]
self.asNondeterministic
)(lambda: self.asNondeterministic()._wrapped())
@@ -303,6 +317,7 @@ class UDFRegistration:
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
]:
raise PySparkTypeError(
errorClass="INVALID_UDF_EVAL_TYPE",
@@ -311,11 +326,18 @@ class UDFRegistration:
"SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_ARROW_UDF, "
"SQL_SCALAR_PANDAS_ITER_UDF,
SQL_SCALAR_ARROW_ITER_UDF, "
"SQL_GROUPED_AGG_PANDAS_UDF,
SQL_GROUPED_AGG_ARROW_UDF, "
- "SQL_GROUPED_AGG_PANDAS_ITER_UDF or
SQL_GROUPED_AGG_ARROW_ITER_UDF"
+ "SQL_GROUPED_AGG_PANDAS_ITER_UDF,
SQL_GROUPED_AGG_ARROW_ITER_UDF "
+ "or SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF"
},
)
self.sparkSession._client.register_udf(
- f.func, f.returnType, name, f.evalType, f.deterministic
+ f.func,
+ f.returnType,
+ name,
+ f.evalType,
+ f.deterministic,
+ # Set for the incremental aggregator (see
pyspark.sql.aggregator).
+ buffer_type=getattr(f, "bufferSchema", None),
)
return f
else:
diff --git a/python/pyspark/sql/pandas/_typing/__init__.pyi
b/python/pyspark/sql/pandas/_typing/__init__.pyi
index f989443a4dd9..442849bcbf0a 100644
--- a/python/pyspark/sql/pandas/_typing/__init__.pyi
+++ b/python/pyspark/sql/pandas/_typing/__init__.pyi
@@ -68,6 +68,8 @@ ArrowScalarIterUDFType = Literal[251]
ArrowGroupedAggUDFType = Literal[252]
ArrowWindowAggUDFType = Literal[253]
ArrowGroupedAggIterUDFType = Literal[254]
+ArrowGroupedAggIncrementalPartialUDFType = Literal[255]
+ArrowGroupedAggIncrementalFinalUDFType = Literal[256]
# Arrow stream types
# A single group of Arrow batches (e.g., one key group in groupBy).
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
b/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
new file mode 100644
index 000000000000..3965d804ae84
--- /dev/null
+++ b/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
@@ -0,0 +1,456 @@
+#
+# 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.
+#
+import unittest
+from decimal import Decimal
+
+from pyspark.errors import AnalysisException, PySparkValueError
+from pyspark.sql import functions as sf
+from pyspark.sql.window import Window
+from pyspark.sql.types import (
+ DecimalType,
+ DoubleType,
+ LongType,
+ StructField,
+ StructType,
+)
+from pyspark.util import PythonEvalType
+from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.utils import (
+ have_pyarrow,
+ pyarrow_requirement_message,
+)
+
+
+if have_pyarrow:
+ from pyspark.sql.aggregator import Aggregator, udaf
+
+ class Mean(Aggregator):
+ @property
+ def bufferSchema(self):
+ return StructType([StructField("sum", DoubleType()),
StructField("count", LongType())])
+
+ @property
+ def outputType(self):
+ return DoubleType()
+
+ def zero(self):
+ return (0.0, 0)
+
+ def reduce(self, buffer, value):
+ (v,) = value
+ if v is None: # ignore nulls, like SQL avg
+ return buffer
+ return (buffer[0] + v, buffer[1] + 1)
+
+ def merge(self, b1, b2):
+ return (b1[0] + b2[0], b1[1] + b2[1])
+
+ def finish(self, buffer):
+ return buffer[0] / buffer[1] if buffer[1] else None
+
+ class DecimalSum(Aggregator):
+ # Non-trivial output/buffer type: the result column and the
intermediate buffer are both
+ # DecimalType, exercising explicit Arrow typing of the emitted arrays
(a bare
+ # ``pa.array([Decimal(...)])`` would infer a decimal type whose
precision/scale need not
+ # match the declared one).
+ @property
+ def bufferSchema(self):
+ return StructType([StructField("total", DecimalType(20, 4))])
+
+ @property
+ def outputType(self):
+ return DecimalType(20, 4)
+
+ def zero(self):
+ return (Decimal(0),)
+
+ def reduce(self, buffer, value):
+ (v,) = value
+ return buffer if v is None else (buffer[0] + Decimal(str(v)),)
+
+ def merge(self, b1, b2):
+ return (b1[0] + b2[0],)
+
+ def finish(self, buffer):
+ return buffer[0]
+
+ class SumSquares(Aggregator):
+ @property
+ def bufferSchema(self):
+ return StructType([StructField("sumsq", DoubleType())])
+
+ @property
+ def outputType(self):
+ return DoubleType()
+
+ def zero(self):
+ return (0.0,)
+
+ def reduce(self, buffer, value):
+ (v,) = value
+ return (buffer[0] + float(v) * float(v),)
+
+ def merge(self, b1, b2):
+ return (b1[0] + b2[0],)
+
+ def finish(self, buffer):
+ return buffer[0]
+
+
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowPythonAggregatorTestsMixin:
+ def _data(self):
+ # 100 rows across 5 keys; repartition so each key is split across
partitions,
+ # exercising map-side PARTIAL combine + post-shuffle FINAL merge.
+ return (
+ self.spark.range(0, 100)
+ .select((sf.col("id") % 5).alias("k"),
sf.col("id").cast("double").alias("v"))
+ .repartition(4, sf.col("v") % 3)
+ )
+
+ def test_incremental_aggregator_matches_builtin_mean(self):
+ df = self._data()
+ result =
df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect()
+ expected =
df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect()
+ got = {r["k"]: r["m"] for r in result}
+ exp = {r["k"]: r["m"] for r in expected}
+ self.assertEqual(got, exp)
+
+ def test_incremental_aggregator_no_group(self):
+ df = self._data()
+ result = df.agg(udaf(Mean())(sf.col("v")).alias("m")).collect()
+ expected = df.agg(sf.avg("v").alias("m")).collect()
+ self.assertAlmostEqual(result[0]["m"], expected[0]["m"], places=6)
+
+ def test_incremental_aggregator_empty_global_input(self):
+ # A global aggregation over empty input must still return one identity
row: finish(zero).
+ empty = self._data().limit(0)
+ result = empty.agg(udaf(Mean())(sf.col("v")).alias("m")).collect()
+ expected = empty.agg(sf.avg("v").alias("m")).collect()
+ self.assertEqual(len(result), 1)
+ self.assertEqual(result[0]["m"], expected[0]["m"])
+
+ def test_incremental_aggregator_custom_buffer(self):
+ df = self._data()
+ result = (
+
df.groupBy("k").agg(udaf(SumSquares())(sf.col("v")).alias("s")).orderBy("k").collect()
+ )
+ expected = (
+ df.groupBy("k").agg(sf.sum(sf.col("v") *
sf.col("v")).alias("s")).orderBy("k").collect()
+ )
+ got = {r["k"]: r["s"] for r in result}
+ exp = {r["k"]: r["s"] for r in expected}
+ for k in exp:
+ self.assertAlmostEqual(got[k], exp[k], places=6)
+
+ def test_incremental_aggregator_decimal_output(self):
+ # Non-trivial output/buffer type (DecimalType), crossing the shuffle
as a decimal buffer
+ # and emitted as a decimal result -- guards the explicit Arrow typing
of both stages.
+ df = self._data()
+ result = (
+
df.groupBy("k").agg(udaf(DecimalSum())(sf.col("v")).alias("s")).orderBy("k").collect()
+ )
+ expected =
df.groupBy("k").agg(sf.sum("v").alias("s")).orderBy("k").collect()
+ got = {r["k"]: r["s"] for r in result}
+ exp = {r["k"]: r["s"] for r in expected}
+ for k in exp:
+ self.assertIsInstance(got[k], Decimal)
+ self.assertAlmostEqual(float(got[k]), exp[k], places=4)
+
+ def test_incremental_aggregator_null_inputs(self):
+ # reduce must tolerate null input values; the null-skipping Mean
should match SQL avg,
+ # including a group whose values are all null (identity buffer ->
finish returns None).
+ df = self.spark.createDataFrame(
+ [("a", 1.0), ("a", None), ("a", 3.0), ("b", None), ("b", None)],
+ "k string, v double",
+ )
+ result =
df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect()
+ expected =
df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect()
+ got = {r["k"]: r["m"] for r in result}
+ exp = {r["k"]: r["m"] for r in expected}
+ self.assertEqual(got, exp)
+
+ def test_multiple_incremental_aggregators(self):
+ # Two aggregators with different buffer schemas over the same input in
one agg call.
+ df = self._data()
+ result = (
+ df.groupBy("k")
+ .agg(
+ udaf(Mean())(sf.col("v")).alias("m"),
+ udaf(SumSquares())(sf.col("v")).alias("s"),
+ )
+ .orderBy("k")
+ .collect()
+ )
+ expected = (
+ df.groupBy("k")
+ .agg(
+ sf.avg("v").alias("m"),
+ sf.sum(sf.col("v") * sf.col("v")).alias("s"),
+ )
+ .orderBy("k")
+ .collect()
+ )
+ got_m = {r["k"]: r["m"] for r in result}
+ got_s = {r["k"]: r["s"] for r in result}
+ for r in expected:
+ self.assertAlmostEqual(got_m[r["k"]], r["m"], places=6)
+ self.assertAlmostEqual(got_s[r["k"]], r["s"], places=6)
+
+ def test_result_independent_of_partition_count(self):
+ # Partial buffers must merge to the same result regardless of how keys
are split.
+ base = self.spark.range(0, 60).select(
+ (sf.col("id") % 3).alias("k"),
sf.col("id").cast("double").alias("v")
+ )
+ results = []
+ for n in (1, 2, 7):
+ rows = (
+ base.repartition(n, sf.col("v"))
+ .groupBy("k")
+ .agg(udaf(Mean())(sf.col("v")).alias("m"))
+ .orderBy("k")
+ .collect()
+ )
+ results.append({r["k"]: r["m"] for r in rows})
+ self.assertEqual(results[0], results[1])
+ self.assertEqual(results[1], results[2])
+
+ def test_sql_registration(self):
+ # Register the aggregator and invoke it from SQL text.
+ df = self._data()
+ df.createOrReplaceTempView("agg_input")
+ self.spark.udf.register("my_mean", udaf(Mean()))
+ result = self.spark.sql(
+ "SELECT k, my_mean(v) AS m FROM agg_input GROUP BY k ORDER BY k"
+ ).collect()
+ expected =
df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect()
+ got = {r["k"]: r["m"] for r in result}
+ exp = {r["k"]: r["m"] for r in expected}
+ self.assertEqual(got, exp)
+
+ def test_named_arguments(self):
+ # A named argument (both DataFrame and SQL forms) must feed the
aggregator's value tuple,
+ # not be silently dropped.
+ df = self._data()
+ result =
df.groupBy("k").agg(udaf(Mean())(v=sf.col("v")).alias("m")).orderBy("k").collect()
+ expected =
df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect()
+ self.assertEqual({r["k"]: r["m"] for r in result}, {r["k"]: r["m"] for
r in expected})
+
+ df.createOrReplaceTempView("agg_input")
+ self.spark.udf.register("my_mean", udaf(Mean()))
+ sql_result = self.spark.sql(
+ "SELECT k, my_mean(v => v) AS m FROM agg_input GROUP BY k ORDER BY
k"
+ ).collect()
+ self.assertEqual({r["k"]: r["m"] for r in sql_result}, {r["k"]: r["m"]
for r in expected})
+
+ def test_distinct_and_filter_rejected(self):
+ # Neither DISTINCT nor FILTER is honored by the two-stage operator, so
both must be
+ # rejected at analysis rather than silently returning the
non-distinct/unfiltered result.
+ df = self._data()
+ df.createOrReplaceTempView("agg_input")
+ self.spark.udf.register("my_mean", udaf(Mean()))
+ with self.assertRaises(AnalysisException):
+ self.spark.sql("SELECT my_mean(DISTINCT v) FROM agg_input GROUP BY
k").collect()
+ with self.assertRaises(AnalysisException):
+ self.spark.sql(
+ "SELECT my_mean(v) FILTER (WHERE v > 0) FROM agg_input GROUP
BY k"
+ ).collect()
+
+ def test_pivot_rejected(self):
+ # ResolvePivot must reject the incremental aggregator (its
null-ignoring fallback rewrite
+ # would produce wrong results), like it already rejects pandas UDAFs.
+ df = self.spark.createDataFrame(
+ [("a", "x", 1.0), ("a", "y", 2.0), ("b", "x", 3.0)],
+ "k string, p string, v double",
+ )
+ with self.assertRaises(AnalysisException):
+ df.groupBy("k").pivot("p").agg(udaf(Mean())(sf.col("v"))).collect()
+
+ def test_window_rejected(self):
+ # The incremental aggregator has no window operator; using it over a
window must fail
+ # analysis rather than crash with an internal error at execution.
+ df = self._data()
+ with self.assertRaises(AnalysisException):
+
df.select(udaf(Mean())(sf.col("v")).over(Window.partitionBy("k"))).collect()
+
+ def test_mixed_with_other_aggregate_rejected(self):
+ # An incremental aggregator mixed with another aggregate in one
Aggregate is unsupported;
+ # the error must be the dedicated (non-pandas) placement error.
+ df = self._data()
+ with self.assertRaises(AnalysisException) as ctx:
+ df.groupBy("k").agg(
+ udaf(Mean())(sf.col("v")).alias("m"), sf.count("*").alias("c")
+ ).collect()
+ self.assertEqual(ctx.exception.getCondition(),
"INVALID_PYTHON_UDF_PLACEMENT")
+
+ def test_duplicate_buffer_field_names_rejected(self):
+ # Duplicate buffer field names silently collapse on the map side and
then fail with an
+ # opaque Arrow error post-shuffle; reject them where the aggregator is
created.
+ class DupBuffer(Mean):
+ @property
+ def bufferSchema(self):
+ return StructType([StructField("x", DoubleType()),
StructField("x", LongType())])
+
+ with self.assertRaises(PySparkValueError) as ctx:
+ udaf(DupBuffer())
+ self.check_error(
+ exception=ctx.exception,
+ errorClass="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT",
+ messageParameters={"field_names": "x"},
+ )
+
+ def test_float_grouping_keys_normalized(self):
+ # 0.0 / -0.0 must fall into one group, and NaN keys (unequal to
themselves) must group
+ # together, matching SQL aggregate semantics -- guards grouping-key
normalization and the
+ # NaN handling of the map-side hash combine.
+ zeros = self.spark.createDataFrame(
+ [(0.0, 1.0), (-0.0, 2.0), (0.0, 3.0)], "k double, v double"
+ )
+ zero_rows =
zeros.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).collect()
+ self.assertEqual(len(zero_rows), 1)
+ self.assertAlmostEqual(zero_rows[0]["m"], 2.0, places=6)
+
+ nans = self.spark.createDataFrame(
+ [(float("nan"), 1.0), (float("nan"), 3.0)], "k double, v double"
+ )
+ nan_rows =
nans.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).collect()
+ self.assertEqual(len(nan_rows), 1)
+ self.assertAlmostEqual(nan_rows[0]["m"], 2.0, places=6)
+
+ def test_complex_grouping_key(self):
+ # A struct grouping key exercises the map-side hash combine's
canonicalization of complex
+ # keys (dict -> hashable) as well as the authoritative FINAL
re-grouping.
+ df = self.spark.createDataFrame(
+ [(1, "a", 1.0), (1, "a", 3.0), (2, "b", 10.0)],
+ "i int, s string, v double",
+ ).select(sf.struct("i", "s").alias("k"), sf.col("v"))
+ result =
df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).collect()
+ got = {(r["k"]["i"], r["k"]["s"]): r["m"] for r in result}
+ self.assertEqual(got, {(1, "a"): 2.0, (2, "b"): 10.0})
+
+ def test_bounded_map_side_combine(self):
+ # A small maxRecordsPerBatch forces the map-side PARTIAL stage to
flush its per-key buffer
+ # in bounded chunks (emitting duplicate keys that the FINAL stage
re-merges authoritatively)
+ # instead of holding every key for the whole partition. The result
must be unchanged --
+ # this guards the cap/flush + chunked-emission path against OOM on
high-cardinality keys.
+ df = self._data()
+ with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch":
2}):
+ result = (
+
df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect()
+ )
+ expected =
df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect()
+ self.assertEqual({r["k"]: r["m"] for r in result}, {r["k"]: r["m"] for
r in expected})
+
+ def test_missing_buffer_schema_rejected(self):
+ # udaf() enforces a struct buffer schema up front, but a low-level
construction (or a
+ # malformed Connect proto) can build an incremental aggregator UDF
without one. That must
+ # surface a classed planner error, not a bare IllegalArgumentException
/ ClassCastException.
+ from pyspark.sql.utils import is_remote
+
+ if is_remote():
+ from pyspark.sql.connect.udf import UserDefinedFunction
+ else:
+ from pyspark.sql.udf import UserDefinedFunction # type:
ignore[assignment]
+
+ bad = UserDefinedFunction(
+ Mean(),
+ returnType=DoubleType(),
+ name="bad_mean",
+
evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ deterministic=True,
+ )._wrapped()
+ df = self._data()
+ with self.assertRaises(AnalysisException) as ctx:
+ df.groupBy("k").agg(bad(sf.col("v"))).collect()
+ self.assertEqual(ctx.exception.getCondition(),
"INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA")
+
+ def test_mixed_pandas_udaf_and_incremental_rejected(self):
+ # Mixing a grouped-agg pandas UDAF with an incremental aggregator in
one Aggregate is
+ # unsupported. It falls through to the dedicated placement error,
which must name BOTH
+ # offending functions rather than dropping the co-offending pandas
UDAF.
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ @pandas_udf("double", PandasUDFType.GROUPED_AGG)
+ def pandas_mean(v):
+ return v.mean()
+
+ df = self._data()
+ with self.assertRaises(AnalysisException) as ctx:
+ df.groupBy("k").agg(
+ udaf(Mean())(sf.col("v")).alias("m"),
+ pandas_mean(sf.col("v")).alias("pm"),
+ ).collect()
+ self.assertEqual(ctx.exception.getCondition(),
"INVALID_PYTHON_UDF_PLACEMENT")
+ message = ctx.exception.getMessage()
+ self.assertIn("Mean", message)
+ self.assertIn("pandas_mean", message)
+
+
+class ArrowPythonAggregatorTests(ArrowPythonAggregatorTestsMixin,
ReusedSQLTestCase):
+ pass
+
+
[email protected](not have_pyarrow, pyarrow_requirement_message)
+class ArrowPythonAggregatorProfilerTests(unittest.TestCase):
+ # Profiling is not supported for incremental aggregators: their ``func``
is an ``Aggregator``
+ # object, so the profiler wrappers would either break the worker (a plain
function has no
+ # ``bufferSchema`` / ``zero`` / ``reduce``) or fail on the driver in
+ # ``inspect.getsourcelines(f.__code__)``. Enabling a profiler must
therefore fall back to the
+ # non-profiled path (with a warning) and still compute the correct result,
not crash. These
+ # confs are set at session creation, so this needs its own session
(classic only).
+ def _run(self, conf_key):
+ import warnings
+
+ from pyspark import SparkConf
+ from pyspark.sql import SparkSession
+
+ conf = SparkConf().set(conf_key, "true")
+ spark = (
+ SparkSession.builder.master("local[4]")
+ .config(conf=conf)
+ .appName(self.__class__.__name__)
+ .getOrCreate()
+ )
+ try:
+ df = spark.range(0, 20).select(
+ (sf.col("id") % 3).alias("k"),
sf.col("id").cast("double").alias("v")
+ )
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ result = (
+
df.groupBy("k").agg(udaf(Mean())(sf.col("v")).alias("m")).orderBy("k").collect()
+ )
+ expected =
df.groupBy("k").agg(sf.avg("v").alias("m")).orderBy("k").collect()
+ self.assertEqual({r["k"]: r["m"] for r in result}, {r["k"]: r["m"]
for r in expected})
+ self.assertTrue(
+ any("incremental Python aggregators" in str(w.message) for w
in caught),
+ "expected an unsupported-profiling warning",
+ )
+ finally:
+ spark.stop()
+
+ def test_cpu_profiler_falls_back(self):
+ self._run("spark.python.profile")
+
+ def test_memory_profiler_falls_back(self):
+ self._run("spark.python.profile.memory")
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git
a/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_aggregator.py
b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_aggregator.py
new file mode 100644
index 000000000000..5b49582e8391
--- /dev/null
+++
b/python/pyspark/sql/tests/connect/arrow/test_parity_arrow_python_aggregator.py
@@ -0,0 +1,29 @@
+#
+# 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.
+#
+
+from pyspark.sql.tests.arrow.test_arrow_python_aggregator import
ArrowPythonAggregatorTestsMixin
+from pyspark.testing.connectutils import ReusedConnectTestCase
+
+
+class ArrowPythonAggregatorParityTests(ArrowPythonAggregatorTestsMixin,
ReusedConnectTestCase):
+ pass
+
+
+if __name__ == "__main__":
+ from pyspark.testing import main
+
+ main()
diff --git a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py
b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py
index 92fb2d0cb06a..d0031545d7b2 100644
--- a/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py
+++ b/python/pyspark/sql/tests/pandas/test_pandas_grouped_map.py
@@ -243,7 +243,8 @@ class ApplyInPandasTestsMixin:
"SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_ARROW_UDF, "
"SQL_SCALAR_PANDAS_ITER_UDF, SQL_SCALAR_ARROW_ITER_UDF, "
"SQL_GROUPED_AGG_PANDAS_UDF, SQL_GROUPED_AGG_ARROW_UDF, "
- "SQL_GROUPED_AGG_PANDAS_ITER_UDF or
SQL_GROUPED_AGG_ARROW_ITER_UDF"
+ "SQL_GROUPED_AGG_PANDAS_ITER_UDF,
SQL_GROUPED_AGG_ARROW_ITER_UDF "
+ "or SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF"
},
)
diff --git a/python/pyspark/sql/udf.py b/python/pyspark/sql/udf.py
index dbfc1b465086..51bdc58791c5 100644
--- a/python/pyspark/sql/udf.py
+++ b/python/pyspark/sql/udf.py
@@ -76,11 +76,17 @@ def _create_udf(
evalType: int,
name: Optional[str] = None,
deterministic: bool = True,
+ bufferSchema: Optional[StructType] = None,
) -> "UserDefinedFunctionLike":
"""Create a regular(non-Arrow-optimized) Python UDF."""
# Set the name of the UserDefinedFunction object to be the name of
function f
udf_obj = UserDefinedFunction(
- f, returnType=returnType, name=name, evalType=evalType,
deterministic=deterministic
+ f,
+ returnType=returnType,
+ name=name,
+ evalType=evalType,
+ deterministic=deterministic,
+ bufferSchema=bufferSchema,
)
return udf_obj._wrapped()
@@ -165,6 +171,7 @@ class UserDefinedFunction:
name: Optional[str] = None,
evalType: int = PythonEvalType.SQL_BATCHED_UDF,
deterministic: bool = True,
+ bufferSchema: Optional[StructType] = None,
):
if not callable(func):
raise PySparkTypeError(
@@ -206,6 +213,12 @@ class UserDefinedFunction:
)
self.evalType = evalType
self.deterministic = deterministic
+ # Schema of the intermediate aggregation buffer, set only for an
incremental Python
+ # aggregator (see :class:`pyspark.sql.aggregator.Aggregator`);
``None`` otherwise. It is a
+ # first-class field so it survives reconstruction paths such as
``_wrapped()``,
+ # ``asNondeterministic()`` and ``spark.udf.register``, and is threaded
to the JVM in
+ # ``_create_judf`` so ``PythonAggregate`` can plan the two-stage
aggregation.
+ self.bufferSchema = bufferSchema
# Extract Python UDF details if transpilation is enabled.
self.transpiled: list = []
self._transpiled_param_names: list[str] = []
@@ -517,6 +530,15 @@ class UserDefinedFunction:
assert sc._jvm is not None
transpiled = self.transpiled if include_transpiled else []
input_categories = self._transpiled_input_categories if
include_transpiled else []
+ # Incremental Python aggregators additionally carry the intermediate
buffer schema, which
+ # the JVM needs at planning time to build the two-stage aggregation
(see PythonAggregate).
+ # Everyone else passes ``None`` here, which Py4J maps to the JVM
``null`` the ``bufferType``
+ # parameter already defaults to.
+ jbuf = (
+ spark._jsparkSession.parseDataType(self.bufferSchema.json())
+ if self.bufferSchema is not None
+ else None
+ )
judf = getattr(sc._jvm,
"org.apache.spark.sql.execution.python.UserDefinedPythonFunction")(
self._name,
wrapped_func,
@@ -525,6 +547,7 @@ class UserDefinedFunction:
self.deterministic,
map(_to_java_column_opt, transpiled),
input_categories,
+ jbuf,
)
return judf
@@ -569,6 +592,19 @@ class UserDefinedFunction:
memory_profiler_enabled = sc._conf.get("spark.python.profile.memory",
"false") == "true"
if profiler_enabled or memory_profiler_enabled:
+ # Profiling is not supported for incremental Python aggregators.
Their ``self.func`` is
+ # an ``Aggregator`` object, not a plain function: the profiler
wrappers below would
+ # replace it with a function the worker cannot drive (it has no
``zero``/``reduce``/
+ # ``bufferSchema``), and the memory profiler's
``inspect.getsourcelines(f.__code__)``
+ # fails on the driver because an ``Aggregator`` instance has no
``__code__``.
+ if self.evalType ==
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF:
+ warnings.warn(
+ "Profiling incremental Python aggregators is not
supported.",
+ UserWarning,
+ )
+ judf = self._judf
+ return Column(judf.apply(_to_seq(sc, jcols)))
+
# Disable profiling Pandas UDFs with iterators as input/output.
if self.evalType in [
PythonEvalType.SQL_SCALAR_PANDAS_ITER_UDF,
@@ -676,6 +712,7 @@ class UserDefinedFunction:
wrapper.returnType = self.returnType # type: ignore[attr-defined]
wrapper.evalType = self.evalType # type: ignore[attr-defined]
wrapper.deterministic = self.deterministic # type:
ignore[attr-defined]
+ wrapper.bufferSchema = self.bufferSchema # type: ignore[attr-defined]
wrapper.asNondeterministic = functools.wraps( # type:
ignore[attr-defined]
self.asNondeterministic
)(lambda: self.asNondeterministic()._wrapped())
@@ -839,6 +876,7 @@ class UDFRegistration:
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
]:
raise PySparkTypeError(
errorClass="INVALID_UDF_EVAL_TYPE",
@@ -847,7 +885,8 @@ class UDFRegistration:
"SQL_SCALAR_PANDAS_UDF, SQL_SCALAR_ARROW_UDF, "
"SQL_SCALAR_PANDAS_ITER_UDF,
SQL_SCALAR_ARROW_ITER_UDF, "
"SQL_GROUPED_AGG_PANDAS_UDF,
SQL_GROUPED_AGG_ARROW_UDF, "
- "SQL_GROUPED_AGG_PANDAS_ITER_UDF or
SQL_GROUPED_AGG_ARROW_ITER_UDF"
+ "SQL_GROUPED_AGG_PANDAS_ITER_UDF,
SQL_GROUPED_AGG_ARROW_ITER_UDF "
+ "or SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF"
},
)
source_udf = _create_udf(
@@ -856,6 +895,8 @@ class UDFRegistration:
name=name,
evalType=f.evalType,
deterministic=f.deterministic,
+ # Preserve the incremental aggregator's buffer schema (None
for other UDFs).
+ bufferSchema=getattr(f, "bufferSchema", None),
)
register_udf = source_udf._unwrapped # type: ignore[attr-defined]
return_udf = register_udf
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 76be9bcc9998..81bea6152f98 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -86,6 +86,8 @@ if typing.TYPE_CHECKING:
ArrowScalarIterUDFType,
ArrowGroupedAggUDFType,
ArrowGroupedAggIterUDFType,
+ ArrowGroupedAggIncrementalPartialUDFType,
+ ArrowGroupedAggIncrementalFinalUDFType,
ArrowWindowAggUDFType,
)
from pyspark.sql._typing import (
@@ -699,6 +701,13 @@ class PythonEvalType:
SQL_WINDOW_AGG_ARROW_UDF: "ArrowWindowAggUDFType" = 253
SQL_GROUPED_AGG_ARROW_ITER_UDF: "ArrowGroupedAggIterUDFType" = 254
+ # Incremental (partial + final) Arrow aggregator. See
``pyspark.sql.aggregator``.
+ # PARTIAL folds input rows into a per-group buffer via
``Aggregator.reduce`` on the map side;
+ # FINAL merges partial buffers via ``Aggregator.merge`` and produces
output via
+ # ``Aggregator.finish`` after the shuffle.
+ SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF:
"ArrowGroupedAggIncrementalPartialUDFType" = 255
+ SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF:
"ArrowGroupedAggIncrementalFinalUDFType" = 256
+
SQL_TABLE_UDF: "SQLTableUDFType" = 300
SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301
SQL_ARROW_UDTF: "SQLArrowUDTFType" = 302
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 5aa781a4c5e5..565323365c33 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -256,6 +256,42 @@ 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_return_type(result: T, expected_type: Type[T]) -> T:
"""
Verify a UDF return value against an expected type.
@@ -637,6 +673,15 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
# The last returnType will be the return type of UDF. Eval types are
grouped below by the
# shape of the value they return.
+ # Incremental Python aggregators: the pickled "function" is the Aggregator
object itself, whose
+ # zero/reduce/merge/finish methods the worker calls directly. Return it
unwrapped (not through
+ # fail_on_stopiteration, which would treat it as a plain callable).
+ if eval_type in (
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ ):
+ return chained_func, args_offsets, kwargs_offsets, return_type
+
# Scalar, aggregation and window UDFs: (func, args_offsets,
kwargs_offsets, return_type).
if eval_type in (
PythonEvalType.SQL_ARROW_BATCHED_UDF,
@@ -2019,6 +2064,8 @@ 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,
):
# NOTE: if timezone is set here, that implies respectSessionTimeZone
is True
if eval_type in (
@@ -2026,6 +2073,10 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF,
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_ARROW_ITER_UDF,
PythonEvalType.SQL_GROUPED_MAP_ARROW_UDF,
PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
@@ -2272,6 +2323,172 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ 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])
+
+ # profiling is not supported for UDF
+ return func, None, ser, 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)
+
+ # profiling is not supported for UDF
+ return grouped_func, None, ser, ser
+
if eval_type == PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF:
import pyarrow as pa
import pandas as pd
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
index 176980e1fafd..af4f185d39aa 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala
@@ -968,7 +968,8 @@ class Analyzer(
// TODO: Support Pandas UDF.
private def checkValidAggregateExpression(expr: Expression): Unit = expr
match {
case a: AggregateExpression =>
- if (a.aggregateFunction.isInstanceOf[PythonUDAF]) {
+ if (a.aggregateFunction.isInstanceOf[PythonUDAF] ||
+ a.aggregateFunction.isInstanceOf[PythonAggregate]) {
throw
QueryCompilationErrors.pandasUDFAggregateNotSupportedInPivotError()
} else {
// OK and leave the argument check to CheckAnalysis.
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala
index 40cf4471c12a..8dac7858724d 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionResolution.scala
@@ -530,8 +530,12 @@ class FunctionResolution(
}
// We get an aggregate function, we need to wrap it in an
AggregateExpression.
case agg: AggregateFunction =>
- // Note: PythonUDAF does not support these advanced clauses.
- if (agg.isInstanceOf[PythonUDAF]) checkUnsupportedAggregateClause(agg,
unresolvedFunc)
+ // Note: neither PythonUDAF nor the incremental PythonAggregate
support these advanced
+ // clauses (DISTINCT / FILTER / ORDER BY / IGNORE NULLS). They have
dedicated physical
+ // operators that do not honor them, so reject rather than silently
drop the clause.
+ if (agg.isInstanceOf[PythonUDAF] || agg.isInstanceOf[PythonAggregate])
{
+ checkUnsupportedAggregateClause(agg, unresolvedFunc)
+ }
// After parse, the functions not set the ordering within group yet.
val newAgg = agg match {
case owg: SupportsOrderingWithinGroup
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
index a477fe4398f6..6c81d2ab7f52 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
@@ -21,6 +21,7 @@ import org.apache.spark.sql.catalyst.expressions.{
AggregateWindowFunction,
CurrentRow,
FrameLessOffsetWindowFunction,
+ PythonAggregate,
RangeFrame,
RankLike,
RowFrame,
@@ -163,6 +164,15 @@ object WindowResolution {
errorClass = "INVALID_WINDOW_SPEC_FOR_AGGREGATION_FUNC",
messageParameters = Map("aggFunc" ->
toSQLExpr(agg.aggregateFunction))
)
+ case AggregateExpression(_: PythonAggregate, _, _, _, _) =>
+ // The incremental Python aggregator has a dedicated two-stage
physical operator and is not
+ // supported as a window function. Reject it here with a clear
analysis error, rather than
+ // letting it be misclassified as a SQL aggregate
(`WindowFunctionType.functionType` only
+ // routes `PythonUDAF` to the Python path) and fail with an internal
error at execution.
+ windowExpression.failAnalysis(
+ errorClass = "UNSUPPORTED_EXPR_FOR_WINDOW",
+ messageParameters = Map("sqlExpr" ->
toSQLExpr(windowExpression.windowFunction))
+ )
case _: AggregateExpression | _: FrameLessOffsetWindowFunction | _:
AggregateWindowFunction =>
case other =>
other.failAnalysis(
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
index c7d9c49b640c..d62131183c49 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
@@ -404,6 +404,55 @@ case class PythonUDAF(
copy(children = newChildren)
}
+/**
+ * A serialized Python aggregator that supports true incremental (partial)
aggregation, the
+ * analog of the Scala typed `org.apache.spark.sql.expressions.Aggregator[IN,
BUF, OUT]`. Unlike
+ * [[PythonUDAF]] (which materializes the whole group and calls Python once),
this is planned as a
+ * two-stage aggregation by
+ * [[org.apache.spark.sql.execution.python.PythonIncrementalAggregateExec]]: a
map-side PARTIAL
+ * stage folds input rows into a per-group buffer via the aggregator's
`reduce`, and a post-shuffle
+ * FINAL stage
+ * merges the partial buffers via `merge` and produces the output via `finish`.
+ *
+ * `bufferSchema` is the schema of the intermediate buffer that crosses the
shuffle between the two
+ * stages (the analog of the Scala aggregator's `bufferEncoder`). It is
exposed here rather than via
+ * [[aggBufferAttributes]] because, like [[PythonUDAF]], this expression is
unevaluable in the JVM;
+ * the physical operator derives the buffer attributes from `bufferSchema`
directly.
+ */
+case class PythonAggregate(
+ name: String,
+ func: PythonFunction,
+ dataType: DataType,
+ children: Seq[Expression],
+ udfDeterministic: Boolean,
+ bufferSchema: StructType,
+ evalType: Int = PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ resultId: ExprId = NamedExpression.newExprId)
+ extends UnevaluableAggregateFunc with PythonFuncExpression {
+
+ override def sql(isDistinct: Boolean): String = {
+ val distinct = if (isDistinct) "DISTINCT " else ""
+ s"$name($distinct${children.mkString(", ")})"
+ }
+
+ override def toAggString(isDistinct: Boolean): String = {
+ val start = if (isDistinct) "(distinct " else "("
+ name + children.mkString(start, ", ", ")") + s"#${resultId.id}$typeSuffix"
+ }
+
+ override lazy val canonicalized: Expression = {
+ val canonicalizedChildren = children.map(_.canonicalized)
+ // `resultId` can be seen as cosmetic variation, as it doesn't affect the
result.
+ this.copy(resultId = ExprId(-1)).withNewChildren(canonicalizedChildren)
+ }
+
+ final override val nodePatterns: Seq[TreePattern] = Seq(PYTHON_UDF)
+
+ override protected def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): PythonAggregate =
+ copy(children = newChildren)
+}
+
abstract class UnevaluableGenerator extends Generator {
final override def eval(input: InternalRow): IterableOnce[InternalRow] =
throw QueryExecutionErrors.cannotEvaluateExpressionError(this)
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala
index ff471cd6f00f..7242c39f07dd 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/NormalizePlan.scala
@@ -111,6 +111,8 @@ object NormalizePlan extends PredicateHelper {
udf.copy(resultId = ExprId(0))
case udaf: PythonUDAF =>
udaf.copy(resultId = ExprId(0))
+ case agg: PythonAggregate =>
+ agg.copy(resultId = ExprId(0))
case a: FunctionTableSubqueryArgumentExpression =>
a.copy(plan = normalizeExprIds(a.plan), exprId = ExprId(0))
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
index 42351b07a688..5db80e86d50c 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
@@ -2436,6 +2436,23 @@ private[sql] object QueryCompilationErrors extends
QueryErrorsBase with Compilat
"functionList" -> groupAggPandasUDFNames.map(toSQLId).mkString(", ")))
}
+ def invalidPythonAggregatePlacementError(
+ pythonAggregateNames: Seq[String]): Throwable = {
+ new AnalysisException(
+ errorClass = "INVALID_PYTHON_UDF_PLACEMENT",
+ messageParameters = Map(
+ "functionList" -> pythonAggregateNames.map(toSQLId).mkString(", ")))
+ }
+
+ def invalidIncrementalPythonAggregatorBufferError(
+ name: String, bufferType: DataType): Throwable = {
+ new AnalysisException(
+ errorClass = "INVALID_PYTHON_AGGREGATOR_BUFFER_SCHEMA",
+ messageParameters = Map(
+ "functionName" -> toSQLId(name),
+ "bufferType" -> Option(bufferType).map(toSQLType).getOrElse("NULL")))
+ }
+
def ambiguousAttributesInSelfJoinError(
ambiguousAttrs: Seq[AttributeReference]): Throwable = {
new AnalysisException(
diff --git
a/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto
b/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto
index 18f8f294f0c0..432bc13e918d 100644
--- a/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/expressions.proto
@@ -475,6 +475,9 @@ message PythonUDF {
string python_ver = 4;
// (Optional) Additional includes for the Python UDF.
repeated string additional_includes = 5;
+ // (Optional) Intermediate buffer schema for an incremental Python aggregator
+ // (see PythonAggregate). Set only for the incremental aggregator eval types.
+ optional DataType buffer_type = 6;
}
message ScalarScalaUDF {
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index 841ce26402cf..78562e8ebf92 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -2201,6 +2201,15 @@ class SparkConnectPlanner(
createUserDefinedPythonFunction(fun)
.builder(fun.getArgumentsList.asScala.map(transformExpression).toSeq)
match {
case udaf: PythonUDAF => udaf.toAggregateExpression()
+ case agg: PythonAggregate =>
+ // The two-stage incremental aggregation operators do not implement
DISTINCT. The SQL path
+ // rejects it in FunctionResolution, but a Connect aggregate is
already resolved and skips
+ // that guard, so reject `is_distinct` here rather than silently
dropping it and returning a
+ // non-distinct result.
+ if (fun.getIsDistinct) {
+ throw
QueryCompilationErrors.functionWithUnsupportedSyntaxError(agg.name, "DISTINCT")
+ }
+ agg.toAggregateExpression()
case other => other
}
}
@@ -2214,7 +2223,9 @@ class SparkConnectPlanner(
func = function,
dataType = transformDataType(udf.getOutputType),
pythonEvalType = udf.getEvalType,
- udfDeterministic = fun.getDeterministic)
+ udfDeterministic = fun.getDeterministic,
+ // Set only for incremental Python aggregators (see PythonAggregate).
+ bufferType = if (udf.hasBufferType) transformDataType(udf.getBufferType)
else null)
}
private def transformPythonFunction(fun: proto.PythonUDF):
SimplePythonFunction = {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index 42c1121f85da..d395b3f986e8 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -579,7 +579,9 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
case PhysicalAggregation(
namedGroupingExpressions, aggregateExpressions,
rewrittenResultExpressions, child) =>
- if
(aggregateExpressions.exists(_.aggregateFunction.isInstanceOf[PythonUDAF])) {
+ if (aggregateExpressions.exists(ae =>
+ ae.aggregateFunction.isInstanceOf[PythonUDAF] ||
+ ae.aggregateFunction.isInstanceOf[PythonAggregate])) {
throw new AnalysisException(
errorClass = "_LEGACY_ERROR_TEMP_3067",
messageParameters = Map.empty)
@@ -719,7 +721,8 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
object Aggregation extends Strategy {
def apply(plan: LogicalPlan): Seq[SparkPlan] = plan match {
case PhysicalAggregation(groupingExpressions, aggExpressions,
resultExpressions, child)
- if
!aggExpressions.exists(_.aggregateFunction.isInstanceOf[PythonUDAF]) =>
+ if !aggExpressions.exists(ae =>
ae.aggregateFunction.isInstanceOf[PythonUDAF] ||
+ ae.aggregateFunction.isInstanceOf[PythonAggregate]) =>
val (functionsWithDistinct, functionsWithoutDistinct) =
aggExpressions.partition(_.isDistinct)
val distinctAggChildSets = functionsWithDistinct.map { ae =>
@@ -798,13 +801,40 @@ abstract class SparkStrategies extends
QueryPlanner[SparkPlan] {
resultExpressions,
planLater(child)))
+ case PhysicalAggregation(groupingExpressions, aggExpressions,
resultExpressions, child)
+ if
aggExpressions.forall(_.aggregateFunction.isInstanceOf[PythonAggregate]) =>
+ // Ideally this should be done in `NormalizeFloatingNumbers`, but we
do it here because
+ // `groupingExpressions` is not extracted during logical phase.
Without this, 0.0/-0.0 (and
+ // distinct NaN bit patterns) would split one logical group across
output rows.
+ val normalizedGroupingExpressions = groupingExpressions.map { e =>
+ NormalizeFloatingNumbers.normalize(e) match {
+ case n: NamedExpression => n
+ // Keep the name of the original expression.
+ case other => Alias(other, e.name)(exprId = e.exprId)
+ }
+ }
+ Seq(execution.python.PythonIncrementalAggregateExec.plan(
+ normalizedGroupingExpressions,
+ aggExpressions,
+ resultExpressions,
+ planLater(child)))
+
case PhysicalAggregation(_, aggExpressions, _, _) =>
- val groupAggPandasUDFNames = aggExpressions
- .map(_.aggregateFunction)
- .filter(_.isInstanceOf[PythonUDAF])
- .map(_.asInstanceOf[PythonUDAF].name)
- // If cannot match the two cases above, then it's an error
- throw
QueryCompilationErrors.invalidPandasUDFPlacementError(groupAggPandasUDFNames.distinct)
+ // Reached when Python aggregate UDFs cannot be planned by the two
cases above -- e.g. a
+ // grouped-agg pandas/arrow UDF or an incremental Python aggregator is
mixed with other
+ // (SQL or differently-typed Python) aggregate functions in the same
Aggregate.
+ val aggFunctions = aggExpressions.map(_.aggregateFunction)
+ val incrementalNames = aggFunctions.collect { case p: PythonAggregate
=> p.name }
+ val pandasUDFNames = aggFunctions.collect { case p: PythonUDAF =>
p.name }
+ if (incrementalNames.nonEmpty) {
+ // The message for pandas UDFs is misleading for an Arrow-based
incremental aggregator, so
+ // report the dedicated, aggregator-neutral error. Name every
offending Python aggregate
+ // function -- both the incremental aggregators and any grouped-agg
pandas/arrow UDAFs
+ // mixed in -- so the diagnostic is complete rather than dropping
the co-offenders.
+ throw QueryCompilationErrors.invalidPythonAggregatePlacementError(
+ (incrementalNames ++ pandasUDFNames).distinct)
+ }
+ throw
QueryCompilationErrors.invalidPandasUDFPlacementError(pandasUDFNames.distinct)
case _ => Nil
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala
new file mode 100644
index 000000000000..9d8931c7f919
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonIncrementalAggregateExec.scala
@@ -0,0 +1,393 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import java.io.File
+
+import scala.collection.mutable.ArrayBuffer
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{JobArtifactSet, SparkEnv, TaskContext}
+import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType}
+import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.plans.physical.{AllTuples,
ClusteredDistribution, Distribution, Partitioning, UnspecifiedDistribution}
+import org.apache.spark.sql.execution.{GroupedIterator, SparkPlan,
UnaryExecNode}
+import org.apache.spark.sql.execution.python.EvalPythonExec.ArgumentMetadata
+import org.apache.spark.sql.types.{DataType, StructField, StructType}
+import org.apache.spark.util.Utils
+
+/**
+ * Execution logic for the post-shuffle FINAL stage of an incremental Python
aggregation (see
+ * [[org.apache.spark.sql.catalyst.expressions.PythonAggregate]]). It groups
the (shuffled) child
+ * rows by the grouping expressions via a local sort + [[GroupedIterator]],
sends each group's
+ * intermediate-buffer columns to the Python worker as Arrow record batches,
and joins the single
+ * result row the worker returns per group back with the grouping key.
+ *
+ * The map-side PARTIAL stage does not extend this base: it hash-combines many
groups per batch
+ * inside the worker and needs no sort (see
[[PythonIncrementalAggregatePartialExec]]). This base is
+ * kept as an abstract class so the per-stage inputs, eval type, output
attributes and final
+ * projection stay explicit hooks.
+ */
+abstract class PythonIncrementalAggregateExecBase extends UnaryExecNode with
PythonSQLMetrics {
+
+ def groupingExpressions: Seq[NamedExpression]
+ def aggExpressions: Seq[AggregateExpression]
+
+ protected val udfExpressions: Seq[PythonAggregate] =
+ aggExpressions.map(_.aggregateFunction.asInstanceOf[PythonAggregate])
+
+ /** The Python eval type for this stage. */
+ protected def evalType: Int
+
+ /** Per-UDF input expressions to project out of the child and send to the
Python worker. */
+ protected def udfInputs: Seq[Seq[Expression]]
+
+ /** Attributes of the row the Python worker returns per group (right side of
the join). */
+ protected def pythonOutputAttributes: Seq[Attribute]
+
+ /** Expressions producing this operator's output from (groupingKey ++
pythonOutput). */
+ protected def outputExpressions: Seq[NamedExpression]
+
+ /** The grouping attributes as seen in the child's output. */
+ protected def groupingAttributes: Seq[Attribute] =
groupingExpressions.map(_.toAttribute)
+
+ /**
+ * Whether to still invoke Python on an empty partition. Only the FINAL
stage of a *global*
+ * (no grouping) aggregation sets this: it must emit the identity row
`finish(zero)` for empty
+ * input, matching SQL aggregate semantics. Everywhere else an empty
partition yields no rows.
+ */
+ protected def emitOnEmptyPartition: Boolean = false
+
+ override def output: Seq[Attribute] = outputExpressions.map(_.toAttribute)
+
+ override def producedAttributes: AttributeSet = AttributeSet(output)
+
+ override def requiredChildOrdering: Seq[Seq[SortOrder]] =
+ Seq(groupingExpressions.map(SortOrder(_, Ascending)))
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val inputRDD = child.execute()
+
+ val sessionLocalTimeZone = conf.sessionLocalTimeZone
+ val largeVarTypes = conf.arrowUseLargeVarTypes
+ val pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf)
+
+ val pyFuncs = udfExpressions.map { u =>
+ (ChainedPythonFunctions(Seq(u.func)), u.resultId.id)
+ }
+
+ // Filter child output attributes down to only those that are UDF inputs,
and eliminate
+ // duplicates, mirroring ArrowAggregatePythonExec.
+ val allInputs = new ArrayBuffer[Expression]
+ val dataTypes = new ArrayBuffer[DataType]
+ val argMetas = PythonIncrementalAggregateExec.buildArgMetas(udfInputs,
allInputs, dataTypes)
+
+ val aggInputSchema = StructType(dataTypes.zipWithIndex.map { case (dt, i)
=>
+ StructField(s"_$i", dt)
+ }.toArray)
+
+ val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
+ val sessionUUID = Option(session).collect {
+ case s if s.sessionState.conf.pythonWorkerLoggingEnabled => s.sessionUUID
+ }
+
+ val groupingExprs = groupingExpressions
+ val childOutput = child.output
+ val joinedAttributes = groupingAttributes ++ pythonOutputAttributes
+ val resultExprs = outputExpressions
+ val localEvalType = evalType
+
+ val emitIdentityOnEmpty = emitOnEmptyPartition
+ inputRDD.mapPartitionsInternal { iter => if (iter.isEmpty &&
!emitIdentityOnEmpty) iter else {
+ val prunedProj = UnsafeProjection.create(allInputs.toSeq, childOutput)
+
+ val groupedItr = if (groupingExprs.isEmpty) {
+ Iterator((new UnsafeRow(), iter))
+ } else {
+ GroupedIterator(iter, groupingExprs, childOutput)
+ }
+
+ // For a global aggregation with empty input, feed one all-null buffer
row so the Python
+ // worker still emits `finish(zero)`. An empty group cannot be sent
through
+ // GroupedPythonArrowInput (it asserts a non-empty batch per group), and
the worker treats a
+ // null partial buffer as contributing nothing to `merge`.
+ lazy val nullInputRow: UnsafeRow =
+ UnsafeProjection.create(aggInputSchema.map(_.dataType).toArray)
+ .apply(new GenericInternalRow(aggInputSchema.length)).copy()
+ val grouped = groupedItr.map { case (key, rows) =>
+ val projected = rows.map(prunedProj)
+ val toSend = if (emitIdentityOnEmpty && groupingExprs.isEmpty &&
!projected.hasNext) {
+ Iterator(nullInputRow)
+ } else {
+ projected
+ }
+ (key, toSend)
+ }
+
+ val context = TaskContext.get()
+
+ // In pipelined mode the queue's add() runs in the writer thread and
remove() in the task
+ // thread; use lock-free mode to skip per-row synchronization (as
ArrowAggregatePythonExec).
+ val pipelined = SparkEnv.get.conf.get(PYTHON_UDF_PIPELINED_EXECUTION)
+ val queue = HybridRowQueue(context.taskMemoryManager(),
+ new File(Utils.getLocalDir(SparkEnv.get.conf)), groupingExprs.length,
lockFree = pipelined)
+ context.addTaskCompletionListener[Unit] { _ => queue.close() }
+
+ val projectedRowIter = grouped.map { case (groupingKey, rows) =>
+ queue.add(groupingKey.asInstanceOf[UnsafeRow])
+ rows
+ }
+
+ val runner = new ArrowPythonWithNamedArgumentRunner(
+ pyFuncs,
+ localEvalType,
+ argMetas,
+ aggInputSchema,
+ sessionLocalTimeZone,
+ largeVarTypes,
+ pythonRunnerConf,
+ pythonMetrics,
+ jobArtifactUUID,
+ sessionUUID) with GroupedPythonArrowInput
+
+ val columnarBatchIter = runner.compute(projectedRowIter,
context.partitionId(), context)
+
+ val joined = new JoinedRow
+ val resultProj = UnsafeProjection.create(resultExprs, joinedAttributes)
+
+ columnarBatchIter.map(_.rowIterator.next()).map { pythonOutputRow =>
+ val leftRow = queue.remove()
+ resultProj(joined(leftRow, pythonOutputRow))
+ }
+ }}
+ }
+}
+
+/**
+ * Map-side PARTIAL stage: hash-combines the input rows of each group into a
per-group intermediate
+ * buffer via the aggregator's `reduce`, inside the Python worker. Unlike the
FINAL stage it needs
+ * neither a clustered distribution nor an ordering: the whole point of the
map-side combine is to
+ * avoid a full pre-shuffle sort. It streams ordinary (multi-group) Arrow
batches to the worker,
+ * which maintains one running buffer per distinct grouping key and emits, at
end of partition, one
+ * row per key -- the grouping key columns followed by one intermediate-buffer
struct column per
+ * aggregator.
+ *
+ * Because keys may be split across partitions (no shuffle here) and the
worker's map-side grouping
+ * is only a combine optimization, correctness does not depend on it being
exhaustive:
+ * [[PythonIncrementalAggregateFinalExec]] re-groups the emitted partial
buffers authoritatively
+ * (by JVM `UnsafeRow` key, after the shuffle) and merges any that share a key.
+ */
+case class PythonIncrementalAggregatePartialExec(
+ groupingExpressions: Seq[NamedExpression],
+ aggExpressions: Seq[AggregateExpression],
+ bufferAttributes: Seq[Attribute],
+ child: SparkPlan) extends UnaryExecNode with PythonSQLMetrics {
+
+ private val udfExpressions: Seq[PythonAggregate] =
+ aggExpressions.map(_.aggregateFunction.asInstanceOf[PythonAggregate])
+
+ private def groupingAttributes: Seq[Attribute] =
groupingExpressions.map(_.toAttribute)
+
+ override def output: Seq[Attribute] = groupingAttributes ++ bufferAttributes
+
+ override def producedAttributes: AttributeSet = AttributeSet(output)
+
+ override def requiredChildDistribution: Seq[Distribution] =
Seq(UnspecifiedDistribution)
+
+ // No ordering: the worker hash-combines rather than relying on grouped
input.
+ override def requiredChildOrdering: Seq[Seq[SortOrder]] = Seq(Nil)
+
+ override def outputPartitioning: Partitioning = child.outputPartitioning
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val inputRDD = child.execute()
+
+ val sessionLocalTimeZone = conf.sessionLocalTimeZone
+ val largeVarTypes = conf.arrowUseLargeVarTypes
+ val pythonRunnerConf = ArrowPythonRunner.getPythonRunnerConfMap(conf)
+
+ val pyFuncs = udfExpressions.map { u =>
+ (ChainedPythonFunctions(Seq(u.func)), u.resultId.id)
+ }
+
+ // The columns sent to Python are the grouping keys first (so the worker
can hash-group by them
+ // and echo them back with each partial buffer), followed by the
deduplicated aggregator input
+ // columns. A UDF input that coincides with a grouping key simply reuses
that leading column.
+ val allInputs = new ArrayBuffer[Expression]
+ val dataTypes = new ArrayBuffer[DataType]
+ groupingExpressions.foreach { g =>
+ allInputs += g
+ dataTypes += g.dataType
+ }
+ val numGroupingKeys = groupingExpressions.length
+
+ val argMetas = PythonIncrementalAggregateExec.buildArgMetas(
+ udfExpressions.map(_.children), allInputs, dataTypes)
+
+ val aggInputSchema = StructType(dataTypes.zipWithIndex.map { case (dt, i)
=>
+ StructField(s"_$i", dt)
+ }.toArray)
+ // The leading `numGroupingKeys` columns are the grouping keys; hand their
schema to the worker.
+ val groupingKeySchemaJson =
StructType(aggInputSchema.fields.take(numGroupingKeys)).json
+
+ val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
+ val sessionUUID = Option(session).collect {
+ case s if s.sessionState.conf.pythonWorkerLoggingEnabled => s.sessionUUID
+ }
+
+ val childOutput = child.output
+ val outputAttrs = output
+
+ inputRDD.mapPartitionsInternal { iter => if (iter.isEmpty) Iterator.empty
else {
+ val prunedProj = UnsafeProjection.create(allInputs.toSeq, childOutput)
+ val projectedRowIter = iter.map(prunedProj)
+
+ val context = TaskContext.get()
+
+ val runner = new ArrowPythonWithNamedArgumentRunner(
+ pyFuncs,
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
+ argMetas,
+ aggInputSchema,
+ sessionLocalTimeZone,
+ largeVarTypes,
+ pythonRunnerConf,
+ pythonMetrics,
+ jobArtifactUUID,
+ sessionUUID) with BatchedPythonArrowInput {
+ // Tell the worker how many leading columns are grouping keys, so it
can hash-group by them
+ // and re-emit them alongside each partial buffer.
+ override protected def evalConf: Map[String, String] =
+ super.evalConf + ("grouping_key_schema" -> groupingKeySchemaJson)
+ }
+
+ val columnarBatchIter = runner.compute(
+ Iterator(projectedRowIter), context.partitionId(), context)
+
+ // Each batch the worker returns holds (grouping key columns ++ one
buffer struct column per
+ // aggregator), i.e. this operator's output columns; copy each row out
as an UnsafeRow.
+ val resultProj = UnsafeProjection.create(outputAttrs, outputAttrs)
+ columnarBatchIter.flatMap(_.rowIterator().asScala).map(resultProj)
+ }}
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ copy(child = newChild)
+}
+
+/**
+ * Post-shuffle FINAL stage: clusters the partial buffers by the grouping key,
merges the buffers
+ * of each group via the aggregator's `merge`, and produces the output via
`finish`. Its input is
+ * the [[PythonIncrementalAggregatePartialExec]] output (grouping key columns
followed by the
+ * intermediate-buffer columns); it sends the buffer columns to Python and
outputs
+ * `resultExpressions`.
+ */
+case class PythonIncrementalAggregateFinalExec(
+ groupingExpressions: Seq[NamedExpression],
+ aggExpressions: Seq[AggregateExpression],
+ bufferAttributes: Seq[Attribute],
+ resultExpressions: Seq[NamedExpression],
+ child: SparkPlan) extends PythonIncrementalAggregateExecBase {
+
+ override protected def evalType: Int =
+ PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF
+
+ // Each aggregator reads its own intermediate-buffer column from the
(shuffled) child.
+ override protected def udfInputs: Seq[Seq[Expression]] =
bufferAttributes.map(Seq(_))
+
+ override protected def pythonOutputAttributes: Seq[Attribute] =
+ aggExpressions.map(_.resultAttribute)
+
+ override protected def outputExpressions: Seq[NamedExpression] =
resultExpressions
+
+ // A global (no-grouping) aggregation must return the identity row even for
empty input. This
+ // stage runs on a single partition (AllTuples), so exactly one identity row
is produced.
+ override protected def emitOnEmptyPartition: Boolean =
groupingExpressions.isEmpty
+
+ override def requiredChildDistribution: Seq[Distribution] = {
+ if (groupingExpressions.isEmpty) {
+ AllTuples :: Nil
+ } else {
+ ClusteredDistribution(groupingExpressions) :: Nil
+ }
+ }
+
+ override def outputPartitioning: Partitioning = child.outputPartitioning
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ copy(child = newChild)
+}
+
+object PythonIncrementalAggregateExec {
+
+ /**
+ * Deduplicates the per-aggregator input expressions into a shared column
list (matched by
+ * [[Expression.semanticEquals]]), returning one [[ArgumentMetadata]] array
per aggregator that
+ * points into that list. `allInputs`/`dataTypes` may be pre-seeded -- the
PARTIAL stage prepends
+ * its grouping-key columns so a UDF input equal to a grouping key reuses
that leading column --
+ * and any newly seen columns are appended to them in place. Shared by both
stages (and mirrors
+ * the same dedup in [[ArrowAggregatePythonExec]]).
+ */
+ private[python] def buildArgMetas(
+ udfInputs: Seq[Seq[Expression]],
+ allInputs: ArrayBuffer[Expression],
+ dataTypes: ArrayBuffer[DataType]): Array[Array[ArgumentMetadata]] = {
+ udfInputs.map { input =>
+ input.map { e =>
+ val (key, value) = e match {
+ case NamedArgumentExpression(key, value) => (Some(key), value)
+ case _ => (None, e)
+ }
+ if (allInputs.exists(_.semanticEquals(value))) {
+ ArgumentMetadata(allInputs.indexWhere(_.semanticEquals(value)), key)
+ } else {
+ allInputs += value
+ dataTypes += value.dataType
+ ArgumentMetadata(allInputs.length - 1, key)
+ }
+ }.toArray
+ }.toArray
+ }
+
+ /**
+ * Builds the two-stage physical plan (PARTIAL -> [Exchange, inserted by
EnsureRequirements] ->
+ * FINAL) for a logical aggregation whose aggregate functions are all
[[PythonAggregate]].
+ */
+ def plan(
+ groupingExpressions: Seq[NamedExpression],
+ aggExpressions: Seq[AggregateExpression],
+ resultExpressions: Seq[NamedExpression],
+ child: SparkPlan): SparkPlan = {
+ // One intermediate-buffer attribute per aggregator, threaded from the
PARTIAL output into the
+ // FINAL inputs (matched by expression id).
+ val bufferAttributes = aggExpressions.map { ae =>
+ val agg = ae.aggregateFunction.asInstanceOf[PythonAggregate]
+ AttributeReference(s"buf_${agg.resultId.id}", agg.bufferSchema, nullable
= true)()
+ }
+ val partial = PythonIncrementalAggregatePartialExec(
+ groupingExpressions, aggExpressions, bufferAttributes, child)
+ // After the PARTIAL stage the grouping expressions are materialized as
plain attributes.
+ val groupingAttributes = groupingExpressions.map(_.toAttribute)
+ PythonIncrementalAggregateFinalExec(
+ groupingAttributes, aggExpressions, bufferAttributes, resultExpressions,
partial)
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
index cfd1fb4ea222..9f53f078b917 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
@@ -28,7 +28,7 @@ import net.razorvine.pickle.Pickler
import org.apache.spark.api.python.{PythonEvalType, PythonFunction,
PythonWorkerUtils, SpecialLengths}
import org.apache.spark.sql.{Column, TableArg}
import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
-import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending,
Descending, Expression, FunctionTableSubqueryArgumentExpression,
NamedArgumentExpression, NullsFirst, NullsLast, PythonUDAF, PythonUDF,
PythonUDTF, PythonUDTFAnalyzeResult, PythonUDTFSelectedExpression, SortOrder,
TranspiledPythonUDF, UnresolvedPolymorphicPythonUDTF, UnresolvedTableArgPlanId}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending,
Descending, Expression, FunctionTableSubqueryArgumentExpression,
NamedArgumentExpression, NullsFirst, NullsLast, PythonAggregate, PythonUDAF,
PythonUDF, PythonUDTF, PythonUDTFAnalyzeResult, PythonUDTFSelectedExpression,
SortOrder, TranspiledPythonUDF, UnresolvedPolymorphicPythonUDTF,
UnresolvedTableArgPlanId}
import org.apache.spark.sql.catalyst.parser.ParserInterface
import org.apache.spark.sql.catalyst.plans.logical.{Generate, LogicalPlan,
NamedParametersSupport, OneRowRelation}
import org.apache.spark.sql.classic.{DataFrame, Dataset, SparkSession}
@@ -56,7 +56,11 @@ case class UserDefinedPythonFunction(
// categories match the bound argument types; when none match, the call
// falls back to the plain Python UDF. `builder` requires the two lists to
// be parallel and skips transpilation otherwise.
- transpiledInputTypes: JList[JList[String]] = Nil.asJava) {
+ transpiledInputTypes: JList[JList[String]] = Nil.asJava,
+ // Schema of the intermediate aggregation buffer, set only for the
incremental Python
+ // aggregator eval types (see [[PythonAggregate]]); `null` otherwise.
Nullable rather than
+ // `Option` so it can be passed positionally from Python over Py4J.
+ bufferType: DataType = null) {
def builder(e: Seq[Expression]): Expression = {
if (pythonEvalType == PythonEvalType.SQL_BATCHED_UDF
@@ -66,7 +70,8 @@ case class UserDefinedPythonFunction(
|| pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_PANDAS_ITER_UDF
|| pythonEvalType == PythonEvalType.SQL_SCALAR_ARROW_UDF
|| pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF
- || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF) {
+ || pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF
+ || pythonEvalType ==
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF) {
/*
* Check if the named arguments:
* - don't have duplicated names
@@ -87,6 +92,21 @@ case class UserDefinedPythonFunction(
|| pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF
|| pythonEvalType == PythonEvalType.SQL_GROUPED_AGG_ARROW_ITER_UDF) {
PythonUDAF(name, func, dataType, e, udfDeterministic, pythonEvalType)
+ } else if (pythonEvalType ==
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF) {
+ // The incremental Python aggregator. `bufferType` (the intermediate
buffer schema) must have
+ // been supplied as a struct when the UDF was created. The single
expression carries the
+ // aggregator for both the PARTIAL and FINAL stages; the physical
operator picks the per-stage
+ // eval type. `udaf()` enforces this up front, but a malformed Connect
proto (eval type set,
+ // `buffer_type` missing or non-struct) or a direct
`UserDefinedFunction(f, evalType=...)` can
+ // reach here without it, so return a classed error rather than a bare
require /
+ // ClassCastException.
+ val bufferStruct = bufferType match {
+ case s: StructType => s
+ case _ =>
+ throw
QueryCompilationErrors.invalidIncrementalPythonAggregatorBufferError(
+ name, bufferType)
+ }
+ PythonAggregate(name, func, dataType, e, udfDeterministic, bufferStruct)
} else {
PythonUDF(name, func, dataType, e, pythonEvalType, udfDeterministic)
}
@@ -172,6 +192,7 @@ case class UserDefinedPythonFunction(
case TranspiledPythonUDF(name, udaf: PythonUDAF, transpiled,
inputCategories) =>
TranspiledPythonUDF(name, udaf.toAggregateExpression(), transpiled,
inputCategories)
case udaf: PythonUDAF => udaf.toAggregateExpression()
+ case agg: PythonAggregate => agg.toAggregateExpression()
case _ => expr
})
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
