windows: Split Aggregations and collectors is own classes (#772)

* windows: Split Aggregations and collectors is own classes

in preparation for multiple aggregations and collectors support

Author: ovv

quixstreams/dataframe/windows/aggregations.py

@@ -0,0 +1,208 @@
+from abc import ABC, abstractmethod
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Hashable,
+    Iterable,
+    Optional,
+    TypeVar,
+    Union,
+)
+
+from typing_extensions import TypeAlias
+
+
+class Aggregator(ABC):
+    """
+    Base class for window aggregation.
+
+    Subclass it to implement custom aggregations.
+
+    An Aggregator reduce incoming items into a single value or group of values. When the window
+    is closed the aggregator produce a result based on the reduced value.
+
+    To store all incoming items without reducing them use a `Collector`.
+    """
+
+    @abstractmethod
+    def initialize(self) -> Any:
+        """
+        This method is triggered once to build the aggregation starting value.
+        It should return the initial value for the aggregation.
+        """
+        ...
+
+    @abstractmethod
+    def agg(self, old: Any, new: Any) -> Any:
+        """
+        This method is trigged when a window is updated with a new value.
+        It should return the updated aggregated value.
+        """
+        ...
+
+    @abstractmethod
+    def result(self, value: Any) -> Any:
+        """
+        This method is triggered when a window is closed.
+        It should return the final aggregation result.
+        """
+        ...
+
+
+V = TypeVar("V", int, float)
+
+
+class ROOT:
+    pass
+
+
+Column: TypeAlias = Union[Hashable, type[ROOT]]
+
+
+class Sum(Aggregator):
+    def __init__(self, column: Column = ROOT) -> None:
+        self.column = column
+
+    def initialize(self) -> int:
+        return 0
+
+    def agg(self, old: V, new: Any) -> V:
+        if self.column is ROOT:
+            return old + new
+        return old + new[self.column]
+
+    def result(self, value: V) -> V:
+        return value
+
+
+class Count(Aggregator):
+    def initialize(self) -> int:
+        return 0
+
+    def agg(self, old: int, new: Any) -> int:
+        return old + 1
+
+    def result(self, value: int) -> int:
+        return value
+
+
+class Mean(Aggregator):
+    def __init__(self, column: Column = ROOT) -> None:
+        self.column = column
+
+    def initialize(self) -> tuple[float, int]:
+        return 0.0, 0
+
+    def agg(self, old: tuple[V, int], new: Any) -> tuple[V, int]:
+        old_sum, old_count = old
+        if self.column is ROOT:
+            return old_sum + new, old_count + 1
+        return old_sum + new[self.column], old_count + 1
+
+    def result(self, value: tuple[Union[int, float], int]) -> float:
+        sum_, count_ = value
+        return sum_ / count_
+
+
+R = TypeVar("R", int, float)
+
+
+class Reduce(Aggregator, Generic[R]):
+    def __init__(
+        self,
+        reducer: Callable[[R, Any], R],
+        initializer: Callable[[Any], R],
+    ) -> None:
+        self._initializer: Callable[[Any], R] = initializer
+        self._reducer: Callable[[R, Any], R] = reducer
+
+    def initialize(self) -> Any:
+        return None
+
+    def agg(self, old: R, new: Any) -> Any:
+        return self._initializer(new) if old is None else self._reducer(old, new)
+
+    def result(self, value: R) -> R:
+        return value
+
+
+class Max(Aggregator):
+    def __init__(self, column: Column = ROOT) -> None:
+        self.column = column
+
+    def initialize(self) -> None:
+        return None
+
+    def agg(self, old: Optional[V], new: Any) -> V:
+        if self.column is not ROOT:
+            new = new[self.column]
+        if old is None:
+            return new
+        return max(old, new)
+
+    def result(self, value: V) -> V:
+        return value
+
+
+class Min(Aggregator):
+    def __init__(self, column: Column = ROOT) -> None:
+        self.column = column
+
+    def initialize(self) -> None:
+        return None
+
+    def agg(self, old: Optional[V], new: Any) -> V:
+        if self.column is not ROOT:
+            new = new[self.column]
+        if old is None:
+            return new
+        return min(old, new)
+
+    def result(self, value: V) -> V:
+        return value
+
+
+I = TypeVar("I")
+
+
+class Collector(ABC, Generic[I]):
+    """
+    Base class for window collections.
+
+    Subclass it to implement custom collections.
+
+    A Collector store incoming items un-modified in an optimized way.
+
+    To reduce incoming items as they come in use an `Aggregator`.
+    """
+
+    @property
+    @abstractmethod
+    def column(self) -> Column:
+        """
+        The column to collect.
+
+        Use `ROOT` to collect the whole message.
+        """
+        ...
+
+    @abstractmethod
+    def result(self, items: Iterable[I]) -> Any:
+        """
+        This method is triggered when a window is closed.
+        It should return the final collection result.
+        """
+        ...
+
+
+class Collect(Collector):
+    def __init__(self, column: Column = ROOT) -> None:
+        self._column = column
+
+    @property
+    def column(self) -> Column:
+        return self._column
+
+    def result(self, items: Iterable[Any]) -> list[Any]:
+        return list(items)

quixstreams/dataframe/windows/base.py

@@ -22,6 +22,8 @@
 from quixstreams.processing import ProcessingContext
 from quixstreams.state import WindowedPartitionTransaction
 
