cloe-nessy 0.2.9__py3-none-any.whl

cloe_nessy/pipeline/actions/read_catalog_table.py

@@ -0,0 +1,68 @@
+from typing import Any
+
+from ...integration.reader import CatalogReader
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class ReadCatalogTableAction(PipelineAction):
+    """Reads a table from Unity Catalog using a specified table identifier and optional reader configurations.
+
+    This function retrieves data from a catalog table using the
+    [`CatalogReader`][cloe_nessy.integration.reader.catalog_reader] identified
+    by either the `table_identifier` parameter or the `table_metadata` from the
+    provided `PipelineContext` of a previous step. The retrieved data is loaded
+    into a DataFrame and returned as part of an updated `PipelineContext`.
+
+    Example:
+    ```yaml
+    Read Sales Table:
+        action: READ_CATALOG_TABLE
+        options:
+            table_identifier: my_catalog.business_schema.sales_table
+            options: <options for the reader>
+    ```
+    """
+
+    name: str = "READ_CATALOG_TABLE"
+
+    @staticmethod
+    def run(
+        context: PipelineContext,
+        *,
+        table_identifier: str | None = None,
+        options: dict[str, str] | None = None,
+        **_: Any,  # define kwargs to match the base class signature
+    ) -> PipelineContext:
+        """Reads a table from Unity Catalog using a specified table identifier and optional reader configurations.
+
+        Args:
+            context: The pipeline's context, which contains
+                metadata and configuration for the action.
+            table_identifier: The identifier of the catalog table to
+                read. If not provided, the function will attempt to use the table
+                identifier from the `table_metadata` in the `context`.
+            options: A dictionary of options for customizing
+                the catalog reader's behavior, such as filters or reading modes. Defaults
+                to None.
+
+        Raises:
+            ValueError: If neither `table_identifier` nor `table_metadata.identifier` in the `context` is provided.
+
+        Returns:
+        An updated pipeline context containing the data read from the catalog table as a DataFrame.
+        """
+        if not options:
+            options = dict()
+
+        if (table_metadata := context.table_metadata) and table_identifier is None:
+            table_identifier = table_metadata.identifier
+        if table_identifier is None:
+            raise ValueError("Table name must be specified or a valid Table object with identifier must be set.")
+
+        table_reader = CatalogReader()
+        df = table_reader.read(
+            table_identifier=table_identifier,
+            **options,
+        )
+        return context.from_existing(data=df)

cloe_nessy/pipeline/actions/read_excel.py

