add aggregation to odfv

HaoXuAI · HaoXuAI · commit a101f0fab950 · 2025-10-14T15:28:05.000-07:00
Signed-off-by: hao-xu5 &lt;hxu44@apple.com&gt;
diff --git a/sdk/python/feast/infra/compute_engines/backends/__init__.py b/sdk/python/feast/infra/compute_engines/backends/__init__.py
diff --git a/sdk/python/feast/infra/compute_engines/backends/base.py b/sdk/python/feast/infra/compute_engines/backends/base.py
@@ -0,0 +1,79 @@
+from abc import ABC, abstractmethod
+from datetime import timedelta
+
+
+class DataFrameBackend(ABC):
+    """
+    Abstract interface for DataFrame operations used by the LocalComputeEngine.
+
+    This interface defines the contract for implementing pluggable DataFrame backends
+    such as Pandas, Polars, or DuckDB. Each backend must support core table operations
+    such as joins, filtering, aggregation, conversion to/from Arrow, and deduplication.
+
+    The purpose of this abstraction is to allow seamless swapping of execution backends
+    without changing DAGNode or ComputeEngine logic. All nodes operate on pyarrow.Table
+    as the standard input/output format, while the backend defines how the computation
+    is actually performed.
+
+    Expected implementations include:
+    - PandasBackend
+    - PolarsBackend
+    - DuckDBBackend (future)
+
+    Methods
+    -------
+    from_arrow(table: pa.Table) -> Any
+        Convert a pyarrow.Table to the backend-native DataFrame format.
+
+    to_arrow(df: Any) -> pa.Table
+        Convert a backend-native DataFrame to pyarrow.Table.
+
+    join(left: Any, right: Any, on: List[str], how: str) -> Any
+        Join two dataframes on specified keys with given join type.
+
+    groupby_agg(df: Any, group_keys: List[str], agg_ops: Dict[str, Tuple[str, str]]) -> Any
+        Group and aggregate the dataframe. `agg_ops` maps output column names
+        to (aggregation function, source column name) pairs.
+
+    filter(df: Any, expr: str) -> Any
+        Apply a filter expression (string-based) to the DataFrame.
+
+    to_timedelta_value(delta: timedelta) -> Any
+        Convert a Python timedelta object to a backend-compatible value
+        that can be subtracted from a timestamp column.
+
+    drop_duplicates(df: Any, keys: List[str], sort_by: List[str], ascending: bool = False) -> Any
+        Deduplicate the DataFrame by key columns, keeping the first row
+        by descending or ascending sort order.
+
+    rename_columns(df: Any, columns: Dict[str, str]) -> Any
+        Rename columns in the DataFrame according to the provided mapping.
+    """
+
+    @abstractmethod
+    def columns(self, df): ...
+
+    @abstractmethod
+    def from_arrow(self, table): ...
+
+    @abstractmethod
+    def join(self, left, right, on, how): ...
+
+    @abstractmethod
+    def groupby_agg(self, df, group_keys, agg_ops): ...
+
+    @abstractmethod
+    def filter(self, df, expr): ...
+
+    @abstractmethod
+    def to_arrow(self, df): ...
+
+    @abstractmethod
+    def to_timedelta_value(self, delta: timedelta): ...
+
+    @abstractmethod
+    def drop_duplicates(self, df, keys, sort_by, ascending: bool = False):
+        pass
+
+    @abstractmethod
+    def rename_columns(self, df, columns: dict[str, str]): ...
diff --git a/sdk/python/feast/infra/compute_engines/backends/factory.py b/sdk/python/feast/infra/compute_engines/backends/factory.py
@@ -0,0 +1,53 @@
+from typing import Optional
+
+import pandas as pd
+import pyarrow
+
+from feast.infra.compute_engines.backends.base import DataFrameBackend
+from feast.infra.compute_engines.backends.pandas_backend import PandasBackend
+
+
+class BackendFactory:
+    """
+    Factory class for constructing DataFrameBackend implementations based on backend name
+    or runtime entity_df type.
+    """
+
+    @staticmethod
+    def from_name(name: str) -> DataFrameBackend:
+        if name == "pandas":
+            return PandasBackend()
+        if name == "polars":
+            return BackendFactory._get_polars_backend()
+        raise ValueError(f"Unsupported backend name: {name}")
+
+    @staticmethod
+    def infer_from_entity_df(entity_df) -> Optional[DataFrameBackend]:
+        if (
+            not entity_df
+            or isinstance(entity_df, pyarrow.Table)
+            or isinstance(entity_df, pd.DataFrame)
+        ):
+            return PandasBackend()
+
+        if BackendFactory._is_polars(entity_df):
+            return BackendFactory._get_polars_backend()
+        return None
+
+    @staticmethod
+    def _is_polars(entity_df) -> bool:
+        try:
+            import polars as pl
+        except ImportError:
+            raise ImportError(
+                "Polars is not installed. Please install it to use Polars backend."
+            )
+        return isinstance(entity_df, pl.DataFrame)
+
+    @staticmethod
+    def _get_polars_backend():
+        from feast.infra.compute_engines.backends.polars_backend import (
+            PolarsBackend,
+        )
+
+        return PolarsBackend()
diff --git a/sdk/python/feast/infra/compute_engines/backends/pandas_backend.py b/sdk/python/feast/infra/compute_engines/backends/pandas_backend.py
@@ -0,0 +1,46 @@
+from datetime import timedelta
+
+import pandas as pd
+import pyarrow as pa
+
+from feast.infra.compute_engines.backends.base import DataFrameBackend
+
+
+class PandasBackend(DataFrameBackend):
+    def columns(self, df):
+        return df.columns.tolist()
+
+    def from_arrow(self, table):
+        return table.to_pandas()
+
+    def join(self, left, right, on, how):
+        return left.merge(right, on=on, how=how)
+
+    def groupby_agg(self, df, group_keys, agg_ops):
+        return (
+            df.groupby(group_keys)
+            .agg(
+                **{
+                    alias: pd.NamedAgg(column=col, aggfunc=func)
+                    for alias, (func, col) in agg_ops.items()
+                }
+            )
+            .reset_index()
+        )
+
+    def filter(self, df, expr):
+        return df.query(expr)
+
+    def to_arrow(self, df):
+        return pa.Table.from_pandas(df)
+
+    def to_timedelta_value(self, delta: timedelta):
+        return pd.to_timedelta(delta)
+
+    def drop_duplicates(self, df, keys, sort_by, ascending: bool = False):
+        return df.sort_values(by=sort_by, ascending=ascending).drop_duplicates(
+            subset=keys
+        )
+
+    def rename_columns(self, df, columns: dict[str, str]):
+        return df.rename(columns=columns)
diff --git a/sdk/python/feast/infra/compute_engines/backends/polars_backend.py b/sdk/python/feast/infra/compute_engines/backends/polars_backend.py
@@ -0,0 +1,47 @@
+from datetime import timedelta
+
+import polars as pl
+import pyarrow as pa
+
+from feast.infra.compute_engines.backends.base import DataFrameBackend
+
+
+class PolarsBackend(DataFrameBackend):
+    def columns(self, df: pl.DataFrame) -> list[str]:
+        return df.columns
+
+    def from_arrow(self, table: pa.Table) -> pl.DataFrame:
+        return pl.from_arrow(table)
+
+    def to_arrow(self, df: pl.DataFrame) -> pa.Table:
+        return df.to_arrow()
+
+    def join(self, left: pl.DataFrame, right: pl.DataFrame, on, how) -> pl.DataFrame:
+        return left.join(right, on=on, how=how)
+
+    def groupby_agg(self, df: pl.DataFrame, group_keys, agg_ops) -> pl.DataFrame:
+        agg_exprs = [
+            getattr(pl.col(col), func)().alias(alias)
+            for alias, (func, col) in agg_ops.items()
+        ]
+        return df.groupby(group_keys).agg(agg_exprs)
+
+    def filter(self, df: pl.DataFrame, expr: str) -> pl.DataFrame:
+        return df.filter(pl.sql_expr(expr))
+
+    def to_timedelta_value(self, delta: timedelta):
+        return pl.duration(milliseconds=delta.total_seconds() * 1000)
+
+    def drop_duplicates(
+        self,
+        df: pl.DataFrame,
+        keys: list[str],
+        sort_by: list[str],
+        ascending: bool = False,
+    ) -> pl.DataFrame:
+        return df.sort(by=sort_by, descending=not ascending).unique(
+            subset=keys, keep="first"
+        )
+
+    def rename_columns(self, df: pl.DataFrame, columns: dict[str, str]) -> pl.DataFrame:
+        return df.rename(columns)
diff --git a/sdk/python/feast/infra/compute_engines/local/compute.py b/sdk/python/feast/infra/compute_engines/local/compute.py
@@ -14,8 +14,8 @@
 from feast.infra.common.retrieval_task import HistoricalRetrievalTask
 from feast.infra.compute_engines.base import ComputeEngine
 from feast.infra.compute_engines.dag.context import ExecutionContext