+from .aggregations import Aggregator, Collector
+
 if TYPE_CHECKING:
     from quixstreams.dataframe.dataframe import StreamingDataFrame
 
@@ -38,7 +40,6 @@ class WindowResult(TypedDict):
 Message: TypeAlias = tuple[WindowResult, Any, int, Any]
 
 WindowAggregateFunc = Callable[[Any, Any], Any]
-WindowMergeFunc = Callable[[Any], Any]
 
 TransformRecordCallbackExpandedWindowed = Callable[
     [Any, Any, int, Any, WindowedPartitionTransaction],
@@ -51,13 +52,26 @@ def __init__(
         self,
         name: str,
         dataframe: "StreamingDataFrame",
+        aggregators: dict[str, Aggregator],
+        collectors: dict[str, Collector],
     ) -> None:
         if not name:
             raise ValueError("Window name must not be empty")
 
         self._name = name
         self._dataframe = dataframe
 
+        self._aggregators = aggregators
+        self._aggregate = len(aggregators) > 0
+
+        self._collectors = collectors
+        self._collect = len(collectors) > 0
+
+        if not self._collect and not self._aggregate:
+            raise ValueError("At least one aggregation or collector must be defined")
+        elif len(collectors) + len(aggregators) > 1:
+            raise ValueError("Only one aggregation or collector can be defined")
+
     @property
     def name(self) -> str:
         return self._name
@@ -260,7 +274,3 @@ def get_window_ranges(
         current_window_start -= step_ms
 
     return window_ranges
-
-
-def default_merge_func(state_value: Any) -> Any:
-    return state_value

quixstreams/dataframe/windows/count_based.py

@@ -3,13 +3,11 @@
 
 from quixstreams.state import WindowedPartitionTransaction
 
+from .aggregations import Aggregator, Collector
 from .base import (
     Window,
-    WindowAggregateFunc,
     WindowKeyResult,
-    WindowMergeFunc,
     WindowResult,
-    default_merge_func,
 )
 
 if TYPE_CHECKING:
@@ -35,22 +33,21 @@ class CountWindow(Window):
 
     def __init__(
         self,
-        name: str,
         count: int,
-        aggregate_func: WindowAggregateFunc,
-        aggregate_default: Any,
-        aggregate_collection: bool,
+        name: str,
         dataframe: "StreamingDataFrame",
-        merge_func: Optional[WindowMergeFunc] = None,
+        aggregators: dict[str, Aggregator],
+        collectors: dict[str, Collector],
         step: Optional[int] = None,
     ):
-        super().__init__(name, dataframe)
+        super().__init__(
+            name=name,
+            dataframe=dataframe,
+            aggregators=aggregators,
+            collectors=collectors,
+        )
 
         self._max_count = count
-        self._aggregate_func = aggregate_func
-        self._aggregate_default = aggregate_default
-        self._aggregate_collection = aggregate_collection
-        self._merge_func = merge_func or default_merge_func
         self._step = step
 
     def process_window(
@@ -86,34 +83,37 @@ def process_window(
         msg_id = None
         if len(data["windows"]) == 0:
             # for new tumbling window, reset the collection id to 0
-            msg_id = 0
+            if self._collect:
+                window_value = msg_id = 0
+            else:
+                window_value = self._aggregators["value"].initialize()
 
             data["windows"].append(
                 CountWindowData(
                     count=0,
                     start=timestamp_ms,
                     end=timestamp_ms,
-                    value=msg_id
-                    if self._aggregate_collection
-                    else self._aggregate_default,
+                    value=window_value,
                 )
             )
         elif self._step is not None and data["windows"][0]["count"] % self._step == 0:
-            if self._aggregate_collection:
-                msg_id = data["windows"][0]["value"] + data["windows"][0]["count"]
+            if self._collect:
+                window_value = msg_id = (
+                    data["windows"][0]["value"] + data["windows"][0]["count"]
+                )
+            else:
+                window_value = self._aggregators["value"].initialize()
 
             data["windows"].append(
                 CountWindowData(
                     count=0,
                     start=timestamp_ms,
                     end=timestamp_ms,
-                    value=msg_id
-                    if self._aggregate_collection
-                    else self._aggregate_default,
+                    value=window_value,
                 )
             )
 
-        if self._aggregate_collection:
+        if self._collect:
             if msg_id is None:
                 msg_id = data["windows"][0]["value"] + data["windows"][0]["count"]
 
@@ -127,7 +127,7 @@ def process_window(
             elif timestamp_ms > window["end"]:
                 window["end"] = timestamp_ms
 
-            if self._aggregate_collection:
+            if self._collect:
                 # window must close
                 if window["count"] >= self._max_count:
                     values = state.get_from_collection(
@@ -141,7 +141,7 @@ def process_window(
                             WindowResult(
                                 start=window["start"],
                                 end=window["end"],
-                                value=self._merge_func(values),
+                                value=self._collectors["value"].result(values),
                             ),
                         )
                     )
@@ -157,14 +157,14 @@ def process_window(
 
                     state.delete_from_collection(end=delete_end, start=delete_start)
             else:
-                window["value"] = self._aggregate_func(window["value"], value)
+                window["value"] = self._aggregators["value"].agg(window["value"], value)
 
                 result = (
                     key,
                     WindowResult(
                         start=window["start"],
                         end=window["end"],
-                        value=self._merge_func(window["value"]),
+                        value=self._aggregators["value"].result(window["value"]),
                     ),
                 )