@@ -0,0 +1,177 @@
+from collections.abc import Callable
+from functools import reduce
+
+from pyspark.sql import DataFrame
+
+from ...file_utilities import get_file_paths
+from ...integration.reader import ExcelDataFrameReader
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class ReadExcelAction(PipelineAction):
+    """Reads data from an Excel file or directory of Excel files and returns a DataFrame.
+
+    The function reads Excel files using the
+    [`ExcelDataFrameReader`][cloe_nessy.integration.reader.excel_reader] either
+    from a single file or a directory path. It can read specific sheets, handle
+    file extensions, and offers various options to customize how the data is
+    read, such as specifying headers, index columns, and handling missing
+    values. The resulting data is returned as a DataFrame, and metadata about
+    the read files can be included in the context.
+
+    Example:
+    ```yaml
+    Read Excel Table:
+        action: READ_EXCEL
+        options:
+            file: excel_file_folder/excel_files_june/interesting_excel_file.xlsx
+            usecols:
+                - key_column
+                - interesting_column
+            options: <more options for the reader>
+    ```
+    """
+
+    name: str = "READ_EXCEL"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        file: str | None = None,
+        path: str | None = None,
+        extension: str = "xlsx",
+        recursive: bool = False,
+        sheet_name: str | int | list = 0,
+        sheet_name_as_column: bool = False,
+        header: int | list[int] = 0,
+        index_col: int | list[int] | None = None,
+        usecols: int | str | list | Callable | None = None,
+        dtype: str | None = None,
+        fillna: str | dict[str, list[str]] | dict[str, str] | None = None,
+        true_values: list | None = None,
+        false_values: list | None = None,
+        nrows: int | None = None,
+        na_values: list[str] | dict[str, list[str]] | None = None,
+        keep_default_na: bool = True,
+        parse_dates: bool | list | dict = False,
+        date_parser: Callable | None = None,
+        thousands: str | None = None,
+        include_index: bool = False,
+        options: dict | None = None,
+        add_metadata_column: bool = True,
+        load_as_strings: bool = False,
+        **_,
+    ) -> PipelineContext:
+        """Reads data from an Excel file or directory of Excel files and returns a DataFrame.
+
+        Args:
+            context: The context in which the action is executed.
+            file: The path to a single Excel file. Either `file` or `path` must be specified.
+            path: The directory path containing multiple Excel files. Either `file` or `path` must be specified.
+            extension: The file extension to look for when reading from a directory.
+            recursive: Whether to include subdirectories when reading from a directory path.
+            sheet_name: The sheet name(s) or index(es) to read from the Excel file.
+            sheet_name_as_column: Whether to add a column with the sheet name to the DataFrame.
+            header: Row number(s) to use as the column labels.
+            index_col: Column(s) to use as the index of the DataFrame.
+            usecols: Subset of columns to parse. Can be an integer, string, list,
+                or function.
+            dtype: Data type for the columns.
+            fillna: Method or value to use to fill NaN values.
+            true_values: Values to consider as True.
+            false_values: Values to consider as False.
+            nrows: Number of rows to parse.
+            na_values: Additional strings to recognize as NaN/NA.
+            keep_default_na: Whether to append default NaN values when custom `na_values` are specified.
+            parse_dates: Options for parsing date columns.
+            date_parser: Function to use for converting strings to datetime objects.
+            thousands: Thousands separator to use when parsing numeric columns.
+            include_index: Whether to include an index column in the output DataFrame.
+            options: Additional options to pass to the DataFrame reader.
+            add_metadata_column: Whether to add a metadata column with file information to the DataFrame.
+            load_as_strings: Whether to load all columns as strings.
+
+        Raises:
+            ValueError: Raised if both `file` and `path` are specified, or if neither is provided.
+
+        Returns:
+            The updated context, with the read data as a DataFrame.
+        """
+        if not options:
+            options = dict()
+
+        if file is not None and path is not None:
+            self._tabular_logger.error("message: Only one of file or path have to be specified.")
+            raise ValueError("Only one of file or path have to be specified.")
+
+        excel_reader = ExcelDataFrameReader()
+        if file is not None:
+            df = excel_reader.read(
+                location=file,
+                sheet_name=sheet_name,
+                sheet_name_as_column=sheet_name_as_column,
+                header=header,
+                index_col=index_col,
+                usecols=usecols,
+                true_values=true_values,
+                false_values=false_values,
+                nrows=nrows,
+                dtype=dtype,
+                fillna=fillna,
+                na_values=na_values,
+                keep_default_na=keep_default_na,
+                parse_dates=parse_dates,
+                date_parser=date_parser,
+                thousands=thousands,
+                include_index=include_index,
+                options=options,
+                add_metadata_column=add_metadata_column,
+                load_as_strings=load_as_strings,
+            )
+        elif path is not None:
+            file_list = get_file_paths(path, extension, recursive)
+            df_dict: dict = {}
+            for path in file_list:
+                df_dict[path] = excel_reader.read(
+                    location=path,
+                    sheet_name=sheet_name,
+                    sheet_name_as_column=sheet_name_as_column,
+                    header=header,
+                    index_col=index_col,
+                    usecols=usecols,
+                    dtype=dtype,
+                    fillna=fillna,
+                    true_values=true_values,
+                    false_values=false_values,
+                    nrows=nrows,
+                    na_values=na_values,
+                    keep_default_na=keep_default_na,
+                    parse_dates=parse_dates,
+                    date_parser=date_parser,
+                    thousands=thousands,
+                    include_index=include_index,
+                    options=options,
+                    add_metadata_column=add_metadata_column,
+                    load_as_strings=load_as_strings,
+                )
+            df = reduce(DataFrame.unionAll, list(df_dict.values()))
+
+        else:
+            self._tabular_logger.error("action_name: READ_EXCEL | message: Either file or path have to be specified.")
+            raise ValueError("Either file or path have to be specified.")
+
+        runtime_info = context.runtime_info
+
+        if add_metadata_column:
+            read_files_list = list(set([x.file_path for x in df.select("__metadata.file_path").collect()]))
+            if runtime_info is None:
+                runtime_info = {"read_files": read_files_list}
+            else:
+                try:
+                    runtime_info["read_files"] = list(set(runtime_info["read_files"] + read_files_list))
+                except KeyError:
+                    runtime_info["read_files"] = read_files_list
+
+        return context.from_existing(data=df)

