feat: Add date_partition_column_format for spark source (feast-dev#5273)

* feat: Add date_partition_column_format for spark source

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

* feat: Add date_partition_column_format for spark source

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

* feat: Add date_partition_column_format for spark source

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

---------

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

Author: joeyutong

protos/feast/core/DataSource.proto

@@ -226,6 +226,9 @@ message DataSource {
 
     // Format of files at `path` (e.g. parquet, avro, etc)
     string file_format = 4;
+
+    // Date Format of date partition column (e.g. %Y-%m-%d)
+    string date_partition_column_format = 5;
   }
 
   // Defines configuration for custom third-party data sources.

sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py

@@ -2,8 +2,9 @@
 import tempfile
 import uuid
 import warnings
+from dataclasses import asdict, dataclass
 from datetime import datetime, timezone
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union, cast
 
 import numpy as np
 import pandas
@@ -55,6 +56,12 @@ class SparkOfflineStoreConfig(FeastConfigBaseModel):
     """ AWS Region if applicable for s3-based staging locations"""
 
 
+@dataclass(frozen=True)
+class SparkFeatureViewQueryContext(offline_utils.FeatureViewQueryContext):
+    min_date_partition: Optional[str]
+    max_date_partition: str
+
+
 class SparkOfflineStore(OfflineStore):
     @staticmethod
     def pull_latest_from_table_or_query(
@@ -101,6 +108,7 @@ def pull_latest_from_table_or_query(
         aliases_as_string = ", ".join(aliases)
 
         date_partition_column = data_source.date_partition_column
+        date_partition_column_format = data_source.date_partition_column_format
 
         start_date_str = _format_datetime(start_date)
         end_date_str = _format_datetime(end_date)
@@ -112,7 +120,7 @@ def pull_latest_from_table_or_query(
                     SELECT {fields_as_string},
                     ROW_NUMBER() OVER({partition_by_join_key_string} ORDER BY {timestamp_desc_string}) AS feast_row_
                     FROM {from_expression} t1
-                    WHERE {timestamp_field} BETWEEN TIMESTAMP('{start_date_str}') AND TIMESTAMP('{end_date_str}'){" AND " + date_partition_column + " >= '" + start_date.strftime("%Y-%m-%d") + "' AND " + date_partition_column + " <= '" + end_date.strftime("%Y-%m-%d") + "' " if date_partition_column != "" and date_partition_column is not None else ""}
+                    WHERE {timestamp_field} BETWEEN TIMESTAMP('{start_date_str}') AND TIMESTAMP('{end_date_str}'){" AND " + date_partition_column + " >= '" + start_date.strftime(date_partition_column_format) + "' AND " + date_partition_column + " <= '" + end_date.strftime(date_partition_column_format) + "' " if date_partition_column != "" and date_partition_column is not None else ""}
                 ) t2
                 WHERE feast_row_ = 1
                 """
@@ -136,8 +144,12 @@ def get_historical_features(
         full_feature_names: bool = False,
     ) -> RetrievalJob:
         assert isinstance(config.offline_store, SparkOfflineStoreConfig)
+        date_partition_column_formats = []
         for fv in feature_views:
             assert isinstance(fv.batch_source, SparkSource)
+            date_partition_column_formats.append(
+                fv.batch_source.date_partition_column_format
+            )
 
         warnings.warn(
             "The spark offline store is an experimental feature in alpha development. "
@@ -186,8 +198,27 @@ def get_historical_features(
             entity_df_event_timestamp_range,
         )
 
+        spark_query_context = [
+            SparkFeatureViewQueryContext(
+                **asdict(context),
+                min_date_partition=datetime.fromisoformat(
+                    context.min_event_timestamp
+                ).strftime(date_format)
+                if context.min_event_timestamp is not None
+                else None,
+                max_date_partition=datetime.fromisoformat(
+                    context.max_event_timestamp
+                ).strftime(date_format),
+            )
+            for date_format, context in zip(
+                date_partition_column_formats, query_context
+            )
+        ]
+
         query = offline_utils.build_point_in_time_query(
-            feature_view_query_contexts=query_context,
+            feature_view_query_contexts=cast(
+                List[offline_utils.FeatureViewQueryContext], spark_query_context
+            ),
             left_table_query_string=tmp_entity_df_table_name,
             entity_df_event_timestamp_col=event_timestamp_col,
             entity_df_columns=entity_schema.keys(),
@@ -651,13 +682,13 @@ def _cast_data_frame(
         FROM {{ featureview.table_subquery }}
         WHERE {{ featureview.timestamp_field }} <= '{{ featureview.max_event_timestamp }}'
         {% if featureview.date_partition_column != "" and featureview.date_partition_column is not none %}
-        AND {{ featureview.date_partition_column }} <= '{{ featureview.max_event_timestamp[:10] }}'
+        AND {{ featureview.date_partition_column }} <= '{{ featureview.max_date_partition }}'
         {% endif %}
 
         {% if featureview.ttl == 0 %}{% else %}
         AND {{ featureview.timestamp_field }} >= '{{ featureview.min_event_timestamp }}'
         {% if featureview.date_partition_column != "" and featureview.date_partition_column is not none %}
-          AND {{ featureview.date_partition_column }} >= '{{ featureview.min_event_timestamp[:10] }}'
+          AND {{ featureview.date_partition_column }} >= '{{ featureview.min_date_partition }}'
         {% endif %}
         {% endif %}
     ),

sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py

@@ -46,6 +46,7 @@ def __init__(
         owner: Optional[str] = "",
         timestamp_field: Optional[str] = None,
         date_partition_column: Optional[str] = None,
+        date_partition_column_format: Optional[str] = "%Y-%m-%d",
     ):
         """Creates a SparkSource object.
 
@@ -97,6 +98,7 @@ def __init__(
             query=query,
             path=path,
             file_format=file_format,
+            date_partition_column_format=date_partition_column_format,
         )
 
     @property
@@ -127,6 +129,13 @@ def file_format(self):
         """
         return self.spark_options.file_format
 
+    @property
+    def date_partition_column_format(self):
+        """
+        Returns the date partition column format of this feature data source.
+        """
+        return self.spark_options.date_partition_column_format
+
     @staticmethod
     def from_proto(data_source: DataSourceProto) -> Any:
         assert data_source.HasField("spark_options")
@@ -139,6 +148,7 @@ def from_proto(data_source: DataSourceProto) -> Any:
             query=spark_options.query,
             path=spark_options.path,
             file_format=spark_options.file_format,
+            date_partition_column_format=spark_options.date_partition_column_format,
             date_partition_column=data_source.date_partition_column,
             timestamp_field=data_source.timestamp_field,
             created_timestamp_column=data_source.created_timestamp_column,
@@ -240,6 +250,7 @@ def __init__(
         query: Optional[str],
         path: Optional[str],
         file_format: Optional[str],
+        date_partition_column_format: Optional[str] = "%Y-%m-%d",
     ):
         # Check that only one of the ways to load a spark dataframe can be used. We have
         # to treat empty string and null the same due to proto (de)serialization.
@@ -261,6 +272,7 @@ def __init__(
         self._query = query
         self._path = path
         self._file_format = file_format
+        self._date_partition_column_format = date_partition_column_format
 
     @property
     def table(self):
@@ -294,6 +306,14 @@ def file_format(self):
     def file_format(self, file_format):
         self._file_format = file_format
 
+    @property
+    def date_partition_column_format(self):
+        return self._date_partition_column_format
+
+    @date_partition_column_format.setter
+    def date_partition_column_format(self, date_partition_column_format):
+        self._date_partition_column_format = date_partition_column_format
+
     @classmethod
     def from_proto(cls, spark_options_proto: DataSourceProto.SparkOptions):
         """
@@ -308,6 +328,7 @@ def from_proto(cls, spark_options_proto: DataSourceProto.SparkOptions):
             query=spark_options_proto.query,
             path=spark_options_proto.path,
             file_format=spark_options_proto.file_format,
+            date_partition_column_format=spark_options_proto.date_partition_column_format,
         )
 
         return spark_options
@@ -323,6 +344,7 @@ def to_proto(self) -> DataSourceProto.SparkOptions:
             query=self.query,
             path=self.path,
             file_format=self.file_format,
+            date_partition_column_format=self.date_partition_column_format,
         )
 
         return spark_options_proto