cloe-nessy 0.3.8__py3-none-any.whl → 0.3.9__py3-none-any.whl

cloe_nessy/integration/reader/api_reader.py

@@ -39,6 +39,7 @@ class APIReader(BaseReader):
 
     def read(
         self,
+        *,
         endpoint: str = "",
         method: str = "GET",
         key: str | None = None,
@@ -66,7 +67,7 @@ class APIReader(BaseReader):
             max_retries: The maximum number of retries for the request.
             options: Additional options for the createDataFrame function.
             add_metadata_column: If set, adds a __metadata column containing metadata about the API response.
-            kwargs: This method does not accept any additional keyword arguments.
+            **kwargs: Additional keyword arguments to maintain compatibility with the base class method.
 
         Returns:
             DataFrame: The Spark DataFrame containing the read data in the json_object column.
@@ -74,7 +75,8 @@ class APIReader(BaseReader):
         Raises:
             RuntimeError: If there is an error with the API request or reading the data.
         """
-        options = options or {}
+        if options is None:
+            options = {}
         try:
             response = self.api_client.request(
                 method=method,

cloe_nessy/integration/reader/catalog_reader.py

@@ -17,12 +17,13 @@ class CatalogReader(BaseReader):
         """Initializes the CatalogReader object."""
         super().__init__()
 
-    def read(self, table_identifier: str = "", **kwargs: Any) -> DataFrame:
+    def read(self, table_identifier: str = "", *, options: dict[str, str] | None = None, **kwargs: Any) -> DataFrame:
         """Reads a table from the Unity Catalog.
 
         Args:
             table_identifier: The table identifier in the Unity Catalog in the format 'catalog.schema.table'.
-            **kwargs: This method does not accept any additional keyword arguments.
+            options: PySpark options for the read table operation (not used in the current implementation).
+            **kwargs: Additional keyword arguments to maintain compatibility with the base class method.
 
         Returns:
             The Spark DataFrame containing the read data.
@@ -31,6 +32,8 @@ class CatalogReader(BaseReader):
             ValueError: If the table_identifier is not provided, is not a string, or is not in the correct format.
             Exception: For any other unexpected errors.
         """
+        if options is None:
+            options = {}
         if not table_identifier:
             raise ValueError("table_identifier is required")
         if not isinstance(table_identifier, str):
@@ -39,7 +42,7 @@ class CatalogReader(BaseReader):
             raise ValueError("table_identifier must be in the format 'catalog.schema.table'")
 
         try:
-            df = self._spark.read.table(table_identifier)
+            df = self._spark.read.table(table_identifier, **options)
             return df
         except AnalysisException as err:
             raise ValueError(f"Table not found: {table_identifier}") from err

cloe_nessy/integration/reader/excel_reader.py

@@ -92,7 +92,7 @@ class ExcelDataFrameReader(BaseReader):
                 pyspark.pandas.read_excel and handed to TextFileReader.
             load_as_strings: If True, converts all columns to string type to avoid datatype conversion errors in Spark.
             add_metadata_column: If True, adds a metadata column containing the file location and sheet name.
-            kwargs: This method does not accept any additional keyword arguments.
+            **kwargs: Additional keyword arguments to maintain compatibility with the base class method.
         """
         if options is None:
             options = {}

cloe_nessy/integration/reader/file_reader.py

@@ -1,7 +1,9 @@
 from typing import Any
 
 import pyspark.sql.functions as F
-from pyspark.sql import DataFrame
+from pyspark.sql import DataFrame, DataFrameReader
+from pyspark.sql.streaming import DataStreamReader
+from pyspark.sql.types import StructType
 
 from ...file_utilities import get_file_paths
 from .reader import BaseReader
@@ -17,9 +19,18 @@ class FileReader(BaseReader):
         """Initializes the FileReader object."""
         super().__init__()
 
+    def _get_reader(self) -> DataFrameReader:
+        """Returns a DataFrameReader."""
+        return self._spark.read
+
+    def _get_stream_reader(self) -> DataStreamReader:
+        """Returns a DataFrameReader."""
+        return self._spark.readStream
+
     def read(
         self,
         location: str,
+        *,
         spark_format: str | None = None,
         extension: str | None = None,
         schema: str | None = None,
@@ -38,7 +49,22 @@ class FileReader(BaseReader):
             search_subdirs: Whether to include files in subdirectories.
             options: Spark DataFrame reader options.
             add_metadata_column: Whether to include __metadata column in the DataFrame.
-            kwargs: This method does not accept any additional keyword arguments.
+            **kwargs: Additional keyword arguments to maintain compatibility with the base class method.
+
+        Raises:
+            ValueError: If neither spark_format nor extension is provided.
+            ValueError: If the provided extension is not supported.
+            Exception: If there is an error while reading the files.
+
+        Note:
+            - The `spark_format` parameter is used to specify the format of the files to be read.
+            - If `spark_format` is not provided, the method will try to infer it from the `extension`.
+            - The `extension` parameter is used to specify the file extension (e.g., 'csv', 'json', etc.).
+            - If both `spark_format` and `extension` are provided, `spark_format` will take precedence.
+            - The method will raise an error if neither `spark_format` nor `extension` is provided.
+
+        Returns:
+            A DataFrame containing the data from the files.
         """
         if options is None:
             options = {}
@@ -67,7 +93,7 @@ class FileReader(BaseReader):
         self._console_logger.debug(f"File paths: {file_paths}")
         assert spark_format is not None
 
-        reader = self._spark.read.format(spark_format)
+        reader = self._get_reader().format(spark_format)
         if schema:
             reader.schema(schema)
         else:
@@ -78,7 +104,7 @@ class FileReader(BaseReader):
 
         try:
             self._console_logger.debug("Loading files into DataFrame")
-            df = reader.load(file_paths)
+            df = reader.load([str(p) for p in file_paths])
             self._console_logger.debug("Successfully loaded files into DataFrame")
             if add_metadata_column:
                 df = self._add_metadata_column(df)
@@ -89,9 +115,56 @@ class FileReader(BaseReader):
             self._console_logger.info(f"Successfully read files from [ '{location}' ]")
             return df
 
+    def read_stream(
+        self,
+        location: str = "",
+        schema: StructType | str | None = None,
+        format: str = "delta",
+        add_metadata_column: bool = False,
+        options: dict[str, Any] | None = None,
+        **_: Any,
+    ) -> DataFrame:
+        """Reads specified location as a stream and returns streaming DataFrame.
+
+        Arguments:
+            location : Location of files to read.
+            format: Format of files to read.
+            schema: Schema of the file.
+            add_metadata_column: Whether to include __metadata column in the DataFrame.
+            options: Spark DataFrame reader options.
+
+        Raises:
+            ValueError: If location is not provided.
+
+        Returns:
+            A Streaming DataFrame
+        """
+        if not location:
+            raise ValueError("Location is required for streaming.")
+        self._console_logger.debug(f"Reading files from [ '{location}' ] ...")
+        try:
+            if options is None:
+                options = {}
+            reader = self._get_stream_reader()
+            reader.format(format)
+            reader.option("rescuedDataColumn", "_rescued_data")
+            if schema is None:
+                options["inferSchema"] = True
+            else:
+                reader.schema(schema)
+            reader.options(**options)
+            df = reader.load(location)
+            if add_metadata_column:
+                df = self._add_metadata_column(df)
+        except Exception as e:
+            self._console_logger.error(f"Failed to read files from [ '{location}' ]: {e}")
+            raise
+        else:
+            self._console_logger.info(f"Successfully read files from [ '{location}' ]")
+            return df
+
     def _add_metadata_column(self, df: DataFrame) -> DataFrame:
         """Add all metadata columns to the DataFrame."""
-        # Extract metadata fields into separate columns
         metadata_columns = df.select("_metadata.*").columns
 
         entries = [(F.lit(field), F.col(f"_metadata.{field}")) for field in metadata_columns]

cloe_nessy/integration/writer/__init__.py

@@ -1,3 +1,10 @@
 from .catalog_writer import CatalogWriter
+from .delta_writer import DeltaAppendWriter, DeltaMergeWriter
+from .file_writer import FileWriter
 
-__all__ = ["CatalogWriter"]
+__all__ = [
+    "CatalogWriter",
+    "DeltaAppendWriter",
+    "DeltaMergeWriter",
+    "FileWriter",
+]

cloe_nessy/integration/writer/delta_writer/__init__.py

@@ -0,0 +1,7 @@
+from .delta_append_writer import DeltaAppendWriter
+from .delta_merge_writer import DeltaMergeWriter
+
+__all__ = [
+    "DeltaAppendWriter",
+    "DeltaMergeWriter",
+]

cloe_nessy/integration/writer/delta_writer/delta_append_writer.py

@@ -0,0 +1,108 @@
+from pyspark.sql import DataFrame
+
+from ....object_manager import table_log_decorator
+from ....session import SessionManager
+from ..file_writer import FileWriter
+from .delta_table_operation_type import DeltaTableOperationType
+from .delta_writer_base import BaseDeltaWriter
+
+
+class DeltaAppendWriter(BaseDeltaWriter):
+    """A class for appending DataFrames to Delta tables."""
+
+    def __init__(self):
+        super().__init__()
+        self._spark = SessionManager.get_spark_session()
+        self._dbutils = SessionManager.get_utils()
+
+    @table_log_decorator(operation="append")
+    def write(
+        self,
+        table_identifier: str,
+        table_location: str,
+        data_frame: DataFrame,
+        ignore_empty_df: bool = False,
+        options: dict[str, str] | None = None,
+    ):
+        """Appends the provided DataFrame to a Delta table.
+
+        Args:
+            table_identifier: The identifier of the Delta table in the format 'catalog.schema.table'.
+            table_location: The location of the Delta table.
+            data_frame: The DataFrame to append to the table.
+            ignore_empty_df: If True, the function returns early without
+                doing anything if the DataFrame is empty.
+            options: Additional keyword arguments that will be passed to the 'write' method of the
+                FileDataFrameWriter instance. These can be any parameters accepted by the 'write'
+                method, which could include options for configuring the write operation, such as
+                'checkpointLocation' for specifying the path where checkpoints will be stored, or
+                'path' for specifying the path where the output data will be written.
+        """
+        if self._empty_dataframe_check(data_frame, ignore_empty_df):
+            return
+        writer = FileWriter()
+        writer.write(
+            data_frame=data_frame,
+            location=table_location,
+            format="DELTA",
+            mode="APPEND",
+            options=options,
+        )
+        self._report_delta_table_operation_metrics(
+            table_identifier=table_identifier, operation_type=DeltaTableOperationType.WRITE
+        )
+
+    @table_log_decorator(operation="stream_append")
+    def write_stream(
+        self,
+        table_identifier: str,
+        table_location: str,
+        data_frame: DataFrame,
+        checkpoint_location: str | None = None,
+        trigger_dict: dict | None = None,
+        options: dict[str, str] | None = None,
+        await_termination: bool = False,
+    ):
+        """Appends the provided DataFrame to a Delta table.
+
+        Args:
+            table_identifier: The identifier of the Delta table in the format 'catalog.schema.table'.
+            table_location: The location of the Delta table.
+            data_frame: The DataFrame to append to the table.
+            checkpoint_location: Location of checkpoint. If None, defaults
+                to the location of the table being written, with '_checkpoint_'
+                added before name. Default None.
+            trigger_dict: A dictionary specifying the trigger configuration for the streaming query.
+                Supported keys include:
+
+                - "processingTime": Specifies a time interval (e.g., "10 seconds") for micro-batch processing.
+                - "once": Processes all available data once and then stops.
+                - "continuous": Specifies a time interval (e.g., "1 second") for continuous processing.
+                - "availableNow": Processes all available data immediately and then stops.
+
+                If nothing is provided, the default is {"availableNow": True}.
+            options: Additional keyword arguments that will be passed to the
+                'write' method of the FileDataFrameWriter instance. These can be
+                any parameters accepted by the 'write' method, which could
+                include options for configuring the write operation.
+            await_termination: If True, the function will wait for the streaming
+                query to finish before returning. This is useful for ensuring that
+                the data has been fully written before proceeding with other
+                operations.
+
+        Returns:
+            None.
+        """
+        writer = FileWriter()
+        writer.write_stream(
+            data_frame=data_frame,
+            location=table_location,
+            format="DELTA",
+            checkpoint_location=checkpoint_location,
+            mode="APPEND",
+            trigger_dict=trigger_dict,
+            options=options,
+        )
+        self._report_delta_table_operation_metrics(
+            table_identifier=table_identifier, operation_type=DeltaTableOperationType.WRITE
+        )

cloe_nessy/integration/writer/delta_writer/delta_merge_writer.py

@@ -0,0 +1,215 @@
+from typing import Any, Self
+
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pyspark.sql import DataFrame
+
+from ....models import Table
+from ....object_manager import table_log_decorator
+from ....session import SessionManager
+from .delta_table_operation_type import DeltaTableOperationType
+from .delta_writer_base import BaseDeltaWriter
+
+
+class DeltaMergeConfig(BaseModel):
+    """Configuration for Merge options.
+
+    Args:
+        dataframe_columns: The columns of the DataFrame.
+        key_columns: List of column names that form the key for the merge
+            operation.
+        when_matched_update: Flag to specify whether to perform an update
+            operation when matching records are found in the target Delta table.
+        when_matched_delete: Flag to specify whether to perform a delete
+            operation when matching records are found in the target Delta table.
+        when_not_matched_insert: Flag to specify whether to perform an insert
+            operation when matching records are not found in the target Delta
+            table.
+        cols_to_exclude_from_update: List of column names to be excluded from
+            the update in the target Delta table.
+        use_partition_pruning: Flag to specify whether to use partition
+            pruning to optimize the performance of the merge operation.
+        partition_by: List of column names to partition by.
+    """
+
+    dataframe_columns: list[str]
+    key_columns: list[str]
+    cols_to_exclude_from_update: list[str] = Field(default_factory=list)
+    when_matched_update: bool = True
+    when_matched_delete: bool = False
+    when_not_matched_insert: bool = True
+    use_partition_pruning: bool = True
+    partition_by: list[str] = Field(default_factory=list)
+    cols_to_merge: list[str] = Field(default_factory=list, alias="_cols_to_merge")
+    cols_to_update: set[str] = Field(default_factory=set, alias="_cols_to_update")
+    cols_to_insert: set[str] = Field(default_factory=set, alias="_cols_to_insert")
+    final_cols_to_update: dict[str, str] = Field(default_factory=dict)
+    final_cols_to_insert: dict[str, str] = Field(default_factory=dict)
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate_update_delete(cls, config: Any):
+        """Update and delete operations must be mutually exclusive."""
+        if config.get("when_matched_update") and config.get("when_matched_delete"):
+            raise ValueError("Update and delete operations cannot be used together.")
+        return config
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate_key_columns(cls, config: Any):
+        """Key columns must exist in the data frame."""
+        key_columns = config.get("key_columns")
+        dataframe_columns = config.get("dataframe_columns")
+        if not set(key_columns).issubset(set(dataframe_columns)):
+            raise ValueError("Key columns must exist in the DataFrame columns.")
+        return config
+
+    @model_validator(mode="before")
+    @classmethod
+    def _derive_merge_columns(cls, config: Any):
+        """Derive update and insert columns from the DataFrame columns."""
+        dataframe_columns = config.get("dataframe_columns", [])
+        config["_cols_to_merge"] = list(set(dataframe_columns))
+        if config.get("cols_to_exclude_from_update"):
+            config["_cols_to_update"] = set(config["_cols_to_merge"]) - set(config["cols_to_exclude_from_update"])
+        else:
+            config["_cols_to_update"] = set(config["_cols_to_merge"])
+
+        config["_cols_to_insert"] = config["_cols_to_merge"]
+        config["final_cols_to_update"] = {col: f"source.{col}" for col in config["_cols_to_update"]}
+        config["final_cols_to_insert"] = {col: f"source.{col}" for col in config["_cols_to_insert"]}
+        return config
+
+    @model_validator(mode="after")
+    @classmethod
+    def _validate_partition_pruning(cls, config: Self):
+        """If partition_pruning is set, the partition by columns must be known."""
+        if config.use_partition_pruning is True and not config.partition_by:
+            raise ValueError("Partition columns must be specified when using partition pruning.")
+        return config
+
+    @model_validator(mode="after")
+    @classmethod
+    def _validate_cols_exist(cls, config: Any):
+        """If partition_pruning is set, the partition by columns must be known."""
+        if any(col not in config.cols_to_merge for col in config.cols_to_update) or any(
+            col not in config.cols_to_merge for col in config.cols_to_insert
+        ):
+            raise ValueError(
+                "You specified column names for UPDATE or INSERT that either don't exist in the dataframe "
+                "or are explicitly excluded from the MERGE.",
+            )
+        return config
+
+
+class DeltaMergeWriter(BaseDeltaWriter):
+    """A class for merging DataFrames to Delta tables."""
+
+    def __init__(self):
+        super().__init__()
+        self._spark = SessionManager.get_spark_session()
+        self._dbutils = SessionManager.get_utils()
+
+    def _validate_table_inputs(
+        self, table: Table | None, table_identifier: str | None, storage_path: str | None
+    ) -> tuple[str, str]:
+        """Validates and retrieves table identifier and storage path."""
+        if table is None and (table_identifier is None or storage_path is None):
+            raise ValueError("Either a Table object or table_identifier and storage_path must be provided.")
+        if table is not None:
+            table_identifier = table.identifier
+            storage_path = str(table.storage_path)
+        if not storage_path:
+            raise ValueError("Storage path must be provided or extracted from the Table object.")
+        assert table_identifier is not None, "Table identifier must be provided."
+        return table_identifier, storage_path
+
+    def _build_match_conditions(self, data_frame: DataFrame, config: DeltaMergeConfig) -> str:
+        """Builds match conditions for the Delta table merge."""
+        match_conditions = self._merge_match_conditions(config.key_columns)
+        if config.use_partition_pruning:
+            match_conditions_list = [match_conditions] + [
+                self._partition_pruning_conditions(data_frame, config.partition_by),
+            ]
+            match_conditions = " AND ".join(match_conditions_list)
+        return match_conditions
+
+    def _build_merge_operations(
+        self, delta_table, data_frame: DataFrame, config: DeltaMergeConfig, match_conditions: str
+    ):
+        """Builds the Delta table merge operations."""
+        delta_table_merge = delta_table.alias("target").merge(
+            source=data_frame.alias("source"),
+            condition=match_conditions,
+        )
+        if config.when_matched_update:
+            delta_table_merge = delta_table_merge.whenMatchedUpdate(set=config.final_cols_to_update)
+        elif config.when_matched_delete:
+            delta_table_merge = delta_table_merge.whenMatchedDelete()
+        if config.when_not_matched_insert:
+            delta_table_merge = delta_table_merge.whenNotMatchedInsert(values=config.final_cols_to_insert)
+        return delta_table_merge
+
+    @table_log_decorator(operation="merge")
+    def write(
+        self,
+        data_frame: DataFrame,
+        table: Table | None = None,
+        table_identifier: str | None = None,
+        storage_path: str | None = None,
+        ignore_empty_df: bool = False,
+        **kwargs: Any,
+    ):
+        """Merges the data in a spark DataFrame into a Delta table.
+
+        This function performs a merge operation between a DataFrame and a Delta
+        table. The function supports update, delete, and insert operations on
+        the target Delta table based on conditions specified by the user. The
+        function also supports partition pruning to optimize the performance of
+        the merge operation.
+
+        Args:
+            table: The Table object representing the Delta table.
+            table_identifier: The identifier of the Delta table in the format
+                'catalog.schema.table'.
+            storage_path: The location of the Delta table.
+            data_frame: The DataFrame to be merged into the Delta table.
+            ignore_empty_df: A flag indicating whether to ignore an empty source
+                dataframe.
+            kwargs: Passed to the
+                [`DeltaMergeConfig`][cloe_nessy.integration.writer.delta_merge_writer.DeltaMergeConfig].
+
+        Raises:
+            ValueError: If both, table and table_identifier or storage_path are provided.
+            EmptyDataframeException: If the source dataframe is empty and
+                ignore_empty_df is False.
+            ValueError: If the specified columns for update or insert do not
+                exist in the DataFrame or are explicitly excluded from the
+                merge operation.
+            ValueError: If partition columns are not specified when using
+                partition pruning.
+        """
+        if self._empty_dataframe_check(data_frame, ignore_empty_df):
+            return
+        table_identifier, storage_path = self._validate_table_inputs(table, table_identifier, storage_path)
+
+        config = DeltaMergeConfig(dataframe_columns=data_frame.columns, **kwargs)
+
+        delta_table = self.table_manager.get_delta_table(location=storage_path, spark=data_frame.sparkSession)
+
+        match_conditions = self._build_match_conditions(data_frame, config)
+
+        delta_table_merge = self._build_merge_operations(delta_table, data_frame, config, match_conditions)
+        delta_table_merge.execute()
+        self._report_delta_table_operation_metrics(
+            table_identifier,
+            operation_type=DeltaTableOperationType.MERGE,
+        )
+
+    @table_log_decorator(operation="stream_merge")
+    def write_stream(self):
+        """Not implemented yet. See docs for more details."""
+        raise NotImplementedError(
+            "Streaming merge is not implemented yet. Please use the `write` method for batch merges."
+        )

cloe_nessy/integration/writer/delta_writer/delta_table_operation_type.py

@@ -0,0 +1,21 @@
+from enum import Enum
+
+
+class DeltaTableOperationType(Enum):
+    """Mapping between Delta table operation types and their operation metric keys available in the Delta table history.
+
+    Values of metric keys included in this mapping are reported using the
+    logging capabilities of the Delta operations of the DeltaManager.
+
+    See https://docs.databricks.com/delta/history.html for a complete list and
+    description of available metrics for each operation type.
+    """
+
+    UPDATE = ["numUpdatedRows"]
+    DELETE = ["numDeletedRows", "numRemovedFiles"]
+    MERGE = ["numSourceRows", "numTargetRowsInserted", "numTargetRowsUpdated", "numTargetRowsDeleted", "numOutputRows"]
+    WRITE = ["numOutputRows"]
+    TRUNCATE = ["numRemovedFiles"]
+    OPTIMIZE = ["numAddedFiles", "numRemovedFiles", "minFileSize", "p50FileSize", "maxFileSize"]
+    VACUUM = ["numDeletedFiles"]
+    STREAMING_UPDATE = ["numRemovedFiles", "numOutputRows", "numOutputBytes", "numAddedFiles"]