feat: Offline store update pull_all_from_table_or_query to make timestampfield optional (#5281)

* Update offline store pull all API date field optional

Signed-off-by: HaoXuAI <[email protected]>

* Update offline store pull all API date field optional

Signed-off-by: HaoXuAI <[email protected]>

* update

Signed-off-by: HaoXuAI <[email protected]>

* update

Signed-off-by: HaoXuAI <[email protected]>

* update

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* update source read node

Signed-off-by: HaoXuAI <[email protected]>

* fix linting

Signed-off-by: HaoXuAI <[email protected]>

* fix linting

Signed-off-by: HaoXuAI <[email protected]>

* fix linting

Signed-off-by: HaoXuAI <[email protected]>

* fix linting

Signed-off-by: HaoXuAI <[email protected]>

* fix linting

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

* fix test

Signed-off-by: HaoXuAI <[email protected]>

---------

Signed-off-by: HaoXuAI <[email protected]>

Author: HaoXuAI

sdk/python/feast/infra/common/retrieval_task.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Union
+from typing import Optional, Union
 
 import pandas as pd
 
@@ -15,5 +15,5 @@ class HistoricalRetrievalTask:
     feature_view: Union[BatchFeatureView, StreamFeatureView]
     full_feature_name: bool
     registry: Registry
-    start_time: datetime
-    end_time: datetime
+    start_time: Optional[datetime] = None
+    end_time: Optional[datetime] = None

sdk/python/feast/infra/compute_engines/feature_builder.py

@@ -72,10 +72,10 @@ def _should_validate(self):
     def build(self) -> ExecutionPlan:
         last_node = self.build_source_node()
 
-        # PIT join entities to the feature data, and perform filtering
-        if isinstance(self.task, HistoricalRetrievalTask):
-            last_node = self.build_join_node(last_node)
+        # Join entity_df with source if needed
+        last_node = self.build_join_node(last_node)
 
+        # PIT filter, TTL, and user-defined filter
         last_node = self.build_filter_node(last_node)
 
         if self._should_aggregate():

sdk/python/feast/infra/compute_engines/local/feature_builder.py

@@ -2,7 +2,6 @@
 
 from feast.infra.common.materialization_job import MaterializationTask
 from feast.infra.common.retrieval_task import HistoricalRetrievalTask
-from feast.infra.compute_engines.dag.plan import ExecutionPlan
 from feast.infra.compute_engines.feature_builder import FeatureBuilder
 from feast.infra.compute_engines.local.backends.base import DataFrameBackend
 from feast.infra.compute_engines.local.nodes import (
@@ -95,25 +94,3 @@ def build_output_nodes(self, input_node):
         node = LocalOutputNode("output")
         node.add_input(input_node)
         self.nodes.append(node)
-
-    def build(self) -> ExecutionPlan:
-        last_node = self.build_source_node()
-
-        if isinstance(self.task, HistoricalRetrievalTask):
-            last_node = self.build_join_node(last_node)
-
-        last_node = self.build_filter_node(last_node)
-
-        if self._should_aggregate():
-            last_node = self.build_aggregation_node(last_node)
-        elif isinstance(self.task, HistoricalRetrievalTask):
-            last_node = self.build_dedup_node(last_node)
-
-        if self._should_transform():
-            last_node = self.build_transformation_node(last_node)
-
-        if self._should_validate():
-            last_node = self.build_validation_node(last_node)
-
-        self.build_output_nodes(last_node)
-        return ExecutionPlan(self.nodes)

sdk/python/feast/infra/compute_engines/local/nodes.py

@@ -1,8 +1,9 @@
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Optional
 
 import pyarrow as pa
 
+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
 from feast.infra.compute_engines.local.backends.base import DataFrameBackend
@@ -15,14 +16,40 @@
 
 
 class LocalSourceReadNode(LocalNode):
-    def __init__(self, name: str, feature_view, task):
+    def __init__(
+        self,
+        name: str,
+        source: DataSource,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+    ):
         super().__init__(name)
-        self.feature_view = feature_view
-        self.task = task
+        self.source = source
+        self.start_time = start_time
+        self.end_time = end_time
 
     def execute(self, context: ExecutionContext) -> ArrowTableValue:
-        # TODO : Implement the logic to read from offline store
-        return ArrowTableValue(data=pa.Table.from_pandas(context.entity_df))
+        offline_store = context.offline_store
+        (
+            join_key_columns,
+            feature_name_columns,
+            timestamp_field,
+            created_timestamp_column,
+        ) = context.column_info
+
+        # 📥 Reuse Feast's robust query resolver
+        retrieval_job = offline_store.pull_all_from_table_or_query(
+            config=context.repo_config,
+            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=self.start_time,
+            end_date=self.end_time,
+        )
+        arrow_table = retrieval_job.to_arrow()
+        return ArrowTableValue(data=arrow_table)
 
 
 class LocalJoinNode(LocalNode):

sdk/python/feast/infra/compute_engines/spark/feature_builder.py

@@ -9,9 +9,8 @@
     SparkAggregationNode,
     SparkDedupNode,
     SparkFilterNode,
-    SparkHistoricalRetrievalReadNode,
     SparkJoinNode,
-    SparkMaterializationReadNode,
+    SparkReadNode,
     SparkTransformationNode,
     SparkWriteNode,
 )
@@ -27,12 +26,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
 

sdk/python/feast/infra/compute_engines/spark/node.py

@@ -1,13 +1,12 @@
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import List, Optional, Union, cast
 
 from pyspark.sql import DataFrame, SparkSession, Window
 from pyspark.sql import functions as F
 
 from feast import BatchFeatureView, StreamFeatureView
 from feast.aggregation import Aggregation
-from feast.infra.common.materialization_job import MaterializationTask
-from feast.infra.common.retrieval_task import HistoricalRetrievalTask
+from feast.data_source import DataSource
 from feast.infra.compute_engines.dag.context import ExecutionContext
 from feast.infra.compute_engines.dag.model import DAGFormat
 from feast.infra.compute_engines.dag.node import DAGNode
@@ -23,7 +22,6 @@
 from feast.infra.offline_stores.offline_utils import (
     infer_event_timestamp_from_entity_df,
 )
-from feast.utils import _get_fields_with_aliases
 
 ENTITY_TS_ALIAS = "__entity_event_timestamp"
 
@@ -49,18 +47,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[datetime] = None,
+        end_time: Optional[datetime] = 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 +70,15 @@ 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 +89,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,
             },
         )
 
