Repository: feast-dev/feast

Issue:
Get_historical_features() Does Not Have Option To Return Distributed Dataframe Like A Spark DF
When doing get_historical_features().to_df on a large training dataset in databricks I am hitting memory full errors. Since to_df is returning the data as a pandas dataframe it is not able to use the full capacity of the databricks cluster and distribute it across the different nodes the way a spark dataframe might.


Discussion:
I believe this is actually possible. See https://github.com/feast-dev/feast/blob/master/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark.py#L252
@adchia but this is only available if you are using a spark offline store correct? for example, you cannot query a snowflake table and return it as a spark df, correct? 
Ah correct. Right now it only allows returning it in arrow batches. Is the intention to be able to materialize through a spark df?   Or to do training e.g. with MLLib.

In either case, might make sense for the SparkOfflineStore (this naming doesn't really make too much sense...) to be able to query from data warehouses (but maybe run the join within the DWH)
@adchia The intention is to just return snowflake data offline as a training dataset with the results being a distributed spark df. Is it possible to use snowflake as the source with the SparkOfflineStore? I found a way to hack around it by manipulating the to_sql() string and issuing that as a spark query using snowflakes spark connector, but wondering if there is a cleaner way
That works as a start and I don't think it's that hacky. Would be happy to check that in. 

An alternative would be to flush the output of Snowflakes query to S3 and read that, but that also seems pretty hacky

Resolving patch:
diff --git a/sdk/python/feast/errors.py b/sdk/python/feast/errors.py
--- a/sdk/python/feast/errors.py
+++ b/sdk/python/feast/errors.py
@@ -56,6 +56,14 @@ def __init__(self, name, project=None):
             super().__init__(f"Feature view {name} does not exist")
 
 
+class InvalidSparkSessionException(Exception):
+    def __init__(self, spark_arg):
+        super().__init__(
+            f" Need Spark Session to convert results to spark data frame\
+               recieved {type(spark_arg)} instead. "
+        )
+
+
 class OnDemandFeatureViewNotFoundException(FeastObjectNotFoundException):
     def __init__(self, name, project=None):
         if project:
diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py
--- a/sdk/python/feast/infra/offline_stores/snowflake.py
+++ b/sdk/python/feast/infra/offline_stores/snowflake.py
@@ -1,7 +1,9 @@
 import contextlib
 import os
 import uuid
+import warnings
 from datetime import datetime
+from functools import reduce
 from pathlib import Path
 from typing import (
     Any,
@@ -21,11 +23,16 @@
 import pyarrow
 from pydantic import Field, StrictStr
 from pydantic.typing import Literal
+from pyspark.sql import DataFrame, SparkSession
 from pytz import utc
 
 from feast import OnDemandFeatureView
 from feast.data_source import DataSource
-from feast.errors import EntitySQLEmptyResults, InvalidEntityType
+from feast.errors import (
+    EntitySQLEmptyResults,
+    InvalidEntityType,
+    InvalidSparkSessionException,
+)
 from feast.feature_logging import LoggingConfig, LoggingSource
 from feast.feature_view import DUMMY_ENTITY_ID, DUMMY_ENTITY_VAL, FeatureView
 from feast.infra.offline_stores import offline_utils
@@ -57,6 +64,8 @@
 
     raise FeastExtrasDependencyImportError("snowflake", str(e))
 
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+
 
 class SnowflakeOfflineStoreConfig(FeastConfigBaseModel):
     """Offline store config for Snowflake"""
@@ -447,6 +456,41 @@ def to_sql(self) -> str:
         with self._query_generator() as query:
             return query
 
+    def to_spark_df(self, spark_session: SparkSession) -> DataFrame:
+        """
+        Method to convert snowflake query results to pyspark data frame.
+
+        Args:
+            spark_session: spark Session variable of current environment.
+
+        Returns:
+            spark_df: A pyspark dataframe.
+        """
+
+        if isinstance(spark_session, SparkSession):
+            with self._query_generator() as query:
+
+                arrow_batches = execute_snowflake_statement(
+                    self.snowflake_conn, query
+                ).fetch_arrow_batches()
+
+                if arrow_batches:
+                    spark_df = reduce(
+                        DataFrame.unionAll,
+                        [
+                            spark_session.createDataFrame(batch.to_pandas())
+                            for batch in arrow_batches
+                        ],
+                    )
+
+                    return spark_df
+
+                else:
+                    raise EntitySQLEmptyResults(query)
+
+        else:
+            raise InvalidSparkSessionException(spark_session)
+
     def persist(self, storage: SavedDatasetStorage, allow_overwrite: bool = False):
         assert isinstance(storage, SavedDatasetSnowflakeStorage)
         self.to_snowflake(table_name=storage.snowflake_options.table)


Test patch: