fix: Update import logic to remove pyspark dependency from Snowflak…

…e Offline Store (#3397)

* fix: move pyspark import inside relevant method

Signed-off-by: Ryan Beauchamp <ryan.beauchamp@ezcater.com>

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

Signed-off-by: Ryan Beauchamp <ryan.beauchamp@ezcater.com>

* fix: import future annotations

Signed-off-by: Ryan Beauchamp <ryan.beauchamp@ezcater.com>

* fix: annotate only the method signature

Signed-off-by: Ryan Beauchamp <ryan.beauchamp@ezcater.com>

Signed-off-by: Ryan Beauchamp <ryan.beauchamp@ezcater.com>

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: