Add file store functionality

Signed-off-by: Kevin Zhang <kzhang@tecton.ai>
feast-dev · adchia · Mar 4, 2022 · Feb 25, 2022 · Feb 26, 2022 · Feb 26, 2022
commit acf1c28c152cec60995d32f1ecc450611f0e4893
diff --git a/sdk/python/feast/infra/offline_stores/third_party/spark_source.py b/sdk/python/feast/infra/offline_stores/third_party/spark_source.py
@@ -1,8 +1,10 @@
 import pickle
+from enum import Enum
 import warnings
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
 from pyspark.sql.utils import AnalysisException
+from pyspark.sql import SparkSession, DataFrame
 
 from feast.data_source import DataSource
 from feast.errors import DataSourceNotFoundException
@@ -14,18 +16,20 @@
 from feast.saved_dataset import SavedDatasetStorage
 from feast.type_map import spark_to_feast_value_type
 from feast.value_type import ValueType
+from feast.infra.offline_stores.offline_utils import get_temp_entity_table_name
 
+class SparkSourceFormat(Enum):
+    csv = "csv"
+    json = "json"
+    parquet = "parquet"
 
 class SparkSource(DataSource):
     def __init__(
         self,
         table: Optional[str] = None,
         query: Optional[str] = None,
-        # TODO support file readers
-        # path: Optional[str] = None,
-        # jdbc=None,
-        # format: Optional[str] = None,
-        # options: Optional[Dict[str, Any]] = None,
+        path: Optional[str] = None,
+        file_format: Optional[str] = None,
         event_timestamp_column: Optional[str] = None,
         created_timestamp_column: Optional[str] = None,
         field_mapping: Optional[Dict[str, str]] = None,
@@ -42,14 +46,26 @@ def __init__(
             "This API is unstable and it could and most probably will be changed in the future.",
             RuntimeWarning,
         )
+        self.allowed_formats = [format.value for format in SparkSourceFormat]
+
+        # Check that only one of the ways to load a spark dataframe can be used.
+        if sum([(arg is not None) for arg in [table, query, path]]) != 1:
+            raise ValueError("Exactly one of params(table, query, path) must be specified.")
+
+        if path is not None:
+            if file_format is None:
+                raise ValueError(
+                    "If 'path' is specified, then 'file_format' is required."
+                )
+            if file_format not in self.allowed_formats:
+                raise ValueError(f"'file_format' should be one of {self.allowed_formats}")
+
 
         self._spark_options = SparkOptions(
             table=table,
             query=query,
-            # path=path,
-            # jdbc=None,
-            # format=format,
-            # options=options,
+            path=path,
+            file_format=file_format,
         )
 
     @property
@@ -80,6 +96,21 @@ def query(self):
         """
         return self._spark_options.query
 
+
+    @property
+    def path(self):
+        """
+        Returns the path of the spark data source file.
+        """
+        return self._spark_options.path
+
+    @property
+    def file_format(self):
+        """
+        Returns the file format of this feature data source.
+        """
+        return self._spark_options.file_format
+
     @staticmethod
     def from_proto(data_source: DataSourceProto) -> Any:
 
@@ -118,7 +149,6 @@ def validate(self, config: RepoConfig):
 
     @staticmethod
     def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
-        # TODO see feast.type_map for examples
         return spark_to_feast_value_type
 
     def get_table_column_names_and_types(
@@ -145,9 +175,19 @@ def get_table_query_string(self) -> str:
         """Returns a string that can directly be used to reference this table in SQL"""
         if self.table:
             return f"`{self.table}`"
-        else:
+        if self.query:
             return f"({self.query})"
 
+        # If both the table query string and the actual query are null, we can load from file.
+        spark_session = SparkSession.getActiveSession()
+        if spark_session is None:
+            raise AssertionError("Could not find an active spark session.")
+        df = spark_session.read.format(self.file_format).load(self.path)
+
+        tmp_table_name = get_temp_entity_table_name()
+        df.createOrReplaceTempView(tmp_table_name)
+
+        return f"`{tmp_table_name}`"
 
 class SparkOptions:
     def __init__(
@@ -213,7 +253,6 @@ def to_proto(self) -> DataSourceProto.CustomSourceOptions:
 
         return spark_options_proto
 
-
 class SavedDatasetSparkStorage(SavedDatasetStorage):
     _proto_attr_name = "spark_storage"
 

@@ -0,0 +1,50 @@
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from pyspark.sql import SparkSession
+from feast.driver_test_data import (
+    create_driver_hourly_stats_df,
+    create_customer_daily_profile_df,
+)
+
+CURRENT_DIR = Path(__file__).parent
+DRIVER_ENTITIES = [1001, 1002, 1003]
+CUSTOMER_ENTITIES = [201, 202, 203]
+START_DATE = datetime.strptime("2022-01-01", "%Y-%m-%d")
+END_DATE = START_DATE + timedelta(days=7)
+
+def bootstrap():
+    # Bootstrap() will automatically be called from the init_repo() during `feast init`
+    generate_example_data(
+        spark_session=SparkSession.builder.getOrCreate(),
+        base_dir=str(CURRENT_DIR),
+    )
+
+
+def example_data_exists(base_dir: str) -> bool:
+    for path in [
+        Path(base_dir) / "data" / "driver_hourly_stats",
+        Path(base_dir) / "data" / "customer_daily_profile",
+    ]:
+        if not path.exists():
+            return False
+    return True
+
+
+def generate_example_data(spark_session: SparkSession, base_dir: str) -> None:
+    spark_session.createDataFrame(
+        data=create_driver_hourly_stats_df(DRIVER_ENTITIES, START_DATE, END_DATE)
+    ).write.parquet(
+        path=str(Path(base_dir) / "data" / "driver_hourly_stats"),
+        mode="overwrite",
+    )
+
+    spark_session.createDataFrame(
+        data=create_customer_daily_profile_df(CUSTOMER_ENTITIES, START_DATE, END_DATE)
+    ).write.parquet(
+        path=str(Path(base_dir) / "data" / "customer_daily_profile"),
+        mode="overwrite",
+    )
+
+if __name__ == "__main__":
+    bootstrap()
@@ -0,0 +1,69 @@
+# # # # # # # # # # # # # # # # # # # # # # # #
+# This is an example feature definition file  #
+# # # # # # # # # # # # # # # # # # # # # # # #
+
+from pathlib import Path
+
+from feast import Entity, Feature, FeatureView, ValueType
+
+from google.protobuf.duration_pb2 import Duration
+
+from feast_spark_offline_store import SparkSource
+
+# Constants related to the generated data sets
+CURRENT_DIR = Path(__file__).parent
+
+
+# Entity definitions
+driver = Entity(
+    name="driver_id",
+    value_type=ValueType.INT64,
+    description="driver id",
+)
+customer = Entity(
+    name="customer_id",
+    value_type=ValueType.INT64,
+    description="customer id",
+)
+
+# Sources
+driver_hourly_stats = SparkSource(
+    path=f"{CURRENT_DIR}/data/driver_hourly_stats",
+    file_format="parquet",
+    event_timestamp_column="event_timestamp",
+    created_timestamp_column="created",
+)
+customer_daily_profile = SparkSource(
+    path=f"{CURRENT_DIR}/data/customer_daily_profile",
+    file_format="parquet",
+    event_timestamp_column="event_timestamp",
+    created_timestamp_column="created",
+)
+
+# Feature Views
+driver_hourly_stats_view = FeatureView(
+    name="driver_hourly_stats",
+    entities=["driver_id"],
+    ttl=Duration(seconds=86400 * 7),  # one week
+    features=[
+        Feature(name="conv_rate", dtype=ValueType.FLOAT),
+        Feature(name="acc_rate", dtype=ValueType.FLOAT),
+        Feature(name="avg_daily_trips", dtype=ValueType.INT64),
+    ],
+    online=True,
+    batch_source=driver_hourly_stats,
+    tags={},
+)
+customer_daily_profile_view = FeatureView(
+    name="customer_daily_profile",
+    entities=["customer_id"],
+    ttl=Duration(seconds=86400 * 7),  # one week
+    features=[
+        Feature(name="current_balance", dtype=ValueType.FLOAT),
+        Feature(name="avg_passenger_count", dtype=ValueType.FLOAT),
+        Feature(name="lifetime_trip_count", dtype=ValueType.INT64),
+    ],
+    online=True,
+    batch_source=customer_daily_profile,
+    tags={},
+)
@@ -0,0 +1,15 @@
+project: my_project
+registry: data/registry.db
+provider: local
+offline_store:
+    type: feast_spark_offline_store.spark.SparkOfflineStore
+    spark_conf:
+        spark.master: "local[*]"
+        spark.ui.enabled: "false"
+        spark.eventLog.enabled: "false"
+        spark.sql.catalogImplementation: "hive"
+        spark.sql.parser.quotedRegexColumnNames: "true"
+        spark.sql.session.timeZone: "UTC"
+online_store:
+    path: data/online_store.db
+## etc: etc....