kedro-sumz-utils 0.0.1__tar.gz

kedro_sumz_utils-0.0.1/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.1
+Name: kedro-sumz-utils
+Version: 0.0.1
+Summary: Utility functions and kedro datasets for Data projects.
+Author-email: Juan Barreto <juan.barreto@sumz.co>
+Project-URL: Homepage, https://github.com/pypa/sampleproject
+Project-URL: Bug Tracker, https://github.com/pypa/sampleproject/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: kedro[spark]<0.19.0,>=0.18.0
+Requires-Dist: kedro-mlflow==0.11.10

kedro_sumz_utils-0.0.1/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kedro-sumz-utils"
+version = "0.0.1"
+authors = [
+  { name="Juan Barreto", email="juan.barreto@sumz.co" },
+]
+description = "Utility functions and kedro datasets for Data projects."
+readme = "README.md"
+requires-python = ">=3.9"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+dependencies = [
+  "kedro[spark]>=0.18.0,<0.19.0",
+  "kedro-mlflow==0.11.10"
+]
+[project.urls]
+"Homepage" = "https://github.com/pypa/sampleproject"
+"Bug Tracker" = "https://github.com/pypa/sampleproject/issues"

kedro_sumz_utils-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

kedro_sumz_utils-0.0.1/src/kedro_sumz/extras/context/__init__.py

@@ -0,0 +1,81 @@
+"""Custom Kedro-Virgo context class."""
+import logging
+from pathlib import Path
+from typing import Any, Dict, Union
+
+from kedro.config import ConfigLoader
+from kedro.framework.context import KedroContext
+from pluggy import PluginManager
+
+logger = logging.getLogger(__name__)
+
+
+class KedroSparkContext(KedroContext):
+    """Custom kedro context which initializes the Spark Session."""
+
+    def __init__(
+        self,
+        package_name: str,
+        project_path: Union[str, Path],
+        config_loader: ConfigLoader,
+        hook_manager: PluginManager,
+        env: str = None,
+        extra_params: Dict[str, Any] = None,
+    ):
+        logger.info("Creating Kedro Context")
+        super().__init__(
+            package_name, project_path, config_loader, hook_manager, env, extra_params
+        )
+        self._spark_session = None
+        self._package_name = (
+            package_name if package_name else Path(__file__).parent.name
+        )
+
+        if not self._init_databricks():
+            self._init_spark()
+
+    def _init_spark(self):
+        from pyspark import SparkConf
+        from pyspark.sql import SparkSession
+        
+        if SparkSession.getActiveSession() is None:
+            logger.info("Initializing Spark session")
+            parameters = self._config_loader["spark"]
+            logger.info(parameters)
+            spark_conf = SparkConf().setAll(list(parameters.items()))
+
+            spark_session_conf = (
+                SparkSession.builder.appName(self._package_name)
+                .master("local[*, 4]")
+                .config(conf=spark_conf)
+            )
+
+            self._spark_session = spark_session_conf.getOrCreate()
+            self._spark_session.sparkContext.setLogLevel("WARN")
+            logger.info("Spark Web URL: %s", self._spark_session.sparkContext.uiWebUrl)
+        else:
+            self._spark_session = SparkSession.getActiveSession()
+
+    def _init_databricks(self) -> bool:
+        try:
+            from databricks.connect import DatabricksSession
+            from pyspark.sql import SparkSession
+
+            if SparkSession.getActiveSession() is None:
+                logger.info("Initializing Databricks session")
+                self._spark_session = DatabricksSession.builder.profile(
+                    "connect"
+                ).getOrCreate()
+            else:
+                self._spark_session = SparkSession.getActiveSession()
+            return True
+        except ImportError:
+            logger.info(
+                "Databricks Connect is not installed. Initializing Spark session instead."
+            )
+            return False
+
+    @property
+    def spark_session(self):
+        """Spark session property."""
+        return self._spark_session

kedro_sumz_utils-0.0.1/src/kedro_sumz/extras/datasets/__init__.py

@@ -0,0 +1,12 @@
+"""Dummy dataset for testing purposes."""
+from typing import Any
+
+from kedro.io import MemoryDataSet
+
+
+class DummyDataset(MemoryDataSet):
+    """Dummy dataset for testing purposes."""
+
+    # pylint: disable=useless-super-delegation,too-few-public-methods
+    def __init__(self, data: Any = True, copy_mode: str = None):
+        super().__init__(data, copy_mode)

kedro_sumz_utils-0.0.1/src/kedro_sumz/extras/datasets/delta.py

