fix: Update import logic to remove pyspark dependency from Snowflake Offline Store (feast-dev#3397)

Ryan Beauchamp · web-flow · commit cf073e68b0b3 · 2022-12-17T20:35:08.000-08:00
* fix: move pyspark import inside relevant method

Signed-off-by: Ryan Beauchamp &lt;ryan.beauchamp@ezcater.com&gt;

* fix: put import inside type checking so that type hints works

Signed-off-by: Ryan Beauchamp &lt;ryan.beauchamp@ezcater.com&gt;

* fix: import future annotations

Signed-off-by: Ryan Beauchamp &lt;ryan.beauchamp@ezcater.com&gt;

* fix: annotate only the method signature

Signed-off-by: Ryan Beauchamp &lt;ryan.beauchamp@ezcater.com&gt;

Signed-off-by: Ryan Beauchamp &lt;ryan.beauchamp@ezcater.com&gt;
diff --git a/sdk/python/feast/infra/offline_stores/snowflake.py b/sdk/python/feast/infra/offline_stores/snowflake.py
@@ -6,6 +6,7 @@
 from functools import reduce
 from pathlib import Path
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     ContextManager,
@@ -63,12 +64,8 @@
 
     raise FeastExtrasDependencyImportError("snowflake", str(e))
 
-try:
+if TYPE_CHECKING:
     from pyspark.sql import DataFrame, SparkSession
-except ImportError as e:
-    from feast.errors import FeastExtrasDependencyImportError
-
-    raise FeastExtrasDependencyImportError("spark", str(e))
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -462,7 +459,7 @@ def to_sql(self) -> str:
         with self._query_generator() as query:
             return query
 
-    def to_spark_df(self, spark_session: SparkSession) -> DataFrame:
+    def to_spark_df(self, spark_session: "SparkSession") -> "DataFrame":
         """
         Method to convert snowflake query results to pyspark data frame.
 
@@ -473,6 +470,13 @@ def to_spark_df(self, spark_session: SparkSession) -> DataFrame:
             spark_df: A pyspark dataframe.
         """
 
+        try:
+            from pyspark.sql import DataFrame, SparkSession
+        except ImportError as e:
+            from feast.errors import FeastExtrasDependencyImportError
+
+            raise FeastExtrasDependencyImportError("spark", str(e))
+
         if isinstance(spark_session, SparkSession):
             with self._query_generator() as query:
 