cloe_nessy/pipeline/actions/read_files.py

@@ -0,0 +1,105 @@
+from typing import Any
+
+from ...integration.reader import FileReader
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class ReadFilesAction(PipelineAction):
+    """Reads files from a specified location.
+
+    If an extension is provided, all files with the given extension will be read
+    using the [`FileReader`][cloe_nessy.integration.reader.file_reader]. If no
+    extension is provided, the `spark_format` must be set, and all files in the
+    location will be read using a DataFrameReader with the specified format.
+
+    Example:
+    ```yaml
+    Read Excel Table:
+        action: READ_FILES
+        options:
+            location: excel_file_folder/excel_files_june/
+            search_subdirs: True
+            spark_format: AVRO
+    ```
+    """
+
+    name: str = "READ_FILES"
+
+    @staticmethod
+    def run(
+        context: PipelineContext,
+        *,
+        location: str | None = None,
+        search_subdirs: bool = False,
+        extension: str | None = None,
+        spark_format: str | None = None,
+        schema: str | None = None,
+        add_metadata_column: bool = True,
+        options: dict[str, str] | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Reads files from a specified location.
+
+        Args:
+            context: The context in which this Action is executed.
+            location: The location from which to read files.
+            search_subdirs: Recursively search subdirectories for files
+                if an extension is provided.
+            extension: The file extension to filter files by.
+            spark_format: The format to use for reading the files.
+            schema: The schema of the data. If None, schema is obtained from
+                the context metadata.
+            add_metadata_column: Whether to include the `__metadata` column with
+                file metadata in the DataFrame.
+            options: Additional options passed to the reader.
+
+        Raises:
+            ValueError: If neither `extension` nor `spark_format` are provided, or if
+                no location is specified.
+
+        Returns:
+            The context after the Action has been executed, containing the read data as a DataFrame.
+        """
+        if not location:
+            raise ValueError("No location provided. Please specify location to read files from.")
+        if not options:
+            options = dict()
+
+        if (metadata := context.table_metadata) and schema is None:
+            schema = metadata.schema
+
+        file_reader = FileReader()
+        if extension:
+            df = file_reader.read(
+                location=location,
+                schema=schema,
+                extension=extension,
+                search_subdirs=search_subdirs,
+                options=options,
+                add_metadata_column=add_metadata_column,
+            )
+        elif spark_format:
+            df = file_reader.read(
+                location=location,
+                schema=schema,
+                spark_format=spark_format,
+                options=options,
+                add_metadata_column=add_metadata_column,
+            )
+        else:
+            raise ValueError("Please provide either the 'extension' or 'spark_format'")
+
+        runtime_info = context.runtime_info
+
+        if add_metadata_column:
+            read_files_list = [x.file_path for x in df.select("__metadata.file_path").drop_duplicates().collect()]
+            if runtime_info is None:
+                runtime_info = {"read_files": read_files_list}
+            else:
+                try:
+                    runtime_info["read_files"] = list(set(runtime_info["read_files"] + read_files_list))
+                except KeyError:
+                    runtime_info["read_files"] = read_files_list
+
+        return context.from_existing(data=df, runtime_info=runtime_info)

cloe_nessy/pipeline/actions/read_metadata_yaml.py

@@ -0,0 +1,66 @@
+import pathlib
+from typing import Any
+
+from ...models import Schema
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class ReadMetadataYAMLAction(PipelineAction):
+    """Reads schema metadata from a yaml file using the [`Schema`][cloe_nessy.models.schema] model.
+
+    Example:
+    ```yaml
+    Read Schema Metadata:
+        action: READ_METADATA_YAML_ACTION
+        options:
+            path: excel_file_folder/excel_files_june/
+            file_name: sales_schema.yml
+            table_name: sales
+    ```
+    """
+
+    name: str = "READ_METADATA_YAML_ACTION"
+
+    @staticmethod
+    def run(
+        context: PipelineContext,
+        *,
+        path: str | None = None,
+        file_name: str | None = None,
+        table_name: str | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Reads schema metadata from a yaml file using the `Schema` model.
+
+        Args:
+            context: The context in which this Action is executed.
+            path: The path to the data contract directory.
+            file_name: The name of the file that defines the schema.
+            table_name: The name of the table for which to retrieve metadata.
+
+        Raises:
+            ValueError: If any issues occur while reading the schema, such as an invalid schema,
+                missing file, or missing path.
+
+        Returns:
+            The context after the execution of this Action, containing the table metadata.
+        """
+        if not path:
+            raise ValueError("No path provided. Please specify path to schema metadata.")
+        if not file_name:
+            raise ValueError("No file_name provided. Please specify file name.")
+        if not table_name:
+            raise ValueError("No table_name provided. Please specify table name.")
+
+        path_obj = pathlib.Path(path)
+
+        schema, errors = Schema.read_instance_from_file(path_obj / file_name)
+        if errors:
+            raise ValueError(f"Errors while reading schema metadata: {errors}")
+        if not schema:
+            raise ValueError("No schema found in metadata.")
+
+        table = schema.get_table_by_name(table_name=table_name)
+
+        return context.from_existing(table_metadata=table)

cloe_nessy/pipeline/actions/transform_change_datatype.py

@@ -0,0 +1,56 @@
+from typing import Any
+
+import pyspark.sql.functions as F
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformChangeDatatypeAction(PipelineAction):
+    """Changes the datatypes of specified columns in the given DataFrame.
+
+    Example:
+    ```yaml
+    Transform Columns:
+        action: TRANSFORM_CHANGE_DATATYPE
+        options:
+            columns:
+                id: string
+                revenue: long
+    ```
+    """
+
+    name: str = "TRANSFORM_CHANGE_DATATYPE"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        columns: dict[str, str] | None = None,
+        **_: Any,  # define kwargs to match the base class signature
+    ) -> PipelineContext:
+        """Changes the datatypes of specified columns in the given DataFrame.
+
+        Args:
+            context: The context in which this Action is executed.
+            columns: A dictionary where the key is the column
+                name and the value is the desired datatype.
+
+        Raises:
+            ValueError: If no columns are provided.
+            ValueError: If the data from context is None.
+
+        Returns:
+            The context after the execution of this Action, containing the DataFrame with updated column datatypes.
+        """
+        if not columns:
+            raise ValueError("No columns provided.")
+
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df = context.data
+        change_columns = {col: F.col(col).cast(dtype) for col, dtype in columns.items()}
+        df = df.withColumns(change_columns)  # type: ignore
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_concat_columns.py

