[WIP][SPARK-54647][PYTHON] Support User-Defined Aggregate Functions (UDAF)

[Yicong-Huang]

What changes were proposed in this pull request?

Add support for User-Defined Aggregate Functions (UDAF) in PySpark. Currently PySpark supports User-Defined Functions (UDF) and User-Defined Table Functions (UDTF), but lacks support for UDAF. Users need to write custom aggregation logic in Scala/Java or use less efficient workarounds.

This change adds UDAF support using a two-stage aggregation pattern with mapInArrow and applyInArrow. The basic idea is to implement aggregation (and partial aggregation) by:

df.selectExpr("rand() as key").mapInArrow(reduce).groupBy(key).applyInArrow(merge)

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.

Aligned with Scala side, the implementation provides a Python Aggregator base class that users can subclass:

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

Users can create UDAF instances using the udaf() function and use them with DataFrame.agg():

sum_udaf = udaf(MySum(), "bigint")
df.agg(sum_udaf(df.value))
df.groupBy("group").agg(sum_udaf(df.value))

Key changes:

• Added pyspark.sql.udaf module with Aggregator base class, UserDefinedAggregateFunction wrapper, and udaf() factory function
• Integrated UDAF support in GroupedData.agg() by detecting UDAF columns via _udaf_func attribute

Why are the changes needed?

Currently PySpark lacks support for User-Defined Aggregate Functions (UDAF), which limits users' ability to express complex aggregation logic directly in Python. Users must either write custom aggregation logic in Scala/Java or use less efficient workarounds. This change adds UDAF support to complement existing UDF and UDTF support in PySpark, aligning with the Scala/Java Aggregator interface in org.apache.spark.sql.expressions.Aggregator.

Does this PR introduce any user-facing change?

Yes. This PR adds a new feature - User-Defined Aggregate Functions (UDAF) support in PySpark. Users can now define custom aggregation logic by subclassing the Aggregator class and using the udaf() function to create UDAF instances that can be used with DataFrame.agg() and GroupedData.agg().

Example:

class MySum(Aggregator):
    def zero(self):
        return 0
    def reduce(self, buffer, value):
        return buffer + value
    def merge(self, buffer1, buffer2):
        return buffer1 + buffer2
    def finish(self, reduction):
        return reduction

sum_udaf = udaf(MySum(), "bigint")
df.agg(sum_udaf(df.value))

How was this patch tested?

Added comprehensive unit tests in python/pyspark/sql/tests/test_udaf.py covering:

• Basic aggregation (sum, average, max)
• Grouped aggregation with groupBy().agg()
• Null value handling
• Empty DataFrame handling
• Large datasets (20000+ rows) distributed across partitions
• Error handling for invalid inputs
• Integration with df.agg() and df.groupBy().agg()

Was this patch authored or co-authored using generative AI tooling?

No.

[allisonwang-db]

Nice feature!

[allisonwang-db]

Can we add some tests for more complicated data structures? like dictionary?

[Yicong-Huang]

added more data types!

[zhengruifeng]

do we necessarily need this class?
I see UDTF doesn't need a base class.

    >>> class TestUDTF:
    ...     def eval(self, *args: Any):
    ...         yield "hello", "world"

[Yicong-Huang]

we could do duck typing if not go with the inheritance. I think it is debatable and we could offer both solutions (with or without a base class).

[zhengruifeng]

why not a dedicated pyhsical plan?

[zhengruifeng]

If we want to support partial aggregation with existing arrow UDFs, I think we should use a modified FlatMapGroupsInArrowExec with requiredChildDistribution = UnspecifiedDistribution.

[zhengruifeng]

MapInArrowExec dosen't requiredChildOrdering, where does it sort the data for partial aggregation?

[Yicong-Huang]

The Sort is now explicitly added in RewritePythonAggregatorUDAF before MapInArrow.

[zhengruifeng]

The basic idea is to implement aggregation (and partial aggregation) by:
df.selectExpr("rand() as key").mapInArrow(reduce).groupBy(key).applyInArrow(merge)

I think there should be a sortWithinPartitions before mapInArrow for partial aggregation.

[zhengruifeng]

The whole approach is based on mapInArrow and applyInArrow, how does it support function registration so that it can be used in SQL?

[zhengruifeng]

group_buffers buffers all the aggregators within a partition, it will cause memory issue if the cardinality is large.
A reasonable physical plan should sort the partition by the key, and then output the partial aggregation result after finishing each group

[zhengruifeng]

it mimic the HashAggregateExec, while SortAggregateExec is more stable

[Yicong-Huang]

thanks for the suggestion. I will take a look on different aggregateExec implementations.

[dtenedor]

@Yicong-Huang please let me help you as a reviewer for this, I implemented remote UDAFs several times for other systems preivously

[Yicong-Huang]

@dtenedor thanks! could you please have a pass on the current implementation?