update

Signed-off-by: HaoXuAI <sduxuhao@gmail.com>
franciscojavierarceo · Apr 20, 2025 · Apr 18, 2025 · Apr 18, 2025 · Apr 18, 2025 · Apr 18, 2025
commit b62fd1a8be2cb8d613d36b11167b0e553bd38871
@@ -9,11 +9,9 @@
     SparkAggregationNode,
     SparkDedupNode,
     SparkFilterNode,
-    SparkHistoricalRetrievalReadNode,
     SparkJoinNode,
-    SparkMaterializationReadNode,
     SparkTransformationNode,
-    SparkWriteNode,
+    SparkWriteNode, SparkReadNode,
 )
 
 
@@ -27,12 +25,10 @@ def __init__(
         self.spark_session = spark_session
 
     def build_source_node(self):
-        if isinstance(self.task, MaterializationTask):
-            node = SparkMaterializationReadNode("source", self.task)
-        else:
-            node = SparkHistoricalRetrievalReadNode(
-                "source", self.task, self.spark_session
-            )
+        source = self.feature_view.batch_source
+        start_time = self.task.start_time
+        end_time = self.task.end_time
+        node = SparkReadNode("source", source, start_time, end_time)
         self.nodes.append(node)
         return node
 

@@ -4,6 +4,7 @@
 from pyspark.sql import DataFrame, SparkSession, Window
 from pyspark.sql import functions as F
 
+from data_source import DataSource
 from feast import BatchFeatureView, StreamFeatureView
 from feast.aggregation import Aggregation
 from feast.infra.common.materialization_job import MaterializationTask
@@ -49,18 +50,21 @@ def rename_entity_ts_column(
     return entity_df
 
 
-class SparkMaterializationReadNode(DAGNode):
+class SparkReadNode(DAGNode):
     def __init__(
-        self, name: str, task: Union[MaterializationTask, HistoricalRetrievalTask]
+        self,
+            name: str,
+            source: DataSource,
+            start_time: Optional[timedelta] = None,
+            end_time: Optional[timedelta] = None,
     ):
         super().__init__(name)
-        self.task = task
+        self.source = source
+        self.start_time = start_time
+        self.end_time = end_time
 
     def execute(self, context: ExecutionContext) -> DAGValue:
         offline_store = context.offline_store
-        start_time = self.task.start_time
-        end_time = self.task.end_time
-
         (
             join_key_columns,
             feature_name_columns,
@@ -69,15 +73,14 @@ def execute(self, context: ExecutionContext) -> DAGValue:
         ) = context.column_info
 
         # 📥 Reuse Feast's robust query resolver
-        retrieval_job = offline_store.pull_latest_from_table_or_query(
+        retrieval_job = offline_store.pull_all_from_table_or_query(
             config=context.repo_config,
-            data_source=self.task.feature_view.batch_source,
+            data_source=self.source,
             join_key_columns=join_key_columns,
             feature_name_columns=feature_name_columns,
             timestamp_field=timestamp_field,
-            created_timestamp_column=created_timestamp_column,
-            start_date=start_time,
-            end_date=end_time,
+            start_date=self.start_time,
+            end_date=self.end_time,
         )
         spark_df = cast(SparkRetrievalJob, retrieval_job).to_spark_df()
 
@@ -88,74 +91,8 @@ def execute(self, context: ExecutionContext) -> DAGValue:
                 "source": "feature_view_batch_source",
                 "timestamp_field": timestamp_field,
                 "created_timestamp_column": created_timestamp_column,
-                "start_date": start_time,
-                "end_date": end_time,
-            },
-        )
-
-
-class SparkHistoricalRetrievalReadNode(DAGNode):
-    def __init__(
-        self, name: str, task: HistoricalRetrievalTask, spark_session: SparkSession
-    ):
-        super().__init__(name)
-        self.task = task
-        self.spark_session = spark_session
-
-    def execute(self, context: ExecutionContext) -> DAGValue:
-        """
-        Read data from the offline store on the Spark engine.
-        TODO: Some functionality is duplicated with SparkMaterializationReadNode and spark get_historical_features.
-        Args:
-            context: SparkExecutionContext
-        Returns: DAGValue
-        """
-        fv = self.task.feature_view
-        source = fv.batch_source
-
-        (
-            join_key_columns,
-            feature_name_columns,
-            timestamp_field,
-            created_timestamp_column,
-        ) = context.column_info
-
-        # TODO: Use pull_all_from_table_or_query when it supports not filtering by timestamp
-        # retrieval_job = offline_store.pull_all_from_table_or_query(
-        #     config=context.repo_config,
-        #     data_source=source,
-        #     join_key_columns=join_key_columns,
-        #     feature_name_columns=feature_name_columns,
-        #     timestamp_field=timestamp_field,
-        #     start_date=min_ts,
-        #     end_date=max_ts,
-        # )
-        # spark_df = cast(SparkRetrievalJob, retrieval_job).to_spark_df()
-
-        columns = join_key_columns + feature_name_columns + [timestamp_field]
-        if created_timestamp_column:
-            columns.append(created_timestamp_column)
-
-        (fields_with_aliases, aliases) = _get_fields_with_aliases(
-            fields=columns,
-            field_mappings=source.field_mapping,
-        )
-        fields_with_alias_string = ", ".join(fields_with_aliases)
-
-        from_expression = source.get_table_query_string()
-
-        query = f"""
-            SELECT {fields_with_alias_string}
-            FROM {from_expression}
-        """
-        spark_df = self.spark_session.sql(query)
-
-        return DAGValue(
-            data=spark_df,
-            format=DAGFormat.SPARK,
-            metadata={
-                "source": "feature_view_batch_source",
-                "timestamp_field": timestamp_field,
+                "start_date": self.start_time,
+                "end_date": self.end_time,
             },
         )
 

@@ -294,8 +294,8 @@ def pull_all_from_table_or_query(
         join_key_columns: List[str],
         feature_name_columns: List[str],
         timestamp_field: str,
-        start_date: Optional[datetime],
-        end_date: Optional[datetime],
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
     ) -> RetrievalJob:
         """
         Extracts all the entity rows (i.e. the combination of join key columns, feature columns, and