fix fv offline

HaoXuAI · HaoXuAI · commit a09b7e5df6d8 · 2025-04-20T23:02:30.000-07:00
Signed-off-by: HaoXuAI &lt;sduxuhao@gmail.com&gt;
diff --git a/protos/feast/core/FeatureView.proto b/protos/feast/core/FeatureView.proto
@@ -74,7 +74,11 @@ message FeatureViewSpec {
     DataSource stream_source = 9;
 
     // Whether these features should be served online or not
+    // This is also used to determine whether the features should be written to the online store
     bool online = 8;
+
+    // Whether these features should be written to the offline store
+    bool offline = 13;
 }
 
 message FeatureViewMeta {
diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py
@@ -236,6 +236,7 @@ def __copy__(self):
             schema=self.schema,
             tags=self.tags,
             online=self.online,
+            offline=self.offline,
         )
 
         # This is deliberately set outside of the FV initialization as we do not have the Entity objects.
@@ -258,6 +259,7 @@ def __eq__(self, other):
             sorted(self.entities) != sorted(other.entities)
             or self.ttl != other.ttl
             or self.online != other.online
+            or self.offline != other.offline
             or self.batch_source != other.batch_source
             or self.stream_source != other.stream_source
             or sorted(self.entity_columns) != sorted(other.entity_columns)
diff --git a/sdk/python/feast/infra/compute_engines/local/nodes.py b/sdk/python/feast/infra/compute_engines/local/nodes.py
@@ -1,8 +1,9 @@
 from datetime import datetime, timedelta
-from typing import Optional
+from typing import Optional, Union
 
 import pyarrow as pa
 
+from feast import BatchFeatureView, StreamFeatureView
 from feast.data_source import DataSource
 from feast.infra.compute_engines.dag.context import ExecutionContext
 from feast.infra.compute_engines.local.arrow_table_value import ArrowTableValue
@@ -208,7 +209,9 @@ def execute(self, context: ExecutionContext) -> ArrowTableValue:
 
 
 class LocalOutputNode(LocalNode):
-    def __init__(self, name: str, feature_view):
+    def __init__(
+        self, name: str, feature_view: Union[BatchFeatureView, StreamFeatureView]
+    ):
         super().__init__(name)
         self.feature_view = feature_view
 
diff --git a/sdk/python/feast/infra/compute_engines/spark/feature_builder.py b/sdk/python/feast/infra/compute_engines/spark/feature_builder.py
@@ -73,7 +73,8 @@ def build_transformation_node(self, input_node):
         return node
 
     def build_output_nodes(self, input_node):
-        node = SparkWriteNode("output", input_node, self.feature_view)
+        node = SparkWriteNode("output", self.feature_view)
+        node.add_input(input_node)
         self.nodes.append(node)
         return node
 
diff --git a/sdk/python/feast/infra/compute_engines/spark/nodes.py b/sdk/python/feast/infra/compute_engines/spark/nodes.py
@@ -271,11 +271,9 @@ class SparkWriteNode(DAGNode):
     def __init__(
         self,
         name: str,
-        input_node: DAGNode,
         feature_view: Union[BatchFeatureView, StreamFeatureView],
     ):
         super().__init__(name)
-        self.add_input(input_node)
         self.feature_view = feature_view
 
     def execute(self, context: ExecutionContext) -> DAGValue:
@@ -286,8 +284,6 @@ def execute(self, context: ExecutionContext) -> DAGValue:
 
         # ✅ 1. Write to online or offline store (if enabled)
         if self.feature_view.online or self.feature_view.offline:
-            print("Spark DF count:", spark_df.count())
-            print("Num partitions:", spark_df.rdd.getNumPartitions())
             spark_df.mapInArrow(
                 lambda x: map_in_arrow(x, serialized_artifacts), spark_df.schema
             ).count()
diff --git a/sdk/python/feast/infra/compute_engines/spark/utils.py b/sdk/python/feast/infra/compute_engines/spark/utils.py
@@ -37,6 +37,7 @@ def map_in_arrow(
             offline_store,
             repo_config,
         ) = serialized_artifacts.unserialize()
