How to Mock a UDAF Aggregate Function in Python

from collections import defaultdict

class MockUDAF:
    """A minimal mock of a User-Defined Aggregate Function.

    Simulates aggregate lifecycle: initialize, update per row,
    and finalize the result.
    """

    def __init__(self):
        self._buffer = defaultdict(int)

    def initialize(self):
        """Reset internal state."""
        self._buffer.clear()

    def update(self, value):
        """Accumulate one input value."""
        if value is not None:
            self._buffer[value] += 1

    def merge(self, other):
        """Merge another MockUDAF (simulates partial aggregates)."""
        for key, count in other._buffer.items():
            self._buffer[key] += count

    def finalize(self):
        """Return the aggregate result."""
        return dict(sorted(self._buffer.items(), key=lambda item: -item[1]))

    @staticmethod
    def run(values):
        """Convenience: full pipeline in one call."""
        udaf = MockUDAF()
        udaf.initialize()
        for v in values:
            udaf.update(v)
        return udaf.finalize()


if __name__ == "__main__":
    data = ["apple", "banana", "apple", "cherry", "banana", "apple", None]
    result = MockUDAF.run(data)
    print(result)