@@ -0,0 +1,155 @@
+"""``AbstractDataSet`` implementation to access DeltaTables using
+``delta-spark``
+"""
+from pathlib import PurePosixPath
+from typing import NoReturn
+
+from delta.tables import DeltaTable
+
+# pylint: disable=unused-import
+from kedro.extras.datasets.spark import SparkHiveDataSet  # noqa: F401
+from kedro.extras.datasets.spark.spark_dataset import (
+    _split_filepath,
+    _strip_dbfs_prefix,
+)
+from kedro.io.core import AbstractDataSet, DataSetError
+from pyspark.sql import SparkSession
+from pyspark.sql.utils import AnalysisException
+
+from kedro_sumz.utils import get_kedro_context
+
+
+class DeltaTableDataSet(AbstractDataSet[None, DeltaTable]):
+    """``DeltaTableDataSet`` loads data into DeltaTable objects.
+
+    Example usage for the
+    `YAML API <https://kedro.readthedocs.io/en/stable/data/\
+    data_catalog.html#use-the-data-catalog-with-the-yaml-api>`_:
+
+    .. code-block:: yaml
+
+        weather@spark:
+          type: spark.SparkDataSet
+          filepath: data/02_intermediate/data.parquet
+          file_format: "delta"
+
+        weather@delta:
+          type: virgo.extras.datasets.DeltaTableDataSet
+          directory: data/02_intermediate
+          table: "data.parquet"
+
+        weather_table@delta:
+          type: virgo.extras.datasets.delta.DeltaTableDataSet
+          directory: intermediate
+          table: data
+
+    Example usage for the
+    `Python API <https://kedro.readthedocs.io/en/stable/data/\
+    data_catalog.html#use-the-data-catalog-with-the-code-api>`_:
+    ::
+
+        >>> from pyspark.sql import SparkSession
+        >>> from pyspark.sql.types import (StructField, StringType,
+        >>>                                IntegerType, StructType)
+        >>>
+        >>> from kedro.extras.datasets.spark import DeltaTableDataSet, SparkDataSet
+        >>>
+        >>> schema = StructType([StructField("name", StringType(), True),
+        >>>                      StructField("age", IntegerType(), True)])
+        >>>
+        >>> data = [('Alex', 31), ('Bob', 12), ('Clarke', 65), ('Dave', 29)]
+        >>>
+        >>> spark_df = SparkSession.builder.getOrCreate().createDataFrame(data, schema)
+        >>>
+        >>> data_set = SparkHiveDataSet(database="test_database", table="test_table",
+        >>>                             write_mode="overwrite")
+        >>> data_set.save(spark_df)
+        >>> delta_dataset = DeltaTableDataSet(directory="test_database", table="test_table")
+        >>> delta_table = delta_dataset.load()
+        >>>
+        >>> delta_table.update()
+    """
+
+    # this dataset cannot be used with ``ParallelRunner``,
+    # therefore it has the attribute ``_SINGLE_PROCESS = True``
+    # for parallelism within a Spark pipeline please consider
+    # using ``ThreadRunner`` instead
+    _SINGLE_PROCESS = True
+
+    def __init__(
+        self,
+        directory: str,
+        table: str,
+        dataset_type: str = None,
+    ) -> None:
+        """Creates a new instance of ``DeltaTableDataSet``.
+
+        Args:
+            directory: Base path to a folder or a database.
+            table: Name of the table or file.
+            dataset_type: Type of the dataset. Can be either "table" or "file".
+        """
+
+        if not dataset_type and ("/" in directory or "/" in table):
+            self._type = "file"
+        elif not dataset_type and "." in directory:
+            self._type = "table"
+        else:
+            self._type = dataset_type
+        if not self._type:
+            raise DataSetError(
+                "Could not determine dataset type. Please specify it explicitly."
+            )
+        if self._type == "table":
+            path = f"{directory}.{table}"
+            self._database = ".".join(path.split(".")[:-1])
+            self._table = table
+            self._full_table_address = f"{self._database}.{self._table}"
+        else:
+            path = f"{directory}/{table}"
+            fs_prefix, filepath = _split_filepath(path)
+            self._fs_prefix = fs_prefix
+            self._filepath = PurePosixPath(filepath)
+            self._full_table_address = (
+                f"delta.`{self._fs_prefix + str(self._filepath)}`"
+            )
+            self._type = "file"
+
+    @staticmethod
+    def _get_spark() -> SparkSession:
+        """Get the active Spark session."""
+        if hasattr(get_kedro_context(), "spark_session"):
+            return get_kedro_context().spark_session
+        else:
+            SparkSession.getActiveSession()
+
+    def _load(self) -> DeltaTable:
+        return DeltaTable.forName(self._get_spark(), self._full_table_address)
+
+    def _save(self, data: None) -> NoReturn:
+        raise DataSetError(f"{self.__class__.__name__} is a read only dataset type")
+
+    def _exists(self) -> bool:
+        # noqa # pylint:disable=protected-access
+        if self._type == "table":
+            return (
+                self._get_spark()
+                ._jsparkSession.catalog()
+                .tableExists(self._database, self._table)
+            )
+
+        load_path = _strip_dbfs_prefix(self._fs_prefix + str(self._filepath))
+
+        try:
+            self._get_spark().read.load(path=load_path, format="delta")
+        except AnalysisException as exception:
+            if "is not a Delta table" in exception.desc:
+                return False
+            raise
+
+        return True
+
+    def _describe(self):
+        return {
+            "full_table_address": self._full_table_address,
+        }

