[jira] [Updated] (SPARK-54647) Support User-Defined Aggregate Functions (UDAF) in PySpark

["ASF GitHub Bot (Jira)" <ji...@apache.org>]

     [ https://issues.apache.org/jira/browse/SPARK-54647?page=com.atlassian.jira.plugin.system.issuetabpanels:all-tabpanel ]

ASF GitHub Bot updated SPARK-54647:
-----------------------------------
    Labels: pull-request-available  (was: )

> Support User-Defined Aggregate Functions (UDAF) in PySpark
> ----------------------------------------------------------
>
>                 Key: SPARK-54647
>                 URL: https://issues.apache.org/jira/browse/SPARK-54647
>             Project: Spark
>          Issue Type: New Feature
>          Components: PySpark
>    Affects Versions: 4.2.0
>            Reporter: Yicong Huang
>            Priority: Major
>              Labels: pull-request-available
>
> Currently PySpark supports User-Defined Functions (UDF) and User-Defined Table Functions (UDTF), but lacks support for User-Defined Aggregate Functions (UDAF). Users need to write custom aggregation logic in Scala/Java or use less efficient workarounds. This limits the ability to express complex aggregation logic directly in Python.
> This change adds UDAF support in PySpark using a two-stage aggregation pattern with mapInArrow and applyInArrow. The basic idea is to implement aggregation (and partial aggregation) by:
> {code}
> df.selectExpr("rand() as key").mapInArrow(func1).groupBy(key).applyInArrow(func2)
> {code}
> Where func1 calls Aggregator.reduce() for partial aggregation within each partition, and func2 calls Aggregator.merge() to combine partial results, then Aggregator.finish() for final results.
> The implementation provides a Python Aggregator base class that users can subclass:
> {code:python}
> class Aggregator:
>     def zero(self) -> Any:
>         """Return zero value for aggregation buffer"""
>         raise NotImplementedError
>     
>     def reduce(self, buffer: Any, value: Any) -> Any:
>         """Combine input value into buffer"""
>         raise NotImplementedError
>     
>     def merge(self, buffer1: Any, buffer2: Any) -> Any:
>         """Merge two intermediate buffers"""
>         raise NotImplementedError
>     
>     def finish(self, reduction: Any) -> Any:
>         """Produce final result from buffer"""
>         raise NotImplementedError
> {code}
> Users can create UDAF instances using the udaf() function and use them with DataFrame.agg():
> {code:python}
> sum_udaf = udaf(MySum(), "bigint")
> df.agg(sum_udaf(df.value))
> df.groupBy("group").agg(sum_udaf(df.value))
> {code}
> The implementation uses iterator-based mapInArrow and applyInArrow APIs to reduce memory footprint, serializes Aggregator instances using cloudpickle for distribution to worker nodes, and handles both grouped and non-grouped aggregations. The random key range is determined by spark.sql.shuffle.partitions or DataFrame partition count to ensure proper data distribution.



--
This message was sent by Atlassian Jira
(v8.20.10#820010)

---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@spark.apache.org
For additional commands, e-mail: issues-help@spark.apache.org