star

venv/lib/python3.9/site-packages/mlflow/data/__init__.py

@@ -0,0 +1,75 @@
+import sys
+from contextlib import suppress
+from typing import Union
+
+from mlflow.data import dataset_registry
+from mlflow.data import sources as mlflow_data_sources
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json, get_registered_sources
+from mlflow.entities import Dataset as DatasetEntity
+from mlflow.entities import DatasetInput
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+with suppress(ImportError):
+    # Suppressing ImportError to pass mlflow-skinny testing.
+    from mlflow.data import meta_dataset  # noqa: F401
+
+
+def get_source(dataset: Union[DatasetEntity, DatasetInput, Dataset]) -> DatasetSource:
+    """Obtains the source of the specified dataset or dataset input.
+
+    Args:
+        dataset:
+            An instance of :py:class:`mlflow.data.dataset.Dataset <mlflow.data.dataset.Dataset>`,
+            :py:class:`mlflow.entities.Dataset`, or :py:class:`mlflow.entities.DatasetInput`.
+
+    Returns:
+        An instance of :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+
+    """
+    if isinstance(dataset, DatasetInput):
+        dataset: DatasetEntity = dataset.dataset
+
+    if isinstance(dataset, DatasetEntity):
+        dataset_source: DatasetSource = get_dataset_source_from_json(
+            source_json=dataset.source,
+            source_type=dataset.source_type,
+        )
+    elif isinstance(dataset, Dataset):
+        dataset_source: DatasetSource = dataset.source
+    else:
+        raise MlflowException(
+            f"Unrecognized dataset type {type(dataset)}. Expected one of: "
+            f"`mlflow.data.dataset.Dataset`,"
+            f" `mlflow.entities.Dataset`, `mlflow.entities.DatasetInput`.",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    return dataset_source
+
+
+__all__ = ["get_source"]
+
+
+def _define_dataset_constructors_in_current_module():
+    data_module = sys.modules[__name__]
+    for (
+        constructor_name,
+        constructor_fn,
+    ) in dataset_registry.get_registered_constructors().items():
+        setattr(data_module, constructor_name, constructor_fn)
+        __all__.append(constructor_name)
+
+
+_define_dataset_constructors_in_current_module()
+
+
+def _define_dataset_sources_in_sources_module():
+    for source in get_registered_sources():
+        setattr(mlflow_data_sources, source.__name__, source)
+        mlflow_data_sources.__all__.append(source.__name__)
+
+
+_define_dataset_sources_in_sources_module()

venv/lib/python3.9/site-packages/mlflow/data/artifact_dataset_sources.py

@@ -0,0 +1,170 @@
+import re
+import warnings
+from pathlib import Path
+from typing import Any, TypeVar
+from urllib.parse import urlparse
+
+from mlflow.artifacts import download_artifacts
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repository_registry import get_registered_artifact_repositories
+from mlflow.utils.uri import is_local_uri
+
+
+def register_artifact_dataset_sources():
+    from mlflow.data.dataset_source_registry import register_dataset_source
+
+    registered_source_schemes = set()
+    artifact_schemes_to_exclude = [
+        "http",
+        "https",
+        "runs",
+        "models",
+        "mlflow-artifacts",
+        # DBFS supports two access patterns: dbfs:/ (URI) and /dbfs (FUSE).
+        # The DBFS artifact repository online supports dbfs:/ (URI). To ensure
+        # a consistent dictionary representation of DBFS datasets across the URI and
+        # FUSE representations, we exclude dbfs from the set of dataset sources
+        # that are autogenerated using artifact repositories and instead define
+        # a separate DBFSDatasetSource elsewhere
+        "dbfs",
+    ]
+    schemes_to_artifact_repos = get_registered_artifact_repositories()
+    for scheme, artifact_repo in schemes_to_artifact_repos.items():
+        if scheme in artifact_schemes_to_exclude or scheme in registered_source_schemes:
+            continue
+
+        if "ArtifactRepository" in artifact_repo.__name__:
+            # Artifact repository name is something like "LocalArtifactRepository",
+            # "S3ArtifactRepository", etc. To preserve capitalization, strip ArtifactRepository
+            # and replace it with ArtifactDatasetSource
+            dataset_source_name = artifact_repo.__name__.replace(
+                "ArtifactRepository", "ArtifactDatasetSource"
+            )
+        else:
+            # Artifact repository name has some other form, e.g. "dbfs_artifact_repo_factory".
+            # In this case, generate the name by capitalizing the first letter of the scheme and
+            # appending ArtifactRepository
+            scheme = str(scheme)
+
+            def camelcase_scheme(scheme):
+                parts = re.split(r"[-_]", scheme)
+                return "".join([part.capitalize() for part in parts])
+
+            source_name_prefix = camelcase_scheme(scheme)
+            dataset_source_name = source_name_prefix + "ArtifactDatasetSource"
+
+        try:
+            registered_source_schemes.add(scheme)
+            dataset_source = _create_dataset_source_for_artifact_repo(
+                scheme=scheme, dataset_source_name=dataset_source_name
+            )
+            register_dataset_source(dataset_source)
+        except Exception as e:
+            warnings.warn(
+                f"Failed to register a dataset source for URIs with scheme '{scheme}': {e}",
+                stacklevel=2,
+            )
+
+
+def _create_dataset_source_for_artifact_repo(scheme: str, dataset_source_name: str):
+    from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
+
+    if scheme in ["", "file"]:
+        source_type = "local"
+        class_docstring = "Represents the source of a dataset stored on the local filesystem."
+    else:
+        source_type = scheme
+        class_docstring = (
+            f"Represents a filesystem-based or blob-storage-based dataset source identified by a"
+            f" URI with scheme '{scheme}'."
+        )
+
+    DatasetForArtifactRepoSourceType = TypeVar(dataset_source_name)
+
+    class ArtifactRepoSource(FileSystemDatasetSource):
+        def __init__(self, uri: str):
+            self._uri = uri
+
+        @property
+        def uri(self):
+            """
+            The URI with scheme '{scheme}' referring to the dataset source filesystem location.
+
+            Returns
+                The URI with scheme '{scheme}' referring to the dataset source filesystem
+                location.
+            """
+            return self._uri
+
+        @staticmethod
+        def _get_source_type() -> str:
+            return source_type
+
+        def load(self, dst_path=None) -> str:
+            """
+            Downloads the dataset source to the local filesystem.
+
+            Args:
+                dst_path: Path of the local filesystem destination directory to which to download
+                    the dataset source. If the directory does not exist, it is created. If
+                    unspecified, the dataset source is downloaded to a new uniquely-named
+                    directory on the local filesystem, unless the dataset source already
+                    exists on the local filesystem, in which case its local path is
+                    returned directly.
+
+            Returns:
+                The path to the downloaded dataset source on the local filesystem.
+            """
+            return download_artifacts(artifact_uri=self.uri, dst_path=dst_path)
+
+        @staticmethod
+        def _can_resolve(raw_source: Any):
+            is_local_source_type = ArtifactRepoSource._get_source_type() == "local"
+
+            if not isinstance(raw_source, str) and (
+                not isinstance(raw_source, Path) and is_local_source_type
+            ):
+                return False
+
+            try:
+                if is_local_source_type:
+                    return is_local_uri(str(raw_source), is_tracking_or_registry_uri=False)
+                else:
+                    parsed_source = urlparse(str(raw_source))
+                    return parsed_source.scheme == scheme
+            except Exception:
+                return False
+
+        @classmethod
+        def _resolve(cls, raw_source: Any) -> DatasetForArtifactRepoSourceType:
+            return cls(str(raw_source))
+
+        def to_dict(self) -> dict[Any, Any]:
+            """
+            Returns:
+                A JSON-compatible dictionary representation of the {dataset_source_name}.
+            """
+            return {
+                "uri": self.uri,
+            }
+
+        @classmethod
+        def from_dict(cls, source_dict: dict[Any, Any]) -> DatasetForArtifactRepoSourceType:
+            uri = source_dict.get("uri")
+            if uri is None:
+                raise MlflowException(
+                    f'Failed to parse {dataset_source_name}. Missing expected key: "uri"',
+                    INVALID_PARAMETER_VALUE,
+                )
+
+            return cls(uri=uri)
+
+    ArtifactRepoSource.__name__ = dataset_source_name
+    ArtifactRepoSource.__qualname__ = dataset_source_name
+    ArtifactRepoSource.__doc__ = class_docstring
+    ArtifactRepoSource.to_dict.__doc__ = ArtifactRepoSource.to_dict.__doc__.format(
+        dataset_source_name=dataset_source_name
+    )
+    ArtifactRepoSource.uri.__doc__ = ArtifactRepoSource.uri.__doc__.format(scheme=scheme)
+    return ArtifactRepoSource

venv/lib/python3.9/site-packages/mlflow/data/code_dataset_source.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from typing_extensions import Self
+
+from mlflow.data.dataset_source import DatasetSource
+
+
+class CodeDatasetSource(DatasetSource):
+    def __init__(
+        self,
+        tags: dict[Any, Any],
+    ):
+        self._tags = tags
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "code"
+
+    def load(self, **kwargs):
+        """
+        Load is not implemented for Code Dataset Source.
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        return False
+
+    @classmethod
+    def _resolve(cls, raw_source: str) -> Self:
+        raise NotImplementedError
+
+    def to_dict(self) -> dict[Any, Any]:
+        return {"tags": self._tags}
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> Self:
+        return cls(
+            tags=source_dict.get("tags"),
+        )

venv/lib/python3.9/site-packages/mlflow/data/dataset.py

@@ -0,0 +1,123 @@
+import json
+from abc import abstractmethod
+from typing import Any, Optional
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.entities import Dataset as DatasetEntity
+
+
+class Dataset:
+    """
+    Represents a dataset for use with MLflow Tracking, including the name, digest (hash),
+    schema, and profile of the dataset as well as source information (e.g. the S3 bucket or
+    managed Delta table from which the dataset was derived). Most datasets expose features
+    and targets for training and evaluation as well.
+    """
+
+    def __init__(
+        self, source: DatasetSource, name: Optional[str] = None, digest: Optional[str] = None
+    ):
+        """
+        Base constructor for a dataset. All subclasses must call this constructor.
+        """
+        self._name = name
+        self._source = source
+        # Note: Subclasses should call super() once they've initialized all of
+        # the class attributes necessary for digest computation
+        self._digest = digest or self._compute_digest()
+
+    @abstractmethod
+    def _compute_digest(self) -> str:
+        """Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+
+        Returns:
+            A string digest for the dataset. We recommend a maximum digest length
+            of 10 characters with an ideal length of 8 characters.
+
+        """
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the dataset.
+
+        Subclasses should override this method to provide additional fields in the config dict,
+        e.g., schema, profile, etc.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type.
+        """
+        return {
+            "name": self.name,
+            "digest": self.digest,
+            "source": self.source.to_json(),
+            "source_type": self.source._get_source_type(),
+        }
+
+    def to_json(self) -> str:
+        """
+        Obtains a JSON string representation of the :py:class:`Dataset
+        <mlflow.data.dataset.Dataset>`.
+
+        Returns:
+            A JSON string representation of the :py:class:`Dataset <mlflow.data.dataset.Dataset>`.
+        """
+
+        return json.dumps(self.to_dict())
+
+    @property
+    def name(self) -> str:
+        """
+        The name of the dataset, e.g. ``"iris_data"``, ``"myschema.mycatalog.mytable@v1"``, etc.
+        """
+        if self._name is not None:
+            return self._name
+        else:
+            return "dataset"
+
+    @property
+    def digest(self) -> str:
+        """
+        A unique hash or fingerprint of the dataset, e.g. ``"498c7496"``.
+        """
+        return self._digest
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        Information about the dataset's source, represented as an instance of
+        :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`. For example, this
+        may be the S3 location or the name of the managed Delta Table from which the dataset
+        was derived.
+        """
+        return self._source
+
+    @property
+    @abstractmethod
+    def profile(self) -> Optional[Any]:
+        """
+        Optional summary statistics for the dataset, such as the number of rows in a table, the
+        mean / median / std of each table column, etc.
+        """
+
+    @property
+    @abstractmethod
+    def schema(self) -> Optional[Any]:
+        """
+        Optional dataset schema, such as an instance of :py:class:`mlflow.types.Schema` representing
+        the features and targets of the dataset.
+        """
+
+    def _to_mlflow_entity(self) -> DatasetEntity:
+        """
+        Returns:
+            A `mlflow.entities.Dataset` instance representing the dataset.
+        """
+        dataset_dict = self.to_dict()
+        return DatasetEntity(
+            name=dataset_dict["name"],
+            digest=dataset_dict["digest"],
+            source_type=dataset_dict["source_type"],
+            source=dataset_dict["source"],
+            schema=dataset_dict.get("schema"),
+            profile=dataset_dict.get("profile"),
+        )

venv/lib/python3.9/site-packages/mlflow/data/dataset_registry.py

@@ -0,0 +1,156 @@
+import inspect
+import warnings
+from contextlib import suppress
+from typing import Callable, Optional
+
+import mlflow.data
+from mlflow.data.dataset import Dataset
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.plugins import get_entry_points
+
+
+class DatasetRegistry:
+    def __init__(self):
+        self.constructors = {}
+
+    def register_constructor(
+        self, constructor_fn: Callable, constructor_name: Optional[str] = None
+    ) -> str:
+        """Registers a dataset constructor.
+
+        Args:
+            constructor_fn: A function that accepts at least the following
+                inputs and returns an instance of a subclass of
+                :py:class:`mlflow.data.dataset.Dataset`:
+
+                - name: Optional. A string dataset name
+                - digest: Optional. A string dataset digest.
+
+            constructor_name: The name of the constructor, e.g.
+                "from_spark". The name must begin with the
+                string "from_" or "load_". If unspecified, the `__name__`
+                attribute of the `constructor_fn` is used instead and must
+                begin with the string "from_" or "load_".
+
+        Returns:
+            The name of the registered constructor, e.g. "from_pandas" or "load_delta".
+        """
+        if constructor_name is None:
+            constructor_name = constructor_fn.__name__
+        DatasetRegistry._validate_constructor(constructor_fn, constructor_name)
+        self.constructors[constructor_name] = constructor_fn
+        return constructor_name
+
+    def register_entrypoints(self):
+        """
+        Registers dataset sources defined as Python entrypoints. For reference, see
+        https://mlflow.org/docs/latest/plugins.html#defining-a-plugin.
+        """
+        for entrypoint in get_entry_points("mlflow.dataset_constructor"):
+            try:
+                self.register_constructor(
+                    constructor_fn=entrypoint.load(), constructor_name=entrypoint.name
+                )
+            except Exception as exc:
+                warnings.warn(
+                    f"Failure attempting to register dataset constructor"
+                    f' "{entrypoint.name}": {exc}.',
+                    stacklevel=2,
+                )
+
+    @staticmethod
+    def _validate_constructor(constructor_fn: Callable, constructor_name: str):
+        if not constructor_name.startswith("load_") and not constructor_name.startswith("from_"):
+            raise MlflowException(
+                f"Invalid dataset constructor name: {constructor_name}."
+                f" Constructor name must start with 'load_' or 'from_'.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        signature = inspect.signature(constructor_fn)
+        parameters = signature.parameters
+        for expected_kwarg in ["name", "digest"]:
+            if expected_kwarg not in parameters or parameters[expected_kwarg].kind not in [
+                inspect.Parameter.KEYWORD_ONLY,
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            ]:
+                raise MlflowException(
+                    f"Invalid dataset constructor function: {constructor_fn.__name__}. Function"
+                    f" must define an optional parameter named '{expected_kwarg}'.",
+                    INVALID_PARAMETER_VALUE,
+                )
+
+        if not issubclass(signature.return_annotation, Dataset):
+            raise MlflowException(
+                f"Invalid dataset constructor function: {constructor_fn.__name__}. Function must"
+                f" have a return type annotation that is a subclass of"
+                f" :py:class:`mlflow.data.dataset.Dataset`.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+
+def register_constructor(constructor_fn: Callable, constructor_name: Optional[str] = None) -> str:
+    """Registers a dataset constructor.
+
+    Args:
+        constructor_fn: A function that accepts at least the following
+            inputs and returns an instance of a subclass of
+            :py:class:`mlflow.data.dataset.Dataset`:
+
+            - name: Optional. A string dataset name
+            - digest: Optional. A string dataset digest.
+
+        constructor_name: The name of the constructor, e.g.
+            "from_spark". The name must begin with the
+            string "from_" or "load_". If unspecified, the `__name__`
+            attribute of the `constructor_fn` is used instead and must
+            begin with the string "from_" or "load_".
+
+    Returns:
+        The name of the registered constructor, e.g. "from_pandas" or "load_delta".
+
+    """
+    registered_constructor_name = _dataset_registry.register_constructor(
+        constructor_fn=constructor_fn, constructor_name=constructor_name
+    )
+    setattr(mlflow.data, registered_constructor_name, constructor_fn)
+    mlflow.data.__all__.append(registered_constructor_name)
+    return registered_constructor_name
+
+
+def get_registered_constructors() -> dict[str, Callable]:
+    """Obtains the registered dataset constructors.
+
+    Returns:
+        A dictionary mapping constructor names to constructor functions.
+
+    """
+    return _dataset_registry.constructors
+
+
+_dataset_registry = DatasetRegistry()
+_dataset_registry.register_entrypoints()
+
+# use contextlib suppress to ignore import errors
+with suppress(ImportError):
+    from mlflow.data.pandas_dataset import from_pandas
+
+    _dataset_registry.register_constructor(from_pandas)
+with suppress(ImportError):
+    from mlflow.data.numpy_dataset import from_numpy
+
+    _dataset_registry.register_constructor(from_numpy)
+with suppress(ImportError):
+    from mlflow.data.huggingface_dataset import from_huggingface
+
+    _dataset_registry.register_constructor(from_huggingface)
+with suppress(ImportError):
+    from mlflow.data.tensorflow_dataset import from_tensorflow
+
+    _dataset_registry.register_constructor(from_tensorflow)
+with suppress(ImportError):
+    from mlflow.data.spark_dataset import from_spark, load_delta
+
+    _dataset_registry.register_constructor(load_delta)
+    _dataset_registry.register_constructor(from_spark)

venv/lib/python3.9/site-packages/mlflow/data/dataset_source.py

@@ -0,0 +1,110 @@
+import json
+from abc import abstractmethod
+from typing import Any
+
+
+class DatasetSource:
+    """
+    Represents the source of a dataset used in MLflow Tracking, providing information such as
+    cloud storage location, delta table name / version, etc.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def _get_source_type() -> str:
+        """Obtains a string representing the source type of the dataset.
+
+        Returns:
+            A string representing the source type of the dataset, e.g. "s3", "delta_table", ...
+
+        """
+
+    @abstractmethod
+    def load(self) -> Any:
+        """
+        Loads files / objects referred to by the DatasetSource. For example, depending on the type
+        of :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`, this may download
+        source CSV files from S3 to the local filesystem, load a source Delta Table as a Spark
+        DataFrame, etc.
+
+        Returns:
+            The downloaded source, e.g. a local filesystem path, a Spark DataFrame, etc.
+
+        """
+
+    @staticmethod
+    @abstractmethod
+    def _can_resolve(raw_source: Any) -> bool:
+        """Determines whether this type of DatasetSource can be resolved from a specified raw source
+        object. For example, an S3DatasetSource can be resolved from an S3 URI like
+        "s3://mybucket/path/to/iris/data" but not from an Azure Blob Storage URI like
+        "wasbs:/account@host.blob.core.windows.net".
+
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data".
+
+        Returns:
+            True if this DatasetSource can resolve the raw source, False otherwise.
+
+        """
+
+    @classmethod
+    @abstractmethod
+    def _resolve(cls, raw_source: Any) -> "DatasetSource":
+        """Constructs an instance of the DatasetSource from a raw source object, such as a
+        string URI like "s3://mybucket/path/to/iris/data" or a delta table identifier
+        like "my.delta.table@2".
+
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data".
+
+        Returns:
+            A DatasetSource instance derived from the raw_source.
+
+        """
+
+    @abstractmethod
+    def to_dict(self) -> dict[str, Any]:
+        """Obtains a JSON-compatible dictionary representation of the DatasetSource.
+
+        Returns:
+            A JSON-compatible dictionary representation of the DatasetSource.
+
+        """
+
+    def to_json(self) -> str:
+        """
+        Obtains a JSON string representation of the
+        :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+
+        Returns:
+            A JSON string representation of the
+            :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+        """
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "DatasetSource":
+        """Constructs an instance of the DatasetSource from a dictionary representation.
+
+        Args:
+            source_dict: A dictionary representation of the DatasetSource.
+
+        Returns:
+            A DatasetSource instance.
+
+        """
+
+    @classmethod
+    def from_json(cls, source_json: str) -> "DatasetSource":
+        """Constructs an instance of the DatasetSource from a JSON string representation.
+
+        Args:
+            source_json: A JSON string representation of the DatasetSource.
+
+        Returns:
+            A DatasetSource instance.
+
+        """
+        return cls.from_dict(json.loads(source_json))

venv/lib/python3.9/site-packages/mlflow/data/dataset_source_registry.py

@@ -0,0 +1,219 @@
+import warnings
+from typing import Any, Optional
+
+from mlflow.data.artifact_dataset_sources import register_artifact_dataset_sources
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.http_dataset_source import HTTPDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
+from mlflow.utils.plugins import get_entry_points
+
+
+class DatasetSourceRegistry:
+    def __init__(self):
+        self.sources = []
+
+    def register(self, source: DatasetSource):
+        """Registers a DatasetSource for use with MLflow Tracking.
+
+        Args:
+            source: The DatasetSource to register.
+        """
+        self.sources.append(source)
+
+    def register_entrypoints(self):
+        """
+        Registers dataset sources defined as Python entrypoints. For reference, see
+        https://mlflow.org/docs/latest/plugins.html#defining-a-plugin.
+        """
+        for entrypoint in get_entry_points("mlflow.dataset_source"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    "Failure attempting to register dataset constructor"
+                    + f' "{entrypoint}": {exc}',
+                    stacklevel=2,
+                )
+
+    def resolve(
+        self, raw_source: Any, candidate_sources: Optional[list[DatasetSource]] = None
+    ) -> DatasetSource:
+        """Resolves a raw source object, such as a string URI, to a DatasetSource for use with
+        MLflow Tracking.
+
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data" or a
+                HuggingFace :py:class:`datasets.Dataset` object.
+            candidate_sources: A list of DatasetSource classes to consider as potential sources
+                when resolving the raw source. Subclasses of the specified candidate sources are
+                also considered. If unspecified, all registered sources are considered.
+
+        Raises:
+            MlflowException: If no DatasetSource class can resolve the raw source.
+
+        Returns:
+            The resolved DatasetSource.
+        """
+        matching_sources = []
+        for source in self.sources:
+            if candidate_sources and not any(
+                issubclass(source, candidate_src) for candidate_src in candidate_sources
+            ):
+                continue
+            try:
+                if source._can_resolve(raw_source):
+                    matching_sources.append(source)
+            except Exception as e:
+                warnings.warn(
+                    f"Failed to determine whether {source.__name__} can resolve source"
+                    f" information for '{raw_source}'. Exception: {e}",
+                    stacklevel=2,
+                )
+                continue
+
+        if len(matching_sources) > 1:
+            source_class_names_str = ", ".join([source.__name__ for source in matching_sources])
+            warnings.warn(
+                f"The specified dataset source can be interpreted in multiple ways:"
+                f" {source_class_names_str}. MLflow will assume that this is a"
+                f" {matching_sources[-1].__name__} source.",
+                stacklevel=2,
+            )
+
+        for matching_source in reversed(matching_sources):
+            try:
+                return matching_source._resolve(raw_source)
+            except Exception as e:
+                warnings.warn(
+                    f"Encountered an unexpected error while using {matching_source.__name__} to"
+                    f" resolve source information for '{raw_source}'. Exception: {e}",
+                    stacklevel=2,
+                )
+                continue
+
+        raise MlflowException(
+            f"Could not find a source information resolver for the specified"
+            f" dataset source: {raw_source}.",
+            RESOURCE_DOES_NOT_EXIST,
+        )
+
+    def get_source_from_json(self, source_json: str, source_type: str) -> DatasetSource:
+        """Parses and returns a DatasetSource object from its JSON representation.
+
+        Args:
+            source_json: The JSON representation of the DatasetSource.
+            source_type: The string type of the DatasetSource, which indicates how to parse the
+                source JSON.
+        """
+        for source in reversed(self.sources):
+            if source._get_source_type() == source_type:
+                return source.from_json(source_json)
+
+        raise MlflowException(
+            f"Could not parse dataset source from JSON due to unrecognized"
+            f" source type: {source_type}.",
+            RESOURCE_DOES_NOT_EXIST,
+        )
+
+
+def register_dataset_source(source: DatasetSource):
+    """Registers a DatasetSource for use with MLflow Tracking.
+
+    Args:
+        source: The DatasetSource to register.
+    """
+    _dataset_source_registry.register(source)
+
+
+def resolve_dataset_source(
+    raw_source: Any, candidate_sources: Optional[list[DatasetSource]] = None
+) -> DatasetSource:
+    """Resolves a raw source object, such as a string URI, to a DatasetSource for use with
+    MLflow Tracking.
+
+    Args:
+        raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data" or a
+            HuggingFace :py:class:`datasets.Dataset` object.
+        candidate_sources: A list of DatasetSource classes to consider as potential sources
+            when resolving the raw source. Subclasses of the specified candidate
+            sources are also considered. If unspecified, all registered sources
+            are considered.
+
+    Raises:
+        MlflowException: If no DatasetSource class can resolve the raw source.
+
+    Returns:
+        The resolved DatasetSource.
+    """
+    return _dataset_source_registry.resolve(
+        raw_source=raw_source, candidate_sources=candidate_sources
+    )
+
+
+def get_dataset_source_from_json(source_json: str, source_type: str) -> DatasetSource:
+    """Parses and returns a DatasetSource object from its JSON representation.
+
+    Args:
+        source_json: The JSON representation of the DatasetSource.
+        source_type: The string type of the DatasetSource, which indicates how to parse the
+            source JSON.
+    """
+    return _dataset_source_registry.get_source_from_json(
+        source_json=source_json, source_type=source_type
+    )
+
+
+def get_registered_sources() -> list[DatasetSource]:
+    """Obtains the registered dataset sources.
+
+    Returns:
+        A list of registered dataset sources.
+
+    """
+    return _dataset_source_registry.sources
+
+
+# NB: The ordering here is important. The last dataset source to be registered takes precedence
+# when resolving dataset information for a raw source (e.g. a string like "s3://mybucket/my/path").
+# Dataset sources derived from artifact repositories are the most generic / provide the most
+# general information about dataset source locations, so they are registered first. More specific
+# source information is provided by specialized dataset platform sources like
+# HuggingFaceDatasetSource, so these sources are registered next. Finally, externally-defined
+# dataset sources are registered last because externally-defined behavior should take precedence
+# over any internally-defined generic behavior
+_dataset_source_registry = DatasetSourceRegistry()
+register_artifact_dataset_sources()
+_dataset_source_registry.register(HTTPDatasetSource)
+_dataset_source_registry.register_entrypoints()
+
+try:
+    from mlflow.data.huggingface_dataset_source import HuggingFaceDatasetSource
+
+    _dataset_source_registry.register(HuggingFaceDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.spark_dataset_source import SparkDatasetSource
+
+    _dataset_source_registry.register(SparkDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.delta_dataset_source import DeltaDatasetSource
+
+    _dataset_source_registry.register(DeltaDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+
+    _dataset_source_registry.register(CodeDatasetSource)
+except ImportError:
+    pass
+try:
+    from mlflow.data.uc_volume_dataset_source import UCVolumeDatasetSource
+
+    _dataset_source_registry.register(UCVolumeDatasetSource)
+except ImportError:
+    pass

venv/lib/python3.9/site-packages/mlflow/data/delta_dataset_source.py

@@ -0,0 +1,167 @@
+import logging
+from typing import Any, Optional
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_managed_catalog_messages_pb2 import (
+    GetTable,
+    GetTableResponse,
+)
+from mlflow.protos.databricks_managed_catalog_service_pb2 import DatabricksUnityCatalogService
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils._unity_catalog_utils import get_full_name_from_sc
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+from mlflow.utils.string_utils import _backtick_quote
+
+DATABRICKS_HIVE_METASTORE_NAME = "hive_metastore"
+# these two catalog names both points to the workspace local default HMS (hive metastore).
+DATABRICKS_LOCAL_METASTORE_NAMES = [DATABRICKS_HIVE_METASTORE_NAME, "spark_catalog"]
+# samples catalog is managed by databricks for hosting public dataset like NYC taxi dataset.
+# it is neither a UC nor local metastore catalog
+DATABRICKS_SAMPLES_CATALOG_NAME = "samples"
+
+_logger = logging.getLogger(__name__)
+
+
+class DeltaDatasetSource(DatasetSource):
+    """
+    Represents the source of a dataset stored at in a delta table.
+    """
+
+    def __init__(
+        self,
+        path: Optional[str] = None,
+        delta_table_name: Optional[str] = None,
+        delta_table_version: Optional[int] = None,
+        delta_table_id: Optional[str] = None,
+    ):
+        if (path, delta_table_name).count(None) != 1:
+            raise MlflowException(
+                'Must specify exactly one of "path" or "table_name"',
+                INVALID_PARAMETER_VALUE,
+            )
+        self._path = path
+        if delta_table_name is not None:
+            self._delta_table_name = get_full_name_from_sc(
+                delta_table_name, _get_active_spark_session()
+            )
+        else:
+            self._delta_table_name = delta_table_name
+        self._delta_table_version = delta_table_version
+        self._delta_table_id = delta_table_id
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "delta_table"
+
+    def load(self, **kwargs):
+        """
+        Loads the dataset source as a Delta Dataset Source.
+
+        Returns:
+            An instance of ``pyspark.sql.DataFrame``.
+        """
+        from pyspark.sql import SparkSession
+
+        spark = SparkSession.builder.getOrCreate()
+
+        spark_read_op = spark.read.format("delta")
+        if self._delta_table_version is not None:
+            spark_read_op = spark_read_op.option("versionAsOf", self._delta_table_version)
+
+        if self._path:
+            return spark_read_op.load(self._path)
+        else:
+            backticked_delta_table_name = ".".join(
+                map(_backtick_quote, self._delta_table_name.split("."))
+            )
+            return spark_read_op.table(backticked_delta_table_name)
+
+    @property
+    def path(self) -> Optional[str]:
+        return self._path
+
+    @property
+    def delta_table_name(self) -> Optional[str]:
+        return self._delta_table_name
+
+    @property
+    def delta_table_id(self) -> Optional[str]:
+        return self._delta_table_id
+
+    @property
+    def delta_table_version(self) -> Optional[int]:
+        return self._delta_table_version
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        return False
+
+    @classmethod
+    def _resolve(cls, raw_source: str) -> "DeltaDatasetSource":
+        raise NotImplementedError
+
+    # check if table is in the Databricks Unity Catalog
+    def _is_databricks_uc_table(self):
+        if self._delta_table_name is not None:
+            catalog_name = self._delta_table_name.split(".", 1)[0]
+            return (
+                catalog_name not in DATABRICKS_LOCAL_METASTORE_NAMES
+                and catalog_name != DATABRICKS_SAMPLES_CATALOG_NAME
+            )
+        else:
+            return False
+
+    def _lookup_table_id(self, table_name):
+        try:
+            req_body = message_to_json(GetTable(full_name_arg=table_name))
+            _METHOD_TO_INFO = extract_api_info_for_service(
+                DatabricksUnityCatalogService, _REST_API_PATH_PREFIX
+            )
+            db_creds = get_databricks_host_creds()
+            endpoint, method = _METHOD_TO_INFO[GetTable]
+            # We need to replace the full_name_arg in the endpoint definition with
+            # the actual table name for the REST API to work.
+            final_endpoint = endpoint.replace("{full_name_arg}", table_name)
+            resp = call_endpoint(
+                host_creds=db_creds,
+                endpoint=final_endpoint,
+                method=method,
+                json_body=req_body,
+                response_proto=GetTableResponse,
+            )
+            return resp.table_id
+        except Exception:
+            return None
+
+    def to_dict(self) -> dict[Any, Any]:
+        info = {}
+        if self._path:
+            info["path"] = self._path
+        if self._delta_table_name:
+            info["delta_table_name"] = self._delta_table_name
+        if self._delta_table_version:
+            info["delta_table_version"] = self._delta_table_version
+        if self._is_databricks_uc_table():
+            info["is_databricks_uc_table"] = True
+            if self._delta_table_id:
+                info["delta_table_id"] = self._delta_table_id
+            else:
+                info["delta_table_id"] = self._lookup_table_id(self._delta_table_name)
+        return info
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "DeltaDatasetSource":
+        return cls(
+            path=source_dict.get("path"),
+            delta_table_name=source_dict.get("delta_table_name"),
+            delta_table_version=source_dict.get("delta_table_version"),
+            delta_table_id=source_dict.get("delta_table_id"),
+        )

venv/lib/python3.9/site-packages/mlflow/data/digest_utils.py

@@ -0,0 +1,108 @@
+import hashlib
+from typing import Any
+
+from packaging.version import Version
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+MAX_ROWS = 10000
+
+
+def compute_pandas_digest(df) -> str:
+    """Computes a digest for the given Pandas DataFrame.
+
+    Args:
+        df: A Pandas DataFrame.
+
+    Returns:
+        A string digest.
+    """
+    import numpy as np
+    import pandas as pd
+
+    # trim to max rows
+    trimmed_df = df.head(MAX_ROWS)
+
+    # keep string and number columns, drop other column types
+    if Version(pd.__version__) >= Version("2.1.0"):
+        string_columns = trimmed_df.columns[(df.map(type) == str).all(0)]
+    else:
+        string_columns = trimmed_df.columns[(df.applymap(type) == str).all(0)]
+    numeric_columns = trimmed_df.select_dtypes(include=[np.number]).columns
+
+    desired_columns = string_columns.union(numeric_columns)
+    trimmed_df = trimmed_df[desired_columns]
+
+    return get_normalized_md5_digest(
+        [
+            pd.util.hash_pandas_object(trimmed_df).values,
+            np.int64(len(df)),
+        ]
+        + [str(x).encode() for x in df.columns]
+    )
+
+
+def compute_numpy_digest(features, targets=None) -> str:
+    """Computes a digest for the given numpy array.
+
+    Args:
+        features: A numpy array containing dataset features.
+        targets: A numpy array containing dataset targets. Optional.
+
+    Returns:
+        A string digest.
+    """
+    import numpy as np
+    import pandas as pd
+
+    hashable_elements = []
+
+    def hash_array(array):
+        flattened_array = array.flatten()
+        trimmed_array = flattened_array[0:MAX_ROWS]
+        try:
+            hashable_elements.append(pd.util.hash_array(trimmed_array))
+        except TypeError:
+            hashable_elements.append(np.int64(trimmed_array.size))
+
+        # hash full array dimensions
+        for x in array.shape:
+            hashable_elements.append(np.int64(x))
+
+    def hash_dict_of_arrays(array_dict):
+        for key in sorted(array_dict.keys()):
+            hash_array(array_dict[key])
+
+    for item in [features, targets]:
+        if item is None:
+            continue
+        if isinstance(item, dict):
+            hash_dict_of_arrays(item)
+        else:
+            hash_array(item)
+
+    return get_normalized_md5_digest(hashable_elements)
+
+
+def get_normalized_md5_digest(elements: list[Any]) -> str:
+    """Computes a normalized digest for a list of hashable elements.
+
+    Args:
+        elements: A list of hashable elements for inclusion in the md5 digest.
+
+    Returns:
+        An 8-character, truncated md5 digest.
+    """
+
+    if not elements:
+        raise MlflowException(
+            "No hashable elements were provided for md5 digest creation",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    md5 = hashlib.md5(usedforsecurity=False)
+    for element in elements:
+        md5.update(element)
+
+    return md5.hexdigest()[:8]

venv/lib/python3.9/site-packages/mlflow/data/evaluation_dataset.py

@@ -0,0 +1,553 @@
+import hashlib
+import json
+import logging
+import math
+import struct
+import sys
+
+from packaging.version import Version
+
+import mlflow
+from mlflow.entities import RunTag
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.string_utils import generate_feature_name_if_not_string
+
+try:
+    # `numpy` and `pandas` are not required for `mlflow-skinny`.
+    import numpy as np
+    import pandas as pd
+except ImportError:
+    pass
+
+_logger = logging.getLogger(__name__)
+
+
+def _hash_uint64_ndarray_as_bytes(array):
+    assert len(array.shape) == 1
+    # see struct pack format string https://docs.python.org/3/library/struct.html#format-strings
+    return struct.pack(f">{array.size}Q", *array)
+
+
+def _is_empty_list_or_array(data):
+    if isinstance(data, list):
+        return len(data) == 0
+    elif isinstance(data, np.ndarray):
+        return data.size == 0
+    return False
+
+
+def _is_array_has_dict(nd_array):
+    if _is_empty_list_or_array(nd_array):
+        return False
+
+    # It is less likely the array or list contains heterogeneous elements, so just checking the
+    # first element to avoid performance overhead.
+    elm = nd_array.item(0)
+    if isinstance(elm, (list, np.ndarray)):
+        return _is_array_has_dict(elm)
+    elif isinstance(elm, dict):
+        return True
+
+    return False
+
+
+def _hash_array_of_dict_as_bytes(data):
+    # NB: If an array or list contains dictionary element, it can't be hashed with
+    # pandas.util.hash_array. Hence we need to manually hash the elements here. This is
+    # particularly for the LLM use case where the input can be a list of dictionary
+    # (chat/completion payloads), so doesn't handle more complex case like nested lists.
+    result = b""
+    for elm in data:
+        if isinstance(elm, (list, np.ndarray)):
+            result += _hash_array_of_dict_as_bytes(elm)
+        elif isinstance(elm, dict):
+            result += _hash_dict_as_bytes(elm)
+        else:
+            result += _hash_data_as_bytes(elm)
+    return result
+
+
+def _hash_ndarray_as_bytes(nd_array):
+    if not isinstance(nd_array, np.ndarray):
+        nd_array = np.array(nd_array)
+
+    if _is_array_has_dict(nd_array):
+        return _hash_array_of_dict_as_bytes(nd_array)
+
+    return _hash_uint64_ndarray_as_bytes(
+        pd.util.hash_array(nd_array.flatten(order="C"))
+    ) + _hash_uint64_ndarray_as_bytes(np.array(nd_array.shape, dtype="uint64"))
+
+
+def _hash_data_as_bytes(data):
+    try:
+        if isinstance(data, (list, np.ndarray)):
+            return _hash_ndarray_as_bytes(data)
+        if isinstance(data, dict):
+            return _hash_dict_as_bytes(data)
+        if np.isscalar(data):
+            return _hash_uint64_ndarray_as_bytes(pd.util.hash_array(np.array([data])))
+    finally:
+        return b""  # Skip unsupported types by returning an empty byte string
+
+
+def _hash_dict_as_bytes(data_dict):
+    result = _hash_ndarray_as_bytes(list(data_dict.keys()))
+    try:
+        result += _hash_ndarray_as_bytes(list(data_dict.values()))
+    # If the values containing non-hashable objects, we will hash the values recursively.
+    except Exception:
+        for value in data_dict.values():
+            result += _hash_data_as_bytes(value)
+    return result
+
+
+def _hash_array_like_obj_as_bytes(data):
+    """
+    Helper method to convert pandas dataframe/numpy array/list into bytes for
+    MD5 calculation purpose.
+    """
+    if isinstance(data, pd.DataFrame):
+        # add checking `'pyspark' in sys.modules` to avoid importing pyspark when user
+        # run code not related to pyspark.
+        if "pyspark" in sys.modules:
+            from pyspark.ml.linalg import Vector as spark_vector_type
+        else:
+            spark_vector_type = None
+
+        def _hash_array_like_element_as_bytes(v):
+            if spark_vector_type is not None:
+                if isinstance(v, spark_vector_type):
+                    return _hash_ndarray_as_bytes(v.toArray())
+            if isinstance(v, (dict, list, np.ndarray)):
+                return _hash_data_as_bytes(v)
+
+            try:
+                # Attempt to hash the value, if it fails, return an empty byte string
+                pd.util.hash_array(np.array([v]))
+                return v
+            except TypeError:
+                return b""  # Skip unhashable types by returning an empty byte string
+
+        if Version(pd.__version__) >= Version("2.1.0"):
+            data = data.map(_hash_array_like_element_as_bytes)
+        else:
+            data = data.applymap(_hash_array_like_element_as_bytes)
+        return _hash_uint64_ndarray_as_bytes(pd.util.hash_pandas_object(data))
+    elif isinstance(data, np.ndarray) and len(data) > 0 and isinstance(data[0], list):
+        # convert numpy array of lists into numpy array of the string representation of the lists
+        # because lists are not hashable
+        hashable = np.array(str(val) for val in data)
+        return _hash_ndarray_as_bytes(hashable)
+    elif isinstance(data, np.ndarray) and len(data) > 0 and isinstance(data[0], np.ndarray):
+        # convert numpy array of numpy arrays into 2d numpy arrays
+        # because numpy array of numpy arrays are not hashable
+        hashable = np.array(data.tolist())
+        return _hash_ndarray_as_bytes(hashable)
+    elif isinstance(data, np.ndarray):
+        return _hash_ndarray_as_bytes(data)
+    elif isinstance(data, list):
+        return _hash_ndarray_as_bytes(np.array(data))
+    else:
+        raise ValueError("Unsupported data type.")
+
+
+def _gen_md5_for_arraylike_obj(md5_gen, data):
+    """
+    Helper method to generate MD5 hash array-like object, the MD5 will calculate over:
+     - array length
+     - first NUM_SAMPLE_ROWS_FOR_HASH rows content
+     - last NUM_SAMPLE_ROWS_FOR_HASH rows content
+    """
+    len_bytes = _hash_uint64_ndarray_as_bytes(np.array([len(data)], dtype="uint64"))
+    md5_gen.update(len_bytes)
+    if len(data) < EvaluationDataset.NUM_SAMPLE_ROWS_FOR_HASH * 2:
+        md5_gen.update(_hash_array_like_obj_as_bytes(data))
+    else:
+        if isinstance(data, pd.DataFrame):
+            # Access rows of pandas Df with iloc
+            head_rows = data.iloc[: EvaluationDataset.NUM_SAMPLE_ROWS_FOR_HASH]
+            tail_rows = data.iloc[-EvaluationDataset.NUM_SAMPLE_ROWS_FOR_HASH :]
+        else:
+            head_rows = data[: EvaluationDataset.NUM_SAMPLE_ROWS_FOR_HASH]
+            tail_rows = data[-EvaluationDataset.NUM_SAMPLE_ROWS_FOR_HASH :]
+        md5_gen.update(_hash_array_like_obj_as_bytes(head_rows))
+        md5_gen.update(_hash_array_like_obj_as_bytes(tail_rows))
+
+
+def convert_data_to_mlflow_dataset(data, targets=None, predictions=None):
+    """Convert input data to mlflow dataset."""
+    supported_dataframe_types = [pd.DataFrame]
+    if "pyspark" in sys.modules:
+        from mlflow.utils.spark_utils import get_spark_dataframe_type
+
+        spark_df_type = get_spark_dataframe_type()
+        supported_dataframe_types.append(spark_df_type)
+
+    if predictions is not None:
+        _validate_dataset_type_supports_predictions(
+            data=data, supported_predictions_dataset_types=supported_dataframe_types
+        )
+
+    if isinstance(data, list):
+        # If the list is flat, we assume each element is an independent sample.
+        if not isinstance(data[0], (list, np.ndarray)):
+            data = [[elm] for elm in data]
+
+        return mlflow.data.from_numpy(
+            np.array(data), targets=np.array(targets) if targets else None
+        )
+    elif isinstance(data, np.ndarray):
+        return mlflow.data.from_numpy(data, targets=targets)
+    elif isinstance(data, pd.DataFrame):
+        return mlflow.data.from_pandas(df=data, targets=targets, predictions=predictions)
+    elif "pyspark" in sys.modules and isinstance(data, spark_df_type):
+        return mlflow.data.from_spark(df=data, targets=targets, predictions=predictions)
+    else:
+        # Cannot convert to mlflow dataset, return original data.
+        _logger.info(
+            "Cannot convert input data to `evaluate()` to an mlflow dataset, input must be a list, "
+            f"a numpy array, a panda Dataframe or a spark Dataframe, but received {type(data)}."
+        )
+        return data
+
+
+def _validate_dataset_type_supports_predictions(data, supported_predictions_dataset_types):
+    """
+    Validate that the dataset type supports a user-specified "predictions" column.
+    """
+    if not any(isinstance(data, sdt) for sdt in supported_predictions_dataset_types):
+        raise MlflowException(
+            message=(
+                "If predictions is specified, data must be one of the following types, or an"
+                " MLflow Dataset that represents one of the following types:"
+                f" {supported_predictions_dataset_types}."
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
+class EvaluationDataset:
+    """
+    An input dataset for model evaluation. This is intended for use with the
+    :py:func:`mlflow.models.evaluate()`
+    API.
+    """
+
+    NUM_SAMPLE_ROWS_FOR_HASH = 5
+    SPARK_DATAFRAME_LIMIT = 10000
+
+    def __init__(
+        self,
+        data,
+        *,
+        targets=None,
+        name=None,
+        path=None,
+        feature_names=None,
+        predictions=None,
+    ):
+        """
+        The values of the constructor arguments comes from the `evaluate` call.
+        """
+        if name is not None and '"' in name:
+            raise MlflowException(
+                message=f'Dataset name cannot include a double quote (") but got {name}',
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        if path is not None and '"' in path:
+            raise MlflowException(
+                message=f'Dataset path cannot include a double quote (") but got {path}',
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        self._user_specified_name = name
+        self._path = path
+        self._hash = None
+        self._supported_dataframe_types = (pd.DataFrame,)
+        self._spark_df_type = None
+        self._labels_data = None
+        self._targets_name = None
+        self._has_targets = False
+        self._predictions_data = None
+        self._predictions_name = None
+        self._has_predictions = predictions is not None
+
+        try:
+            # add checking `'pyspark' in sys.modules` to avoid importing pyspark when user
+            # run code not related to pyspark.
+            if "pyspark" in sys.modules:
+                from mlflow.utils.spark_utils import get_spark_dataframe_type
+
+                spark_df_type = get_spark_dataframe_type()
+                self._supported_dataframe_types = (pd.DataFrame, spark_df_type)
+                self._spark_df_type = spark_df_type
+        except ImportError:
+            pass
+
+        if feature_names is not None and len(set(feature_names)) < len(list(feature_names)):
+            raise MlflowException(
+                message="`feature_names` argument must be a list containing unique feature names.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        if self._has_predictions:
+            _validate_dataset_type_supports_predictions(
+                data=data,
+                supported_predictions_dataset_types=self._supported_dataframe_types,
+            )
+
+        has_targets = targets is not None
+        if has_targets:
+            self._has_targets = True
+        if isinstance(data, (np.ndarray, list)):
+            if has_targets and not isinstance(targets, (np.ndarray, list)):
+                raise MlflowException(
+                    message="If data is a numpy array or list of evaluation features, "
+                    "`targets` argument must be a numpy array or list of evaluation labels.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+            shape_message = (
+                "If the `data` argument is a numpy array, it must be a 2-dimensional "
+                "array, with the second dimension representing the number of features. If the "
+                "`data` argument is a list, each of its elements must be a feature array of "
+                "the numpy array or list, and all elements must have the same length."
+            )
+
+            if isinstance(data, list):
+                try:
+                    data = np.array(data)
+                except ValueError as e:
+                    raise MlflowException(
+                        message=shape_message, error_code=INVALID_PARAMETER_VALUE
+                    ) from e
+
+            if len(data.shape) != 2:
+                raise MlflowException(
+                    message=shape_message,
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+            self._features_data = data
+            if has_targets:
+                self._labels_data = (
+                    targets if isinstance(targets, np.ndarray) else np.array(targets)
+                )
+
+                if len(self._features_data) != len(self._labels_data):
+                    raise MlflowException(
+                        message="The input features example rows must be the same length "
+                        "with labels array.",
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+
+            num_features = data.shape[1]
+
+            if feature_names is not None:
+                feature_names = list(feature_names)
+                if num_features != len(feature_names):
+                    raise MlflowException(
+                        message="feature name list must be the same length with feature data.",
+                        error_code=INVALID_PARAMETER_VALUE,
+                    )
+                self._feature_names = feature_names
+            else:
+                self._feature_names = [
+                    f"feature_{str(i + 1).zfill(math.ceil(math.log10(num_features + 1)))}"
+                    for i in range(num_features)
+                ]
+        elif isinstance(data, self._supported_dataframe_types):
+            if has_targets and not isinstance(targets, str):
+                raise MlflowException(
+                    message="If data is a Pandas DataFrame or Spark DataFrame, `targets` argument "
+                    "must be the name of the column which contains evaluation labels in the `data` "
+                    "dataframe.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+            if self._spark_df_type and isinstance(data, self._spark_df_type):
+                if data.count() > EvaluationDataset.SPARK_DATAFRAME_LIMIT:
+                    _logger.warning(
+                        "Specified Spark DataFrame is too large for model evaluation. Only "
+                        f"the first {EvaluationDataset.SPARK_DATAFRAME_LIMIT} rows will be used. "
+                        "If you want evaluate on the whole spark dataframe, please manually call "
+                        "`spark_dataframe.toPandas()`."
+                    )
+                data = data.limit(EvaluationDataset.SPARK_DATAFRAME_LIMIT).toPandas()
+
+            if has_targets:
+                self._labels_data = data[targets].to_numpy()
+                self._targets_name = targets
+
+            if self._has_predictions:
+                self._predictions_data = data[predictions].to_numpy()
+                self._predictions_name = predictions
+
+            if feature_names is not None:
+                self._features_data = data[list(feature_names)]
+                self._feature_names = feature_names
+            else:
+                features_data = data
+
+                if has_targets:
+                    features_data = features_data.drop(targets, axis=1, inplace=False)
+
+                if self._has_predictions:
+                    features_data = features_data.drop(predictions, axis=1, inplace=False)
+
+                self._features_data = features_data
+                self._feature_names = [
+                    generate_feature_name_if_not_string(c) for c in self._features_data.columns
+                ]
+        else:
+            raise MlflowException(
+                message="The data argument must be a numpy array, a list or a Pandas DataFrame, or "
+                "spark DataFrame if pyspark package installed.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        # generate dataset hash
+        md5_gen = hashlib.md5(usedforsecurity=False)
+        _gen_md5_for_arraylike_obj(md5_gen, self._features_data)
+        if self._labels_data is not None:
+            _gen_md5_for_arraylike_obj(md5_gen, self._labels_data)
+        if self._predictions_data is not None:
+            _gen_md5_for_arraylike_obj(md5_gen, self._predictions_data)
+        md5_gen.update(",".join(list(map(str, self._feature_names))).encode("UTF-8"))
+
+        self._hash = md5_gen.hexdigest()
+
+    @property
+    def feature_names(self):
+        return self._feature_names
+
+    @property
+    def features_data(self):
+        """
+        return features data as a numpy array or a pandas DataFrame.
+        """
+        return self._features_data
+
+    @property
+    def labels_data(self):
+        """
+        return labels data as a numpy array
+        """
+        return self._labels_data
+
+    @property
+    def has_targets(self):
+        """
+        Returns True if the dataset has targets, False otherwise.
+        """
+        return self._has_targets
+
+    @property
+    def targets_name(self):
+        """
+        return targets name
+        """
+        return self._targets_name
+
+    @property
+    def predictions_data(self):
+        """
+        return labels data as a numpy array
+        """
+        return self._predictions_data
+
+    @property
+    def has_predictions(self):
+        """
+        Returns True if the dataset has targets, False otherwise.
+        """
+        return self._has_predictions
+
+    @property
+    def predictions_name(self):
+        """
+        return predictions name
+        """
+        return self._predictions_name
+
+    @property
+    def name(self):
+        """
+        Dataset name, which is specified dataset name or the dataset hash if user don't specify
+        name.
+        """
+        return self._user_specified_name if self._user_specified_name is not None else self.hash
+
+    @property
+    def path(self):
+        """
+        Dataset path
+        """
+        return self._path
+
+    @property
+    def hash(self):
+        """
+        Dataset hash, includes hash on first 20 rows and last 20 rows.
+        """
+        return self._hash
+
+    @property
+    def _metadata(self):
+        """
+        Return dataset metadata containing name, hash, and optional path.
+        """
+        metadata = {
+            "name": self.name,
+            "hash": self.hash,
+        }
+        if self.path is not None:
+            metadata["path"] = self.path
+        return metadata
+
+    def _log_dataset_tag(self, client, run_id, model_uuid):
+        """
+        Log dataset metadata as a tag "mlflow.datasets", if the tag already exists, it will
+        append current dataset metadata into existing tag content.
+        """
+        existing_dataset_metadata_str = client.get_run(run_id).data.tags.get(
+            "mlflow.datasets", "[]"
+        )
+        dataset_metadata_list = json.loads(existing_dataset_metadata_str)
+
+        for metadata in dataset_metadata_list:
+            if (
+                metadata["hash"] == self.hash
+                and metadata["name"] == self.name
+                and metadata["model"] == model_uuid
+            ):
+                break
+        else:
+            dataset_metadata_list.append({**self._metadata, "model": model_uuid})
+
+        dataset_metadata_str = json.dumps(dataset_metadata_list, separators=(",", ":"))
+        client.log_batch(
+            run_id,
+            tags=[RunTag("mlflow.datasets", dataset_metadata_str)],
+        )
+
+    def __hash__(self):
+        return hash(self.hash)
+
+    def __eq__(self, other):
+        if not isinstance(other, EvaluationDataset):
+            return False
+
+        if isinstance(self._features_data, np.ndarray):
+            is_features_data_equal = np.array_equal(self._features_data, other._features_data)
+        else:
+            is_features_data_equal = self._features_data.equals(other._features_data)
+
+        return (
+            is_features_data_equal
+            and np.array_equal(self._labels_data, other._labels_data)
+            and self.name == other.name
+            and self.path == other.path
+            and self._feature_names == other._feature_names
+        )

venv/lib/python3.9/site-packages/mlflow/data/filesystem_dataset_source.py

@@ -0,0 +1,81 @@
+from abc import abstractmethod
+from typing import Any
+
+from mlflow.data.dataset_source import DatasetSource
+
+
+class FileSystemDatasetSource(DatasetSource):
+    """
+    Represents the source of a dataset stored on a filesystem, e.g. a local UNIX filesystem,
+    blob storage services like S3, etc.
+    """
+
+    @property
+    @abstractmethod
+    def uri(self):
+        """The URI referring to the dataset source filesystem location.
+
+        Returns:
+            The URI referring to the dataset source filesystem location,
+            e.g "s3://mybucket/path/to/mydataset", "/tmp/path/to/my/dataset" etc.
+
+        """
+
+    @staticmethod
+    @abstractmethod
+    def _get_source_type() -> str:
+        """
+        Returns:
+            A string describing the filesystem containing the dataset, e.g. "local", "s3", ...
+        """
+
+    @abstractmethod
+    def load(self, dst_path=None) -> str:
+        """Downloads the dataset source to the local filesystem.
+
+        Args:
+            dst_path: Path of the local filesystem destination directory to which to download the
+                dataset source. If the directory does not exist, it is created. If
+                unspecified, the dataset source is downloaded to a new uniquely-named
+                directory on the local filesystem, unless the dataset source already
+                exists on the local filesystem, in which case its local path is returned
+                directly.
+
+        Returns:
+            The path to the downloaded dataset source on the local filesystem.
+
+        """
+
+    @staticmethod
+    @abstractmethod
+    def _can_resolve(raw_source: Any) -> bool:
+        """
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data".
+
+        Returns:
+            True if this DatasetSource can resolve the raw source, False otherwise.
+        """
+
+    @classmethod
+    @abstractmethod
+    def _resolve(cls, raw_source: Any) -> "FileSystemDatasetSource":
+        """
+        Args:
+            raw_source: The raw source, e.g. a string like "s3://mybucket/path/to/iris/data".
+        """
+
+    @abstractmethod
+    def to_dict(self) -> dict[Any, Any]:
+        """
+        Returns:
+            A JSON-compatible dictionary representation of the FileSystemDatasetSource.
+        """
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "FileSystemDatasetSource":
+        """
+        Args:
+            source_dict: A dictionary representation of the FileSystemDatasetSource.
+        """

venv/lib/python3.9/site-packages/mlflow/data/http_dataset_source.py

@@ -0,0 +1,145 @@
+import os
+import re
+from typing import Any
+from urllib.parse import urlparse
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.file_utils import create_tmp_dir
+from mlflow.utils.rest_utils import augmented_raise_for_status, cloud_storage_http_request
+
+
+def _is_path(filename: str) -> bool:
+    """
+    Return True if `filename` is a path, False otherwise. For example,
+    "foo/bar" is a path, but "bar" is not.
+    """
+    return os.path.basename(filename) != filename
+
+
+class HTTPDatasetSource(DatasetSource):
+    """
+    Represents the source of a dataset stored at a web location and referred to
+    by an HTTP or HTTPS URL.
+    """
+
+    def __init__(self, url):
+        self._url = url
+
+    @property
+    def url(self):
+        """The HTTP/S URL referring to the dataset source location.
+
+        Returns:
+            The HTTP/S URL referring to the dataset source location.
+
+        """
+        return self._url
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "http"
+
+    def _extract_filename(self, response) -> str:
+        """
+        Extracts a filename from the Content-Disposition header or the URL's path.
+        """
+        if content_disposition := response.headers.get("Content-Disposition"):
+            for match in re.finditer(r"filename=(.+)", content_disposition):
+                filename = match[1].strip("'\"")
+                if _is_path(filename):
+                    raise MlflowException.invalid_parameter_value(
+                        f"Invalid filename in Content-Disposition header: {filename}. "
+                        "It must be a file name, not a path."
+                    )
+                return filename
+
+        # Extract basename from URL if no valid filename in Content-Disposition
+        return os.path.basename(urlparse(self.url).path)
+
+    def load(self, dst_path=None) -> str:
+        """Downloads the dataset source to the local filesystem.
+
+        Args:
+            dst_path: Path of the local filesystem destination directory to which to download the
+                dataset source. If the directory does not exist, it is created. If
+                unspecified, the dataset source is downloaded to a new uniquely-named
+                directory on the local filesystem.
+
+        Returns:
+            The path to the downloaded dataset source on the local filesystem.
+
+        """
+        resp = cloud_storage_http_request(
+            method="GET",
+            url=self.url,
+            stream=True,
+        )
+        augmented_raise_for_status(resp)
+
+        basename = self._extract_filename(resp)
+
+        if not basename:
+            basename = "dataset_source"
+
+        if dst_path is None:
+            dst_path = create_tmp_dir()
+
+        dst_path = os.path.join(dst_path, basename)
+        with open(dst_path, "wb") as f:
+            chunk_size = 1024 * 1024  # 1 MB
+            for chunk in resp.iter_content(chunk_size=chunk_size):
+                f.write(chunk)
+
+        return dst_path
+
+    @staticmethod
+    def _can_resolve(raw_source: Any) -> bool:
+        """
+        Args:
+            raw_source: The raw source, e.g. a string like "http://mysite/mydata.tar.gz".
+
+        Returns:
+            True if this DatasetSource can resolve the raw source, False otherwise.
+        """
+        if not isinstance(raw_source, str):
+            return False
+
+        try:
+            parsed_source = urlparse(str(raw_source))
+            return parsed_source.scheme in ["http", "https"]
+        except Exception:
+            return False
+
+    @classmethod
+    def _resolve(cls, raw_source: Any) -> "HTTPDatasetSource":
+        """
+        Args:
+            raw_source: The raw source, e.g. a string like "http://mysite/mydata.tar.gz".
+        """
+        return HTTPDatasetSource(raw_source)
+
+    def to_dict(self) -> dict[Any, Any]:
+        """
+        Returns:
+            A JSON-compatible dictionary representation of the HTTPDatasetSource.
+        """
+        return {
+            "url": self.url,
+        }
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "HTTPDatasetSource":
+        """
+        Args:
+            source_dict: A dictionary representation of the HTTPDatasetSource.
+        """
+        url = source_dict.get("url")
+        if url is None:
+            raise MlflowException(
+                'Failed to parse HTTPDatasetSource. Missing expected key: "url"',
+                INVALID_PARAMETER_VALUE,
+            )
+
+        return cls(url=url)

venv/lib/python3.9/site-packages/mlflow/data/huggingface_dataset.py

@@ -0,0 +1,245 @@
+import json
+import logging
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence, Union
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.digest_utils import compute_pandas_digest
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.huggingface_dataset_source import HuggingFaceDatasetSource
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+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
+
+_logger = logging.getLogger(__name__)
+
+_MAX_ROWS_FOR_DIGEST_COMPUTATION_AND_SCHEMA_INFERENCE = 10000
+
+if TYPE_CHECKING:
+    import datasets
+
+
+class HuggingFaceDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a HuggingFace dataset for use with MLflow Tracking.
+    """
+
+    def __init__(  # noqa: D417
+        self,
+        ds: "datasets.Dataset",
+        source: HuggingFaceDatasetSource,
+        targets: Optional[str] = None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+    ):
+        """
+        Args:
+            ds: A Hugging Face dataset. Must be an instance of `datasets.Dataset`.
+                Other types, such as :py:class:`datasets.DatasetDict`, are not supported.
+            source: The source of the Hugging Face dataset.
+            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.
+        """
+        if targets is not None and targets not in ds.column_names:
+            raise MlflowException(
+                f"The specified Hugging Face dataset does not contain the specified targets column"
+                f" '{targets}'.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._ds = ds
+        self._targets = targets
+        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.
+        """
+        df = next(
+            self._ds.to_pandas(
+                batch_size=_MAX_ROWS_FOR_DIGEST_COMPUTATION_AND_SCHEMA_INFERENCE, batched=True
+            )
+        )
+        return compute_pandas_digest(df)
+
+    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 ds(self) -> "datasets.Dataset":
+        """The Hugging Face ``datasets.Dataset`` instance.
+
+        Returns:
+            The Hugging Face ``datasets.Dataset`` instance.
+
+        """
+        return self._ds
+
+    @property
+    def targets(self) -> Optional[str]:
+        """
+        The name of the Hugging Face dataset column containing targets (labels) for supervised
+        learning.
+
+        Returns:
+            The string name of the Hugging Face dataset column containing targets.
+        """
+        return self._targets
+
+    @property
+    def source(self) -> HuggingFaceDatasetSource:
+        """Hugging Face dataset source information.
+
+        Returns:
+            A :py:class:`mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource`
+        """
+        return self._source
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        Summary statistics for the Hugging Face dataset, including the number of rows,
+        size, and size in bytes.
+        """
+        return {
+            "num_rows": self._ds.num_rows,
+            "dataset_size": self._ds.dataset_size,
+            "size_in_bytes": self._ds.size_in_bytes,
+        }
+
+    @cached_property
+    def schema(self) -> Optional[Schema]:
+        """
+        The MLflow ColSpec schema of the Hugging Face dataset.
+        """
+        try:
+            df = next(
+                self._ds.to_pandas(
+                    batch_size=_MAX_ROWS_FOR_DIGEST_COMPUTATION_AND_SCHEMA_INFERENCE, batched=True
+                )
+            )
+            return _infer_schema(df)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for Hugging Face dataset. Exception: %s", e)
+            return None
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        df = self._ds.to_pandas()
+        if self._targets is not None:
+            if self._targets not in df.columns:
+                raise MlflowException(
+                    f"Failed to convert Hugging Face dataset to pyfunc inputs and outputs because"
+                    f" the pandas representation of the Hugging Face 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._ds.to_pandas(),
+            targets=self._targets,
+            path=path,
+            feature_names=feature_names,
+        )
+
+
+def from_huggingface(
+    ds,
+    path: Optional[str] = None,
+    targets: Optional[str] = None,
+    data_dir: Optional[str] = None,
+    data_files: Optional[Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]] = None,
+    revision=None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+    trust_remote_code: Optional[bool] = None,
+) -> HuggingFaceDataset:
+    """
+    Create a `mlflow.data.huggingface_dataset.HuggingFaceDataset` from a Hugging Face dataset.
+
+    Args:
+        ds:
+            A Hugging Face dataset. Must be an instance of `datasets.Dataset`. Other types, such as
+            `datasets.DatasetDict`, are not supported.
+        path: The path of the Hugging Face dataset used to construct the source. This is the same
+            argument as `path` in `datasets.load_dataset()` function. To be able to reload the
+            dataset via MLflow, `path` must match the path of the dataset on the hub, e.g.,
+            "databricks/databricks-dolly-15k". If no path is specified, a `CodeDatasetSource` is,
+            used which will source information from the run context.
+        targets: The name of the Hugging Face `dataset.Dataset` column containing targets (labels)
+            for supervised learning.
+        data_dir: The `data_dir` of the Hugging Face dataset configuration. This is used by the
+            `datasets.load_dataset()` function to reload the dataset upon request via
+            :py:func:`HuggingFaceDataset.source.load()
+            <mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource.load>`.
+        data_files: Paths to source data file(s) for the Hugging Face dataset configuration.
+            This is used by the `datasets.load_dataset()` function to reload the
+            dataset upon request via :py:func:`HuggingFaceDataset.source.load()
+            <mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource.load>`.
+        revision: Version of the dataset script to load. This is used by the
+            `datasets.load_dataset()` function to reload the dataset upon request via
+            :py:func:`HuggingFaceDataset.source.load()
+            <mlflow.data.huggingface_dataset_source.HuggingFaceDatasetSource.load>`.
+        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.
+        trust_remote_code: Whether to trust remote code from the dataset repo.
+    """
+    import datasets
+
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.tracking.context import registry
+
+    if not isinstance(ds, datasets.Dataset):
+        raise MlflowException(
+            f"The specified Hugging Face dataset must be an instance of `datasets.Dataset`."
+            f" Instead, found an instance of: {type(ds)}",
+            INVALID_PARAMETER_VALUE,
+        )
+
+    # Set the source to a `HuggingFaceDatasetSource` if a path is specified, otherwise set it to a
+    # `CodeDatasetSource`.
+    if path is not None:
+        source = HuggingFaceDatasetSource(
+            path=path,
+            config_name=ds.config_name,
+            data_dir=data_dir,
+            data_files=data_files,
+            split=ds.split,
+            revision=revision,
+            trust_remote_code=trust_remote_code,
+        )
+    else:
+        context_tags = registry.resolve_tags()
+        source = CodeDatasetSource(tags=context_tags)
+    return HuggingFaceDataset(ds=ds, targets=targets, source=source, name=name, digest=digest)

venv/lib/python3.9/site-packages/mlflow/data/huggingface_dataset_source.py

@@ -0,0 +1,118 @@
+from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence, Union
+
+from mlflow.data.dataset_source import DatasetSource
+
+if TYPE_CHECKING:
+    import datasets
+
+
+class HuggingFaceDatasetSource(DatasetSource):
+    """Represents the source of a Hugging Face dataset used in MLflow Tracking."""
+
+    def __init__(
+        self,
+        path: str,
+        config_name: Optional[str] = None,
+        data_dir: Optional[str] = None,
+        data_files: Optional[
+            Union[str, Sequence[str], Mapping[str, Union[str, Sequence[str]]]]
+        ] = None,
+        split: Optional[Union[str, "datasets.Split"]] = None,
+        revision: Optional[Union[str, "datasets.Version"]] = None,
+        trust_remote_code: Optional[bool] = None,
+    ):
+        """Create a `HuggingFaceDatasetSource` instance.
+
+        Arguments in `__init__` match arguments of the same name in
+        `datasets.load_dataset() <https://huggingface.co/docs/datasets/v2.14.5/en/package_reference/loading_methods#datasets.load_dataset>`_.
+        The only exception is `config_name` matches `name` in `datasets.load_dataset()`, because
+        we need to differentiate from `mlflow.data.Dataset` `name` attribute.
+
+        Args:
+            path: The path of the Hugging Face dataset, if it is a dataset from HuggingFace hub,
+                `path` must match the hub path, e.g., "databricks/databricks-dolly-15k".
+            config_name: The name of of the Hugging Face dataset configuration.
+            data_dir: The `data_dir` of the Hugging Face dataset configuration.
+            data_files: Paths to source data file(s) for the Hugging Face dataset configuration.
+            split: Which split of the data to load.
+            revision: Version of the dataset script to load.
+            trust_remote_code: Whether to trust remote code from the dataset repo.
+        """
+        self.path = path
+        self.config_name = config_name
+        self.data_dir = data_dir
+        self.data_files = data_files
+        self.split = split
+        self.revision = revision
+        self.trust_remote_code = trust_remote_code
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "hugging_face"
+
+    def load(self, **kwargs):
+        """Load the Hugging Face dataset based on `HuggingFaceDatasetSource`.
+
+        Args:
+            kwargs: Additional keyword arguments used for loading the dataset with the Hugging Face
+                `datasets.load_dataset()` method.
+
+        Returns:
+            An instance of `datasets.Dataset`.
+        """
+        import datasets
+        from packaging.version import Version
+
+        load_kwargs = {
+            "path": self.path,
+            "name": self.config_name,
+            "data_dir": self.data_dir,
+            "data_files": self.data_files,
+            "split": self.split,
+            "revision": self.revision,
+        }
+
+        # this argument only exists in >= 2.16.0
+        if Version(datasets.__version__) >= Version("2.16.0"):
+            load_kwargs["trust_remote_code"] = self.trust_remote_code
+
+        intersecting_keys = set(load_kwargs.keys()) & set(kwargs.keys())
+        if intersecting_keys:
+            raise KeyError(
+                f"Found duplicated arguments in `HuggingFaceDatasetSource` and "
+                f"`kwargs`: {intersecting_keys}. Please remove them from `kwargs`."
+            )
+            load_kwargs.update(kwargs)
+        return datasets.load_dataset(**load_kwargs)
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        # NB: Initially, we expect that Hugging Face dataset sources will only be used with
+        # Hugging Face datasets constructed by from_huggingface_dataset, which can create
+        # an instance of HuggingFaceDatasetSource directly without the need for resolution
+        return False
+
+    @classmethod
+    def _resolve(cls, raw_source: str) -> "HuggingFaceDatasetSource":
+        raise NotImplementedError
+
+    def to_dict(self) -> dict[Any, Any]:
+        return {
+            "path": self.path,
+            "config_name": self.config_name,
+            "data_dir": self.data_dir,
+            "data_files": self.data_files,
+            "split": str(self.split),
+            "revision": self.revision,
+        }
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "HuggingFaceDatasetSource":
+        return cls(
+            path=source_dict.get("path"),
+            config_name=source_dict.get("config_name"),
+            data_dir=source_dict.get("data_dir"),
+            data_files=source_dict.get("data_files"),
+            split=source_dict.get("split"),
+            revision=source_dict.get("revision"),
+        )

venv/lib/python3.9/site-packages/mlflow/data/meta_dataset.py

@@ -0,0 +1,106 @@
+import hashlib
+import json
+from typing import Any, Optional
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.types import Schema
+from mlflow.utils.annotations import experimental
+
+
+@experimental
+class MetaDataset(Dataset):
+    """Dataset that only contains metadata.
+
+    This class is used to represent a dataset that only contains metadata, which is useful when
+    users only want to log metadata to MLflow without logging the actual data. For example, users
+    build a custom dataset from a text file publicly hosted in the Internet, and they want to log
+    the text file's URL to MLflow for future tracking instead of the dataset itself.
+
+    Args:
+        source: dataset source of type `DatasetSource`, indicates where the data is from.
+        name: name of the dataset. If not specified, a name is automatically generated.
+        digest: digest (hash, fingerprint) of the dataset. If not specified, a digest is
+            automatically computed.
+        schame: schema of the dataset.
+
+    .. code-block:: python
+        :caption: Create a MetaDataset
+
+        import mlflow
+
+        mlflow.set_experiment("/test-mlflow-meta-dataset")
+
+        source = mlflow.data.http_dataset_source.HTTPDatasetSource(
+            url="https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz"
+        )
+        ds = mlflow.data.meta_dataset.MetaDataset(source)
+
+        with mlflow.start_run() as run:
+            mlflow.log_input(ds)
+
+    .. code-block:: python
+        :caption: Create a MetaDataset with schema
+
+        import mlflow
+
+        mlflow.set_experiment("/test-mlflow-meta-dataset")
+
+        source = mlflow.data.http_dataset_source.HTTPDatasetSource(
+            url="https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz"
+        )
+        schema = Schema(
+            [
+                ColSpec(type=mlflow.types.DataType.string, name="text"),
+                ColSpec(type=mlflow.types.DataType.integer, name="label"),
+            ]
+        )
+        ds = mlflow.data.meta_dataset.MetaDataset(source, schema=schema)
+
+        with mlflow.start_run() as run:
+            mlflow.log_input(ds)
+    """
+
+    def __init__(
+        self,
+        source: DatasetSource,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+        schema: Optional[Schema] = None,
+    ):
+        # Set `self._schema` before calling the superclass constructor because
+        # `self._compute_digest` depends on `self._schema`.
+        self._schema = schema
+        super().__init__(source=source, name=name, digest=digest)
+
+    def _compute_digest(self) -> str:
+        """Computes a digest for the dataset.
+
+        The digest computation of `MetaDataset` is based on the dataset's name, source, source type,
+        and schema instead of the actual data. Basically we compute the sha256 hash of the config
+        dict.
+        """
+        config = {
+            "name": self.name,
+            "source": self.source.to_json(),
+            "source_type": self.source._get_source_type(),
+            "schema": self.schema.to_dict() if self.schema else "",
+        }
+        return hashlib.sha256(json.dumps(config).encode("utf-8")).hexdigest()[:8]
+
+    @property
+    def schema(self) -> Optional[Any]:
+        """Returns the schema of the dataset."""
+        return self._schema
+
+    def to_dict(self) -> dict[str, str]:
+        """Create config dictionary for the MetaDataset.
+
+        Returns a string dictionary containing the following fields: name, digest, source, source
+        type, schema, and profile.
+        """
+        config = super().to_dict()
+        if self.schema:
+            schema = json.dumps({"mlflow_colspec": self.schema.to_dict()}) if self.schema else None
+            config["schema"] = schema
+        return config

venv/lib/python3.9/site-packages/mlflow/data/numpy_dataset.py

@@ -0,0 +1,221 @@
+import json
+import logging
+from functools import cached_property
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.digest_utils import compute_numpy_digest
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.data.schema import TensorDatasetSchema
+from mlflow.types.utils import _infer_schema
+
+_logger = logging.getLogger(__name__)
+
+
+class NumpyDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a NumPy dataset for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        features: Union[np.ndarray, dict[str, np.ndarray]],
+        source: DatasetSource,
+        targets: Union[np.ndarray, dict[str, np.ndarray]] = None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+    ):
+        """
+        Args:
+            features: A numpy array or dictionary of numpy arrays containing dataset features.
+            source: The source of the numpy dataset.
+            targets: A numpy array or dictionary of numpy arrays containing dataset targets.
+                Optional.
+            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.
+        """
+        self._features = features
+        self._targets = targets
+        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.
+        """
+        return compute_numpy_digest(self._features, self._targets)
+
+    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(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 source(self) -> DatasetSource:
+        """
+        The source of the dataset.
+        """
+        return self._source
+
+    @property
+    def features(self) -> Union[np.ndarray, dict[str, np.ndarray]]:
+        """
+        The features of the dataset.
+        """
+        return self._features
+
+    @property
+    def targets(self) -> Optional[Union[np.ndarray, dict[str, np.ndarray]]]:
+        """
+        The targets of the dataset. May be ``None`` if no targets are available.
+        """
+        return self._targets
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be ``None`` if a profile cannot be computed.
+        """
+
+        def get_profile_attribute(numpy_data, attr_name):
+            if isinstance(numpy_data, dict):
+                return {key: getattr(array, attr_name) for key, array in numpy_data.items()}
+            else:
+                return getattr(numpy_data, attr_name)
+
+        profile = {
+            "features_shape": get_profile_attribute(self._features, "shape"),
+            "features_size": get_profile_attribute(self._features, "size"),
+            "features_nbytes": get_profile_attribute(self._features, "nbytes"),
+        }
+        if self._targets is not None:
+            profile.update(
+                {
+                    "targets_shape": get_profile_attribute(self._targets, "shape"),
+                    "targets_size": get_profile_attribute(self._targets, "size"),
+                    "targets_nbytes": get_profile_attribute(self._targets, "nbytes"),
+                }
+            )
+
+        return profile
+
+    @cached_property
+    def schema(self) -> Optional[TensorDatasetSchema]:
+        """
+        MLflow TensorSpec schema representing the dataset features and targets (optional).
+        """
+        try:
+            features_schema = _infer_schema(self._features)
+            targets_schema = None
+            if self._targets is not None:
+                targets_schema = _infer_schema(self._targets)
+            return TensorDatasetSchema(features=features_schema, targets=targets_schema)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for NumPy dataset. Exception: %s", e)
+            return None
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        """
+        return PyFuncInputsOutputs(self._features, self._targets)
+
+    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.sklearn.evaluate().
+        """
+        return EvaluationDataset(
+            data=self._features,
+            targets=self._targets,
+            path=path,
+            feature_names=feature_names,
+        )
+
+
+def from_numpy(
+    features: Union[np.ndarray, dict[str, np.ndarray]],
+    source: Union[str, DatasetSource] = None,
+    targets: Union[np.ndarray, dict[str, np.ndarray]] = None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+) -> NumpyDataset:
+    """
+    Constructs a :py:class:`NumpyDataset <mlflow.data.numpy_dataset.NumpyDataset>` object from
+    NumPy features, optional targets, and source. If the source is path like, then this will
+    construct a DatasetSource object from the source path. Otherwise, the source is assumed to
+    be a DatasetSource object.
+
+    Args:
+        features: NumPy features, represented as an np.ndarray or dictionary of named np.ndarrays.
+        source: The source from which the numpy data was derived, e.g. a filesystem path, an S3 URI,
+            an HTTPS URL, a delta table name with version, or spark table etc. ``source`` may be
+            specified as a URI, a path-like string, or an instance of
+            :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`. If unspecified,
+            the source is assumed to be the code location (e.g. notebook cell, script, etc.) where
+            :py:func:`from_numpy <mlflow.data.from_numpy>` is being called.
+        targets: Optional NumPy targets, represented as an np.ndarray or dictionary of named
+            np.ndarrays.
+        name: The name of the dataset. If unspecified, a name is generated.
+        digest: The dataset digest (hash). If unspecified, a digest is computed automatically.
+
+    .. code-block:: python
+        :test:
+        :caption: Basic Example
+
+        import mlflow
+        import numpy as np
+
+        x = np.random.uniform(size=[2, 5, 4])
+        y = np.random.randint(2, size=[2])
+        dataset = mlflow.data.from_numpy(x, targets=y)
+
+    .. code-block:: python
+        :test:
+        :caption: Dict Example
+
+        import mlflow
+        import numpy as np
+
+        x = {
+            "feature_1": np.random.uniform(size=[2, 5, 4]),
+            "feature_2": np.random.uniform(size=[2, 5, 4]),
+        }
+        y = np.random.randint(2, size=[2])
+        dataset = mlflow.data.from_numpy(x, targets=y)
+    """
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        if isinstance(source, DatasetSource):
+            resolved_source = source
+        else:
+            resolved_source = resolve_dataset_source(
+                source,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        resolved_source = CodeDatasetSource(tags=context_tags)
+    return NumpyDataset(
+        features=features, source=resolved_source, targets=targets, name=name, digest=digest
+    )

venv/lib/python3.9/site-packages/mlflow/data/pandas_dataset.py

@@ -0,0 +1,229 @@
+import json
+import logging
+from functools import cached_property
+from typing import Any, Optional, Union
+
+import pandas as pd
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.digest_utils import compute_pandas_digest
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.types import Schema
+from mlflow.types.utils import _infer_schema
+
+_logger = logging.getLogger(__name__)
+
+
+class PandasDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a Pandas DataFrame for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        source: DatasetSource,
+        targets: Optional[str] = None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+        predictions: Optional[str] = None,
+    ):
+        """
+        Args:
+            df: A pandas DataFrame.
+            source: The source of the pandas DataFrame.
+            targets: The name of the target column. Optional.
+            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``).
+        """
+        if targets is not None and targets not in df.columns:
+            raise MlflowException(
+                f"The specified pandas DataFrame 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 pandas DataFrame 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.
+        """
+        return compute_pandas_digest(self._df)
+
+    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) -> pd.DataFrame:
+        """
+        The underlying pandas DataFrame.
+        """
+        return self._df
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        The source of the dataset.
+        """
+        return self._source
+
+    @property
+    def targets(self) -> Optional[str]:
+        """
+        The name of the target column. May be ``None`` if no target column is available.
+        """
+        return self._targets
+
+    @property
+    def predictions(self) -> Optional[str]:
+        """
+        The name of the predictions column. May be ``None`` if no predictions column is available.
+        """
+        return self._predictions
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be ``None`` if a profile cannot be computed.
+        """
+        return {
+            "num_rows": len(self._df),
+            "num_elements": int(self._df.size),
+        }
+
+    @cached_property
+    def schema(self) -> Optional[Schema]:
+        """
+        An instance of :py:class:`mlflow.types.Schema` representing the tabular dataset. May be
+        ``None`` if the schema cannot be inferred from the dataset.
+        """
+        try:
+            return _infer_schema(self._df)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for Pandas dataset. Exception: %s", e)
+            return None
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        """
+        if self._targets:
+            inputs = self._df.drop(columns=[self._targets])
+            outputs = self._df[self._targets]
+            return PyFuncInputsOutputs(inputs, outputs)
+        else:
+            return PyFuncInputsOutputs(self._df)
+
+    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,
+            targets=self._targets,
+            path=path,
+            feature_names=feature_names,
+            predictions=self._predictions,
+        )
+
+
+def from_pandas(
+    df: pd.DataFrame,
+    source: Union[str, DatasetSource] = None,
+    targets: Optional[str] = None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+    predictions: Optional[str] = None,
+) -> PandasDataset:
+    """
+    Constructs a :py:class:`PandasDataset <mlflow.data.pandas_dataset.PandasDataset>` instance from
+    a Pandas DataFrame, optional targets, optional predictions, and source.
+
+    Args:
+        df: A Pandas DataFrame.
+        source: The source from which the DataFrame was derived, e.g. a filesystem
+            path, an S3 URI, an HTTPS URL, a delta table name with version, or
+            spark table etc. ``source`` may be specified as a URI, a path-like string,
+            or an instance of
+            :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
+            If unspecified, the source is assumed to be the code location
+            (e.g. notebook cell, script, etc.) where
+            :py:func:`from_pandas <mlflow.data.from_pandas>` is being called.
+        targets: An optional target column name for supervised training. This column
+            must be present in the dataframe (``df``).
+        name: The name of the dataset. If unspecified, a name is generated.
+        digest: The dataset digest (hash). If unspecified, a digest is computed
+            automatically.
+        predictions: An optional predictions column name for model evaluation. This column
+            must be present in the dataframe (``df``).
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow
+        import pandas as pd
+
+        x = pd.DataFrame(
+            [["tom", 10, 1, 1], ["nick", 15, 0, 1], ["july", 14, 1, 1]],
+            columns=["Name", "Age", "Label", "ModelOutput"],
+        )
+        dataset = mlflow.data.from_pandas(x, targets="Label", predictions="ModelOutput")
+    """
+
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        if isinstance(source, DatasetSource):
+            resolved_source = source
+        else:
+            resolved_source = resolve_dataset_source(
+                source,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        resolved_source = CodeDatasetSource(tags=context_tags)
+    return PandasDataset(
+        df=df,
+        source=resolved_source,
+        targets=targets,
+        name=name,
+        digest=digest,
+        predictions=predictions,
+    )

venv/lib/python3.9/site-packages/mlflow/data/pyfunc_dataset_mixin.py

@@ -0,0 +1,29 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.models.utils import PyFuncInput, PyFuncOutput
+
+
+@dataclass
+class PyFuncInputsOutputs:
+    inputs: list[PyFuncInput]
+    outputs: Optional[list[PyFuncOutput]] = None
+
+
+class PyFuncConvertibleDatasetMixin:
+    @abstractmethod
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        May not be implemented by all datasets.
+        """
+
+    @abstractmethod
+    def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
+        """
+        Converts the dataset to an EvaluationDataset for model evaluation.
+        May not be implemented by all datasets.
+        """

venv/lib/python3.9/site-packages/mlflow/data/schema.py

@@ -0,0 +1,76 @@
+from typing import Any
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.types.schema import Schema
+
+
+class TensorDatasetSchema:
+    """
+    Represents the schema of a dataset with tensor features and targets.
+    """
+
+    def __init__(self, features: Schema, targets: Schema = None):
+        if not isinstance(features, Schema):
+            raise MlflowException(
+                f"features must be mlflow.types.Schema, got '{type(features)}'",
+                INVALID_PARAMETER_VALUE,
+            )
+        if targets is not None and not isinstance(targets, Schema):
+            raise MlflowException(
+                f"targets must be either None or mlflow.types.Schema, got '{type(features)}'",
+                INVALID_PARAMETER_VALUE,
+            )
+        self.features = features
+        self.targets = targets
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize into a 'jsonable' dictionary.
+
+        Returns:
+            dictionary representation of the schema's features and targets (if defined).
+
+        """
+
+        return {
+            "mlflow_tensorspec": {
+                "features": self.features.to_json(),
+                "targets": self.targets.to_json() if self.targets is not None else None,
+            },
+        }
+
+    @classmethod
+    def from_dict(cls, schema_dict: dict[str, Any]):
+        """Deserialize from dictionary representation.
+
+        Args:
+            schema_dict: Dictionary representation of model signature. Expected dictionary format:
+                `{'features': <json string>, 'targets': <json string>" }`
+
+        Returns:
+            TensorDatasetSchema populated with the data from the dictionary.
+
+        """
+        if "mlflow_tensorspec" not in schema_dict:
+            raise MlflowException(
+                "TensorDatasetSchema dictionary is missing expected key 'mlflow_tensorspec'",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        schema_dict = schema_dict["mlflow_tensorspec"]
+        features = Schema.from_json(schema_dict["features"])
+        if "targets" in schema_dict and schema_dict["targets"] is not None:
+            targets = Schema.from_json(schema_dict["targets"])
+            return cls(features, targets)
+        else:
+            return cls(features)
+
+    def __eq__(self, other) -> bool:
+        return (
+            isinstance(other, TensorDatasetSchema)
+            and self.features == other.features
+            and self.targets == other.targets
+        )
+
+    def __repr__(self) -> str:
+        return f"features:\n  {self.features!r}\ntargets:\n  {self.targets!r}\n"

venv/lib/python3.9/site-packages/mlflow/data/sources.py

@@ -0,0 +1 @@
+__all__ = []

venv/lib/python3.9/site-packages/mlflow/data/spark_dataset.py

@@ -0,0 +1,404 @@
+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,
+        )
+
+
+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,
+    )