@@ -0,0 +1,88 @@
+from typing import Any
+
+import pyspark.sql.functions as F
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformConcatColumnsAction(PipelineAction):
+    """Concatenates the specified columns in the given DataFrame.
+
+    Example:
+    ```yaml
+    Concat Columns:
+        action: TRANSFORM_CONCAT_COLUMNS
+        options:
+            name: address
+            columns:
+                - street
+                - postcode
+                - country
+            separator: ', '
+    ```
+    """
+
+    name: str = "TRANSFORM_CONCAT_COLUMNS"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        name: str = "",
+        columns: list[str] | None = None,
+        separator: str | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Concatenates the specified columns in the given DataFrame.
+
+        !!!warning
+
+            # Null Handling Behavior
+
+            The behavior of null handling differs based on whether a `separator` is provided:
+
+            - **When `separator` is specified**: The function uses Spark's
+                `concat_ws`, which **ignores `NULL` values**. In this case, `NULL`
+                values are treated as empty strings (`""`) and are excluded from the
+                final concatenated result.
+            - **When `separator` is not specified**: The function defaults to
+                using Spark's `concat`, which **returns `NULL` if any of the
+                concatenated values is `NULL`**. This means the presence of a `NULL`
+                in any input will make the entire output `NULL`.
+
+        Args:
+            context: The context in which this Action is executed.
+            name: The name of the new concatenated column.
+            columns: A list of columns to be concatenated.
+            separator: The separator used between concatenated column values.
+
+        Raises:
+            ValueError: If no name is provided.
+            ValueError: If no columns are provided.
+            ValueError: If the data from context is None.
+            ValueError: If 'columns' is not a list.
+
+        Returns:
+            The context after the execution of this Action, containing the
+                DataFrame with the concatenated column.
+        """
+        if not name:
+            raise ValueError("No name provided.")
+        if not columns:
+            raise ValueError("No columns provided.")
+
+        if context.data is None:
+            raise ValueError("The data from context is required for the operation.")
+
+        df = context.data
+
+        if isinstance(columns, list):
+            if separator:
+                df = df.withColumn(name, F.concat_ws(separator, *columns))  # type: ignore
+            else:
+                df = df.withColumn(name, F.concat(*columns))  # type: ignore
+        else:
+            raise ValueError("'columns' should be a list, like ['col1', 'col2',]")
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_decode.py

@@ -0,0 +1,102 @@
+from typing import Any
+
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import col, from_json, schema_of_json, unbase64
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformDecodeAction(PipelineAction):
+    """Decodes values of a specified column in the DataFrame based on the given format.
+
+    Example:
+    ```yaml
+    Decode Columns:
+        action: TRANSFORM_DECODE
+        options:
+            column: configurations
+            input_format: json
+    ```
+    """
+
+    name: str = "TRANSFORM_DECODE"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        column: str | None = None,
+        input_format: str | None = None,
+        schema: str | None = None,
+        **_: Any,  # define kwargs to match the base class signature
+    ) -> PipelineContext:
+        """Decodes values of a specified column in the DataFrame based on the given format.
+
+        Args:
+            context: The context in which this Action is executed.
+            column: The name of the column that should be decoded.
+            input_format: The format from which the column should be decoded.
+                Currently supported formats are 'base64' and 'json'.
+            schema: For JSON input, the schema of the JSON object. If empty,
+                the schema is inferred from the first row of the DataFrame. For base64 input,
+                the data type to which the column is cast.
+
+        Raises:
+            ValueError: If no column is specified.
+            ValueError: If no input_format is specified.
+            ValueError: If the data from context is None.
+            ValueError: If an invalid input_format is provided.
+
+        Returns:
+            The context after the execution of this Action, containing the DataFrame with the decoded column(s).
+        """
+        if not column:
+            raise ValueError("No column specified.")
+        if not input_format:
+            raise ValueError("No input_format specified")
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df = context.data
+        match input_format.lower():
+            case "base64":
+                df = self._decode_base64(df, column, schema)  # type: ignore
+            case "json":
+                df = self._decode_json(df, column, schema)  # type: ignore
+            case _:
+                raise ValueError(
+                    f"Invalid input_format: [ '{input_format}' ]. Please specify a valid format to decode.",
+                )
+
+        return context.from_existing(data=df)  # type: ignore
+
+    def _decode_base64(self, df: DataFrame, column: str, base64_schema: str | None):
+        """Decode base64 column."""
+        df_decoded = df.withColumn(column, unbase64(col(column)))
+        if base64_schema:
+            df_decoded = df_decoded.withColumn(column, col(column).cast(base64_schema))
+        return df_decoded
+
+    def _decode_json(self, df: DataFrame, column: str, json_schema: str | None):
+        """Decode json column."""
+        distinct_schemas = (
+            df.select(column)
+            .withColumn("json_schema", schema_of_json(col(column)))
+            .select("json_schema")
+            .dropDuplicates()
+        )
+        if not (json_schema or distinct_schemas.count() > 0):
+            raise RuntimeError("Cannot infer schema from empty DataFrame.")
+
+        elif distinct_schemas.count() > 1:
+            raise RuntimeError(f"There is more than one JSON schema in column {column}.")
+
+        if json_schema is None:
+            final_json_schema = distinct_schemas.collect()[0].json_schema
+        else:
+            final_json_schema = json_schema  # type: ignore
+
+        df_decoded = df.withColumn(column, from_json(col(column), final_json_schema)).select(*df.columns, f"{column}.*")
+
+        return df_decoded