From bc19cc4babc1ff032620f6aac68ac6d27933357f Mon Sep 17 00:00:00 2001 From: Lokesh Rangineni Date: Wed, 10 Apr 2024 16:01:47 -0400 Subject: [PATCH] Trying to import pyspark lazily to avoid the dependency on the library. This change will fix the issue - https://github.com/feast-dev/feast/issues/3872 Picked up changes from the PR - https://github.com/feast-dev/feast/pull/3873 Signed-off-by: Lokesh Rangineni --- .../contrib/spark_offline_store/spark_source.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) 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 index 0809043a01..4eb020ebd3 100644 --- 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(".")])