[SPARK-58736][SQL][PYTHON][CONNECT] Support incremental Python aggregators via Arrow

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

dev/sparktestsupport/modules.py

@@ -610,6 +610,7 @@ def __hash__(self):
         "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",
@@ -1253,6 +1254,7 @@ def __hash__(self):
         "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",

python/pyspark/sql/aggregator.py

@@ -0,0 +1,202 @@
+#
+# 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
+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, since the framework
+    may combine partial buffers in any order.
+
+    .. 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`."""
+        ...
+
+    @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__,
+            },
+        )
+
+    udf_obj = UserDefinedFunction(
+        agg,
+        returnType=agg.outputType,
+        name=agg.__class__.__name__,
+        evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+        deterministic=True,
+    )
+    # Threaded to the JVM in UserDefinedFunction._create_judf so PythonAggregate can plan the
+    # two-stage aggregation. Set on both the UDF and its wrapper so it survives
+    # ``spark.udf.register`` (which reconstructs the UDF from the wrapper).
+    udf_obj.bufferSchema = agg.bufferSchema  # type: ignore[attr-defined]
+    wrapped = udf_obj._wrapped()
+    wrapped.bufferSchema = agg.bufferSchema  # type: ignore[attr-defined]
+    return wrapped

python/pyspark/sql/connect/client/core.py

@@ -1081,6 +1081,7 @@ def register_udf(
         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 @@ def register_udf(
             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

python/pyspark/sql/connect/expressions.py

@@ -741,13 +741,16 @@ def __init__(
         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
         )
         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 @@ def to_plan(self, session: "SparkConnectClient") -> proto.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:

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
 

python/pyspark/sql/connect/udf.py

@@ -210,6 +210,8 @@ def to_expr(col: "ColumnOrName") -> Expression:
             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=getattr(self, "bufferSchema", None),
         )
         return CommonInlineUserDefinedFunction(
             function_name=self._name,
@@ -303,6 +305,7 @@ def register(
                 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 +314,18 @@ def register(
                         "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:

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