-from feast.infra.compute_engines.local.backends.base import DataFrameBackend
-from feast.infra.compute_engines.local.backends.factory import BackendFactory
+from feast.infra.compute_engines.backends.base import DataFrameBackend
+from feast.infra.compute_engines.backends.factory import BackendFactory
 from feast.infra.compute_engines.local.feature_builder import LocalFeatureBuilder
 from feast.infra.compute_engines.local.job import (
     LocalMaterializationJob,
diff --git a/sdk/python/feast/infra/compute_engines/local/feature_builder.py b/sdk/python/feast/infra/compute_engines/local/feature_builder.py
@@ -3,7 +3,7 @@
 from feast.infra.common.materialization_job import MaterializationTask
 from feast.infra.common.retrieval_task import HistoricalRetrievalTask
 from feast.infra.compute_engines.feature_builder import FeatureBuilder
-from feast.infra.compute_engines.local.backends.base import DataFrameBackend
+from feast.infra.compute_engines.backends.base import DataFrameBackend
 from feast.infra.compute_engines.local.nodes import (
     LocalAggregationNode,
     LocalDedupNode,
diff --git a/sdk/python/feast/infra/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py
@@ -9,7 +9,7 @@
 from feast.infra.compute_engines.dag.model import DAGFormat
 from feast.infra.compute_engines.dag.node import DAGNode
 from feast.infra.compute_engines.local.arrow_table_value import ArrowTableValue
-from feast.infra.compute_engines.local.backends.base import DataFrameBackend
+from feast.infra.compute_engines.backends.base import DataFrameBackend
 from feast.infra.compute_engines.local.local_node import LocalNode
 from feast.infra.compute_engines.utils import (
     create_offline_store_retrieval_job,
diff --git a/sdk/python/feast/utils.py b/sdk/python/feast/utils.py
@@ -33,6 +33,7 @@
     RequestDataNotFoundInEntityRowsException,
 )
 from feast.field import Field
+from feast.infra.compute_engines.backends.pandas_backend import PandasBackend
 from feast.infra.key_encoding_utils import deserialize_entity_key
 from feast.protos.feast.serving.ServingService_pb2 import (
     FieldStatus,
@@ -561,6 +562,72 @@ def construct_response_feature_vector(
     )
 
 
+def _get_aggregate_operations(agg_specs) -> dict:
+    """
+    Convert Aggregation specs to agg_ops format for PandasBackend.
+
+    Reused from LocalFeatureBuilder logic.
+    TODO: This logic is duplicated from LocalFeatureBuilder._get_aggregate_operations().
+    Consider refactoring to a shared utility module in the future.
+    """
+    agg_ops = {}
+    for agg in agg_specs:
+        if agg.time_window is not None:
+            raise ValueError(
+                "Time window aggregation is not supported in online serving."
+            )
+        alias = f"{agg.function}_{agg.column}"
+        agg_ops[alias] = (agg.function, agg.column)
+    return agg_ops
+
+
+def _apply_aggregations_to_response(
+    response_data: Union[pyarrow.Table, Dict[str, List[Any]]],
+    aggregations,
+    group_keys: List[str],
+    mode: str,
+) -> Union[pyarrow.Table, Dict[str, List[Any]]]:
+    """
+    Apply aggregations using PandasBackend.
+
+    Args:
+        response_data: Either a pyarrow.Table or dict of lists containing the data
+        aggregations: List of Aggregation objects to apply
+        group_keys: List of column names to group by
+        mode: Transformation mode ("python", "pandas", or "substrait")
+
+    Returns:
+        Aggregated data in the same format as input
+
+    TODO: Consider refactoring to support backends other than pandas in the future.
+    """
+    if not aggregations:
+        return response_data
+
+    backend = PandasBackend()
+
+    # Convert to pandas DataFrame
+    if isinstance(response_data, dict):
+        df = pd.DataFrame(response_data)
+    else:  # pyarrow.Table
+        df = backend.from_arrow(response_data)
+
+    if df.empty:
+        return response_data
+
+    # Convert aggregations to agg_ops format
+    agg_ops = _get_aggregate_operations(aggregations)
+
+    # Apply aggregations using PandasBackend
+    result_df = backend.groupby_agg(df, group_keys, agg_ops)
+
+    # Convert back to original format
+    if mode == "python":
+        return {col: result_df[col].tolist() for col in result_df.columns}
+    else:  # pandas or substrait
+        return backend.to_arrow(result_df)
+
+
 def _augment_response_with_on_demand_transforms(
     online_features_response: GetOnlineFeaturesResponse,
     feature_refs: List[str],
@@ -605,6 +672,28 @@ def _augment_response_with_on_demand_transforms(
     for odfv_name, _feature_refs in odfv_feature_refs.items():
         odfv = requested_odfv_map[odfv_name]
         if not odfv.write_to_online_store:
+            # Apply aggregations BEFORE transformation if defined
+            if odfv.aggregations:
+                if odfv.mode == "python":
+                    if initial_response_dict is None:
+                        initial_response_dict = initial_response.to_dict()
+                    initial_response_dict = _apply_aggregations_to_response(
+                        initial_response_dict,
+                        odfv.aggregations,
+                        odfv.entities,
+                        odfv.mode,
+                    )
+                elif odfv.mode in {"pandas", "substrait"}:
+                    if initial_response_arrow is None:
+                        initial_response_arrow = initial_response.to_arrow()
+                    initial_response_arrow = _apply_aggregations_to_response(
+                        initial_response_arrow,
+                        odfv.aggregations,
+                        odfv.entities,
+                        odfv.mode,
+                    )
+
+            # Apply transformation
             if odfv.mode == "python":
                 if initial_response_dict is None:
                     initial_response_dict = initial_response.to_dict()
