cloe-nessy 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

cloe_nessy/pipeline/actions/__init__.py

@@ -7,6 +7,7 @@ from .read_excel import ReadExcelAction
 from .read_files import ReadFilesAction
 from .read_metadata_yaml import ReadMetadataYAMLAction
 from .transform_change_datatype import TransformChangeDatatypeAction
+from .transform_clean_column_names import TransformCleanColumnNamesAction
 from .transform_concat_columns import TransformConcatColumnsAction
 from .transform_decode import TransformDecodeAction
 from .transform_distinct import TransformDistinctAction
@@ -39,6 +40,7 @@ __all__ = [
     "TransformFilterAction",
     "TransformUnionAction",
     "TransformChangeDatatypeAction",
+    "TransformCleanColumnNamesAction",
     "TransformConcatColumnsAction",
     "TransformDecodeAction",
     "TransformDistinctAction",

cloe_nessy/pipeline/actions/transform_clean_column_names.py

@@ -0,0 +1,77 @@
+import json
+import re
+from typing import Any
+
+import pyspark.sql.functions as F
+import pyspark.sql.types as T
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class TransformCleanColumnNamesAction(PipelineAction):
+    """Fixes column names in the DataFrame to be valid.
+
+    Removes invalid characters from the column names, including the fields of a struct and
+    replaces a single leading underscore by a double underscore.
+
+    Example:
+        ```yaml
+        Clean Column Names:
+            action: TRANSFORM_CLEAN_COLUMN_NAMES
+        ```
+    """
+
+    name: str = "TRANSFORM_CLEAN_COLUMN_NAMES"
+
+    def run(
+        self,
+        context: PipelineContext,
+        **_: Any,
+    ) -> PipelineContext:
+        """Fixes column names in the DataFrame to be valid.
+
+        Removes invalid characters from the column names, including the fields of a struct and
+        replaces a single leading underscore by a double underscore.
+
+        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 cleaned column names.
+        """
+        if context.data is None:
+            raise ValueError("Data from the context is required for the operation.")
+
+        with_columns_renamed = {}
+        with_columns_casted: dict[str, T.StructType | T.ArrayType | T.MapType] = {}
+
+        single_underscrore_at_beginning = r"^_(?=[^_])"
+
+        for c in context.data.schema:
+            old_name = c.name
+            new_name = re.sub(single_underscrore_at_beginning, "__", re.sub("\W", "_", old_name))
+            with_columns_renamed[old_name] = new_name
+
+            if isinstance(c.dataType, (T.StructType | T.ArrayType | T.MapType)):
+                old_column_schema = c.dataType.json()
+                new_column_schema = re.sub(
+                    r'(?<="name":")[^"]+',
+                    lambda m: re.sub("\W", "_", str(m.group())),
+                    old_column_schema,
+                )
+                if isinstance(c.dataType, T.StructType):
+                    with_columns_casted[new_name] = T.StructType.fromJson(json.loads(new_column_schema))
+                elif isinstance(c.dataType, T.ArrayType):
+                    with_columns_casted[new_name] = T.ArrayType.fromJson(json.loads(new_column_schema))
+                elif isinstance(c.dataType, T.MapType):
+                    with_columns_casted[new_name] = T.MapType.fromJson(json.loads(new_column_schema))
+
+        df = context.data.withColumnsRenamed(with_columns_renamed)
+        for c_name, c_type in with_columns_casted.items():
+            df = df.withColumn(c_name, F.col(c_name).cast(c_type))
+
+        return context.from_existing(data=df)  # type: ignore

cloe_nessy/pipeline/actions/transform_distinct.py

@@ -7,10 +7,16 @@ from ..pipeline_context import PipelineContext
 class TransformDistinctAction(PipelineAction):
     """Selects distinct rows from the DataFrame in the given context.
 
+    If a subset is given these columns are used for duplicate comparison. If no subset is given all columns are used.
+
     Example:
     ```yaml
     Decode Columns:
         action: TRANSFORM_DISTINCT
+        options:
+            subset:
+                - first_name
+                - last_name
     ```
     """
 
@@ -19,12 +25,15 @@ class TransformDistinctAction(PipelineAction):
     def run(
         self,
         context: PipelineContext,
+        *,
+        subset: list[str] | None = None,
         **_: Any,
     ) -> PipelineContext:
         """Selects distinct rows from the DataFrame in the given context.
 
         Args:
             context: The context in which this Action is executed.
+            subset: List of column names to use for duplicate comparison (default All columns).
 
         Raises:
             ValueError: If the data from the context is None.
@@ -35,6 +44,14 @@ class TransformDistinctAction(PipelineAction):
         if context.data is None:
             raise ValueError("Data from the context is required for the operation.")
 
-        df = context.data.distinct()
+        # check if all columns that are part of the subset are actually part of the dataframe.
+        if subset is not None:
+            subset_columns_not_in_dataframe = set(subset) - set(context.data.columns)
+            if len(subset_columns_not_in_dataframe) != 0:
+                raise ValueError(
+                    f"The following subset columns are not part of the dataframe: {subset_columns_not_in_dataframe}"
+                )
+
+        df = context.data.dropDuplicates(subset=subset)
 
         return context.from_existing(data=df)  # type: ignore

cloe_nessy/session/session_manager.py

@@ -63,13 +63,13 @@ class SessionManager:
     @classmethod
     def get_utils(
         cls,
-    ) -> Any:  # return type should be Union[DBUtils, MsSparkUtils], but can't import locally.
-        """Get or create a DBUtils or MsSparkUtils instance, depending on the context.
+    ) -> Any:  # return type should be Union[DBUtils, MsSparkUtils, RemoteDbUtils].
+        """Get or create a DBUtils, RemoteDbUtils or MsSparkUtils instance, depending on the context.
 
-        In Databricks this will return DBUtils, while in Fabric it will return MsSparkUtils.
+        In Databricks this will return DBUtils, when using Databricks-Connect it returns RemoteDbUtils, and in Fabric it will return MsSparkUtils.
 
         Returns:
-            utils: The DBUtils or MsSparkUtils instance.
+            utils: The DBUtils, RemoteDbUtils or MsSparkUtils instance.
 
         Raises:
             RuntimeError: If the instance cannot be created.
@@ -88,19 +88,25 @@ class SessionManager:
         }
 
         try:
-            cls._utils = utils_function[cls._env](cls)  # type: ignore
+            cls._utils = utils_function[cls._env]()  # type: ignore
         except Exception as e:
             raise RuntimeError(f"Cannot create utils instance. Error: {e}") from e
 
         return cls._utils
 
+    @classmethod
     def _get_dbutils(cls):
+        if cls._env == cls.Environment.DATABRICKS_CONNECT:
+            from databricks.sdk import WorkspaceClient
+
+            return WorkspaceClient().dbutils
+
         from pyspark.dbutils import DBUtils
 
         cls.get_spark_session()
-        utils = DBUtils(cls._spark)
-        return utils
+        return DBUtils(cls._spark)
 
+    @classmethod
     def _get_mssparkutils(cls):
         from notebookutils import mssparkutils  # type: ignore
 

{cloe_nessy-0.3.2.dist-info → cloe_nessy-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: cloe-nessy
-Version: 0.3.2
+Version: 0.3.3
 Summary: Your friendly datalake monster.
 Home-page: https://initions.com/
 Author: initions

{cloe_nessy-0.3.2.dist-info → cloe_nessy-0.3.3.dist-info}/RECORD

@@ -47,17 +47,18 @@ cloe_nessy/pipeline/pipeline_config.py,sha256=BN3ZSbr6bC-X9edoh-n5vRfPHFMbgtAU7m
 cloe_nessy/pipeline/pipeline_context.py,sha256=csElDc6BsynDUtRXgQOSCH7ONc_b-ag0YEg0zlQTz58,1874
 cloe_nessy/pipeline/pipeline_parsing_service.py,sha256=c_nAsgw81QYBM9AFiTxGgqRhNXABkDKplbeoCJPtbpE,6434
 cloe_nessy/pipeline/pipeline_step.py,sha256=UlnmpS6gm_dZ7m9dD1mZvye7mvUF_DA7HjOZo0oGYDU,1977
-cloe_nessy/pipeline/actions/__init__.py,sha256=Psksv49DVhWHR2D1OuMxvYClF1Vjh5shiyy9yBdWnb0,2160
+cloe_nessy/pipeline/actions/__init__.py,sha256=LwKctXy4Jun52BnCVGvWa8nnKVjTSov4GT58j6Zy8zg,2273
 cloe_nessy/pipeline/actions/read_api.py,sha256=wGyPZdeh3Cam_BQBilltWBWCIdD9I_kv4lunEhE39Tg,6625
 cloe_nessy/pipeline/actions/read_catalog_table.py,sha256=aZy4sJLLE8ZQ_SPXGSDoHYaBJTz8s7xQDVn5eYrYHvE,2689
 cloe_nessy/pipeline/actions/read_excel.py,sha256=EgHbK1wO6dkDo0KErYDhK_2sNIkIoa-6As9oo9dNFsE,7708
 cloe_nessy/pipeline/actions/read_files.py,sha256=8twjprqKYEmVu5QITEGe4no45TfhgzZosTFVQ89vV6g,3861
 cloe_nessy/pipeline/actions/read_metadata_yaml.py,sha256=aZtkstf9jBYYN2MGnazz63BG_hJ7mIgAfKiNqUpc26E,2235
 cloe_nessy/pipeline/actions/transform_change_datatype.py,sha256=Nz3Ncr-Zd-wy8g9-aN5XcvpWAHLyWs70RpZ7KqKqIaU,1788
+cloe_nessy/pipeline/actions/transform_clean_column_names.py,sha256=XuVAVEbp-UiF8PO6wAEJyl1TYgBD7MSnuOGhuEvXKv4,2881
 cloe_nessy/pipeline/actions/transform_concat_columns.py,sha256=V0TzeQFpBYur_T1Nv0nRpOU02nKQ2iypo2CCcV2rBtk,3083
 cloe_nessy/pipeline/actions/transform_decode.py,sha256=DmT-29dIqbz_xTj4GSCfnbgYRCiUrWzKvGrRYy1frNw,4004
 cloe_nessy/pipeline/actions/transform_deduplication.py,sha256=2VN5_wza7sD7fERyG6ElGh_Yo-W-Mxw-QBmtDXs1MGQ,5063
-cloe_nessy/pipeline/actions/transform_distinct.py,sha256=sdCElXCM77AQ0m6Zzg_h7cyavBOxo7W9K1NrsvNLufA,1105
+cloe_nessy/pipeline/actions/transform_distinct.py,sha256=R0Wv_YnWOw198r0rPR_72fgH5sp7upgjZzfOPTZ1oPA,1942
 cloe_nessy/pipeline/actions/transform_filter.py,sha256=vOAxKtNWCABLb6G6Xz98NK7fEfgn6QJia31S7IvoUTg,1428
 cloe_nessy/pipeline/actions/transform_generic_sql.py,sha256=cli59HCERFge7f0RB8yXw2oDtHSbMCWQMdeCeqhbdg8,2355
 cloe_nessy/pipeline/actions/transform_group_aggregate.py,sha256=HcY4sqb2yNBCz90jQtxGA8fZPuQXfJuaDmv8lWuoTqg,4050
@@ -69,12 +70,12 @@ cloe_nessy/pipeline/actions/transform_select_columns.py,sha256=Kez8puDK7cRfhleBE
 cloe_nessy/pipeline/actions/transform_union.py,sha256=TDER06IABzxvIez4bGLKCLaDA4eScpTzYRbfUzwv_RQ,2342
 cloe_nessy/pipeline/actions/write_catalog_table.py,sha256=6yAHTX5kZviumgBW_NYVGAUin6U2nDzmic9of6wA8FY,2590
 cloe_nessy/session/__init__.py,sha256=t7_YjUhJYW3km_FrucaUdbIl1boQtwkyhw_8yE10qzc,74
-cloe_nessy/session/session_manager.py,sha256=rd33lSafzomuyGf1BzhyjIWuy9sXgFjr-ca7A7Sw8eo,6490
+cloe_nessy/session/session_manager.py,sha256=PK7awMc6fmot7f9FMmvIUbIzKFgjcy2o2bZS9kjVs10,6733
 cloe_nessy/settings/__init__.py,sha256=ZbkneO3WaKOxon7qHFHnou7EnBOSnBFyKMDZblIEvzM,101
 cloe_nessy/settings/settings.py,sha256=I4n129lrujriW-d8q4as2Kb4_kI932ModfZ5Ow_UpVM,3653
 cloe_nessy/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cloe_nessy/utils/file_and_directory_handler.py,sha256=r2EVt9xG81p6ScaJCwETC5an6pMT6WseB0jMOR-JlpU,602
-cloe_nessy-0.3.2.dist-info/METADATA,sha256=7w0f9JC9rm0tmEBYvkPSuTPTKIoGaHFytY7eYf1GRkU,1837
-cloe_nessy-0.3.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-cloe_nessy-0.3.2.dist-info/top_level.txt,sha256=Z7izn8HmQpg2wBUb-0jzaKlYKMU7Ypzuc9__9vPtW_I,11
-cloe_nessy-0.3.2.dist-info/RECORD,,
+cloe_nessy-0.3.3.dist-info/METADATA,sha256=StCfl2I5dItaMzO10u3gQw6WxfjZUZHRIodEvKuQu_s,1837
+cloe_nessy-0.3.3.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+cloe_nessy-0.3.3.dist-info/top_level.txt,sha256=Z7izn8HmQpg2wBUb-0jzaKlYKMU7Ypzuc9__9vPtW_I,11
+cloe_nessy-0.3.3.dist-info/RECORD,,

{cloe_nessy-0.3.2.dist-info → cloe_nessy-0.3.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (75.8.2)
 Root-Is-Purelib: true
 Tag: py3-none-any