HyukjinKwon commented on code in PR #57952: URL: https://github.com/apache/spark/pull/57952#discussion_r3810285727
########## python/pyspark/sql/aggregator.py: ########## @@ -0,0 +1,211 @@ +# +# 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, 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`.""" Review Comment: Good call -- documented in 82dd5d8. Spelled out the identity law on `zero` (`merge(buffer, zero()) == buffer` and `merge(zero(), buffer) == buffer`), noting that a fresh `zero()` seeds every partition and every early-flushed map-side chunk, so associativity/commutativity alone don't give the partition-independent result. Also referenced it from the class docstring. Thanks for the review and the approval! -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
