feat(datasets): Add support for databricks-connect>=13.0

kedro-datasets/RELEASE.md

@@ -1,5 +1,7 @@
 # Upcoming Release
 ## Major features and improvements
+* Spark and Databricks based datasets now support [databricks-connect>=13.0](https://docs.databricks.com/en/dev-tools/databricks-connect-ref.html)
+
 ## Bug fixes and other changes
 ## Upcoming deprecations for Kedro-Datasets 2.0.0
 * Renamed dataset and error classes, in accordance with the [Kedro lexicon](https://github.com/kedro-org/kedro/wiki/Kedro-documentation-style-guide#kedro-lexicon). Dataset classes ending with "DataSet" are deprecated and will be removed in 2.0.0.

kedro-datasets/kedro_datasets/databricks/managed_table_dataset.py

@@ -9,12 +9,13 @@
 
 import pandas as pd
 from kedro.io.core import Version, VersionNotFoundError
-from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import DataFrame
 from pyspark.sql.types import StructType
 from pyspark.sql.utils import AnalysisException, ParseException
 
 from kedro_datasets import KedroDeprecationWarning
 from kedro_datasets._io import AbstractVersionedDataset, DatasetError
+from kedro_datasets.spark.spark_dataset import _get_spark
 
 logger = logging.getLogger(__name__)
 
@@ -263,10 +264,6 @@ def __init__(  # noqa: PLR0913
             exists_function=self._exists,
         )
 
-    @staticmethod
-    def _get_spark() -> SparkSession:
-        return SparkSession.builder.getOrCreate()
-
     def _load(self) -> Union[DataFrame, pd.DataFrame]:
         """Loads the version of data in the format defined in the init
         (spark|pandas dataframe)
@@ -282,15 +279,15 @@ def _load(self) -> Union[DataFrame, pd.DataFrame]:
         if self._version and self._version.load >= 0:
             try:
                 data = (
-                    self._get_spark()
+                    _get_spark()
                     .read.format("delta")
                     .option("versionAsOf", self._version.load)
                     .table(self._table.full_table_location())
                 )
             except Exception as exc:
                 raise VersionNotFoundError(self._version.load) from exc
         else:
-            data = self._get_spark().table(self._table.full_table_location())
+            data = _get_spark().table(self._table.full_table_location())
         if self._table.dataframe_type == "pandas":
             data = data.toPandas()
         return data
@@ -328,7 +325,7 @@ def _save_upsert(self, update_data: DataFrame) -> None:
             update_data (DataFrame): the Spark dataframe to upsert
         """
         if self._exists():
-            base_data = self._get_spark().table(self._table.full_table_location())
+            base_data = _get_spark().table(self._table.full_table_location())
             base_columns = base_data.columns
             update_columns = update_data.columns
 
@@ -351,13 +348,13 @@ def _save_upsert(self, update_data: DataFrame) -> None:
                 )
 
             update_data.createOrReplaceTempView("update")
-            self._get_spark().conf.set(
+            _get_spark().conf.set(
                 "fullTableAddress", self._table.full_table_location()
             )
-            self._get_spark().conf.set("whereExpr", where_expr)
+            _get_spark().conf.set("whereExpr", where_expr)
             upsert_sql = """MERGE INTO ${fullTableAddress} base USING update ON ${whereExpr}
                 WHEN MATCHED THEN UPDATE SET * WHEN NOT MATCHED THEN INSERT *"""
-            self._get_spark().sql(upsert_sql)
+            _get_spark().sql(upsert_sql)
         else:
             self._save_append(update_data)
 
@@ -379,13 +376,13 @@ def _save(self, data: Union[DataFrame, pd.DataFrame]) -> None:
         if self._table.schema():
             cols = self._table.schema().fieldNames()
             if self._table.dataframe_type == "pandas":
-                data = self._get_spark().createDataFrame(
+                data = _get_spark().createDataFrame(
                     data.loc[:, cols], schema=self._table.schema()
                 )
             else:
                 data = data.select(*cols)
         elif self._table.dataframe_type == "pandas":
-            data = self._get_spark().createDataFrame(data)
+            data = _get_spark().createDataFrame(data)
         if self._table.write_mode == "overwrite":
             self._save_overwrite(data)
         elif self._table.write_mode == "upsert":
@@ -420,7 +417,7 @@ def _exists(self) -> bool:
         """
         if self._table.catalog:
             try:
-                self._get_spark().sql(f"USE CATALOG `{self._table.catalog}`")
+                _get_spark().sql(f"USE CATALOG `{self._table.catalog}`")
             except (ParseException, AnalysisException) as exc:
                 logger.warning(
                     "catalog %s not found or unity not enabled. Error message: %s",
@@ -429,7 +426,7 @@ def _exists(self) -> bool:
                 )
         try:
             return (
-                self._get_spark()
+                _get_spark()
                 .sql(f"SHOW TABLES IN `{self._table.database}`")
                 .filter(f"tableName = '{self._table.table}'")
                 .count()

kedro-datasets/kedro_datasets/spark/deltatable_dataset.py

@@ -11,7 +11,7 @@
 
 from kedro_datasets import KedroDeprecationWarning
 from kedro_datasets._io import AbstractDataset, DatasetError
-from kedro_datasets.spark.spark_dataset import _split_filepath, _strip_dbfs_prefix
+from kedro_datasets.spark.spark_dataset import _get_spark, _split_filepath, _strip_dbfs_prefix
 
 
 class DeltaTableDataset(AbstractDataset[None, DeltaTable]):
@@ -81,13 +81,9 @@ def __init__(self, filepath: str, metadata: Dict[str, Any] = None) -> None:
         self._filepath = PurePosixPath(filepath)
         self.metadata = metadata
 
-    @staticmethod
-    def _get_spark():
-        return SparkSession.builder.getOrCreate()
-
     def _load(self) -> DeltaTable:
         load_path = self._fs_prefix + str(self._filepath)
-        return DeltaTable.forPath(self._get_spark(), load_path)
+        return DeltaTable.forPath(_get_spark(), load_path)
 
     def _save(self, data: None) -> NoReturn:
         raise DatasetError(f"{self.__class__.__name__} is a read only dataset type")
@@ -96,7 +92,7 @@ def _exists(self) -> bool:
         load_path = _strip_dbfs_prefix(self._fs_prefix + str(self._filepath))
 
         try:
-            self._get_spark().read.load(path=load_path, format="delta")
+            _get_spark().read.load(path=load_path, format="delta")
         except AnalysisException as exception:
             # `AnalysisException.desc` is deprecated with pyspark >= 3.4
             message = exception.desc if hasattr(exception, "desc") else str(exception)

kedro-datasets/kedro_datasets/spark/spark_dataset.py

@@ -13,6 +13,7 @@
 from warnings import warn
 
 import fsspec
+import pyspark
 from hdfs import HdfsError, InsecureClient
 from kedro.io.core import Version, get_filepath_str, get_protocol_and_path
 from pyspark.sql import DataFrame, SparkSession
@@ -26,6 +27,27 @@
 logger = logging.getLogger(__name__)
 
 
+def _get_spark() -> Any:
+    """
+    Returns the SparkSession, we need this wrapper because the SparkSession
+    is retrieved differently in databricks-connect vs a Notebook in web UI
+    """
+    try:
+        # When using databricks-connect >= 13.0.0 (a.k.a databricks-connect-v2)
+        # the remote session is instantiated using the databricks module 
+        # If the databricks-connect module is installed, we use a remote session 
+        from databricks.connect import DatabricksSession
+
+        spark = DatabricksSession.builder.getOrCreate()
+
+    except ImportError:
+        # For "normal" spak sessions that don't use databricks-connect
+        # we get spark normally
+        spark = SparkSession.builder.getOrCreate()
+
+    return spark
+
+
 def _parse_glob_pattern(pattern: str) -> str:
     special = ("*", "?", "[")
     clean = []
@@ -315,7 +337,7 @@ def __init__(  # noqa: PLR0913
         elif filepath.startswith("/dbfs/"):
             # dbfs add prefix to Spark path by default
             # See https://github.com/kedro-org/kedro-plugins/issues/117
-            dbutils = _get_dbutils(self._get_spark())
+            dbutils = _get_dbutils(_get_spark())
             if dbutils:
                 glob_function = partial(_dbfs_glob, dbutils=dbutils)
                 exists_function = partial(_dbfs_exists, dbutils=dbutils)
@@ -383,13 +405,9 @@ def _describe(self) -> Dict[str, Any]:
             "version": self._version,
         }
 
-    @staticmethod
-    def _get_spark():
-        return SparkSession.builder.getOrCreate()
-
     def _load(self) -> DataFrame:
         load_path = _strip_dbfs_prefix(self._fs_prefix + str(self._get_load_path()))
-        read_obj = self._get_spark().read
+        read_obj = _get_spark().read
 
         # Pass schema if defined
         if self._schema:
@@ -405,7 +423,7 @@ def _exists(self) -> bool:
         load_path = _strip_dbfs_prefix(self._fs_prefix + str(self._get_load_path()))
 
         try:
-            self._get_spark().read.load(load_path, self._file_format)
+            _get_spark().read.load(load_path, self._file_format)
         except AnalysisException as exception:
             # `AnalysisException.desc` is deprecated with pyspark >= 3.4
             message = exception.desc if hasattr(exception, "desc") else str(exception)

kedro-datasets/kedro_datasets/spark/spark_hive_dataset.py

@@ -6,11 +6,12 @@
 from copy import deepcopy
 from typing import Any, Dict, List
 
-from pyspark.sql import DataFrame, SparkSession, Window
+from pyspark.sql import DataFrame, Window
 from pyspark.sql.functions import col, lit, row_number
 
 from kedro_datasets import KedroDeprecationWarning
 from kedro_datasets._io import AbstractDataset, DatasetError
+from kedro_datasets.spark.spark_dataset import _get_spark
 
 
 class SparkHiveDataset(AbstractDataset[DataFrame, DataFrame]):
@@ -135,20 +136,6 @@ def _describe(self) -> Dict[str, Any]:
             "format": self._format,
         }
 
-    @staticmethod
-    def _get_spark() -> SparkSession:
-        """
-        This method should only be used to get an existing SparkSession
-        with valid Hive configuration.
-        Configuration for Hive is read from hive-site.xml on the classpath.
-        It supports running both SQL and HiveQL commands.
-        Additionally, if users are leveraging the `upsert` functionality,
-        then a `checkpoint` directory must be set, e.g. using
-        `spark.sparkContext.setCheckpointDir("/path/to/dir")`
-        """
-        _spark = SparkSession.builder.getOrCreate()
-        return _spark
-
     def _create_hive_table(self, data: DataFrame, mode: str = None):
         _mode: str = mode or self._write_mode
         data.write.saveAsTable(
@@ -159,7 +146,7 @@ def _create_hive_table(self, data: DataFrame, mode: str = None):
         )
 
     def _load(self) -> DataFrame:
-        return self._get_spark().read.table(self._full_table_address)
+        return _get_spark().read.table(self._full_table_address)
 
     def _save(self, data: DataFrame) -> None:
         self._validate_save(data)
@@ -211,7 +198,7 @@ def _validate_save(self, data: DataFrame):
 
     def _exists(self) -> bool:
         return (
-            self._get_spark()
+            _get_spark()
             ._jsparkSession.catalog()
             .tableExists(self._database, self._table)
         )

kedro-datasets/kedro_datasets/spark/spark_jdbc_dataset.py

@@ -3,10 +3,11 @@
 from copy import deepcopy
 from typing import Any, Dict
 
-from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import DataFrame
 
 from kedro_datasets import KedroDeprecationWarning
 from kedro_datasets._io import AbstractDataset, DatasetError
+from kedro_datasets.spark.spark_dataset import _get_spark
 
 
 class SparkJDBCDataset(AbstractDataset[DataFrame, DataFrame]):
@@ -166,12 +167,8 @@ def _describe(self) -> Dict[str, Any]:
             "save_args": save_args,
         }
 
-    @staticmethod
-    def _get_spark():  # pragma: no cover
-        return SparkSession.builder.getOrCreate()
-
     def _load(self) -> DataFrame:
-        return self._get_spark().read.jdbc(self._url, self._table, **self._load_args)
+        return _get_spark().read.jdbc(self._url, self._table, **self._load_args)
 
     def _save(self, data: DataFrame) -> None:
         return data.write.jdbc(self._url, self._table, **self._save_args)

kedro-datasets/kedro_datasets/spark/spark_streaming_dataset.py

@@ -4,13 +4,14 @@
 from pathlib import PurePosixPath
 from typing import Any, Dict
 
-from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import DataFrame
 from pyspark.sql.utils import AnalysisException
 
 from kedro_datasets import KedroDeprecationWarning
 from kedro_datasets._io import AbstractDataset
 from kedro_datasets.spark.spark_dataset import (
     SparkDataset,
+    _get_spark,
     _split_filepath,
     _strip_dbfs_prefix,
 )
@@ -104,10 +105,6 @@ def _describe(self) -> Dict[str, Any]:
             "save_args": self._save_args,
         }
 
-    @staticmethod
-    def _get_spark():
-        return SparkSession.builder.getOrCreate()
-
     def _load(self) -> DataFrame:
         """Loads data from filepath.
         If the connector type is kafka then no file_path is required, schema needs to be
@@ -117,7 +114,7 @@ def _load(self) -> DataFrame:
         """
         load_path = _strip_dbfs_prefix(self._fs_prefix + str(self._filepath))
         data_stream_reader = (
-            self._get_spark()
+            _get_spark()
             .readStream.schema(self._schema)
             .format(self._file_format)
             .options(**self._load_args)
@@ -146,7 +143,7 @@ def _exists(self) -> bool:
         load_path = _strip_dbfs_prefix(self._fs_prefix + str(self._filepath))
 
         try:
-            self._get_spark().readStream.schema(self._schema).load(
+            _get_spark().readStream.schema(self._schema).load(
                 load_path, self._file_format
             )
         except AnalysisException as exception: