cloe-nessy 0.2.9__py3-none-any.whl

cloe_nessy/pipeline/actions/transform_distinct.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformDistinctAction(PipelineAction):
+    """Selects distinct rows from the DataFrame in the given context.
+
+    Example:
+    ```yaml
+    Decode Columns:
+        action: TRANSFORM_DISTINCT
+    ```
+    """
+
+    name: str = "TRANSFORM_DISTINCT"
+
+    def run(
+        self,
+        context: PipelineContext,
+        **_: Any,
+    ) -> PipelineContext:
+        """Selects distinct rows from the DataFrame in the given context.
+
+        Args:
+            context: The context in which this Action is executed.
+
+        Raises:
+            ValueError: If the data from the context is None.
+
+        Returns:
+            The context after the execution of this Action, containing the DataFrame with distinct rows.
+        """
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df = context.data.distinct()
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_filter.py

@@ -0,0 +1,51 @@
+from typing import Any
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformFilterAction(PipelineAction):
+    """Filters the DataFrame in the given context based on a specified condition.
+
+    Example:
+    ```yaml
+    Decode Columns:
+        action: TRANSFORM_FILTER
+        options:
+            condition: where city="Hamburg"
+    ```
+    """
+
+    name: str = "TRANSFORM_FILTER"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        condition: str = "",
+        **_: Any,
+    ) -> PipelineContext:
+        """Filters the DataFrame in the given context based on a specified condition.
+
+        Args:
+            context: Context in which this Action is executed.
+            condition: A SQL-like expression used to filter the DataFrame.
+
+        Raises:
+            ValueError: If no condition is provided.
+            ValueError: If the data from the context is None.
+
+        Returns:
+            Context after the execution of this Action, containing the filtered DataFrame.
+        """
+        if not condition:
+            raise ValueError("No condition provided.")
+
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df = context.data
+
+        df_filtered = df.filter(condition=condition)
+
+        return context.from_existing(data=df_filtered)  # type: ignore

cloe_nessy/pipeline/actions/transform_generic_sql.py

@@ -0,0 +1,66 @@
+import uuid
+from typing import Any
+
+from ...session import SessionManager
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformSqlAction(PipelineAction):
+    """Executes a SQL statement on a DataFrame within the provided context.
+
+    A temporary view is created from the current DataFrame, and the SQL
+    statement is executed on that view. The resulting DataFrame is returned.
+
+    Example:
+    ```yaml
+    SQL Transform:
+        action: TRANSFORM_SQL
+        options:
+            sql_statement: select city, revenue, firm from {DATA_FRAME} where product="Databricks"
+    ```
+    """
+
+    name: str = "TRANSFORM_SQL"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        sql_statement: str = "",
+        **kwargs: Any,
+    ) -> PipelineContext:
+        """Executes a SQL statement on a DataFrame within the provided context.
+
+        Args:
+            context: Context in which this Action is executed.
+            sql_statement: A string containing the SQL statement to be
+                executed. The source table should be referred to as "{DATA_FRAME}".
+            **kwargs: Additional keyword arguments are passed as placeholders to the
+                SQL statement.
+
+        Raises:
+            ValueError: If "{DATA_FRAME}" is not included in the SQL statement.
+            ValueError: If no SQL statement is provided.
+            ValueError: If the data from the context is None.
+
+        Returns:
+            Context after the execution of this Action, containing the DataFrame resulting from the SQL statement.
+        """
+        if not sql_statement:
+            raise ValueError("No SQL statement provided.")
+
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        _spark = SessionManager.get_spark_session()
+
+        temp_view_name = str(uuid.uuid1()).replace("-", "_")
+        context.data.createTempView(temp_view_name)
+
+        if "FROM {DATA_FRAME}".casefold() not in sql_statement.casefold():
+            raise ValueError("Please use 'FROM {DATA_FRAME}' in your SQL statement.")
+
+        df = _spark.sql(sql_statement.format(DATA_FRAME=temp_view_name, **kwargs))
+
+        return context.from_existing(data=df)

cloe_nessy/pipeline/actions/transform_join.py

@@ -0,0 +1,81 @@
+from typing import Any
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+from ..pipeline_step import PipelineStep
+
+
+class TransformJoinAction(PipelineAction):
+    """Joins the current DataFrame with another DataFrame defined in joined_data.
+
+    The join operation is performed based on specified columns and the type of join
+    indicated by the `how` parameter.
+
+    Example:
+    ```yaml
+    Join Tables:
+        action: TRANSFORM_JOIN
+        options:
+            joined_data: ((step:Transform First Table))
+            join_on: id
+            how: anti
+    ```
+    """
+
+    name: str = "TRANSFORM_JOIN"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        joined_data: PipelineStep | None = None,
+        join_on: list[str] | str | dict[str, str] | None = None,
+        how: str = "inner",
+        **_: Any,
+    ) -> PipelineContext:
+        """Joins the current DataFrame with another DataFrame defined in joined_data.
+
+        Args:
+            context: Context in which this Action is executed.
+            joined_data: The PipelineStep context defining the DataFrame
+                to join with as the right side of the join.
+            join_on: A string for the join column
+                name, a list of column names, or a dictionary mapping columns from the
+                left DataFrame to the right DataFrame. This defines the condition for the
+                join operation.
+            how: The type of join to perform. Must be one of: inner, cross, outer,
+                full, fullouter, left, leftouter, right, rightouter, semi, anti, etc.
+
+        Raises:
+            ValueError: If no joined_data is provided.
+            ValueError: If no join_on is provided.
+            ValueError: If the data from context is None.
+            ValueError: If the data from the joined_data is None.
+
+        Returns:
+            Context after the execution of this Action, containing the result of the join operation.
+        """
+        if joined_data is None or joined_data.result is None or joined_data.result.data is None:
+            raise ValueError("No joined_data provided.")
+        if not join_on:
+            raise ValueError("No join_on provided.")
+
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df_right = joined_data.result.data.alias("right")  # type: ignore
+        df_left = context.data.alias("left")  # type: ignore
+
+        if isinstance(join_on, str):
+            join_condition = [join_on]
+        elif isinstance(join_on, list):
+            join_condition = join_on
+        else:
+            join_condition = [
+                df_left[left_column] == df_right[right_column]  # type: ignore
+                for left_column, right_column in join_on.items()
+            ]
+
+        df = df_left.join(df_right, on=join_condition, how=how)  # type: ignore
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_json_normalize.py

@@ -0,0 +1,106 @@
+from typing import Any, cast
+
+import pyspark.sql.functions as F
+
+from cloe_nessy.pipeline.pipeline_action import PipelineAction
+from cloe_nessy.pipeline.pipeline_context import PipelineContext
+
+
+class TransformJsonNormalize(PipelineAction):
+    """Normalizes and flattens the DataFrame by exploding array columns and flattening struct columns.
+
+    The method performs recursive normalization on the DataFrame present in the context,
+    ensuring that the order of columns is retained and new columns created by flattening
+    structs are appended after existing columns.
+
+    Example:
+    ```yaml
+    Normalize Tables:
+        action: TRANSFORM_JSON_NORMALIZE
+        options:
+            exclude_columns: coordinates
+    ```
+    """
+
+    name: str = "TRANSFORM_JSON_NORMALIZE"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        exclude_columns: list[str] | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Executes the normalization process on the DataFrame present in the context.
+
+        Please note that columns retain their relative order during the
+        normalization process, and new columns created by flattening structs are
+        appended after the existing columns.
+
+        Args:
+            context: The pipeline context that contains the DataFrame to be normalized.
+            exclude_columns: A list of column names to exclude from the normalization process.
+                    These columns will not be exploded or flattened.
+            **_: Additional keyword arguments (not used).
+
+        Returns:
+            A new pipeline context with the normalized DataFrame.
+
+        Raises:
+            ValueError: If the DataFrame in the context is `None`.
+        """
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        if not exclude_columns:
+            exclude_columns = []
+        df = TransformJsonNormalize._normalize(context.data, exclude_columns=cast(list, exclude_columns))
+        return context.from_existing(data=df)
+
+    @staticmethod
+    def _normalize(df, exclude_columns):
+        """Recursively normalizes the given DataFrame by exploding arrays and flattening structs.
+
+        This method performs two primary operations:
+        1. Explodes any array columns, unless they are in the list of excluded columns.
+        2. Flattens any struct columns, renaming nested fields and appending them to the top-level DataFrame.
+
+        The method continues these operations in a loop until there are no array or struct columns left.
+
+        Args:
+            df: The input DataFrame to normalize.
+            exclude_columns: A list of column names to exclude from the normalization process. These columns
+                                         will not be exploded or flattened.
+
+        Returns:
+            pyspark.sql.DataFrame: The normalized DataFrame with no array or struct columns.
+        """
+
+        def explode_arrays(df, exclude_columns):
+            array_present = False
+            for col in df.columns:
+                if df.schema[col].dataType.typeName() == "array" and col not in exclude_columns:
+                    df = df.withColumn(col, F.explode(col))
+                    array_present = True
+            return df, array_present
+
+        def flatten_structs(df):
+            struct_present = False
+            struct_columns = [col for col in df.columns if df.schema[col].dataType.typeName() == "struct"]
+            for col in struct_columns:
+                df = df.select(F.col("*"), F.col(col + ".*"))
+                nested_columns = df.select(F.col(col + ".*")).schema.names
+                for nested_col in nested_columns:
+                    df = df.withColumnRenamed(nested_col, f"{col}_{nested_col}")
+                df = df.drop(col)
+                struct_present = True
+            return df, struct_present
+
+        array_present = True
+        struct_present = True
+
+        while array_present or struct_present:
+            df, array_present = explode_arrays(df, exclude_columns)
+            df, struct_present = flatten_structs(df)
+
+        return df

cloe_nessy/pipeline/actions/transform_rename_columns.py

@@ -0,0 +1,60 @@
+from typing import Any
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformRenameColumnsAction(PipelineAction):
+    """Renames the specified columns in the DataFrame.
+
+    This method updates the DataFrame in the provided context by renaming columns according
+    to the mapping defined in the `columns` dictionary, where each key represents an old column
+    name and its corresponding value represents the new column name.
+
+    Example:
+    ```yaml
+    Rename Column:
+        action: TRANSFORM_RENAME_COLUMNS
+        options:
+            columns:
+                a_very_long_column_name: shortname
+    ```
+    """
+
+    name: str = "TRANSFORM_RENAME_COLUMNS"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        columns: dict[str, str] | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Renames the specified columns in the DataFrame.
+
+        Args:
+            context: Context in which this Action is executed.
+            columns: A dictionary where the key is the old column name
+                and the value is the new column name.
+
+        Raises:
+            ValueError: If no columns are provided.
+            ValueError: If the data from context is None.
+
+        Returns:
+            Context after the execution of this Action.
+        """
+        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
+
+        if isinstance(columns, dict):
+            df = df.withColumnsRenamed(columns)
+        else:
+            raise ValueError("'columns' should be a dict, like {'old_name_1':'new_name_1', 'old_name_2':'new_name_2'}")
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_replace_values.py

@@ -0,0 +1,59 @@
+from typing import Any
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformReplaceValuesAction(PipelineAction):
+    """Replaces specified values in the given DataFrame.
+
+    This method iterates over the specified `replace` dictionary, where each key is a column name
+    and each value is another dictionary containing old values as keys and new values as the corresponding
+    values. The method updates the DataFrame by replacing occurrences of the old values with the new ones
+    in the specified columns.
+
+    Example:
+    ```yaml
+    Replace Values:
+        action: TRANSFORM_REPLACE_VALUES
+        options:
+            replace:
+                empl_function:
+                    sales_employee: seller
+    ```
+    """
+
+    name: str = "TRANSFORM_REPLACE_VALUES"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        replace: dict[str, dict[str, str]] | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Replaces specified values in the given DataFrame.
+
+        Args:
+            context: Context in which this Action is executed.
+            replace: A dictionary where each key is the column name
+                and the corresponding value is another dictionary mapping old values to new values.
+
+        Raises:
+            ValueError: If no replace values are provided.
+            ValueError: If the data from context is None.
+
+        Returns:
+            Context after the execution of this Action.
+        """
+        if not replace:
+            raise ValueError("No replace values provided.")
+
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df = context.data
+        for column, to_replace in replace.items():
+            df = df.replace(to_replace=to_replace, subset=[column])  # type: ignore
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_select_columns.py

@@ -0,0 +1,83 @@
+from typing import Any
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformSelectColumnsAction(PipelineAction):
+    """Selects specified columns from the given DataFrame.
+
+    This method allows you to include or exclude specific columns from the
+    DataFrame. If `include_columns` is provided, only those columns will be
+    selected. If `exclude_columns` is provided, all columns except those will be
+    selected. The method ensures that the specified columns exist in the
+    DataFrame before performing the selection.
+
+    Example:
+    ```yaml
+    Select Columns:
+        action: TRANSFORM_SELECT_COLUMNS
+        options:
+            include_columns:
+                - id
+                - city
+                - product
+    ```
+    """
+
+    name: str = "TRANSFORM_SELECT_COLUMNS"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        include_columns: list[str] | None = None,
+        exclude_columns: list[str] | None = None,
+        raise_on_non_existing_columns: bool = True,
+        **_: Any,
+    ) -> PipelineContext:
+        """Selects specified columns from the given DataFrame.
+
+        Args:
+            context: Context in which this Action is executed.
+            include_columns: A list of column names that should be included.
+                If provided, only these columns will be selected.
+            exclude_columns: A list of column names that should be excluded.
+                If provided, all columns except these will be selected.
+            raise_on_non_existing_columns: If True, raise an error if a specified
+                column is not found in the DataFrame. If False, ignore the column
+                and continue with the selection.
+
+        Raises:
+            ValueError: If a specified column is not found in the DataFrame.
+            ValueError: If neither include_columns nor exclude_columns are provided,
+                or if both are provided.
+
+        Returns:
+            Context after the execution of this Action.
+        """
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        df = context.data
+
+        if (not include_columns and not exclude_columns) or (include_columns and exclude_columns):
+            raise ValueError("Please define either 'include_columns' or 'exclude_columns'.")
+
+        def check_missing_columns(df, columns, raise_on_non_existing_columns):
+            if raise_on_non_existing_columns:
+                missing_columns = [col for col in columns if col not in df.columns]
+                if missing_columns:
+                    raise ValueError(f"Columns not found in DataFrame: {missing_columns}")
+
+        try:
+            if include_columns:
+                check_missing_columns(df, include_columns, raise_on_non_existing_columns)
+                df_selected = df.select(*include_columns)
+            elif exclude_columns:
+                check_missing_columns(df, exclude_columns, raise_on_non_existing_columns)
+                df_selected = df.drop(*exclude_columns)
+        except Exception as e:
+            raise ValueError(f"Column selection error: {e}") from e
+
+        return context.from_existing(data=df_selected)  # type: ignore

cloe_nessy/pipeline/actions/transform_union.py

@@ -0,0 +1,71 @@
+from functools import reduce
+from typing import Any
+
+from pyspark.sql.dataframe import DataFrame
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+from ..pipeline_step import PipelineStep
+
+
+class TransformUnionAction(PipelineAction):
+    """Unions multiple DataFrames together.
+
+    This method takes the current DataFrame from the context and unites it with
+    additional DataFrames specified in the `union_data` argument. All DataFrames
+    must have the same schema. If any DataFrame in `union_data` is None or
+    empty, a ValueError will be raised.
+
+    Example:
+    ```yaml
+    Union Tables:
+        action: TRANSFORM_UNION
+        options:
+            union_data:
+                - ((step: Filter First Table))
+                - ((step: SQL Transform Second Table))
+    ```
+    """
+
+    name: str = "TRANSFORM_UNION"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        union_data: list[PipelineStep] | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Unions multiple DataFrames together.
+
+        Args:
+            context: Context in which this Action is executed.
+            union_data: A list of PipelineSteps that define the DataFrames
+                to union with the current context.
+
+        Raises:
+            ValueError: If no union_data is provided.
+            ValueError: If the data from context is None.
+            ValueError: If the data from any of the union_data is None.
+
+        Returns:
+            Context after the execution of this Action.
+        """
+        if not union_data:
+            raise ValueError("No union_data provided.")
+
+        # Check that all union_data contexts have valid data
+        result_contexts = []
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        for ctx in union_data:
+            if ctx.result is None or ctx.result.data is None:
+                raise ValueError(f"Data from the context of step '{ctx.name}' is required for the operation.")
+            result_contexts.append(ctx.result.data)
+
+        # Union all DataFrames
+        union_dfs = [context.data] + result_contexts
+        df = reduce(DataFrame.unionAll, union_dfs)  # type: ignore
+
+        return context.from_existing(data=df)  # type: ignore