venv/lib/python3.9/site-packages/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"),
+        )

venv/lib/python3.9/site-packages/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)

venv/lib/python3.9/site-packages/mlflow/data/tensorflow_dataset.py

@@ -0,0 +1,348 @@
+import json
+import logging
+from functools import cached_property
+from typing import Any, Optional, Union
+
+import numpy as np
+
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.data.digest_utils import (
+    MAX_ROWS,
+    compute_numpy_digest,
+    get_normalized_md5_digest,
+)
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin, PyFuncInputsOutputs
+from mlflow.data.schema import TensorDatasetSchema
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, INVALID_PARAMETER_VALUE
+from mlflow.types.schema import Schema
+from mlflow.types.utils import _infer_schema
+
+_logger = logging.getLogger(__name__)
+
+
+class TensorFlowDataset(Dataset, PyFuncConvertibleDatasetMixin):
+    """
+    Represents a TensorFlow dataset for use with MLflow Tracking.
+    """
+
+    def __init__(
+        self,
+        features,
+        source: DatasetSource,
+        targets=None,
+        name: Optional[str] = None,
+        digest: Optional[str] = None,
+    ):
+        """
+        Args:
+            features: A TensorFlow dataset or tensor of features.
+            source: The source of the TensorFlow dataset.
+            targets: A TensorFlow dataset or tensor of targets. Optional.
+            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.
+        """
+        import tensorflow as tf
+
+        if not isinstance(features, tf.data.Dataset) and not tf.is_tensor(features):
+            raise MlflowException(
+                f"'features' must be an instance of tf.data.Dataset or a TensorFlow Tensor."
+                f" Found: {type(features)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        if tf.is_tensor(features) and targets is not None and not tf.is_tensor(targets):
+            raise MlflowException(
+                f"If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
+                f" Tensor. Found: {type(targets)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        if (
+            isinstance(features, tf.data.Dataset)
+            and targets is not None
+            and not isinstance(targets, tf.data.Dataset)
+        ):
+            raise MlflowException(
+                "If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
+                f" instance of tf.data.Dataset. Found: {type(targets)}.",
+                INVALID_PARAMETER_VALUE,
+            )
+
+        self._features = features
+        self._targets = targets
+        super().__init__(source=source, name=name, digest=digest)
+
+    def _compute_tensorflow_dataset_digest(  # noqa: D417
+        self,
+        dataset,
+        targets=None,
+    ) -> str:
+        """Computes a digest for the given Tensorflow dataset.
+
+        Args:
+            dataset: A Tensorflow dataset.
+
+        Returns:
+            A string digest.
+        """
+        import pandas as pd
+        import tensorflow as tf
+
+        hashable_elements = []
+
+        def hash_tf_dataset_iterator_element(element):
+            if element is None:
+                return
+            flat_element = tf.nest.flatten(element)
+            flattened_array = np.concatenate([x.flatten() for x in flat_element])
+            trimmed_array = flattened_array[0:MAX_ROWS]
+            try:
+                hashable_elements.append(pd.util.hash_array(trimmed_array))
+            except TypeError:
+                hashable_elements.append(np.int64(trimmed_array.size))
+
+        for element in dataset.as_numpy_iterator():
+            hash_tf_dataset_iterator_element(element)
+        if targets is not None:
+            for element in targets.as_numpy_iterator():
+                hash_tf_dataset_iterator_element(element)
+
+        return get_normalized_md5_digest(hashable_elements)
+
+    def _compute_tensor_digest(
+        self,
+        tensor_data,
+        tensor_targets,
+    ) -> str:
+        """Computes a digest for the given Tensorflow tensor.
+
+        Args:
+            tensor_data: A Tensorflow tensor, representing the features.
+            tensor_targets: A Tensorflow tensor, representing the targets. Optional.
+
+        Returns:
+            A string digest.
+        """
+        if tensor_targets is None:
+            return compute_numpy_digest(tensor_data.numpy())
+        else:
+            return compute_numpy_digest(tensor_data.numpy(), tensor_targets.numpy())
+
+    def _compute_digest(self) -> str:
+        """
+        Computes a digest for the dataset. Called if the user doesn't supply
+        a digest when constructing the dataset.
+        """
+        import tensorflow as tf
+
+        if isinstance(self._features, tf.data.Dataset):
+            return self._compute_tensorflow_dataset_digest(self._features, self._targets)
+        return self._compute_tensor_digest(self._features, self._targets)
+
+    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(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 data(self):
+        """
+        The underlying TensorFlow data.
+        """
+        return self._features
+
+    @property
+    def source(self) -> DatasetSource:
+        """
+        The source of the dataset.
+        """
+        return self._source
+
+    @property
+    def targets(self):
+        """
+        The targets of the dataset.
+        """
+        return self._targets
+
+    @property
+    def profile(self) -> Optional[Any]:
+        """
+        A profile of the dataset. May be None if no profile is available.
+        """
+        import tensorflow as tf
+
+        profile = {
+            "features_cardinality": int(self._features.cardinality().numpy())
+            if isinstance(self._features, tf.data.Dataset)
+            else int(tf.size(self._features).numpy()),
+        }
+        if self._targets is not None:
+            profile.update(
+                {
+                    "targets_cardinality": int(self._targets.cardinality().numpy())
+                    if isinstance(self._targets, tf.data.Dataset)
+                    else int(tf.size(self._targets).numpy()),
+                }
+            )
+        return profile
+
+    @cached_property
+    def schema(self) -> Optional[TensorDatasetSchema]:
+        """
+        An MLflow TensorSpec schema representing the tensor dataset
+        """
+        try:
+            features_schema = TensorFlowDataset._get_tf_object_schema(self._features)
+            targets_schema = None
+            if self._targets is not None:
+                targets_schema = TensorFlowDataset._get_tf_object_schema(self._targets)
+            return TensorDatasetSchema(features=features_schema, targets=targets_schema)
+        except Exception as e:
+            _logger.warning("Failed to infer schema for TensorFlow dataset. Exception: %s", e)
+            return None
+
+    @staticmethod
+    def _get_tf_object_schema(tf_object) -> Schema:
+        import tensorflow as tf
+
+        if isinstance(tf_object, tf.data.Dataset):
+            numpy_data = next(tf_object.as_numpy_iterator())
+            if isinstance(numpy_data, np.ndarray):
+                return _infer_schema(numpy_data)
+            elif isinstance(numpy_data, dict):
+                return TensorFlowDataset._get_schema_from_tf_dataset_dict_numpy_data(numpy_data)
+            elif isinstance(numpy_data, tuple):
+                return TensorFlowDataset._get_schema_from_tf_dataset_tuple_numpy_data(numpy_data)
+            else:
+                raise MlflowException(
+                    f"Failed to infer schema for tf.data.Dataset due to unrecognized numpy iterator"
+                    f" data type. Numpy iterator data types 'np.ndarray', 'dict', and 'tuple' are"
+                    f" supported. Found: {type(numpy_data)}.",
+                    INVALID_PARAMETER_VALUE,
+                )
+        elif tf.is_tensor(tf_object):
+            return _infer_schema(tf_object.numpy())
+        else:
+            raise MlflowException(
+                f"Cannot infer schema of an object that is not an instance of tf.data.Dataset or"
+                f" a TensorFlow Tensor. Found: {type(tf_object)}",
+                INTERNAL_ERROR,
+            )
+
+    @staticmethod
+    def _get_schema_from_tf_dataset_dict_numpy_data(numpy_data: dict[Any, Any]) -> Schema:
+        if not all(isinstance(data_element, np.ndarray) for data_element in numpy_data.values()):
+            raise MlflowException(
+                "Failed to infer schema for tf.data.Dataset. Schemas can only be inferred"
+                " if the dataset consists of tensors. Ragged tensors, tensor arrays, and"
+                " other types are not supported. Additionally, datasets with nested tensors"
+                " are not supported.",
+                INVALID_PARAMETER_VALUE,
+            )
+        return _infer_schema(numpy_data)
+
+    @staticmethod
+    def _get_schema_from_tf_dataset_tuple_numpy_data(numpy_data: tuple[Any]) -> Schema:
+        if not all(isinstance(data_element, np.ndarray) for data_element in numpy_data):
+            raise MlflowException(
+                "Failed to infer schema for tf.data.Dataset. Schemas can only be inferred"
+                " if the dataset consists of tensors. Ragged tensors, tensor arrays, and"
+                " other types are not supported. Additionally, datasets with nested tensors"
+                " are not supported.",
+                INVALID_PARAMETER_VALUE,
+            )
+        return _infer_schema(
+            {
+                # MLflow Schemas currently require each tensor to have a name, if more than
+                # one tensor is defined. Accordingly, use the index as the name
+                str(i): data_element
+                for i, data_element in enumerate(numpy_data)
+            }
+        )
+
+    def to_pyfunc(self) -> PyFuncInputsOutputs:
+        """
+        Converts the dataset to a collection of pyfunc inputs and outputs for model
+        evaluation. Required for use with mlflow.evaluate().
+        """
+        return PyFuncInputsOutputs(self._features, self._targets)
+
+    def to_evaluation_dataset(self, path=None, feature_names=None) -> EvaluationDataset:
+        """
+        Converts the dataset to an EvaluationDataset for model evaluation. Only supported if the
+        dataset is a Tensor. Required for use with mlflow.evaluate().
+        """
+        import tensorflow as tf
+
+        # check that data and targets are Tensors
+        if not tf.is_tensor(self._features):
+            raise MlflowException("Data must be a Tensor to convert to an EvaluationDataset.")
+        if self._targets is not None and not tf.is_tensor(self._targets):
+            raise MlflowException("Targets must be a Tensor to convert to an EvaluationDataset.")
+        return EvaluationDataset(
+            data=self._features.numpy(),
+            targets=self._targets.numpy() if self._targets is not None else None,
+            path=path,
+            feature_names=feature_names,
+        )
+
+
+def from_tensorflow(
+    features,
+    source: Optional[Union[str, DatasetSource]] = None,
+    targets=None,
+    name: Optional[str] = None,
+    digest: Optional[str] = None,
+) -> TensorFlowDataset:
+    """Constructs a TensorFlowDataset object from TensorFlow data, optional targets, and source.
+
+    If the source is path like, then this will construct a DatasetSource object from the source
+    path. Otherwise, the source is assumed to be a DatasetSource object.
+
+    Args:
+        features: A TensorFlow dataset or tensor of features.
+        source: The source from which the data was derived, e.g. a filesystem
+            path, an S3 URI, an HTTPS URL, a delta table name with version, or
+            spark table etc. If source is not a path like string,
+            pass in a DatasetSource object directly. If no source is specified,
+            a CodeDatasetSource is used, which will source information from the run
+            context.
+        targets: A TensorFlow dataset or tensor of targets. Optional.
+        name: The name of the dataset. If unspecified, a name is generated.
+        digest: A dataset digest (hash). If unspecified, a digest is computed
+            automatically.
+    """
+    from mlflow.data.code_dataset_source import CodeDatasetSource
+    from mlflow.data.dataset_source_registry import resolve_dataset_source
+    from mlflow.tracking.context import registry
+
+    if source is not None:
+        if isinstance(source, DatasetSource):
+            resolved_source = source
+        else:
+            resolved_source = resolve_dataset_source(
+                source,
+            )
+    else:
+        context_tags = registry.resolve_tags()
+        resolved_source = CodeDatasetSource(tags=context_tags)
+    return TensorFlowDataset(
+        features=features, source=resolved_source, targets=targets, name=name, digest=digest
+    )

venv/lib/python3.9/site-packages/mlflow/data/uc_volume_dataset_source.py

@@ -0,0 +1,81 @@
+import logging
+from typing import Any
+
+from mlflow.data.dataset_source import DatasetSource
+from mlflow.exceptions import MlflowException
+
+_logger = logging.getLogger(__name__)
+
+
+class UCVolumeDatasetSource(DatasetSource):
+    """Represents the source of a dataset stored in Databricks Unified Catalog Volume.
+
+    If you are using a delta table, please use `mlflow.data.delta_dataset_source.DeltaDatasetSource`
+    instead. This `UCVolumeDatasetSource` does not provide loading function, and is mostly useful
+    when you are logging a `mlflow.data.meta_dataset.MetaDataset` to MLflow, i.e., you want
+    to log the source of dataset to MLflow without loading the dataset.
+
+    Args:
+        path: the UC path of your data. It should be a valid UC path following the pattern
+            "/Volumes/{catalog}/{schema}/{volume}/{file_path}". For example,
+            "/Volumes/MyCatalog/MySchema/MyVolume/MyFile.json".
+    """
+
+    def __init__(self, path: str):
+        self.path = path
+        self._verify_uc_path_is_valid()
+
+    def _verify_uc_path_is_valid(self):
+        """Verify if the path exists in Databricks Unified Catalog."""
+        try:
+            from databricks.sdk import WorkspaceClient
+
+            w = WorkspaceClient()
+        except ImportError:
+            _logger.warning(
+                "Cannot verify the path of `UCVolumeDatasetSource` because of missing"
+                "`databricks-sdk`. Please install `databricks-sdk` via "
+                "`pip install -U databricks-sdk`. This does not block creating "
+                "`UCVolumeDatasetSource`, but your `UCVolumeDatasetSource` might be invalid."
+            )
+            return
+        except Exception:
+            _logger.warning(
+                "Cannot verify the path of `UCVolumeDatasetSource` due to a connection failure "
+                "with Databricks workspace. Please run `mlflow.login()` to log in to Databricks. "
+                "This does not block creating `UCVolumeDatasetSource`, but your "
+                "`UCVolumeDatasetSource` might be invalid."
+            )
+            return
+
+        try:
+            # Check if `self.path` points to a valid UC file.
+            w.files.get_metadata(self.path)
+        except Exception:
+            try:
+                # Check if `self.path` points to a valid UC directory.
+                w.files.get_directory_metadata(self.path)
+                # Append a slash to `self.path` to indicate it's a directory.
+                self.path += "/" if not self.path.endswith("/") else ""
+            except Exception:
+                # Neither file nor directory exists, we throw an exception.
+                raise MlflowException(f"{self.path} does not exist in Databricks Unified Catalog.")
+
+    @staticmethod
+    def _get_source_type() -> str:
+        return "uc_volume"
+
+    @staticmethod
+    def _can_resolve(raw_source: Any):
+        raise NotImplementedError
+
+    @classmethod
+    def _resolve(cls, raw_source: str):
+        raise NotImplementedError
+
+    def to_dict(self) -> dict[Any, Any]:
+        return {"path": self.path}
+
+    @classmethod
+    def from_dict(cls, source_dict: dict[Any, Any]) -> "UCVolumeDatasetSource":
+        return cls(**source_dict)