@@ -227,7 +162,12 @@ def execute(self, context: ExecutionContext) -> DAGValue:
         feature_df: DataFrame = feature_value.data
 
         entity_df = context.entity_df
-        assert entity_df is not None, "entity_df must be set in ExecutionContext"
+        if entity_df is None:
+            return DAGValue(
+                data=feature_df,
+                format=DAGFormat.SPARK,
+                metadata={"joined_on": None},
+            )
 
         # Get timestamp fields from feature view
         join_keys, feature_cols, ts_col, created_ts_col = context.column_info
@@ -272,13 +212,13 @@ def execute(self, context: ExecutionContext) -> DAGValue:
         if ENTITY_TS_ALIAS in input_df.columns:
             filtered_df = filtered_df.filter(F.col(ts_col) <= F.col(ENTITY_TS_ALIAS))
 
-        # Optional TTL filter: feature.ts >= entity.event_timestamp - ttl
-        if self.ttl:
-            ttl_seconds = int(self.ttl.total_seconds())
-            lower_bound = F.col(ENTITY_TS_ALIAS) - F.expr(
-                f"INTERVAL {ttl_seconds} seconds"
-            )
-            filtered_df = filtered_df.filter(F.col(ts_col) >= lower_bound)
+            # Optional TTL filter: feature.ts >= entity.event_timestamp - ttl
+            if self.ttl:
+                ttl_seconds = int(self.ttl.total_seconds())
+                lower_bound = F.col(ENTITY_TS_ALIAS) - F.expr(
+                    f"INTERVAL {ttl_seconds} seconds"
+                )
+                filtered_df = filtered_df.filter(F.col(ts_col) >= lower_bound)
 
         # Optional custom filter condition
         if self.filter_condition:

sdk/python/feast/infra/offline_stores/bigquery.py

@@ -52,6 +52,7 @@
     BigQuerySource,
     SavedDatasetBigQueryStorage,
 )
+from .offline_utils import get_timestamp_filter_sql
 
 try:
     from google.api_core import client_info as http_client_info
@@ -188,8 +189,9 @@ def pull_all_from_table_or_query(
         join_key_columns: List[str],
         feature_name_columns: List[str],
         timestamp_field: str,
-        start_date: datetime,
-        end_date: datetime,
+        created_timestamp_column: Optional[str] = None,
+        start_date: Optional[datetime] = None,
+        end_date: Optional[datetime] = None,
     ) -> RetrievalJob:
         assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)
         assert isinstance(data_source, BigQuerySource)
@@ -201,15 +203,26 @@ def pull_all_from_table_or_query(
             project=project_id,
             location=config.offline_store.location,
         )
+
+        timestamp_fields = [timestamp_field]
+        if created_timestamp_column:
+            timestamp_fields.append(created_timestamp_column)
         field_string = ", ".join(
             BigQueryOfflineStore._escape_query_columns(join_key_columns)
             + BigQueryOfflineStore._escape_query_columns(feature_name_columns)
-            + [timestamp_field]
+            + timestamp_fields
+        )
+        timestamp_filter = get_timestamp_filter_sql(
+            start_date,
+            end_date,
+            timestamp_field,
+            quote_fields=False,
+            cast_style="timestamp_func",
         )
         query = f"""
             SELECT {field_string}
             FROM {from_expression}
-            WHERE {timestamp_field} BETWEEN TIMESTAMP('{start_date}') AND TIMESTAMP('{end_date}')
+            WHERE {timestamp_filter}
         """
         return BigQueryRetrievalJob(
             query=query,