genesis-flow 1.0.0__py3-none-any.whl

mlflow/data/spark_dataset.py

@@ -0,0 +1,406 @@
+import json
+import logging
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from packaging.version import Version
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.delta_dataset_source import DeltaDatasetSource
+from mlflow.data.digest_utils import get_normalized_md5_digest
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.data.spark_dataset_source import SparkDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, INVALID_PARAMETER_VALUE
+from mlflow.types import Schema
+from mlflow.types.utils import _infer_schema
+
+if TYPE_CHECKING:
+    import pyspark
+
+_logger = logging.getLogger(__name__)
+
+
+class SparkDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a Spark dataset (e.g. data derived from a Spark Table / file directory or Delta
+    Table) for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        df: "pyspark.sql.DataFrame",
+        source: DatasetSource,
+        targets: Optional[str] = None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+        predictions: Optional[str] = None,
+    ):
+        if targets is not None and targets not in df.columns:
+            raise MlflowException(
+                f"The specified Spark dataset does not contain the specified targets column"
+                f" '{targets}'.",
+                INVALID_PARAMETER_VALUE,
+            )
+        if predictions is not None and predictions not in df.columns:
+            raise MlflowException(
+                f"The specified Spark dataset does not contain the specified predictions column"
+                f" '{predictions}'.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._df = df
+        self._targets = targets
+        self._predictions = predictions
+        super().__init__(source=source, name=name, digest=digest)
+
+    def _compute_digest(self) -> str:
+        """
+        Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+        """
+        # Retrieve a semantic hash of the DataFrame's logical plan, which is much more efficient
+        # and deterministic than hashing DataFrame records
+        import numpy as np
+        import pyspark
+
+        # Spark 3.1.0+ has a semanticHash() method on DataFrame
+        if Version(pyspark.__version__) >= Version("3.1.0"):
+            semantic_hash = self._df.semanticHash()
+        else:
+            semantic_hash = self._df._jdf.queryExecution().analyzed().semanticHash()
+        return get_normalized_md5_digest([np.int64(semantic_hash)])
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the dataset.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type, schema, and profile.
+        """
+        schema = json.dumps({"mlflow_colspec": self.schema.to_dict()}) if self.schema else None
+        config = super().to_dict()
+        config.update(
+            {
+                "schema": schema,
+                "profile": json.dumps(self.profile),
+            }
+        )
+        return config
+
+    @property
+    def df(self):
+        """The Spark DataFrame instance.
+
+        Returns:
+            The Spark DataFrame instance.
+
+        """
+        return self._df
+
+    @property
+    def targets(self) -> Optional[str]:
+        """The name of the Spark DataFrame column containing targets (labels) for supervised
+        learning.
+
+        Returns:
+            The string name of the Spark DataFrame column containing targets.
+        """
+        return self._targets
+
+    @property
+    def predictions(self) -> Optional[str]:
+        """
+        The name of the predictions column. May be ``None`` if no predictions column
+        was specified when the dataset was created.
+        """
+        return self._predictions
+
+    @property
+    def source(self) -> Union[SparkDatasetSource, DeltaDatasetSource]:
+        """
+        Spark dataset source information.
+
+        Returns:
+            An instance of
+            :py:class:`SparkDatasetSource <mlflow.data.spark_dataset_source.SparkDatasetSource>` or
+            :py:class:`DeltaDatasetSource <mlflow.data.delta_dataset_source.DeltaDatasetSource>`.
+        """
+        return self._source
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be None if no profile is available.
+        """
+        try:
+            from pyspark.rdd import BoundedFloat
+
+            # Use Spark RDD countApprox to get approximate count since count() may be expensive.
+            # Note that we call the Scala RDD API because the PySpark API does not respect the
+            # specified timeout. Reference code:
+            # https://spark.apache.org/docs/3.4.0/api/python/_modules/pyspark/rdd.html
+            # #RDD.countApprox. This is confirmed to work in all Spark 3.x versions
+            py_rdd = self.df.rdd
+            drdd = py_rdd.mapPartitions(lambda it: [float(sum(1 for i in it))])
+            jrdd = drdd.mapPartitions(lambda it: [float(sum(it))])._to_java_object_rdd()
+            jdrdd = drdd.ctx._jvm.JavaDoubleRDD.fromRDD(jrdd.rdd())
+            timeout_millis = 5000
+            confidence = 0.9
+            approx_count_operation = jdrdd.sumApprox(timeout_millis, confidence)
+            approx_count_result = approx_count_operation.initialValue()
+            approx_count_float = BoundedFloat(
+                mean=approx_count_result.mean(),
+                confidence=approx_count_result.confidence(),
+                low=approx_count_result.low(),
+                high=approx_count_result.high(),
+            )
+            approx_count = int(approx_count_float)
+            if approx_count <= 0:
+                # An approximate count of zero likely indicates that the count timed
+                # out before an estimate could be made. In this case, we use the value
+                # "unknown" so that users don't think the dataset is empty
+                approx_count = "unknown"
+
+            return {
+                "approx_count": approx_count,
+            }
+        except Exception as e:
+            _logger.warning(
+                "Encountered an unexpected exception while computing Spark dataset profile."
+                " Exception: %s",
+                e,
+            )
+
+    @cached_property
+    def schema(self) -> Optional[Schema]:
+        """
+        The MLflow ColSpec schema of the Spark dataset.
+        """
+        try:
+            return _infer_schema(self._df)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for Spark dataset. Exception: %s", e)
+            return None
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the Spark DataFrame to pandas and splits the resulting
+        :py:class:`pandas.DataFrame` into: 1. a :py:class:`pandas.DataFrame` of features and
+        2. a :py:class:`pandas.Series` of targets.
+
+        To avoid overuse of driver memory, only the first 10,000 DataFrame rows are selected.
+        """
+        df = self._df.limit(10000).toPandas()
+        if self._targets is not None:
+            if self._targets not in df.columns:
+                raise MlflowException(
+                    f"Failed to convert Spark dataset to pyfunc inputs and outputs because"
+                    f" the pandas representation of the Spark dataset does not contain the"
+                    f" specified targets column '{self._targets}'.",
+                    # This is an internal error because we should have validated the presence of
+                    # the target column in the Hugging Face dataset at construction time
+                    INTERNAL_ERROR,
+                )
+            inputs = df.drop(columns=self._targets)
+            outputs = df[self._targets]
+            return PyFuncInputsOutputs(inputs=inputs, outputs=outputs)
+        else:
+            return PyFuncInputsOutputs(inputs=df, outputs=None)
+
+    def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
+        """
+        Converts the dataset to an EvaluationDataset for model evaluation. Required
+        for use with mlflow.evaluate().
+        """
+        return EvaluationDataset(
+            data=self._df.limit(10000).toPandas(),
+            targets=self._targets,
+            path=path,
+            feature_names=feature_names,
+            predictions=self._predictions,
+            name=self.name,
+            digest=self.digest,
+        )
+
+
+def load_delta(
+    path: Optional[str] = None,
+    table_name: Optional[str] = None,
+    version: Optional[str] = None,
+    targets: Optional[str] = None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+) -> SparkDataset:
+    """
+    Loads a :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>` from a Delta table
+    for use with MLflow Tracking.
+
+    Args:
+        path: The path to the Delta table. Either ``path`` or ``table_name`` must be specified.
+        table_name: The name of the Delta table. Either ``path`` or ``table_name`` must be
+            specified.
+        version: The Delta table version. If not specified, the version will be inferred.
+        targets: Optional. The name of the Delta table column containing targets (labels) for
+            supervised learning.
+        name: The name of the dataset. E.g. "wiki_train". If unspecified, a name is
+            automatically generated.
+        digest: The digest (hash, fingerprint) of the dataset. If unspecified, a digest
+            is automatically computed.
+
+    Returns:
+        An instance of :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>`.
+    """
+    from mlflow.data.spark_delta_utils import (
+        _try_get_delta_table_latest_version_from_path,
+        _try_get_delta_table_latest_version_from_table_name,
+    )
+
+    if (path, table_name).count(None) != 1:
+        raise MlflowException(
+            "Must specify exactly one of `table_name` or `path`.",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    if version is None:
+        if path is not None:
+            version = _try_get_delta_table_latest_version_from_path(path)
+        else:
+            version = _try_get_delta_table_latest_version_from_table_name(table_name)
+
+    if name is None and table_name is not None:
+        name = table_name + (f"@v{version}" if version is not None else "")
+
+    source = DeltaDatasetSource(path=path, delta_table_name=table_name, delta_table_version=version)
+    df = source.load()
+
+    return SparkDataset(
+        df=df,
+        source=source,
+        targets=targets,
+        name=name,
+        digest=digest,
+    )
+
+
+def from_spark(
+    df: "pyspark.sql.DataFrame",
+    path: Optional[str] = None,
+    table_name: Optional[str] = None,
+    version: Optional[str] = None,
+    sql: Optional[str] = None,
+    targets: Optional[str] = None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+    predictions: Optional[str] = None,
+) -> SparkDataset:
+    """
+    Given a Spark DataFrame, constructs a
+    :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>` object for use with
+    MLflow Tracking.
+
+    Args:
+        df: The Spark DataFrame from which to construct a SparkDataset.
+        path: The path of the Spark or Delta source that the DataFrame originally came from. Note
+            that the path does not have to match the DataFrame exactly, since the DataFrame may have
+            been modified by Spark operations. This is used to reload the dataset upon request via
+            :py:func:`SparkDataset.source.load()
+            <mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. If none of ``path``,
+            ``table_name``, or ``sql`` are specified, a CodeDatasetSource is used, which will source
+            information from the run context.
+        table_name: The name of the Spark or Delta table that the DataFrame originally came from.
+            Note that the table does not have to match the DataFrame exactly, since the DataFrame
+            may have been modified by Spark operations. This is used to reload the dataset upon
+            request via :py:func:`SparkDataset.source.load()
+            <mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. If none of ``path``,
+            ``table_name``, or ``sql`` are specified, a CodeDatasetSource is used, which will source
+            information from the run context.
+        version: If the DataFrame originally came from a Delta table, specifies the version of the
+            Delta table. This is used to reload the dataset upon request via
+            :py:func:`SparkDataset.source.load()
+            <mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. ``version`` cannot be
+            specified if ``sql`` is specified.
+        sql: The Spark SQL statement that was originally used to construct the DataFrame. Note that
+            the Spark SQL statement does not have to match the DataFrame exactly, since the
+            DataFrame may have been modified by Spark operations. This is used to reload the dataset
+            upon request via :py:func:`SparkDataset.source.load()
+            <mlflow.data.spark_dataset_source.SparkDatasetSource.load>`. If none of ``path``,
+            ``table_name``, or ``sql`` are specified, a CodeDatasetSource is used, which will source
+            information from the run context.
+        targets: Optional. The name of the Data Frame column containing targets (labels) for
+            supervised learning.
+        name: The name of the dataset. E.g. "wiki_train". If unspecified, a name is automatically
+            generated.
+        digest: The digest (hash, fingerprint) of the dataset. If unspecified, a digest is
+            automatically computed.
+        predictions: Optional. The name of the column containing model predictions,
+            if the dataset contains model predictions. If specified, this column
+            must be present in the dataframe (``df``).
+
+    Returns:
+        An instance of :py:class:`SparkDataset <mlflow.data.spark_dataset.SparkDataset>`.
+    """
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.spark_delta_utils import (
+        _is_delta_table,
+        _is_delta_table_path,
+        _try_get_delta_table_latest_version_from_path,
+        _try_get_delta_table_latest_version_from_table_name,
+    )
+    from mlflow.tracking.context import registry
+
+    if (path, table_name, sql).count(None) < 2:
+        raise MlflowException(
+            "Must specify at most one of `path`, `table_name`, or `sql`.",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    if (sql, version).count(None) == 0:
+        raise MlflowException(
+            "`version` may not be specified when `sql` is specified. `version` may only be"
+            " specified when `table_name` or `path` is specified.",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    if sql is not None:
+        source = SparkDatasetSource(sql=sql)
+    elif path is not None:
+        if _is_delta_table_path(path):
+            version = version or _try_get_delta_table_latest_version_from_path(path)
+            source = DeltaDatasetSource(path=path, delta_table_version=version)
+        elif version is None:
+            source = SparkDatasetSource(path=path)
+        else:
+            raise MlflowException(
+                f"Version '{version}' was specified, but the path '{path}' does not refer"
+                f" to a Delta table.",
+                INVALID_PARAMETER_VALUE,
+            )
+    elif table_name is not None:
+        if _is_delta_table(table_name):
+            version = version or _try_get_delta_table_latest_version_from_table_name(table_name)
+            source = DeltaDatasetSource(
+                delta_table_name=table_name,
+                delta_table_version=version,
+            )
+        elif version is None:
+            source = SparkDatasetSource(table_name=table_name)
+        else:
+            raise MlflowException(
+                f"Version '{version}' was specified, but could not find a Delta table with name"
+                f" '{table_name}'.",
+                INVALID_PARAMETER_VALUE,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        source = CodeDatasetSource(tags=context_tags)
+
+    return SparkDataset(
+        df=df,
+        source=source,
+        targets=targets,
+        name=name,
+        digest=digest,
+        predictions=predictions,
+    )

mlflow/data/spark_dataset_source.py

@@ -0,0 +1,74 @@
+from typing import Any, Optional
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+
+class SparkDatasetSource(DatasetSource):
+    """
+    Represents the source of a dataset stored in a spark table.
+    """
+
+    def __init__(
+        self,
+        path: Optional[str] = None,
+        table_name: Optional[str] = None,
+        sql: Optional[str] = None,
+    ):
+        if (path, table_name, sql).count(None) != 2:
+            raise MlflowException(
+                'Must specify exactly one of "path", "table_name", or "sql"',
+                INVALID_PARAMETER_VALUE,
+            )
+        self._path = path
+        self._table_name = table_name
+        self._sql = sql
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "spark"
+
+    def load(self, **kwargs):
+        """Loads the dataset source as a Spark Dataset Source.
+
+        Returns:
+            An instance of ``pyspark.sql.DataFrame``.
+
+        """
+        from pyspark.sql import SparkSession
+
+        spark = SparkSession.builder.getOrCreate()
+
+        if self._path:
+            return spark.read.parquet(self._path)
+        if self._table_name:
+            return spark.read.table(self._table_name)
+        if self._sql:
+            return spark.sql(self._sql)
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        return False
+
+    @classmethod
+    def _resolve(cls, raw_source: str) -> "SparkDatasetSource":
+        raise NotImplementedError
+
+    def to_dict(self) -> dict[Any, Any]:
+        info = {}
+        if self._path is not None:
+            info["path"] = self._path
+        elif self._table_name is not None:
+            info["table_name"] = self._table_name
+        elif self._sql is not None:
+            info["sql"] = self._sql
+        return info
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "SparkDatasetSource":
+        return cls(
+            path=source_dict.get("path"),
+            table_name=source_dict.get("table_name"),
+            sql=source_dict.get("sql"),
+        )

mlflow/data/spark_delta_utils.py

@@ -0,0 +1,118 @@
+import logging
+import os
+from typing import Optional
+
+from mlflow.utils.string_utils import _backtick_quote
+
+_logger = logging.getLogger(__name__)
+
+
+def _is_delta_table(table_name: str) -> bool:
+    """Checks if a Delta table exists with the specified table name.
+
+    Returns:
+        True if a Delta table exists with the specified table name. False otherwise.
+
+    """
+    from pyspark.sql import SparkSession
+    from pyspark.sql.utils import AnalysisException
+
+    spark = SparkSession.builder.getOrCreate()
+
+    try:
+        # use DESCRIBE DETAIL to check if the table is a Delta table
+        # https://docs.databricks.com/delta/delta-utility.html#describe-detail
+        # format will be `delta` for delta tables
+        spark.sql(f"DESCRIBE DETAIL {table_name}").filter("format = 'delta'").count()
+        return True
+    except AnalysisException:
+        return False
+
+
+def _is_delta_table_path(path: str) -> bool:
+    """Checks if the specified filesystem path is a Delta table.
+
+    Returns:
+        True if the specified path is a Delta table. False otherwise.
+    """
+    if os.path.exists(path) and os.path.isdir(path) and "_delta_log" in os.listdir(path):
+        return True
+    from mlflow.utils.uri import dbfs_hdfs_uri_to_fuse_path
+
+    try:
+        dbfs_path = dbfs_hdfs_uri_to_fuse_path(path)
+        return os.path.exists(dbfs_path) and "_delta_log" in os.listdir(dbfs_path)
+    except Exception:
+        return False
+
+
+def _try_get_delta_table_latest_version_from_path(path: str) -> Optional[int]:
+    """Gets the latest version of the Delta table located at the specified path.
+
+    Args:
+        path: The path to the Delta table.
+
+    Returns:
+        The version of the Delta table, or None if it cannot be resolved (e.g. because the
+        Delta core library is not installed or the specified path does not refer to a Delta
+        table).
+
+    """
+    from pyspark.sql import SparkSession
+
+    try:
+        spark = SparkSession.builder.getOrCreate()
+        j_delta_table = spark._jvm.io.delta.tables.DeltaTable.forPath(spark._jsparkSession, path)
+        return _get_delta_table_latest_version(j_delta_table)
+    except Exception as e:
+        _logger.warning(
+            "Failed to obtain version information for Delta table at path '%s'. Version information"
+            " may not be included in the dataset source for MLflow Tracking. Exception: %s",
+            path,
+            e,
+        )
+
+
+def _try_get_delta_table_latest_version_from_table_name(table_name: str) -> Optional[int]:
+    """Gets the latest version of the Delta table with the specified name.
+
+    Args:
+        table_name: The name of the Delta table.
+
+    Returns:
+        The version of the Delta table, or None if it cannot be resolved (e.g. because the
+        Delta core library is not installed or no such table exists).
+    """
+    from pyspark.sql import SparkSession
+
+    try:
+        spark = SparkSession.builder.getOrCreate()
+        backticked_table_name = ".".join(map(_backtick_quote, table_name.split(".")))
+        j_delta_table = spark._jvm.io.delta.tables.DeltaTable.forName(
+            spark._jsparkSession, backticked_table_name
+        )
+        return _get_delta_table_latest_version(j_delta_table)
+    except Exception as e:
+        _logger.warning(
+            "Failed to obtain version information for Delta table with name '%s'. Version"
+            " information may not be included in the dataset source for MLflow Tracking."
+            " Exception: %s",
+            table_name,
+            e,
+        )
+
+
+def _get_delta_table_latest_version(j_delta_table) -> int:
+    """Obtains the latest version of the specified Delta table Java class.
+
+    Args:
+        j_delta_table: A Java DeltaTable class instance.
+
+    Returns:
+        The version of the Delta table.
+
+    """
+    latest_commit_jdf = j_delta_table.history(1)
+    latest_commit_row = latest_commit_jdf.head()
+    version_field_idx = latest_commit_row.fieldIndex("version")
+    return latest_commit_row.get(version_field_idx)