+        print("write_feature_view", feature_view)
 
         if feature_view.online:
             join_key_to_value_type = {
@@ -55,6 +56,7 @@ def map_in_arrow(
                 progress=lambda x: None,
             )
         if feature_view.offline:
+            print("offline_to_write", table)
             offline_store.offline_write_batch(
                 config=repo_config,
                 feature_view=feature_view,
diff --git a/sdk/python/tests/integration/compute_engines/spark/test_compute.py b/sdk/python/tests/integration/compute_engines/spark/test_compute.py
@@ -33,7 +33,6 @@
 from tests.integration.feature_repos.universal.online_store.redis import (
     RedisOnlineStoreCreator,
 )
-from tests.utils.e2e_test_validation import _check_offline_and_online_features
 
 now = datetime.now()
 today = datetime.today()
@@ -189,6 +188,20 @@ def transform_feature(df: DataFrame) -> DataFrame:
 
 @pytest.mark.integration
 def test_spark_compute_engine_materialize():
+    """
+    Test the SparkComputeEngine materialize method.
+    For the current feature view driver_hourly_stats, The below execution plan:
+    1. feature data from create_feature_dataset
+    2. filter by start_time and end_time, that is, the last 2 days
+        for the driver_id 1001, the data left is row 0
+        for the driver_id 1002, the data left is row 2
+    3. apply the transform_feature function to the data
+        for all features, the value is multiplied by 2
+    4. write the data to the online store and offline store
+
+    Returns:
+
+    """
     spark_environment = create_spark_environment()
     fs = spark_environment.feature_store
     registry = fs.registry
@@ -213,7 +226,7 @@ def transform_feature(df: DataFrame) -> DataFrame:
             Field(name="driver_id", dtype=Int32),
         ],
         online=True,
-        offline=False,
+        offline=True,
         source=data_source,
     )
 
@@ -244,18 +257,62 @@ def tqdm_builder(length):
 
         assert spark_materialize_job.status() == MaterializationJobStatus.SUCCEEDED
 
-        _check_offline_and_online_features(
+        _check_online_features(
             fs=fs,
-            fv=driver_stats_fv,
-            driver_id=1,
-            event_timestamp=now,
-            expected_value=0.3,
+            driver_id=1001,
+            feature="driver_hourly_stats:conv_rate",
+            expected_value=1.6,
             full_feature_names=True,
-            check_offline_store=True,
+        )
+
+        entity_df = create_entity_df()
+
+        _check_offline_features(
+            fs=fs,
+            feature="driver_hourly_stats:conv_rate",
+            entity_df=entity_df,
+            expected_value=1.6,
         )
     finally:
         spark_environment.teardown()
 
 
+def _check_online_features(
+    fs,
+    driver_id,
+    feature,
+    expected_value,
+    full_feature_names: bool = True,
+):
+    online_response = fs.get_online_features(
+        features=[feature],
+        entity_rows=[{"driver_id": driver_id}],
+        full_feature_names=full_feature_names,
+    ).to_dict()
+
+    feature_ref = "__".join(feature.split(":"))
+
+    assert len(online_response["driver_id"]) == 1
+    assert online_response["driver_id"][0] == driver_id
+    assert abs(online_response[feature_ref][0] - expected_value < 1e-6), (
+        "Transformed result"
+    )
+
+
+def _check_offline_features(
+    fs,
+    feature,
+    entity_df,
+    expected_value,
+):
+    offline_df = fs.get_historical_features(
+        entity_df=entity_df,
+        features=[feature],
+    ).to_df()
+
+    assert len(offline_df) == 2
+    assert offline_df["driver_id"].to_list() == [1001, 1002]
+
+
 if __name__ == "__main__":
     test_spark_compute_engine_get_historical_features()
diff --git a/sdk/python/tests/unit/infra/compute_engines/local/test_nodes.py b/sdk/python/tests/unit/infra/compute_engines/local/test_nodes.py
@@ -194,7 +194,7 @@ def test_local_output_node():
     context = create_context(
         node_outputs={"source": ArrowTableValue(pa.Table.from_pandas(sample_df))}
     )
-    node = LocalOutputNode("output")
+    node = LocalOutputNode("output", MagicMock())
     node.add_input(MagicMock())
     node.inputs[0].name = "source"
     result = node.execute(context)

Original file line number	Diff line number	Diff line change
`@@ -194,7 +194,7 @@ def test_local_output_node():`
`194`	`194`	`context = create_context(`
`195`	`195`	`node_outputs={"source": ArrowTableValue(pa.Table.from_pandas(sample_df))}`
`196`	`196`	`)`
`197`		`- node = LocalOutputNode("output")`
	`197`	`+ node = LocalOutputNode("output", MagicMock())`
`198`	`198`	`node.add_input(MagicMock())`
`199`	`199`	`node.inputs[0].name = "source"`
`200`	`200`	`result = node.execute(context)`