feat: Feast Spark Offline Store (feast-dev#2349)

* State of feast Signed-off-by: Kevin Zhang <[email protected]> * Clean up changes Signed-off-by: Kevin Zhang <[email protected]> * Fix random incorrect changes Signed-off-by: Kevin Zhang <[email protected]> * Fix lint Signed-off-by: Kevin Zhang <[email protected]> * Fix build errors Signed-off-by: Kevin Zhang <[email protected]> * Fix lint Signed-off-by: Kevin Zhang <[email protected]> * Add spark offline store components to test against current integration tests Signed-off-by: Kevin Zhang <[email protected]> * Fix lint Signed-off-by: Kevin Zhang <[email protected]> * Rename to pass checks Signed-off-by: Kevin Zhang <[email protected]> * Fix issues Signed-off-by: Kevin Zhang <[email protected]> * Fix type checking issues Signed-off-by: Kevin Zhang <[email protected]> * Fix lint Signed-off-by: Kevin Zhang <[email protected]> * Clean up print statements for first review Signed-off-by: Kevin Zhang <[email protected]> * Fix lint Signed-off-by: Kevin Zhang <[email protected]> * Fix flake 8 lint tests Signed-off-by: Kevin Zhang <[email protected]> * Add warnings for alpha version release Signed-off-by: Kevin Zhang <[email protected]> * Format Signed-off-by: Kevin Zhang <[email protected]> * Address review Signed-off-by: Kevin Zhang <[email protected]> * Address review Signed-off-by: Kevin Zhang <[email protected]> * Fix lint Signed-off-by: Kevin Zhang <[email protected]> * Add file store functionality Signed-off-by: Kevin Zhang <[email protected]> * lint Signed-off-by: Kevin Zhang <[email protected]> * Add example feature repo Signed-off-by: Kevin Zhang <[email protected]> * Update data source creator Signed-off-by: Kevin Zhang <[email protected]> * Make cli work for feast init with spark Signed-off-by: Kevin Zhang <[email protected]> * Update the docs Signed-off-by: Kevin Zhang <[email protected]> * Clean up code Signed-off-by: Kevin Zhang <[email protected]> * Clean up more code Signed-off-by: Kevin Zhang <[email protected]> * Uncomment repo configs Signed-off-by: Kevin Zhang <[email protected]> * Fix setup.py Signed-off-by: Kevin Zhang <[email protected]> * Update dependencies Signed-off-by: Kevin Zhang <[email protected]> * Fix ci dependencies Signed-off-by: Kevin Zhang <[email protected]> * Screwed up rebase Signed-off-by: Kevin Zhang <[email protected]> * Screwed up rebase Signed-off-by: Kevin Zhang <[email protected]> * Screwed up rebase Signed-off-by: Kevin Zhang <[email protected]> * Realign with master Signed-off-by: Kevin Zhang <[email protected]> * Fix accidental changes Signed-off-by: Kevin Zhang <[email protected]> * Make type map change cleaner Signed-off-by: Kevin Zhang <[email protected]> * Address review comments Signed-off-by: Kevin Zhang <[email protected]> * Fix tests accidentally broken Signed-off-by: Kevin Zhang <[email protected]> * Add comments Signed-off-by: Kevin Zhang <[email protected]> * Reformat Signed-off-by: Kevin Zhang <[email protected]> * Fix logger Signed-off-by: Kevin Zhang <[email protected]> * Remove unused imports Signed-off-by: Kevin Zhang <[email protected]> * Fix imports Signed-off-by: Kevin Zhang <[email protected]> * Fix CI dependencies Signed-off-by: Danny Chiao <[email protected]> * Prefix destinations with project name Signed-off-by: Kevin Zhang <[email protected]> * Update comment Signed-off-by: Kevin Zhang <[email protected]> * Fix 3.8 Signed-off-by: Kevin Zhang <[email protected]> * temporary fix Signed-off-by: Kevin Zhang <[email protected]> * rollback Signed-off-by: Kevin Zhang <[email protected]> * update Signed-off-by: Kevin Zhang <[email protected]> * Update ci? Signed-off-by: Kevin Zhang <[email protected]> * Move third party to contrib Signed-off-by: Kevin Zhang <[email protected]> * Fix imports Signed-off-by: Kevin Zhang <[email protected]> * Remove third_party refactor Signed-off-by: Kevin Zhang <[email protected]> * Revert ci requirements and update comment in type map Signed-off-by: Kevin Zhang <[email protected]> * Revert 3.8-requirements Signed-off-by: Kevin Zhang <[email protected]> Co-authored-by: Danny Chiao <[email protected]>
achals · Mar 4, 2022 · 98b8d8d · 98b8d8d
1 parent 74f887f
commit 98b8d8d
Show file tree

Hide file tree

Showing 21 changed files with 1,401 additions and 208 deletions.
diff --git a/docs/reference/data-sources/README.md b/docs/reference/data-sources/README.md
@@ -9,3 +9,5 @@ Please see [Data Source](../../getting-started/concepts/feature-view.md#data-sou
 {% page-ref page="bigquery.md" %}
 
 {% page-ref page="redshift.md" %}
+
+{% page-ref page="spark.md" %}
diff --git a/docs/reference/data-sources/spark.md b/docs/reference/data-sources/spark.md
@@ -0,0 +1,45 @@
+# Spark
+
+## Description
+
+**NOTE**: Spark data source api is currently in alpha development and the API is not completely stable. The API may change or update in the future.
+
+The spark data source API allows for the retrieval of historical feature values from file/database sources for building training datasets as well as materializing features into an online store.
+
+* Either a table name, a SQL query, or a file path can be provided.
+
+## Examples
+
+Using a table reference from SparkSession(for example, either in memory or a Hive Metastore)
+
+```python
+from feast import SparkSource
+
+my_spark_source = SparkSource(
+    table="FEATURE_TABLE",
+)
+```
+
+Using a query
+
+```python
+from feast import SparkSource
+
+my_spark_source = SparkSource(
+    query="SELECT timestamp as ts, created, f1, f2 "
+          "FROM spark_table",
+)
+```
+
+Using a file reference
+
+```python
+from feast import SparkSource
+
+my_spark_source = SparkSource(
+    path=f"{CURRENT_DIR}/data/driver_hourly_stats",
+    file_format="parquet",
+    event_timestamp_column="event_timestamp",
+    created_timestamp_column="created",
+)
+```
diff --git a/docs/reference/offline-stores/README.md b/docs/reference/offline-stores/README.md
@@ -9,3 +9,5 @@ Please see [Offline Store](../../getting-started/architecture-and-components/off
 {% page-ref page="bigquery.md" %}
 
 {% page-ref page="redshift.md" %}
+
+{% page-ref page="spark.md" %}
diff --git a/docs/reference/offline-stores/spark.md b/docs/reference/offline-stores/spark.md
@@ -0,0 +1,38 @@
+# Spark
+
+## Description
+
+The Spark offline store is an offline store currently in alpha development that provides support for reading [SparkSources](../data-sources/spark.md).
+
+## Disclaimer
+
+This Spark offline store still does not achieve full test coverage and continues to fail some integration tests when integrating with the feast universal test suite. Please do NOT assume complete stability of the API.
+
+* Spark tables and views are allowed as sources that are loaded in from some Spark store(e.g in Hive or in memory).
+* Entity dataframes can be provided as a SQL query or can be provided as a Pandas dataframe. Pandas dataframes will be converted to a Spark dataframe and processed as a temporary view.
+* A `SparkRetrievalJob` is returned when calling `get_historical_features()`.
+  * This allows you to call
+     * `to_df` to retrieve the pandas dataframe.
+     * `to_arrow` to retrieve the dataframe as a pyarrow Table.
+     * `to_spark_df` to retrieve the dataframe the spark.
+
+## Example
+
+{% code title="feature_store.yaml" %}
+```yaml
+project: my_project
+registry: data/registry.db
+provider: local
+offline_store:
+    type: spark
+    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
+```
+{% endcode %}
diff --git a/sdk/python/feast/__init__.py b/sdk/python/feast/__init__.py
@@ -3,6 +3,9 @@
 from pkg_resources import DistributionNotFound, get_distribution
 
 from feast.infra.offline_stores.bigquery_source import BigQuerySource
+from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import (
+    SparkSource,
+)
 from feast.infra.offline_stores.file_source import FileSource
 from feast.infra.offline_stores.redshift_source import RedshiftSource
 from feast.infra.offline_stores.snowflake_source import SnowflakeSource
@@ -47,4 +50,5 @@
     "RedshiftSource",
     "RequestFeatureView",
     "SnowflakeSource",
+    "SparkSource",
 ]
diff --git a/sdk/python/feast/cli.py b/sdk/python/feast/cli.py
@@ -477,7 +477,9 @@ def materialize_incremental_command(ctx: click.Context, end_ts: str, views: List
 @click.option(
     "--template",
     "-t",
-    type=click.Choice(["local", "gcp", "aws", "snowflake"], case_sensitive=False),
+    type=click.Choice(
+        ["local", "gcp", "aws", "snowflake", "spark"], case_sensitive=False
+    ),
     help="Specify a template for the created project",
     default="local",
 )

diff --git a/sdk/python/feast/inference.py b/sdk/python/feast/inference.py
@@ -8,6 +8,7 @@
     FileSource,
     RedshiftSource,
     SnowflakeSource,
+    SparkSource,
 )
 from feast.data_source import DataSource
 from feast.errors import RegistryInferenceFailure
@@ -84,7 +85,9 @@ def update_data_sources_with_inferred_event_timestamp_col(
         ):
             # prepare right match pattern for data source
             ts_column_type_regex_pattern = ""
-            if isinstance(data_source, FileSource):
+            if isinstance(data_source, FileSource) or isinstance(
+                data_source, SparkSource
+            ):
                 ts_column_type_regex_pattern = r"^timestamp"
             elif isinstance(data_source, BigQuerySource):
                 ts_column_type_regex_pattern = "TIMESTAMP|DATETIME"
@@ -97,7 +100,7 @@ def update_data_sources_with_inferred_event_timestamp_col(
                     "DataSource",
                     """
                     DataSource inferencing of event_timestamp_column is currently only supported
-                    for FileSource and BigQuerySource.
+                    for FileSource, SparkSource, BigQuerySource, RedshiftSource, and SnowflakeSource.
                     """,
                 )
             #  for informing the type checker

diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/__init__.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/__init__.py
Original file line number	Diff line number	Diff line change
Expand Up		@@ -9,3 +9,5 @@ Please see [Data Source](../../getting-started/concepts/feature-view.md#data-sou
		{% page-ref page="bigquery.md" %}

		{% page-ref page="redshift.md" %}

		{% page-ref page="spark.md" %}
Original file line number	Diff line number	Diff line change
Expand Up		@@ -9,3 +9,5 @@ Please see [Offline Store](../../getting-started/architecture-and-components/off
		{% page-ref page="bigquery.md" %}

		{% page-ref page="redshift.md" %}

		{% page-ref page="spark.md" %}