fix: Trying to import pyspark lazily to avoid the dependency on the library (#4091)

lokeshrangineni · web-flow · commit a05cdbcd38d8 · 2024-04-15T00:35:42.000-07:00
diff --git a/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py b/sdk/python/feast/infra/offline_stores/contrib/spark_offline_store/spark_source.py
@@ -4,8 +4,6 @@
 from enum import Enum
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
-from pyspark.sql import SparkSession
-
 from feast import flags_helper
 from feast.data_source import DataSource
 from feast.errors import DataSourceNoNameException, DataSourceNotFoundException
@@ -162,6 +160,13 @@ def get_table_column_names_and_types(
 
     def get_table_query_string(self) -> str:
         """Returns a string that can directly be used to reference this table in SQL"""
+        try:
+            from pyspark.sql import SparkSession
+        except ImportError as e:
+            from feast.errors import FeastExtrasDependencyImportError
+
+            raise FeastExtrasDependencyImportError("spark", str(e))
+
         if self.table:
             # Backticks make sure that spark sql knows this a table reference.
             table = ".".join([f"`{x}`" for x in self.table.split(".")])
