add compute engine

Signed-off-by: HaoXuAI <sduxuhao@gmail.com>
franciscojavierarceo · Apr 16, 2025 · Apr 4, 2025 · Apr 4, 2025 · Apr 4, 2025 · Apr 4, 2025
commit 1183a9a314cbdebf2c0d45b8b33f5d2add4a1162
@@ -55,6 +55,7 @@ class BatchFeatureView(FeatureView):
     entity_columns: List[Field]
     features: List[Field]
     online: bool
+    offline: bool
     description: str
     tags: Dict[str, str]
     owner: str
@@ -74,6 +75,7 @@ def __init__(
         ttl: Optional[timedelta] = None,
         tags: Optional[Dict[str, str]] = None,
         online: bool = True,
+        offline: bool = True,
         description: str = "",
         owner: str = "",
         schema: Optional[List[Field]] = None,
@@ -110,6 +112,7 @@ def __init__(
             ttl=ttl,
             tags=tags,
             online=online,
+            offline=offline,
             description=description,
             owner=owner,
             schema=schema,

@@ -93,6 +93,7 @@ class FeatureView(BaseFeatureView):
     entity_columns: List[Field]
     features: List[Field]
     online: bool
+    offline: bool
     description: str
     tags: Dict[str, str]
     owner: str
@@ -107,6 +108,7 @@ def __init__(
         entities: Optional[List[Entity]] = None,
         ttl: Optional[timedelta] = timedelta(days=0),
         online: bool = True,
+        offline: bool = False,
         description: str = "",
         tags: Optional[Dict[str, str]] = None,
         owner: str = "",
@@ -127,6 +129,8 @@ def __init__(
                 can result in extremely computationally intensive queries.
             online (optional): A boolean indicating whether online retrieval is enabled for
                 this feature view.
+            offline (optional): A boolean indicating whether write to offline store is enabled for
+                this feature view.
             description (optional): A human-readable description.
             tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
             owner (optional): The owner of the feature view, typically the email of the

@@ -0,0 +1,49 @@
+from abc import ABC
+from dataclasses import dataclass
+from typing import Union, List
+
+import pandas as pd
+import pyarrow as pa
+
+from feast import RepoConfig, BatchFeatureView, StreamFeatureView
+from feast.infra.materialization.batch_materialization_engine import MaterializationTask, MaterializationJob
+from feast.infra.offline_stores.offline_store import OfflineStore
+from feast.infra.online_stores.online_store import OnlineStore
+from feast.infra.registry.registry import Registry
+
+
+@dataclass
+class HistoricalRetrievalTask:
+    entity_df: Union[pd.DataFrame, str]
+    feature_views: List[Union[BatchFeatureView, StreamFeatureView]]
+    full_feature_names: bool
+    registry: Registry
+    config: RepoConfig
+
+
+class ComputeEngine(ABC):
+    """
+    The interface that Feast uses to control the compute system that handles materialization and get_historical_features.
+    """
+
+    def __init__(
+            self,
+            *,
+            registry: Registry,
+            repo_config: RepoConfig,
+            offline_store: OfflineStore,
+            online_store: OnlineStore,
+            **kwargs,
+    ):
+        self.registry = registry
+        self.repo_config = repo_config
+        self.offline_store = offline_store
+        self.online_store = online_store
+
+    def materialize(self,
+                    task: MaterializationTask) -> MaterializationJob:
+        raise NotImplementedError
+
+    def get_historical_features(self,
+                                task: HistoricalRetrievalTask) -> pa.Table:
+        raise NotImplementedError
diff --git a/sdk/python/feast/infra/compute_engines/dag/builder.py b/sdk/python/feast/infra/compute_engines/dag/builder.py
@@ -0,0 +1,70 @@
+from abc import ABC, abstractmethod
+from typing import Union
+
+from feast import BatchFeatureView, StreamFeatureView
+from feast.infra.compute_engines.dag.plan import ExecutionPlan
+from feast.infra.compute_engines.base import HistoricalRetrievalTask
+from feast.infra.materialization.batch_materialization_engine import MaterializationTask
+
+
+class DAGBuilder(ABC):
+    def __init__(self,
+                 feature_view: Union[BatchFeatureView, StreamFeatureView],
+                 task: Union[MaterializationTask, HistoricalRetrievalTask]
+                 ):
+        self.feature_view = feature_view
+        self.task = task
+        self.nodes = []
+
+    @abstractmethod
+    def build_source_node(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_aggregation_node(self,
+                               input_node):
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_join_node(self,
+                        input_node):
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_transformation_node(self,
+                                  input_node):
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_output_nodes(self,
+                           input_node):
+        raise NotImplementedError
+
+    @abstractmethod
+    def build_validation_node(self,
+                              input_node):
+        raise
+
+    def build(self) -> ExecutionPlan:
+        last_node = self.build_source_node()
+
+        if getattr(self.feature_view.transformation, "requires_aggregation", False):
+            last_node = self.build_aggregation_node(last_node)
+
+        if self._should_join():
+            last_node = self.build_join_node(last_node)
+
+        if self.feature_view.transformation:
+            last_node = self.build_transformation_node(last_node)
+
+        if getattr(self.feature_view, "enable_validation", False):
+            last_node = self.build_validation_node(last_node)
+
+        self.build_output_nodes(last_node)
+        return ExecutionPlan(self.nodes)
+
+    def _should_join(self):
+        return (
+                self.feature_view.compute_config.join_strategy == "engine"
+                or self.task.config.compute_engine.get("point_in_time_join") == "engine"
+        )
@@ -0,0 +1,27 @@
+from enum import Enum
+from dataclasses import dataclass, field
+import pandas as pd
+from typing import Union, List, Dict
+
+from feast.entity import Entity
+from feast.infra.offline_stores.offline_store import OfflineStore
+from feast.infra.online_stores.online_store import OnlineStore
+from feast.repo_config import RepoConfig
+from feast.infra.compute_engines.dag.value import DAGValue
+
+
+class DAGFormat(str, Enum):
+    SPARK = "spark"
+    PANDAS = "pandas"
+    ARROW = "arrow"
+
+
+@dataclass
+class ExecutionContext:
+    project: str
+    repo_config: RepoConfig
+    offline_store: OfflineStore
+    online_store: OnlineStore
+    entity_defs: List[Entity]
+    entity_df: Union[pd.DataFrame, None] = None
+    node_outputs: Dict[str, DAGValue] = field(default_factory=dict)
@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+from feast.infra.compute_engines.dag.model import ExecutionContext
+from infra.compute_engines.dag.value import DAGValue
+
+
+class DAGNode(ABC):
+    name: str
+    inputs: List["DAGNode"]
+    outputs: List["DAGNode"]
+
+    def __init__(self,
+                 name: str):
+        self.name = name
+        self.inputs = []
+        self.outputs = []
+
+    def add_input(self,
+                  node: "DAGNode"):
+        if node in self.inputs:
+            raise ValueError(f"Input node {node.name} already added to {self.name}")
+        self.inputs.append(node)
+        node.outputs.append(self)
+
+    @abstractmethod
+    def execute(self,
+                context: ExecutionContext) -> DAGValue:
+        ...
@@ -0,0 +1,63 @@
+from typing import List
+
+from feast.infra.compute_engines.dag.node import DAGNode
+from feast.infra.compute_engines.dag.value import DAGValue
+from feast.infra.compute_engines.dag.model import ExecutionContext
+
+
+class ExecutionPlan:
+    """
+    ExecutionPlan represents an ordered sequence of DAGNodes that together define
+    a data processing pipeline for feature materialization or historical retrieval.
+
+    This plan is constructed as a topological sort of the DAG — meaning that each
+    node appears after all its input dependencies. The plan is executed in order,
+    caching intermediate results (`DAGValue`) so that each node can reuse outputs
+    from upstream nodes without recomputation.
+
+    Key Concepts:
+    - DAGNode: Each node performs a specific logical step (e.g., read, aggregate, join).
+    - DAGValue: Output of a node, includes data (e.g., Spark DataFrame) and metadata.
+    - ExecutionContext: Contains runtime information (config, registry, stores, entity_df).
+    - node_outputs: A cache of intermediate results keyed by node name.
+
+    Usage:
+        plan = ExecutionPlan(dag_nodes)
+        result = plan.execute(context)
+
+    This design enables modular compute backends (e.g., Spark, Pandas, Arrow), where
+    each node defines its execution logic independently while benefiting from shared
+    execution orchestration, caching, and context injection.
+
+    Example:
+        DAG:
+            ReadNode -> AggregateNode -> JoinNode -> TransformNode -> WriteNode
+
+        Execution proceeds step by step, passing intermediate DAGValues through
+        the plan while respecting node dependencies and formats.
+
+    This approach is inspired by execution DAGs in systems like Apache Spark,
+    Apache Beam, and Dask — but specialized for Feast’s feature computation domain.
+    """
+    def __init__(self, nodes: List[DAGNode]):
+        self.nodes = nodes
+
+    def execute(self, context: ExecutionContext) -> DAGValue:
+        node_outputs: dict[str, DAGValue] = {}
+
+        for node in self.nodes:
+            # Gather input values
+            for input_node in node.inputs:
+                if input_node.name not in node_outputs:
+                    node_outputs[input_node.name] = input_node.execute(context)
+
+            # Execute this node
+            output = node.execute(context)
+            node_outputs[node.name] = output
+
+        # Inject into context for downstream access
+        context.node_outputs = node_outputs
+
+        # Return output of final node
+        return node_outputs[self.nodes[-1].name]
+
@@ -0,0 +1,18 @@
+from typing import Any, Optional
+
+from feast.infra.compute_engines.dag.model import DAGFormat
+
+
+class DAGValue:
+    def __init__(self,
+                 data: Any,
+                 format: DAGFormat,
+                 metadata: Optional[dict] = None):
+        self.data = data
+        self.format = format
+        self.metadata = metadata or {}
+
+    def assert_format(self,
+                      expected: DAGFormat):
+        if self.format != expected:
+            raise ValueError(f"Expected format {expected}, but got {self.format}")
@@ -0,0 +1,81 @@
+import pyarrow as pa
+from feast.infra.compute_engines.base import ComputeEngine, HistoricalRetrievalTask
+from feast.infra.compute_engines.spark.spark_dag_builder import SparkDAGBuilder
+from feast.infra.materialization.batch_materialization_engine import MaterializationTask, MaterializationJob, \
+    MaterializationJobStatus
+from feast.infra.materialization.contrib.spark.spark_materialization_engine import SparkMaterializationJob
+from feast.infra.compute_engines.dag.model import ExecutionContext
+
+
+class SparkComputeEngine(ComputeEngine):
+    def materialize(self, task: MaterializationTask) -> MaterializationJob:
+        job_id = f"{task.feature_view.name}-{task.start_time}-{task.end_time}"
+
+        try:
+            # ✅ 1. Build typed execution context
+            entities = []
+            for entity_name in task.feature_view.entities:
+                entities.append(self.registry.get_entity(entity_name, task.project))
+
+            context = ExecutionContext(
+                project=task.project,
+                repo_config=self.repo_config,
+                offline_store=self.offline_store,
+                online_store=self.online_store,
+                entity_defs=entities
+            )
+
+            # ✅ 2. Construct DAG and run it
+            builder = SparkDAGBuilder(
+                feature_view=task.feature_view,
+                task=task,
+            )
+            plan = builder.build()
+            plan.execute(context)
+
+            # ✅ 3. Report success
+            return SparkMaterializationJob(
+                job_id=job_id,
+                status=MaterializationJobStatus.SUCCEEDED
+            )
+
+        except Exception as e:
+            # 🛑 Handle failure
+            return SparkMaterializationJob(
+                job_id=job_id,
+                status=MaterializationJobStatus.ERROR,
+                error=e
+            )
+
+    def get_historical_features(self, task: HistoricalRetrievalTask) -> pa.Table:
+        # ✅ 1. Validate input
+        assert len(task.feature_views) == 1, "Multi-view support not yet implemented"
+        feature_view = task.feature_views[0]
+
+        if isinstance(task.entity_df, str):
+            raise NotImplementedError("SQL-based entity_df is not yet supported in DAG")
+
+        # ✅ 2. Build typed execution context
+        entity_defs = [
+            task.registry.get_entity(name, task.config.project)
+            for name in feature_view.entities
+        ]
+
+        context = ExecutionContext(
+            project=task.config.project,
+            repo_config=task.config,
+            offline_store=self.offline_store,
+            online_store=self.online_store,
+            entity_defs=entity_defs,
+            entity_df=task.entity_df,
+        )
+
+        # ✅ 3. Construct and execute DAG
+        builder = SparkDAGBuilder(feature_view=feature_view, task=task)
+        plan = builder.build()
+
+        result = plan.execute(context=context)
+        spark_df = result.data  # should be a Spark DataFrame
+
+        # ✅ 4. Return as Arrow
+        return spark_df.toPandas().to_arrow()