kedro_sumz_utils-0.0.1/src/kedro_sumz/extras/datasets/mlflow.py

@@ -0,0 +1,277 @@
+from functools import partial
+from typing import Any, Dict, Optional
+from copy import deepcopy
+import mlflow
+from kedro.io import AbstractDataSet, DataSetError
+from mlflow.tracking import MlflowClient
+from kedro_mlflow.io.models.mlflow_model_logger_dataset import (
+    MlflowModelLoggerDataSet
+)
+from kedro_sumz.utils.mlflow import ModelWithConfig
+
+
+class MlflowModelWithConfigDataSet(MlflowModelLoggerDataSet):
+    """
+    Wrapper for saving, logging and loading for all MLFlow model flavor
+    including its configuration
+    """
+    def __init__(
+        self,
+        flavor: str,
+        run_id: Optional[str] = None,
+        artifact_path: Optional[str] = "model",
+        pyfunc_workflow: Optional[str] = None,
+        load_args: Optional[Dict[str, Any]] = None,
+        save_args: Optional[Dict[str, Any]] = None
+    ) -> None:
+        """Initializes MlflowModelWithConfigDataSet
+
+        Args:
+            flavor (str): Flavor of the model
+            run_id (Optional[str], optional): Mlflow run id. Defaults to None.
+            artifact_path (Optional[str], optional): Path of the model artifact. Defaults to "model".
+            pyfunc_workflow (Optional[str], optional): Name of the pyfunc model. Defaults to None.
+            load_args (Optional[Dict[str, Any]], optional): Load args used in the 
+            mlflow load_model function. Defaults to None.
+            save_args (Optional[Dict[str, Any]], optional): Save args used in the
+            mlflow log_model function. Defaults to None.
+        """
+        super().__init__(flavor, run_id, artifact_path, pyfunc_workflow, load_args, save_args)
+
+    def _save_model_in_run(self, model: ModelWithConfig):
+        """
+        Saves a model in a run as an artifact
+
+        Args:
+            model (ModelWithConfig): Object which stores the model and its configuration
+
+        Raises:
+            DataSetError: When the object passed is not an instance of the ModelWithConfig class
+        """
+        if not isinstance(model, ModelWithConfig):
+            raise DataSetError(
+                "The model to be logged must be an instance of the ModelWithConfig class"
+            )
+
+        if not self._logging_activated:
+            return
+        
+        save_args = self._prepare_save_args(
+            model_with_config=model
+        )
+
+        if self._flavor == "mlflow.pyfunc":
+            # PyFunc models utilise either `python_model` or `loader_module`
+            # workflow. We we assign the passed `model` object to one of those keys
+            # depending on the chosen `pyfunc_workflow`.
+            save_args[self._pyfunc_workflow] = model.model
+            self._mlflow_model_module.log_model(
+                **save_args
+            )
+        else:
+            # Otherwise we save using the common workflow
+            self._mlflow_model_module.log_model(
+                model.model,
+                **save_args
+            )
+    
+    def _prepare_save_args(
+            self,
+            model_with_config: ModelWithConfig
+    ) -> Dict[str, Any]:
+        save_args = deepcopy(self._save_args)
+
+        save_args["signature"] = self._save_arg_value(
+            catalog_save_arg=self._save_args.get("signature"),
+            model_with_config_arg=model_with_config.signature
+        )
+
+        save_args["conda_env"] = self._save_arg_value(
+            catalog_save_arg=self._save_args.get("conda_env"),
+            model_with_config_arg=model_with_config.conda_env
+        )
+
+        save_args["pip_requirements"] = self._save_arg_value(
+            catalog_save_arg=self._save_args.get("pip_requirements"),
+            model_with_config_arg=model_with_config.pip_requirements
+        )
+
+        save_args["extra_pip_requirements"] = self._save_arg_value(
+            catalog_save_arg=self._save_args.get("extra_pip_requirements"),
+            model_with_config_arg=model_with_config.extra_pip_requirements
+        )
+
+        save_args["artifact_path"] = self._artifact_path
+
+        return save_args
+    
+    def _save_arg_value(
+            self,
+            catalog_save_arg: Any,
+            model_with_config_arg: Any
+    )->Any:
+        if catalog_save_arg is None and model_with_config_arg is None:
+            return None
+        elif catalog_save_arg is None:
+            return model_with_config_arg
+        else:
+            return catalog_save_arg
+
+
+class MlflowMetricsDictDataSet(AbstractDataSet):
+    """This class represent MLflow metrics dictionary dataset."""
+
+    def __init__(
+        self,
+        prefix: str = None,
+        run_id: str = None,
+    ):
+        """Initialise MlflowMetricsDictDataSet.
+
+        Args:
+            prefix (str, optional): Prefix for metrics logged in MLflow. Defaults to None
+            run_id (str, optional): ID of MLflow run. Defaults to None
+        """
+        self._prefix = prefix
+        self.run_id = run_id
+        self._logging_activated = True  # by default, logging is activated!
+
+    @property
+    def run_id(self):
+        """Get run id.
+
+        If active run is not found, tries to find last experiment.
+
+        Raise `DataSetError` exception if run id can't be found.
+
+        Returns:
+            str: String contains run_id.
+        """
+        if self._run_id is not None:
+            return self._run_id
+        run = mlflow.active_run()
+        if run:
+            return run.info.run_id
+        raise DataSetError("Cannot find run id.")
+
+    @run_id.setter
+    def run_id(self, run_id):
+        self._run_id = run_id
+
+    # we want to be able to turn logging off for an entire pipeline run
+    # To avoid that a single call to a dataset in the catalog creates a new run automatically
+    # we want to be able to turn everything off
+    @property
+    def _logging_activated(self):
+        return self.__logging_activated
+
+    @_logging_activated.setter
+    def _logging_activated(self, flag):
+        if not isinstance(flag, bool):
+            raise ValueError(f"_logging_activated must be a boolean, got {type(flag)}")
+        self.__logging_activated = flag
+
+    def _load(self) -> Dict[str, float]:
+        """Load MlflowMetricsDictDataSet.
+
+        Returns:
+            Dict[str, float]: Dictionary with MLflow metrics.
+        """
+        all_metrics = self._get_all_metrics()
+
+        dataset_metrics = self._filter_dataset_metrics(all_metrics)
+
+        if not dataset_metrics:
+            raise DataSetError(
+            "Tried to load an empty metrics dictionary"
+            f"be sure that metrics with the prefix '{self._prefix}' exists"
+            "in the current run"
+            )
+            
+        return dataset_metrics
+
+    def _save(self, data: Dict[str, float]) -> None:
+        """
+        Save given MLflow metrics dictionary dataset 
+        and log it in MLflow as metrics.
+
+        Args:
+            data (Dict[str, float]): Metrics dictionary.
+        """
+        if not data:
+            raise DataSetError(
+                "Metrics dictionary should have at least one metric"
+            )
+
+        client = MlflowClient()
+        try:
+            run_id = self.run_id
+        except DataSetError:
+            # If run_id can't be found log_metric would create new run.
+            run_id = None
+
+        log_metric = (
+            partial(client.log_metric, run_id)
+            if run_id is not None
+            else mlflow.log_metric
+        )
+
+        if not all([isinstance(value, float) for value in data.values()]):
+            raise DataSetError(f"All metric values should be of type `float`")
+
+        metrics = {f"{self._prefix}.{key}": value for key, value in data.items()}
+
+        if self._logging_activated:
+            for k, v in metrics.items():
+                log_metric(k, v)
+
+    def _exists(self) -> bool:
+        """Check if MLflow metrics dataset exists.
+
+        Returns:
+            bool: Is MLflow metrics dataset exists?
+        """
+        all_metrics = self._get_all_metrics()
+        dataset_metrics = self._filter_dataset_metrics(all_metrics)
+        return len(dataset_metrics.keys())>0
+
+    def _describe(self) -> Dict[str, Any]:
+        """Describe MLflow metrics dataset.
+
+        Returns:
+            Dict[str, Any]: Dictionary with MLflow metrics 
+            dictionary dataset description.
+        """
+        return {
+            "run_id": self._run_id,
+            "prefix": self._prefix,
+        }
+
+    def _filter_dataset_metrics(self, metrics: Dict[str, float]) -> Dict[str, float]:
+        """
+        Filter a metric dictionary by keeping only those that 
+        belongs to the dataset
+
+        Args:
+            metric (Dict[str, float]): Metrics dictionary.
+        Returns:
+            Dict[str, float]: Dictionary with the metrics that belongs 
+            to the dataset
+        """
+        dataset_metrics = {
+            k: v for k, v in metrics.items()
+            if k.startswith(self._prefix)
+        }
+
+        return dataset_metrics
+
+    def _get_all_metrics(self)->Dict[str, float]:
+        """Retrieve all metrics logged in the current run
+
+        Returns:
+            Dict[str, float]: Dictionary with all the metrics logged
+            in the current run
+        """
+        client = MlflowClient()
+        run = client.get_run(self.run_id)
+        return run.data.metrics