chore: import upstream snapshot with attribution

tests/data/test_artifact_dataset_sources.py

@@ -0,0 +1,137 @@
+import json
+import os
+from unittest import mock
+
+import pytest
+
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json, resolve_dataset_source
+from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
+from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
+
+
+@pytest.mark.parametrize(
+    ("source_uri", "source_type", "source_class_name"),
+    [
+        ("/tmp/path/to/my/local/file.txt", "local", "LocalArtifactDatasetSource"),
+        ("file:///tmp/path/to/my/local/directory", "local", "LocalArtifactDatasetSource"),
+        ("s3://mybucket/path/to/my/file.txt", "s3", "S3ArtifactDatasetSource"),
+        ("gs://mybucket/path/to/my/dir", "gs", "GCSArtifactDatasetSource"),
+        ("wasbs://user@host.blob.core.windows.net/dir", "wasbs", "AzureBlobArtifactDatasetSource"),
+        ("ftp://mysite.com/path/to/my/file.txt", "ftp", "FTPArtifactDatasetSource"),
+        ("sftp://mysite.com/path/to/my/dir", "sftp", "SFTPArtifactDatasetSource"),
+        ("hdfs://host_name:8020/hdfs/path/to/my/file.txt", "hdfs", "HdfsArtifactDatasetSource"),
+        ("viewfs://host_name:8020/path/to/my/dir", "viewfs", "HdfsArtifactDatasetSource"),
+    ],
+)
+def test_expected_artifact_dataset_sources_are_registered_and_resolvable(
+    source_uri, source_type, source_class_name
+):
+    dataset_source = resolve_dataset_source(source_uri)
+    assert isinstance(dataset_source, FileSystemDatasetSource)
+    assert dataset_source._get_source_type() == source_type
+    assert type(dataset_source).__name__ == source_class_name
+    assert type(dataset_source).__qualname__ == source_class_name
+    assert dataset_source.uri == source_uri
+
+
+@pytest.mark.parametrize(
+    ("source_uri", "source_type"),
+    [
+        ("/tmp/path/to/my/local/file.txt", "local"),
+        ("file:///tmp/path/to/my/local/directory", "local"),
+        ("s3://mybucket/path/to/my/file.txt", "s3"),
+        ("gs://mybucket/path/to/my/dir", "gs"),
+        ("wasbs://user@host.blob.core.windows.net/dir", "wasbs"),
+        ("ftp://mysite.com/path/to/my/file.txt", "ftp"),
+        ("sftp://mysite.com/path/to/my/dir", "sftp"),
+        ("hdfs://host_name:8020/hdfs/path/to/my/file.txt", "hdfs"),
+        ("viewfs://host_name:8020/path/to/my/dir", "viewfs"),
+    ],
+)
+def test_to_and_from_json(source_uri, source_type):
+    dataset_source = resolve_dataset_source(source_uri)
+    assert dataset_source._get_source_type() == source_type
+    source_json = dataset_source.to_json()
+
+    parsed_source_json = json.loads(source_json)
+    assert parsed_source_json["uri"] == source_uri
+
+    reloaded_source = get_dataset_source_from_json(
+        source_json, source_type=dataset_source._get_source_type()
+    )
+    assert isinstance(reloaded_source, FileSystemDatasetSource)
+    assert type(dataset_source) == type(reloaded_source)
+    assert reloaded_source.uri == dataset_source.uri
+
+
+@pytest.mark.parametrize(
+    ("source_uri", "source_type"),
+    [
+        ("/tmp/path/to/my/local/file.txt", "local"),
+        ("file:///tmp/path/to/my/local/directory", "local"),
+        ("s3://mybucket/path/to/my/file.txt", "s3"),
+        ("gs://mybucket/path/to/my/dir", "gs"),
+        ("wasbs://user@host.blob.core.windows.net/dir", "wasbs"),
+        ("ftp://mysite.com/path/to/my/file.txt", "ftp"),
+        ("sftp://mysite.com/path/to/my/dir", "sftp"),
+        ("hdfs://host_name:8020/hdfs/path/to/my/file.txt", "hdfs"),
+        ("viewfs://host_name:8020/path/to/my/dir", "viewfs"),
+    ],
+)
+def test_load_makes_expected_mlflow_artifacts_download_call(source_uri, source_type, tmp_path):
+    dataset_source = resolve_dataset_source(source_uri)
+    assert dataset_source._get_source_type() == source_type
+
+    with mock.patch("mlflow.artifacts.download_artifacts") as download_imp_mock:
+        dataset_source.load()
+        download_imp_mock.assert_called_once_with(artifact_uri=source_uri, dst_path=None)
+
+    with mock.patch("mlflow.artifacts.download_artifacts") as download_imp_mock:
+        dataset_source.load(dst_path=str(tmp_path))
+        download_imp_mock.assert_called_once_with(artifact_uri=source_uri, dst_path=str(tmp_path))
+
+
+@pytest.mark.parametrize("dst_path", [None, "dst"])
+def test_local_load(dst_path, tmp_path):
+    if dst_path is not None:
+        dst_path = str(tmp_path / dst_path)
+
+    # Test string file paths
+    file_path = str(tmp_path / "myfile.txt")
+    with open(file_path, "w") as f:
+        f.write("text")
+
+    file_dataset_source = resolve_dataset_source(file_path)
+    assert file_dataset_source._get_source_type() == "local"
+    assert file_dataset_source.load(dst_path=dst_path) == dst_path or file_path
+    with open(file_path) as f:
+        assert f.read() == "text"
+
+    # Test directory paths with pathlib.Path
+    dir_path = tmp_path / "mydir"
+    os.makedirs(dir_path)
+
+    dir_dataset_source = resolve_dataset_source(dir_path)
+    assert file_dataset_source._get_source_type() == "local"
+    assert dir_dataset_source.load() == dst_path or str(dir_path)
+
+
+@pytest.mark.parametrize("dst_path", [None, "dst"])
+def test_s3_load(mock_s3_bucket, dst_path, tmp_path):
+    if dst_path is not None:
+        dst_path = str(tmp_path / dst_path)
+
+    file_path = str(tmp_path / "myfile.txt")
+    with open(file_path, "w") as f:
+        f.write("text")
+
+    S3ArtifactRepository(f"s3://{mock_s3_bucket}").log_artifact(file_path)
+
+    s3_source_uri = f"s3://{mock_s3_bucket}/myfile.txt"
+    s3_dataset_source = resolve_dataset_source(s3_source_uri)
+    assert s3_dataset_source._get_source_type() == "s3"
+    downloaded_source = s3_dataset_source.load(dst_path=dst_path)
+    if dst_path is not None:
+        assert downloaded_source == os.path.join(dst_path, "myfile.txt")
+    with open(downloaded_source) as f:
+        assert f.read() == "text"

tests/data/test_code_dataset_source.py

@@ -0,0 +1,16 @@
+from mlflow.data.code_dataset_source import CodeDatasetSource
+
+
+def test_code_dataset_source_from_path():
+    tags = {
+        "mlflow_source_type": "NOTEBOOK",
+        "mlflow_source_name": "some_random_notebook_path",
+    }
+    code_datasource = CodeDatasetSource(tags)
+    assert code_datasource.to_dict() == {
+        "tags": tags,
+    }
+
+
+def test_code_datasource_type():
+    assert CodeDatasetSource._get_source_type() == "code"

tests/data/test_dataset.py

@@ -0,0 +1,42 @@
+import json
+
+from mlflow.types.schema import Schema
+
+from tests.resources.data.dataset import SampleDataset
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+def test_conversion_to_json():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = SampleDataset(data_list=[1, 2, 3], source=source, name="testname")
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+
+    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
+    assert Schema.from_json(schema_json) == dataset.schema
+
+
+def test_digest_property_has_expected_value():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = SampleDataset(data_list=[1, 2, 3], source=source, name="testname")
+    assert dataset.digest == dataset._compute_digest()
+
+
+def test_expected_name_is_used():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+
+    dataset_without_name = SampleDataset(data_list=[1, 2, 3], source=source)
+    assert dataset_without_name.name == "dataset"
+
+    dataset_with_name = SampleDataset(data_list=[1, 2, 3], source=source, name="testname")
+    assert dataset_with_name.name == "testname"

tests/data/test_dataset_registry.py

@@ -0,0 +1,151 @@
+from unittest import mock
+
+import pytest
+
+import mlflow.data
+from mlflow.data.dataset import Dataset
+from mlflow.data.dataset_registry import DatasetRegistry, register_constructor
+from mlflow.data.dataset_source_registry import DatasetSourceRegistry, resolve_dataset_source
+from mlflow.exceptions import MlflowException
+
+from tests.resources.data.dataset import SampleDataset
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+@pytest.fixture
+def dataset_source_registry():
+    registry = DatasetSourceRegistry()
+    with mock.patch("mlflow.data.dataset_source_registry._dataset_source_registry", wraps=registry):
+        yield registry
+
+
+@pytest.fixture
+def dataset_registry():
+    registry = DatasetRegistry()
+    with mock.patch("mlflow.data.dataset_registry._dataset_registry", wraps=registry):
+        yield registry
+
+
+def test_register_constructor_function_performs_validation():
+    registry = DatasetRegistry()
+
+    def from_good_function(
+        path: str,
+        name: str | None = None,
+        digest: str | None = None,
+    ) -> Dataset:
+        pass
+
+    registry.register_constructor(from_good_function)
+
+    def bad_name_fn(
+        name: str | None = None,
+        digest: str | None = None,
+    ) -> Dataset:
+        pass
+
+    with pytest.raises(MlflowException, match="Constructor name must start with"):
+        registry.register_constructor(bad_name_fn)
+
+    with pytest.raises(MlflowException, match="Constructor name must start with"):
+        registry.register_constructor(
+            constructor_fn=from_good_function, constructor_name="bad_name"
+        )
+
+    def from_no_name_fn(
+        digest: str | None = None,
+    ) -> Dataset:
+        pass
+
+    with pytest.raises(MlflowException, match="must define an optional parameter named 'name'"):
+        registry.register_constructor(from_no_name_fn)
+
+    def from_no_digest_fn(
+        name: str | None = None,
+    ) -> Dataset:
+        pass
+
+    with pytest.raises(MlflowException, match="must define an optional parameter named 'digest'"):
+        registry.register_constructor(from_no_digest_fn)
+
+    def from_bad_return_type_fn(
+        path: str,
+        name: str | None = None,
+        digest: str | None = None,
+    ) -> str:
+        pass
+
+    with pytest.raises(MlflowException, match="must have a return type annotation.*Dataset"):
+        registry.register_constructor(from_bad_return_type_fn)
+
+    def from_no_return_type_fn(
+        path: str,
+        name: str | None = None,
+        digest: str | None = None,
+    ):
+        pass
+
+    with pytest.raises(MlflowException, match="must have a return type annotation.*Dataset"):
+        registry.register_constructor(from_no_return_type_fn)
+
+
+def test_register_constructor_from_entrypoints_and_call(dataset_registry, tmp_path):
+    from mlflow_test_plugin.dummy_dataset import DummyDataset
+
+    dataset_registry.register_entrypoints()
+
+    dataset = mlflow.data.from_dummy(
+        data_list=[1, 2, 3],
+        # Use a DummyDatasetSource URI from mlflow_test_plugin.dummy_dataset_source, which
+        # is registered as an entrypoint whenever mlflow-test-plugin is installed
+        source="dummy:" + str(tmp_path),
+        name="dataset_name",
+        digest="foo",
+    )
+    assert isinstance(dataset, DummyDataset)
+    assert dataset.data_list == [1, 2, 3]
+    assert dataset.name == "dataset_name"
+    assert dataset.digest == "foo"
+
+
+def test_register_constructor_and_call(dataset_registry, dataset_source_registry, tmp_path):
+    dataset_source_registry.register(SampleDatasetSource)
+
+    def from_test(data_list, source, name=None, digest=None) -> SampleDataset:
+        resolved_source: SampleDatasetSource = resolve_dataset_source(
+            source, candidate_sources=[SampleDatasetSource]
+        )
+        return SampleDataset(data_list=data_list, source=resolved_source, name=name, digest=digest)
+
+    register_constructor(constructor_fn=from_test)
+    register_constructor(constructor_name="from_test_2", constructor_fn=from_test)
+
+    dataset1 = mlflow.data.from_test(
+        data_list=[1, 2, 3],
+        # Use a SampleDatasetSourceURI
+        source="test:" + str(tmp_path),
+        name="name1",
+        digest="digest1",
+    )
+    assert isinstance(dataset1, SampleDataset)
+    assert dataset1.data_list == [1, 2, 3]
+    assert dataset1.name == "name1"
+    assert dataset1.digest == "digest1"
+
+    dataset2 = mlflow.data.from_test_2(
+        data_list=[4, 5, 6],
+        # Use a SampleDatasetSourceURI
+        source="test:" + str(tmp_path),
+        name="name2",
+        digest="digest2",
+    )
+    assert isinstance(dataset2, SampleDataset)
+    assert dataset2.data_list == [4, 5, 6]
+    assert dataset2.name == "name2"
+    assert dataset2.digest == "digest2"
+
+
+def test_dataset_source_registration_failure(dataset_source_registry):
+    with mock.patch.object(dataset_source_registry, "register", side_effect=ImportError("Error")):
+        with pytest.warns(UserWarning, match="Failure attempting to register dataset constructor"):
+            dataset_source_registry.register_entrypoints()

tests/data/test_dataset_source.py

@@ -0,0 +1,73 @@
+import json
+
+import pandas as pd
+import pytest
+
+import mlflow.data
+from mlflow.exceptions import MlflowException
+
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+def test_load(tmp_path):
+    assert SampleDatasetSource("test:" + str(tmp_path)).load() == str(tmp_path)
+
+
+def test_conversion_to_json_and_back():
+    uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(uri)
+    source_json = source.to_json()
+    assert json.loads(source_json)["uri"] == uri
+    reloaded_source = SampleDatasetSource.from_json(source_json)
+    assert reloaded_source.uri == source.uri
+
+
+def test_get_source_obtains_expected_file_source(tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    path = tmp_path / "temp.csv"
+    df.to_csv(path)
+    pandas_ds = mlflow.data.from_pandas(df, source=path)
+
+    source1 = mlflow.data.get_source(pandas_ds)
+    assert json.loads(source1.to_json()) == json.loads(pandas_ds.source.to_json())
+
+    with mlflow.start_run() as r:
+        mlflow.log_input(pandas_ds)
+
+    run = mlflow.get_run(r.info.run_id)
+
+    ds_input = run.inputs.dataset_inputs[0]
+    source2 = mlflow.data.get_source(ds_input)
+    assert json.loads(source2.to_json()) == json.loads(pandas_ds.source.to_json())
+
+    ds_entity = run.inputs.dataset_inputs[0].dataset
+    source3 = mlflow.data.get_source(ds_entity)
+    assert json.loads(source3.to_json()) == json.loads(pandas_ds.source.to_json())
+
+    assert source1.load() == source2.load() == source3.load() == str(path)
+
+
+def test_get_source_obtains_expected_code_source():
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    pandas_ds = mlflow.data.from_pandas(df)
+
+    source1 = mlflow.data.get_source(pandas_ds)
+    assert json.loads(source1.to_json()) == json.loads(pandas_ds.source.to_json())
+
+    with mlflow.start_run() as r:
+        mlflow.log_input(pandas_ds)
+
+    run = mlflow.get_run(r.info.run_id)
+
+    ds_input = run.inputs.dataset_inputs[0]
+    source2 = mlflow.data.get_source(ds_input)
+    assert json.loads(source2.to_json()) == json.loads(pandas_ds.source.to_json())
+
+    ds_entity = run.inputs.dataset_inputs[0].dataset
+    source3 = mlflow.data.get_source(ds_entity)
+    assert json.loads(source3.to_json()) == json.loads(pandas_ds.source.to_json())
+
+
+def test_get_source_throws_for_invalid_input(tmp_path):
+    with pytest.raises(MlflowException, match="Unrecognized dataset type.*str"):
+        mlflow.data.get_source(str(tmp_path))

tests/data/test_dataset_source_registry.py

@@ -0,0 +1,178 @@
+from typing import Any
+from unittest import mock
+
+import pytest
+
+from mlflow.data.dataset_source_registry import DatasetSourceRegistry
+from mlflow.exceptions import MlflowException
+
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+def test_register_entrypoints_and_resolve(tmp_path):
+    from mlflow_test_plugin.dummy_dataset_source import DummyDatasetSource
+
+    registry = DatasetSourceRegistry()
+    registry.register_entrypoints()
+
+    uri = "dummy:" + str(tmp_path)
+    resolved_source = registry.resolve(uri)
+    assert isinstance(resolved_source, DummyDatasetSource)
+    # Verify that the DummyDatasetSource is constructed with the correct URI
+    assert resolved_source.uri == uri
+
+
+def test_register_dataset_source_and_resolve(tmp_path):
+    registry = DatasetSourceRegistry()
+    registry.register(SampleDatasetSource)
+
+    uri = "test:" + str(tmp_path)
+    resolved_source = registry.resolve(uri)
+    assert isinstance(resolved_source, SampleDatasetSource)
+    # Verify that the SampleDatasetSource is constructed with the correct URI
+    assert resolved_source.uri == uri
+
+
+def test_register_dataset_source_and_load_from_json(tmp_path):
+    registry = DatasetSourceRegistry()
+    registry.register(SampleDatasetSource)
+    resolved_source = registry.resolve("test:" + str(tmp_path))
+    resolved_source_json = resolved_source.to_json()
+    source_from_json = registry.get_source_from_json(
+        source_json=resolved_source_json, source_type="test"
+    )
+    assert source_from_json.uri == resolved_source.uri
+
+
+def test_load_from_json_throws_for_unrecognized_source_type(tmp_path):
+    registry = DatasetSourceRegistry()
+    registry.register(SampleDatasetSource)
+
+    with pytest.raises(MlflowException, match="unrecognized source type: foo"):
+        registry.get_source_from_json(source_json='{"bar": "123"}', source_type="foo")
+
+    class CandidateDatasetSource1(SampleDatasetSource):
+        @staticmethod
+        def _get_source_type() -> str:
+            return "candidate1"
+
+        @staticmethod
+        def _can_resolve(raw_source: Any) -> bool:
+            return raw_source.startswith("candidate1")
+
+    class CandidateDatasetSource2(CandidateDatasetSource1):
+        @staticmethod
+        def _get_source_type() -> str:
+            return "candidate2"
+
+        @staticmethod
+        def _can_resolve(raw_source: Any) -> bool:
+            return raw_source.startswith("candidate2")
+
+    registry = DatasetSourceRegistry()
+    registry.register(SampleDatasetSource)
+    registry.register(CandidateDatasetSource1)
+    registry.register(CandidateDatasetSource2)
+
+    registry.resolve("test:" + str(tmp_path))
+    registry.resolve("test:" + str(tmp_path), candidate_sources=[SampleDatasetSource])
+    with pytest.raises(MlflowException, match="Could not find a source information resolver"):
+        # SampleDatasetSource is the only source that can resolve raw sources with scheme "test",
+        # and SampleDatasetSource is not a subclass of CandidateDatasetSource1
+        registry.resolve("test:" + str(tmp_path), candidate_sources=[CandidateDatasetSource1])
+
+    registry.resolve("candidate1:" + str(tmp_path))
+    registry.resolve("candidate1:" + str(tmp_path), candidate_sources=[CandidateDatasetSource1])
+    # CandidateDatasetSource1 is a subclass of SampleDatasetSource and is therefore considered
+    # as a candidate for resolution
+    registry.resolve("candidate1:" + str(tmp_path), candidate_sources=[SampleDatasetSource])
+    with pytest.raises(MlflowException, match="Could not find a source information resolver"):
+        # CandidateDatasetSource2 is not a superclass of CandidateDatasetSource1 or
+        # SampleDatasetSource and cannot resolve raw sources with scheme "candidate1"
+        registry.resolve("candidate1:" + str(tmp_path), candidate_sources=[CandidateDatasetSource2])
+
+
+def test_resolve_dataset_source_maintains_consistent_order_and_uses_last_registered_match(tmp_path):
+    from mlflow_test_plugin.dummy_dataset_source import DummyDatasetSource
+
+    class SampleDatasetSourceCopy1(SampleDatasetSource):
+        pass
+
+    class SampleDatasetSourceCopy2(SampleDatasetSource):
+        pass
+
+    registry1 = DatasetSourceRegistry()
+    registry1.register(SampleDatasetSource)
+    registry1.register(SampleDatasetSourceCopy1)
+    registry1.register(SampleDatasetSourceCopy2)
+    source1 = registry1.resolve("test:/" + str(tmp_path))
+    assert isinstance(source1, SampleDatasetSourceCopy2)
+
+    registry2 = DatasetSourceRegistry()
+    registry2.register(SampleDatasetSource)
+    registry2.register(SampleDatasetSourceCopy2)
+    registry2.register(SampleDatasetSourceCopy1)
+    source2 = registry2.resolve("test:/" + str(tmp_path))
+    assert isinstance(source2, SampleDatasetSourceCopy1)
+
+    # Verify that a different matching dataset source can still be resolved via `candidates`
+    source3 = registry2.resolve(
+        "test:/" + str(tmp_path), candidate_sources=[SampleDatasetSourceCopy2]
+    )
+    assert isinstance(source3, SampleDatasetSourceCopy2)
+
+    # Verify that last registered order applies to entrypoints too
+    class DummyDatasetSourceCopy(DummyDatasetSource):
+        pass
+
+    registry3 = DatasetSourceRegistry()
+    registry3.register(DummyDatasetSourceCopy)
+    source4 = registry3.resolve("dummy:/" + str(tmp_path))
+    assert isinstance(source4, DummyDatasetSourceCopy)
+    registry3.register_entrypoints()
+    source5 = registry3.resolve("dummy:/" + str(tmp_path))
+    assert isinstance(source5, DummyDatasetSource)
+
+
+def test_resolve_dataset_source_warns_when_multiple_matching_sources_found(tmp_path):
+    class SampleDatasetSourceCopy1(SampleDatasetSource):
+        pass
+
+    class SampleDatasetSourceCopy2(SampleDatasetSource):
+        pass
+
+    registry1 = DatasetSourceRegistry()
+    registry1.register(SampleDatasetSource)
+    registry1.register(SampleDatasetSourceCopy1)
+    registry1.register(SampleDatasetSourceCopy2)
+
+    with mock.patch("mlflow.data.dataset_source_registry.warnings.warn") as mock_warn:
+        registry1.resolve("test:/" + str(tmp_path))
+        mock_warn.assert_called_once()
+        call_args, _ = mock_warn.call_args
+        multiple_match_msg = call_args[0]
+        assert (
+            "The specified dataset source can be interpreted in multiple ways" in multiple_match_msg
+        )
+        assert (
+            "SampleDatasetSource, SampleDatasetSourceCopy1, SampleDatasetSourceCopy2"
+            in multiple_match_msg
+        )
+        assert (
+            "MLflow will assume that this is a SampleDatasetSourceCopy2 source"
+            in multiple_match_msg
+        )
+
+
+def test_dataset_sources_are_importable_from_sources_module(tmp_path):
+    from mlflow.data.sources import LocalArtifactDatasetSource
+
+    src = LocalArtifactDatasetSource(tmp_path)
+    assert src._get_source_type() == "local"
+    assert src.uri == tmp_path
+
+    from mlflow.data.sources import DeltaDatasetSource
+
+    src = DeltaDatasetSource(path=tmp_path)
+    assert src._get_source_type() == "delta_table"
+    assert src.path == tmp_path

tests/data/test_delta_dataset_source.py

@@ -0,0 +1,248 @@
+import json
+from unittest import mock
+
+import pandas as pd
+import pytest
+
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json
+from mlflow.data.delta_dataset_source import DeltaDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_managed_catalog_messages_pb2 import GetTable, GetTableResponse
+from mlflow.utils.proto_json_utils import message_to_json
+
+
+@pytest.fixture(scope="module")
+def spark_session():
+    from pyspark.sql import SparkSession
+
+    with (
+        SparkSession.builder
+        .master("local[*]")
+        .config("spark.jars.packages", "io.delta:delta-spark_2.13:4.0.0")
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog"
+        )
+        .getOrCreate()
+    ) as session:
+        yield session
+
+
+def test_delta_dataset_source_from_path(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.delta")
+    df_spark.write.format("delta").mode("overwrite").save(path)
+
+    delta_datasource = DeltaDatasetSource(path=path)
+    loaded_df_spark = delta_datasource.load()
+    assert loaded_df_spark.count() == df_spark.count()
+    assert delta_datasource.to_dict()["path"] == path
+
+    reloaded_source = get_dataset_source_from_json(
+        delta_datasource.to_json(), source_type=delta_datasource._get_source_type()
+    )
+    assert isinstance(reloaded_source, DeltaDatasetSource)
+    assert type(delta_datasource) == type(reloaded_source)
+    assert reloaded_source.to_json() == delta_datasource.to_json()
+
+
+def test_delta_dataset_source_from_table(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.format("delta").mode("overwrite").saveAsTable(
+        "default.temp_delta", path=tmp_path
+    )
+
+    delta_datasource = DeltaDatasetSource(delta_table_name="temp_delta")
+    loaded_df_spark = delta_datasource.load()
+    assert loaded_df_spark.count() == df_spark.count()
+    assert delta_datasource.to_dict()["delta_table_name"] == "temp_delta"
+
+    reloaded_source = get_dataset_source_from_json(
+        delta_datasource.to_json(), source_type=delta_datasource._get_source_type()
+    )
+    assert isinstance(reloaded_source, DeltaDatasetSource)
+    assert type(delta_datasource) == type(reloaded_source)
+    assert reloaded_source.to_json() == delta_datasource.to_json()
+
+
+def test_delta_dataset_source_from_table_versioned(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.format("delta").mode("overwrite").saveAsTable(
+        "default.temp_delta_versioned", path=tmp_path
+    )
+
+    df2 = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"])
+    df2_spark = spark_session.createDataFrame(df2)
+    df2_spark.write.format("delta").mode("overwrite").saveAsTable(
+        "default.temp_delta_versioned", path=tmp_path
+    )
+
+    delta_datasource = DeltaDatasetSource(
+        delta_table_name="temp_delta_versioned", delta_table_version=1
+    )
+    loaded_df_spark = delta_datasource.load()
+    assert loaded_df_spark.count() == df2_spark.count()
+    config = delta_datasource.to_dict()
+    assert config["delta_table_name"] == "temp_delta_versioned"
+    assert config["delta_table_version"] == 1
+
+    reloaded_source = get_dataset_source_from_json(
+        delta_datasource.to_json(), source_type=delta_datasource._get_source_type()
+    )
+    assert isinstance(reloaded_source, DeltaDatasetSource)
+    assert type(delta_datasource) == type(reloaded_source)
+    assert reloaded_source.to_json() == delta_datasource.to_json()
+
+
+def test_delta_dataset_source_too_many_inputs(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.format("delta").mode("overwrite").saveAsTable(
+        "default.temp_delta_too_many_inputs", path=tmp_path
+    )
+
+    with pytest.raises(MlflowException, match='Must specify exactly one of "path" or "table_name"'):
+        DeltaDatasetSource(path=tmp_path, delta_table_name="temp_delta_too_many_inputs")
+
+
+def test_uc_table_id_retrieval_works(spark_session, tmp_path):
+    def mock_resolve_table_name(table_name, spark):
+        if table_name == "temp_delta_versioned_with_id":
+            return "default.temp_delta_versioned_with_id"
+        return table_name
+
+    def mock_lookup_table_id(table_name):
+        if table_name == "default.temp_delta_versioned_with_id":
+            return "uc_table_id_1"
+        return None
+
+    with (
+        mock.patch(
+            "mlflow.data.delta_dataset_source.get_full_name_from_sc",
+            side_effect=mock_resolve_table_name,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source.DeltaDatasetSource._lookup_table_id",
+            side_effect=mock_lookup_table_id,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source._get_active_spark_session",
+            return_value=None,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source.DeltaDatasetSource._is_databricks_uc_table",
+            return_value=True,
+        ),
+    ):
+        df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+        df_spark = spark_session.createDataFrame(df)
+        df_spark.write.format("delta").mode("overwrite").saveAsTable(
+            "default.temp_delta_versioned_with_id", path=tmp_path
+        )
+
+        df2 = pd.DataFrame([[1, 2, 3]], columns=["a", "b", "c"])
+        df2_spark = spark_session.createDataFrame(df2)
+        df2_spark.write.format("delta").mode("overwrite").saveAsTable(
+            "default.temp_delta_versioned_with_id", path=tmp_path
+        )
+
+        delta_datasource = DeltaDatasetSource(
+            delta_table_name="temp_delta_versioned_with_id", delta_table_version=1
+        )
+        loaded_df_spark = delta_datasource.load()
+        assert loaded_df_spark.count() == df2_spark.count()
+        assert delta_datasource.to_json() == json.dumps({
+            "delta_table_name": "default.temp_delta_versioned_with_id",
+            "delta_table_version": 1,
+            "is_databricks_uc_table": True,
+            "delta_table_id": "uc_table_id_1",
+        })
+
+
+def _args(endpoint, json_body):
+    return {
+        "host_creds": None,
+        "endpoint": f"/api/2.0/unity-catalog/tables/{endpoint}",
+        "method": "GET",
+        "json_body": json_body,
+        "response_proto": GetTableResponse,
+    }
+
+
+@pytest.mark.parametrize(
+    ("call_endpoint_response", "expected_lookup_response", "test_table_name"),
+    [
+        (None, None, "delta_table_1"),
+        (Exception("Exception from call_endpoint"), None, "delta_table_2"),
+        (GetTableResponse(table_id="uc_table_id_1"), "uc_table_id_1", "delta_table_3"),
+    ],
+)
+def test_lookup_table_id(
+    call_endpoint_response, expected_lookup_response, test_table_name, tmp_path
+):
+    def mock_resolve_table_name(table_name, spark):
+        if table_name == test_table_name:
+            return f"default.{test_table_name}"
+        return table_name
+
+    def mock_call_endpoint(host_creds, endpoint, method, json_body, response_proto):
+        if isinstance(call_endpoint_response, Exception):
+            raise call_endpoint_response
+        return call_endpoint_response
+
+    with (
+        mock.patch(
+            "mlflow.data.delta_dataset_source.get_full_name_from_sc",
+            side_effect=mock_resolve_table_name,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source._get_active_spark_session",
+            return_value=None,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source.get_databricks_host_creds",
+            return_value=None,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source.DeltaDatasetSource._is_databricks_uc_table",
+            return_value=True,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source.call_endpoint",
+            side_effect=mock_call_endpoint,
+        ) as mock_endpoint,
+    ):
+        delta_datasource = DeltaDatasetSource(
+            delta_table_name=test_table_name, delta_table_version=1
+        )
+        assert delta_datasource._lookup_table_id(test_table_name) == expected_lookup_response
+        req_body = message_to_json(GetTable(full_name_arg=test_table_name))
+        call_args = _args(test_table_name, req_body)
+        mock_endpoint.assert_any_call(**call_args)
+
+
+@pytest.mark.parametrize(
+    ("table_name", "expected_result"),
+    [
+        ("default.test", True),
+        ("hive_metastore.test", False),
+        ("spark_catalog.test", False),
+        ("samples.test", False),
+    ],
+)
+def test_is_databricks_uc_table(table_name, expected_result):
+    with (
+        mock.patch(
+            "mlflow.data.delta_dataset_source.get_full_name_from_sc",
+            return_value=table_name,
+        ),
+        mock.patch(
+            "mlflow.data.delta_dataset_source._get_active_spark_session",
+            return_value=None,
+        ),
+    ):
+        delta_datasource = DeltaDatasetSource(delta_table_name=table_name, delta_table_version=1)
+        assert delta_datasource._is_databricks_uc_table() == expected_result

tests/data/test_http_dataset_source.py

@@ -0,0 +1,188 @@
+import json
+import os
+from unittest import mock
+
+import pandas as pd
+import pytest
+
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json, resolve_dataset_source
+from mlflow.data.http_dataset_source import HTTPDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.utils.os import is_windows
+from mlflow.utils.rest_utils import cloud_storage_http_request
+
+
+def test_source_to_and_from_json():
+    url = "http://mywebsite.com/path/to/my/dataset.txt"
+    source = HTTPDatasetSource(url)
+    assert source.to_json() == json.dumps({"url": url})
+
+    reloaded_source = get_dataset_source_from_json(
+        source.to_json(), source_type=source._get_source_type()
+    )
+    assert isinstance(reloaded_source, HTTPDatasetSource)
+    assert type(source) == type(reloaded_source)
+    assert source.url == reloaded_source.url == url
+
+
+def test_http_dataset_source_is_registered_and_resolvable():
+    source1 = resolve_dataset_source(
+        "http://mywebsite.com/path/to/my/dataset.txt", candidate_sources=[HTTPDatasetSource]
+    )
+    assert isinstance(source1, HTTPDatasetSource)
+    assert source1.url == "http://mywebsite.com/path/to/my/dataset.txt"
+
+    source2 = resolve_dataset_source(
+        "https://otherwebsite.net", candidate_sources=[HTTPDatasetSource]
+    )
+    assert isinstance(source2, HTTPDatasetSource)
+    assert source2.url == "https://otherwebsite.net"
+
+    with pytest.raises(MlflowException, match="Could not find a source information resolver"):
+        resolve_dataset_source("s3://mybucket", candidate_sources=[HTTPDatasetSource])
+
+    with pytest.raises(MlflowException, match="Could not find a source information resolver"):
+        resolve_dataset_source("otherscheme://mybucket", candidate_sources=[HTTPDatasetSource])
+
+    with pytest.raises(MlflowException, match="Could not find a source information resolver"):
+        resolve_dataset_source("htp://mybucket", candidate_sources=[HTTPDatasetSource])
+
+
+def test_source_load(tmp_path):
+    source1 = HTTPDatasetSource(
+        "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
+    )
+
+    loaded1 = source1.load()
+    parsed1 = pd.read_csv(loaded1, sep=";")
+    # Verify that the expected data was downloaded by checking for an expected column and asserting
+    # that several rows are present
+    assert "fixed acidity" in parsed1.columns
+    assert len(parsed1) > 10
+
+    loaded2 = source1.load(dst_path=tmp_path)
+    assert loaded2 == str(tmp_path / "winequality-red.csv")
+    parsed2 = pd.read_csv(loaded2, sep=";")
+    # Verify that the expected data was downloaded by checking for an expected column and asserting
+    # that several rows are present
+    assert "fixed acidity" in parsed2.columns
+    assert len(parsed1) > 10
+
+    source2 = HTTPDatasetSource(
+        "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv#foo?query=param"
+    )
+    loaded3 = source2.load(dst_path=tmp_path)
+    assert loaded3 == str(tmp_path / "winequality-red.csv")
+    parsed3 = pd.read_csv(loaded3, sep=";")
+    assert "fixed acidity" in parsed3.columns
+    assert len(parsed1) > 10
+
+    source3 = HTTPDatasetSource("https://github.com/")
+    loaded4 = source3.load()
+    assert os.path.exists(loaded4)
+    assert os.path.basename(loaded4) == "dataset_source"
+
+    source4 = HTTPDatasetSource("https://github.com")
+    loaded5 = source4.load()
+    assert os.path.exists(loaded5)
+    assert os.path.basename(loaded5) == "dataset_source"
+
+    def cloud_storage_http_request_with_fast_fail(*args, **kwargs):
+        kwargs["max_retries"] = 1
+        kwargs["timeout"] = 5
+        return cloud_storage_http_request(*args, **kwargs)
+
+    source5 = HTTPDatasetSource("https://nonexistentwebsitebuiltbythemlflowteam112312.com")
+    with (
+        mock.patch(
+            "mlflow.data.http_dataset_source.cloud_storage_http_request",
+            side_effect=cloud_storage_http_request_with_fast_fail,
+        ),
+        pytest.raises(Exception, match="Max retries exceeded with url"),
+    ):
+        source5.load()
+
+
+@pytest.mark.parametrize(
+    ("attachment_filename", "expected_filename"),
+    [
+        ("testfile.txt", "testfile.txt"),
+        ('"testfile.txt"', "testfile.txt"),
+        ("'testfile.txt'", "testfile.txt"),
+        (None, "winequality-red.csv"),
+    ],
+)
+def test_source_load_with_content_disposition_header(attachment_filename, expected_filename):
+    def download_with_mock_content_disposition_headers(*args, **kwargs):
+        response = cloud_storage_http_request(*args, **kwargs)
+        if attachment_filename is not None:
+            response.headers["Content-Disposition"] = f"attachment; filename={attachment_filename}"
+        else:
+            response.headers["Content-Disposition"] = "attachment"
+        return response
+
+    with mock.patch(
+        "mlflow.data.http_dataset_source.cloud_storage_http_request",
+        side_effect=download_with_mock_content_disposition_headers,
+    ):
+        source = HTTPDatasetSource(
+            "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
+        )
+        source.load()
+        loaded = source.load()
+        assert os.path.exists(loaded)
+        assert os.path.basename(loaded) == expected_filename
+
+
+@pytest.mark.parametrize(
+    "filename",
+    [
+        "/foo/bar.txt",
+        "./foo/bar.txt",
+        "../foo/bar.txt",
+        "foo/bar.txt",
+    ],
+)
+def test_source_load_with_content_disposition_header_invalid_filename(filename):
+    def download_with_mock_content_disposition_headers(*args, **kwargs):
+        response = cloud_storage_http_request(*args, **kwargs)
+        response.headers["Content-Disposition"] = f"attachment; filename={filename}"
+        return response
+
+    with mock.patch(
+        "mlflow.data.http_dataset_source.cloud_storage_http_request",
+        side_effect=download_with_mock_content_disposition_headers,
+    ):
+        source = HTTPDatasetSource(
+            "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
+        )
+
+        with pytest.raises(MlflowException, match="Invalid filename in Content-Disposition header"):
+            source.load()
+
+
+@pytest.mark.skipif(not is_windows(), reason="This test only passes on Windows")
+@pytest.mark.parametrize(
+    "filename",
+    [
+        r"..\..\poc.txt",
+        r"Users\User\poc.txt",
+    ],
+)
+def test_source_load_with_content_disposition_header_invalid_filename_windows(filename):
+    def download_with_mock_content_disposition_headers(*args, **kwargs):
+        response = cloud_storage_http_request(*args, **kwargs)
+        response.headers = {"Content-Disposition": f"attachment; filename={filename}"}
+        return response
+
+    with mock.patch(
+        "mlflow.data.http_dataset_source.cloud_storage_http_request",
+        side_effect=download_with_mock_content_disposition_headers,
+    ):
+        source = HTTPDatasetSource(
+            "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-red.csv"
+        )
+
+        # Expect an MlflowException for invalid filenames
+        with pytest.raises(MlflowException, match="Invalid filename in Content-Disposition header"):
+            source.load()

tests/data/test_huggingface_dataset_and_source.py

@@ -0,0 +1,292 @@
+import json
+import os
+
+import datasets
+import pandas as pd
+import pytest
+from huggingface_hub.errors import HfHubHTTPError
+
+import mlflow.data
+import mlflow.data.huggingface_dataset
+from mlflow.data.code_dataset_source import CodeDatasetSource
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.huggingface_dataset import HuggingFaceDataset
+from mlflow.data.huggingface_dataset_source import HuggingFaceDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.types.schema import Schema
+from mlflow.types.utils import _infer_schema
+
+from tests.helper_functions import skip_if_hf_hub_unhealthy
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+pytestmark = skip_if_hf_hub_unhealthy()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def prefetch_huggingface_datasets():
+    """Pre-warm the HF cache so individual tests don't hit the Hub and risk HTTP 429s."""
+
+    try:
+        datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
+        datasets.load_dataset(
+            "cornell-movie-review-data/rotten_tomatoes",
+            split="train",
+            revision="aa13bc287fa6fcab6daf52f0dfb9994269ffea28",
+            trust_remote_code=True,
+        )
+        datasets.load_dataset(
+            "cornell-movie-review-data/rotten_tomatoes",
+            split="train",
+            revision="c33cbf965006dba64f134f7bef69c53d5d0d285d",
+        )
+        datasets.load_dataset(
+            "fka/awesome-chatgpt-prompts",
+            data_files={"train": "prompts.csv"},
+            split="train",
+        )
+    except HfHubHTTPError as e:
+        if e.response is not None and e.response.status_code == 429:
+            pytest.skip(f"HF Hub returned 429 while pre-warming the cache: {e}")
+        raise
+
+
+def test_from_huggingface_dataset_constructs_expected_dataset():
+    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
+    mlflow_ds = mlflow.data.from_huggingface(ds, path="cornell-movie-review-data/rotten_tomatoes")
+
+    assert isinstance(mlflow_ds, HuggingFaceDataset)
+    assert mlflow_ds.ds == ds
+    assert mlflow_ds.schema == _infer_schema(ds.to_pandas())
+    assert mlflow_ds.profile == {
+        "num_rows": ds.num_rows,
+        "dataset_size": ds.dataset_size,
+        "size_in_bytes": ds.size_in_bytes,
+    }
+
+    assert isinstance(mlflow_ds.source, HuggingFaceDatasetSource)
+
+    with pytest.raises(KeyError, match="Found duplicated arguments*"):
+        # Test that we raise an error if the same key is specified in both
+        # `HuggingFaceDatasetSource` and `kwargs`.
+        mlflow_ds.source.load(path="dummy_path")
+
+    reloaded_ds = mlflow_ds.source.load()
+    assert reloaded_ds.builder_name == ds.builder_name
+    assert reloaded_ds.config_name == ds.config_name
+    assert reloaded_ds.split == ds.split == "train"
+    assert reloaded_ds.num_rows == ds.num_rows
+
+    reloaded_mlflow_ds = mlflow.data.from_huggingface(
+        reloaded_ds, path="cornell-movie-review-data/rotten_tomatoes"
+    )
+    assert reloaded_mlflow_ds.digest == mlflow_ds.digest
+
+
+def test_from_huggingface_dataset_constructs_expected_dataset_with_revision():
+    # Load this revision:
+    # https://huggingface.co/datasets/cornell-movie-review-data/rotten_tomatoes/commit/aa13bc287fa6fcab6daf52f0dfb9994269ffea28
+    revision = "aa13bc287fa6fcab6daf52f0dfb9994269ffea28"
+    ds = datasets.load_dataset(
+        "cornell-movie-review-data/rotten_tomatoes",
+        split="train",
+        revision=revision,
+        trust_remote_code=True,
+    )
+
+    mlflow_ds_new = mlflow.data.from_huggingface(
+        ds,
+        path="cornell-movie-review-data/rotten_tomatoes",
+        revision=revision,
+        trust_remote_code=True,
+    )
+
+    mlflow_ds_new.source.load()
+    assert mlflow_ds_new.source.revision == revision
+
+
+def test_from_huggingface_dataset_constructs_expected_dataset_with_data_files():
+    data_files = {"train": "prompts.csv"}
+    ds = datasets.load_dataset("fka/awesome-chatgpt-prompts", data_files=data_files, split="train")
+    mlflow_ds = mlflow.data.from_huggingface(
+        ds, path="fka/awesome-chatgpt-prompts", data_files=data_files
+    )
+
+    assert isinstance(mlflow_ds, HuggingFaceDataset)
+    assert mlflow_ds.ds == ds
+    assert mlflow_ds.schema == _infer_schema(ds.to_pandas())
+    assert mlflow_ds.profile == {
+        "num_rows": ds.num_rows,
+        "dataset_size": ds.dataset_size,
+        "size_in_bytes": ds.size_in_bytes,
+    }
+
+    assert isinstance(mlflow_ds.source, HuggingFaceDatasetSource)
+    reloaded_ds = mlflow_ds.source.load()
+    assert reloaded_ds.builder_name == ds.builder_name
+    assert reloaded_ds.config_name == ds.config_name
+    assert reloaded_ds.split == ds.split == "train"
+    assert reloaded_ds.num_rows == ds.num_rows
+
+    reloaded_mlflow_ds = mlflow.data.from_huggingface(
+        reloaded_ds, path="fka/awesome-chatgpt-prompts", data_files=data_files
+    )
+    assert reloaded_mlflow_ds.digest == mlflow_ds.digest
+
+
+def test_from_huggingface_dataset_constructs_expected_dataset_with_data_dir(tmp_path):
+    df = pd.DataFrame.from_dict({"a": [1, 2, 3], "b": [4, 5, 6]})
+    data_dir = "data"
+    os.makedirs(tmp_path / data_dir)
+    df.to_csv(tmp_path / data_dir / "my_data.csv")
+    ds = datasets.load_dataset(str(tmp_path), data_dir=data_dir, name="default", split="train")
+    mlflow_ds = mlflow.data.from_huggingface(ds, path=str(tmp_path), data_dir=data_dir)
+
+    assert mlflow_ds.ds == ds
+    assert mlflow_ds.schema == _infer_schema(ds.to_pandas())
+    assert mlflow_ds.profile == {
+        "num_rows": ds.num_rows,
+        "dataset_size": ds.dataset_size,
+        "size_in_bytes": ds.size_in_bytes,
+    }
+
+    assert isinstance(mlflow_ds.source, HuggingFaceDatasetSource)
+    reloaded_ds = mlflow_ds.source.load()
+    assert reloaded_ds.builder_name == ds.builder_name
+    assert reloaded_ds.config_name == ds.config_name
+    assert reloaded_ds.split == ds.split == "train"
+    assert reloaded_ds.num_rows == ds.num_rows
+
+    reloaded_mlflow_ds = mlflow.data.from_huggingface(
+        reloaded_ds, path=str(tmp_path), data_dir=data_dir
+    )
+    assert reloaded_mlflow_ds.digest == mlflow_ds.digest
+
+
+def test_from_huggingface_dataset_respects_user_specified_name_and_digest():
+    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
+    mlflow_ds = mlflow.data.from_huggingface(
+        ds, path="cornell-movie-review-data/rotten_tomatoes", name="myname", digest="mydigest"
+    )
+    assert mlflow_ds.name == "myname"
+    assert mlflow_ds.digest == "mydigest"
+
+
+def test_from_huggingface_dataset_digest_is_consistent_for_large_ordered_datasets(tmp_path):
+    assert (
+        mlflow.data.huggingface_dataset._MAX_ROWS_FOR_DIGEST_COMPUTATION_AND_SCHEMA_INFERENCE
+        < 200000
+    )
+
+    df = pd.DataFrame.from_dict({
+        "a": list(range(200000)),
+        "b": list(range(200000)),
+    })
+    data_dir = "data"
+    os.makedirs(tmp_path / data_dir)
+    df.to_csv(tmp_path / data_dir / "my_data.csv")
+
+    ds = datasets.load_dataset(str(tmp_path), data_dir=data_dir, name="default", split="train")
+    mlflow_ds = mlflow.data.from_huggingface(ds, path=str(tmp_path), data_dir=data_dir)
+    assert mlflow_ds.digest == "1dda4ce8"
+
+
+def test_from_huggingface_dataset_throws_for_dataset_dict():
+    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes")
+    assert isinstance(ds, datasets.DatasetDict)
+
+    with pytest.raises(
+        MlflowException, match="must be an instance of `datasets.Dataset`.*DatasetDict"
+    ):
+        mlflow.data.from_huggingface(ds, path="cornell-movie-review-data/rotten_tomatoes")
+
+
+def test_from_huggingface_dataset_no_source_specified():
+    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
+    mlflow_ds = mlflow.data.from_huggingface(ds)
+
+    assert isinstance(mlflow_ds, HuggingFaceDataset)
+
+    assert isinstance(mlflow_ds.source, CodeDatasetSource)
+    assert "mlflow.source.name" in mlflow_ds.source.to_json()
+
+
+def test_dataset_conversion_to_json():
+    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
+    mlflow_ds = mlflow.data.from_huggingface(ds, path="cornell-movie-review-data/rotten_tomatoes")
+
+    dataset_json = mlflow_ds.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == mlflow_ds.name
+    assert parsed_json["digest"] == mlflow_ds.digest
+    assert parsed_json["source"] == mlflow_ds.source.to_json()
+    assert parsed_json["source_type"] == mlflow_ds.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(mlflow_ds.profile)
+
+    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
+    assert Schema.from_json(schema_json) == mlflow_ds.schema
+
+
+def test_dataset_source_conversion_to_json():
+    ds = datasets.load_dataset(
+        "cornell-movie-review-data/rotten_tomatoes",
+        split="train",
+        revision="c33cbf965006dba64f134f7bef69c53d5d0d285d",
+    )
+    mlflow_ds = mlflow.data.from_huggingface(
+        ds,
+        path="cornell-movie-review-data/rotten_tomatoes",
+        revision="c33cbf965006dba64f134f7bef69c53d5d0d285d",
+    )
+    source = mlflow_ds.source
+
+    source_json = source.to_json()
+    parsed_source = json.loads(source_json)
+    assert parsed_source["revision"] == "c33cbf965006dba64f134f7bef69c53d5d0d285d"
+    assert parsed_source["split"] == "train"
+    assert parsed_source["config_name"] == "default"
+    assert parsed_source["path"] == "cornell-movie-review-data/rotten_tomatoes"
+    assert not parsed_source["data_dir"]
+    assert not parsed_source["data_files"]
+
+    reloaded_source = HuggingFaceDatasetSource.from_json(source_json)
+    assert json.loads(reloaded_source.to_json()) == parsed_source
+
+    reloaded_source = get_dataset_source_from_json(
+        source_json, source_type=source._get_source_type()
+    )
+    assert isinstance(reloaded_source, HuggingFaceDatasetSource)
+    assert type(source) == type(reloaded_source)
+    assert reloaded_source.to_json() == source.to_json()
+
+
+def test_to_evaluation_dataset():
+    import numpy as np
+
+    ds = datasets.load_dataset("cornell-movie-review-data/rotten_tomatoes", split="train")
+    dataset = mlflow.data.from_huggingface(
+        ds, path="cornell-movie-review-data/rotten_tomatoes", targets="label"
+    )
+
+    evaluation_dataset = dataset.to_evaluation_dataset()
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert evaluation_dataset.features_data.equals(dataset.ds.to_pandas().drop("label", axis=1))
+    assert np.array_equal(evaluation_dataset.labels_data, dataset.ds.to_pandas()["label"].values)
+
+
+def test_from_huggingface_dataset_with_sample_source():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+
+    data = {"text": ["This is a sample text.", "Another sample text."], "label": [0, 1]}
+    dataset = datasets.Dataset.from_dict(data)
+
+    train_dataset = mlflow.data.from_huggingface(
+        dataset,
+        name="sample-text-dataset",
+        source=source,
+    )
+
+    assert isinstance(train_dataset, HuggingFaceDataset)
+    assert train_dataset.source == source

tests/data/test_meta_dataset.py

@@ -0,0 +1,119 @@
+import json
+from unittest.mock import patch
+
+import pytest
+
+pd = pytest.importorskip("pandas")
+
+from mlflow.data.delta_dataset_source import DeltaDatasetSource
+from mlflow.data.http_dataset_source import HTTPDatasetSource
+from mlflow.data.huggingface_dataset_source import HuggingFaceDatasetSource
+from mlflow.data.meta_dataset import MetaDataset
+from mlflow.data.pandas_dataset import from_pandas
+from mlflow.data.uc_volume_dataset_source import UCVolumeDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.types import DataType
+from mlflow.types.schema import ColSpec, Schema
+
+
+@pytest.mark.parametrize(
+    ("dataset_source_class", "path"),
+    [
+        (HTTPDatasetSource, "test:/my/test/uri"),
+        (DeltaDatasetSource, "fake/path/to/delta"),
+        (HuggingFaceDatasetSource, "databricks/databricks-dolly-15k"),
+    ],
+)
+def test_create_meta_dataset_from_source(dataset_source_class, path):
+    source = dataset_source_class(path)
+    dataset = MetaDataset(source=source)
+
+    json_str = dataset.to_json()
+    parsed_json = json.loads(json_str)
+
+    assert parsed_json["digest"] is not None
+    assert path in parsed_json["source"]
+    assert parsed_json["source_type"] == dataset_source_class._get_source_type()
+
+
+@pytest.mark.parametrize(
+    ("dataset_source_class", "path"),
+    [
+        (HTTPDatasetSource, "test:/my/test/uri"),
+        (DeltaDatasetSource, "fake/path/to/delta"),
+        (HuggingFaceDatasetSource, "databricks/databricks-dolly-15k"),
+    ],
+)
+def test_create_meta_dataset_from_source_with_schema(dataset_source_class, path):
+    source = dataset_source_class(path)
+    schema = Schema([
+        ColSpec(type=DataType.long, name="foo"),
+        ColSpec(type=DataType.integer, name="bar"),
+    ])
+    dataset = MetaDataset(source=source, schema=schema)
+
+    json_str = dataset.to_json()
+    parsed_json = json.loads(json_str)
+
+    assert parsed_json["digest"] is not None
+    assert path in parsed_json["source"]
+    assert parsed_json["source_type"] == dataset_source_class._get_source_type()
+    assert json.loads(parsed_json["schema"])["mlflow_colspec"] == schema.to_dict()
+
+
+def test_meta_dataset_digest():
+    http_source = HTTPDatasetSource("test:/my/test/uri")
+    dataset1 = MetaDataset(source=http_source)
+    schema = Schema([
+        ColSpec(type=DataType.long, name="foo"),
+        ColSpec(type=DataType.integer, name="bar"),
+    ])
+    dataset2 = MetaDataset(source=http_source, schema=schema)
+
+    assert dataset1.digest != dataset2.digest
+
+    delta_source = DeltaDatasetSource("fake/path/to/delta")
+    dataset3 = MetaDataset(source=delta_source)
+    assert dataset1.digest != dataset3.digest
+
+
+def test_meta_dataset_with_uc_source():
+    path = "/Volumes/dummy_catalog/dummy_schema/dummy_volume/tmp.yaml"
+
+    with (
+        patch(
+            "mlflow.data.uc_volume_dataset_source.UCVolumeDatasetSource._verify_uc_path_is_valid",
+            side_effect=MlflowException(f"{path} does not exist in Databricks Unified Catalog."),
+        ),
+        pytest.raises(
+            MlflowException, match=f"{path} does not exist in Databricks Unified Catalog."
+        ),
+    ):
+        uc_volume_source = UCVolumeDatasetSource(path)
+
+    with patch(
+        "mlflow.data.uc_volume_dataset_source.UCVolumeDatasetSource._verify_uc_path_is_valid",
+    ):
+        uc_volume_source = UCVolumeDatasetSource(path)
+        dataset = MetaDataset(source=uc_volume_source)
+        json_str = dataset.to_json()
+        parsed_json = json.loads(json_str)
+
+        assert parsed_json["digest"] is not None
+        assert path in parsed_json["source"]
+        assert parsed_json["source_type"] == "uc_volume"
+
+
+def test_create_meta_dataset_from_dataset():
+    pandas_dataset = from_pandas(
+        df=pd.DataFrame({"a": [1, 2, 3]}),
+        source="/tmp/test.csv",
+    )
+
+    meta_dataset = MetaDataset(source=pandas_dataset)
+
+    parsed_json = json.loads(meta_dataset.to_json())
+
+    assert parsed_json["source_type"] == pandas_dataset._get_source_type()
+    dataset_json = json.loads(parsed_json["source"])
+    assert dataset_json["source_type"] == pandas_dataset.source._get_source_type()

tests/data/test_numpy_dataset.py

@@ -0,0 +1,257 @@
+import json
+
+import numpy as np
+import pandas as pd
+import pytest
+
+import mlflow.data
+from mlflow.data.code_dataset_source import CodeDatasetSource
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
+from mlflow.data.numpy_dataset import NumpyDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncInputsOutputs
+from mlflow.data.schema import TensorDatasetSchema
+from mlflow.types.utils import _infer_schema
+
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+def test_conversion_to_json():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = NumpyDataset(features=np.array([1, 2, 3]), source=source, name="testname")
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+
+    parsed_schema = json.loads(parsed_json["schema"])
+    assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
+
+
+@pytest.mark.parametrize(
+    ("features", "targets"),
+    [
+        (
+            {
+                "a": np.array([1, 2, 3]),
+                "b": np.array([[4, 5]]),
+            },
+            {
+                "c": np.array([1]),
+                "d": np.array([[[2]]]),
+            },
+        ),
+        (
+            np.array([1, 2, 3]),
+            {
+                "c": np.array([1]),
+                "d": np.array([[[2]]]),
+            },
+        ),
+        (
+            {
+                "a": np.array([1, 2, 3]),
+                "b": np.array([[4, 5]]),
+            },
+            np.array([1, 2, 3]),
+        ),
+    ],
+)
+def test_conversion_to_json_with_multi_tensor_features_and_targets(features, targets):
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = NumpyDataset(features=features, targets=targets, source=source)
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+    parsed_schema = json.loads(parsed_json["schema"])
+    assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
+
+
+@pytest.mark.parametrize(
+    ("features", "targets"),
+    [
+        (
+            {
+                "a": np.array([1, 2, 3]),
+                "b": np.array([[4, 5]]),
+            },
+            {
+                "c": np.array([1]),
+                "d": np.array([[[2]]]),
+            },
+        ),
+        (
+            np.array([1, 2, 3]),
+            {
+                "c": np.array([1]),
+                "d": np.array([[[2]]]),
+            },
+        ),
+        (
+            {
+                "a": np.array([1, 2, 3]),
+                "b": np.array([[4, 5]]),
+            },
+            np.array([1, 2, 3]),
+        ),
+    ],
+)
+def test_schema_and_profile_with_multi_tensor_features_and_targets(features, targets):
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = NumpyDataset(features=features, targets=targets, source=source)
+
+    assert isinstance(dataset.schema, TensorDatasetSchema)
+    assert dataset.schema.features == _infer_schema(features)
+    assert dataset.schema.targets == _infer_schema(targets)
+
+    if isinstance(features, dict):
+        assert {
+            "features_shape": {key: array.shape for key, array in features.items()},
+            "features_size": {key: array.size for key, array in features.items()},
+            "features_nbytes": {key: array.nbytes for key, array in features.items()},
+        }.items() <= dataset.profile.items()
+    else:
+        assert {
+            "features_shape": features.shape,
+            "features_size": features.size,
+            "features_nbytes": features.nbytes,
+        }.items() <= dataset.profile.items()
+
+    if isinstance(targets, dict):
+        assert {
+            "targets_shape": {key: array.shape for key, array in targets.items()},
+            "targets_size": {key: array.size for key, array in targets.items()},
+            "targets_nbytes": {key: array.nbytes for key, array in targets.items()},
+        }.items() <= dataset.profile.items()
+    else:
+        assert {
+            "targets_shape": targets.shape,
+            "targets_size": targets.size,
+            "targets_nbytes": targets.nbytes,
+        }.items() <= dataset.profile.items()
+
+
+def test_digest_property_has_expected_value():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    features = np.array([1, 2, 3])
+    targets = np.array([4, 5, 6])
+    dataset_with_features = NumpyDataset(features=features, source=source, name="testname")
+    assert dataset_with_features.digest == dataset_with_features._compute_digest()
+    assert dataset_with_features.digest == "fdf1765f"
+    dataset_with_features_and_targets = NumpyDataset(
+        features=features, targets=targets, source=source, name="testname"
+    )
+    assert (
+        dataset_with_features_and_targets.digest
+        == dataset_with_features_and_targets._compute_digest()
+    )
+    assert dataset_with_features_and_targets.digest == "1387de76"
+
+
+def test_features_property():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    features = np.array([1, 2, 3])
+    dataset = NumpyDataset(features=features, source=source, name="testname")
+    assert np.array_equal(dataset.features, features)
+
+
+def test_targets_property():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    features = np.array([1, 2, 3])
+    targets = np.array([4, 5, 6])
+    dataset_with_targets = NumpyDataset(
+        features=features, targets=targets, source=source, name="testname"
+    )
+    assert np.array_equal(dataset_with_targets.targets, targets)
+    dataset_without_targets = NumpyDataset(features=features, source=source, name="testname")
+    assert dataset_without_targets.targets is None
+
+
+def test_to_pyfunc():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    features = np.array([1, 2, 3])
+    dataset = NumpyDataset(features=features, source=source, name="testname")
+    assert isinstance(dataset.to_pyfunc(), PyFuncInputsOutputs)
+
+
+def test_to_evaluation_dataset():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    features = np.array([[1, 2], [3, 4]])
+    targets = np.array([0, 1])
+    dataset = NumpyDataset(features=features, targets=targets, source=source, name="testname")
+    evaluation_dataset = dataset.to_evaluation_dataset()
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert np.array_equal(evaluation_dataset.features_data, features)
+    assert np.array_equal(evaluation_dataset.labels_data, targets)
+
+
+def test_from_numpy_features_only(tmp_path):
+    features = np.array([1, 2, 3])
+    path = tmp_path / "temp.csv"
+    pd.DataFrame(features).to_csv(path)
+    mlflow_features = mlflow.data.from_numpy(features, source=path)
+
+    assert isinstance(mlflow_features, NumpyDataset)
+    assert np.array_equal(mlflow_features.features, features)
+    assert mlflow_features.schema == TensorDatasetSchema(features=_infer_schema(features))
+    assert mlflow_features.profile == {
+        "features_shape": features.shape,
+        "features_size": features.size,
+        "features_nbytes": features.nbytes,
+    }
+
+    assert isinstance(mlflow_features.source, FileSystemDatasetSource)
+
+
+def test_from_numpy_features_and_targets(tmp_path):
+    features = np.array([[1, 2, 3], [3, 2, 1], [2, 3, 1]])
+    targets = np.array([4, 5, 6])
+    path = tmp_path / "temp.csv"
+    pd.DataFrame(features).to_csv(path)
+    mlflow_ds = mlflow.data.from_numpy(features, targets=targets, source=path)
+
+    assert isinstance(mlflow_ds, NumpyDataset)
+    assert np.array_equal(mlflow_ds.features, features)
+    assert np.array_equal(mlflow_ds.targets, targets)
+    assert mlflow_ds.schema == TensorDatasetSchema(
+        features=_infer_schema(features), targets=_infer_schema(targets)
+    )
+    assert mlflow_ds.profile == {
+        "features_shape": features.shape,
+        "features_size": features.size,
+        "features_nbytes": features.nbytes,
+        "targets_shape": targets.shape,
+        "targets_size": targets.size,
+        "targets_nbytes": targets.nbytes,
+    }
+
+    assert isinstance(mlflow_ds.source, FileSystemDatasetSource)
+
+
+def test_from_numpy_no_source_specified():
+    features = np.array([1, 2, 3])
+    mlflow_features = mlflow.data.from_numpy(features)
+
+    assert isinstance(mlflow_features, NumpyDataset)
+
+    assert isinstance(mlflow_features.source, CodeDatasetSource)
+    assert "mlflow.source.name" in mlflow_features.source.to_json()

tests/data/test_pandas_dataset.py

@@ -0,0 +1,283 @@
+import json
+
+import pandas as pd
+import pytest
+
+import mlflow.data
+from mlflow.data.code_dataset_source import CodeDatasetSource
+from mlflow.data.delta_dataset_source import DeltaDatasetSource
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
+from mlflow.data.pandas_dataset import PandasDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncInputsOutputs
+from mlflow.data.spark_dataset_source import SparkDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.types.schema import Schema
+from mlflow.types.utils import _infer_schema
+
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+@pytest.fixture(scope="module")
+def spark_session():
+    from pyspark.sql import SparkSession
+
+    with (
+        SparkSession.builder
+        .master("local[*]")
+        .config("spark.jars.packages", "io.delta:delta-spark_2.12:3.0.0")
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog"
+        )
+        .getOrCreate()
+    ) as session:
+        yield session
+
+
+def test_conversion_to_json():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+
+    dataset = PandasDataset(
+        df=pd.DataFrame([1, 2, 3], columns=["Numbers"]),
+        source=source,
+        name="testname",
+    )
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+
+    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
+    assert Schema.from_json(schema_json) == dataset.schema
+
+
+def test_digest_property_has_expected_value():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = PandasDataset(
+        df=pd.DataFrame([1, 2, 3], columns=["Numbers"]),
+        source=source,
+        name="testname",
+    )
+    assert dataset.digest == dataset._compute_digest()
+    assert dataset.digest == "31ccce44"
+
+
+def test_df_property():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    df = pd.DataFrame([1, 2, 3], columns=["Numbers"])
+    dataset = PandasDataset(
+        df=df,
+        source=source,
+        name="testname",
+    )
+    assert dataset.df.equals(df)
+
+
+def test_targets_property():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    df_no_targets = pd.DataFrame([1, 2, 3], columns=["Numbers"])
+    dataset_no_targets = PandasDataset(
+        df=df_no_targets,
+        source=source,
+        name="testname",
+    )
+    assert dataset_no_targets._targets is None
+    df_with_targets = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    dataset_with_targets = PandasDataset(
+        df=df_with_targets,
+        source=source,
+        targets="c",
+        name="testname",
+    )
+    assert dataset_with_targets._targets == "c"
+
+
+def test_with_invalid_targets():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    with pytest.raises(
+        MlflowException,
+        match="The specified pandas DataFrame does not contain the specified targets column 'd'.",
+    ):
+        PandasDataset(
+            df=df,
+            source=source,
+            targets="d",
+            name="testname",
+        )
+
+
+def test_to_pyfunc():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    df = pd.DataFrame([1, 2, 3], columns=["Numbers"])
+    dataset = PandasDataset(
+        df=df,
+        source=source,
+        name="testname",
+    )
+    assert isinstance(dataset.to_pyfunc(), PyFuncInputsOutputs)
+
+
+def test_to_pyfunc_with_outputs():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    dataset = PandasDataset(
+        df=df,
+        source=source,
+        targets="c",
+        name="testname",
+    )
+    input_outputs = dataset.to_pyfunc()
+    assert isinstance(input_outputs, PyFuncInputsOutputs)
+    assert input_outputs.inputs.equals(pd.DataFrame([[1, 2], [1, 2]], columns=["a", "b"]))
+    assert input_outputs.outputs.equals(pd.Series([3, 3], name="c"))
+
+
+def test_from_pandas_with_targets(tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    path = tmp_path / "temp.csv"
+    df.to_csv(path)
+    dataset = mlflow.data.from_pandas(df, targets="c", source=path)
+    input_outputs = dataset.to_pyfunc()
+    assert isinstance(input_outputs, PyFuncInputsOutputs)
+    assert input_outputs.inputs.equals(pd.DataFrame([[1, 2], [1, 2]], columns=["a", "b"]))
+    assert input_outputs.outputs.equals(pd.Series([3, 3], name="c"))
+
+
+def test_from_pandas_file_system_datasource(tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    path = tmp_path / "temp.csv"
+    df.to_csv(path)
+    mlflow_df = mlflow.data.from_pandas(df, source=path)
+
+    assert isinstance(mlflow_df, PandasDataset)
+    assert mlflow_df.df.equals(df)
+    assert mlflow_df.schema == _infer_schema(df)
+    assert mlflow_df.profile == {
+        "num_rows": len(df),
+        "num_elements": df.size,
+    }
+
+    assert isinstance(mlflow_df.source, FileSystemDatasetSource)
+
+
+def test_from_pandas_spark_datasource(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    spark_datasource = SparkDatasetSource(path=path)
+    mlflow_df = mlflow.data.from_pandas(df, source=spark_datasource)
+
+    assert isinstance(mlflow_df, PandasDataset)
+    assert mlflow_df.df.equals(df)
+    assert mlflow_df.schema == _infer_schema(df)
+    assert mlflow_df.profile == {
+        "num_rows": len(df),
+        "num_elements": df.size,
+    }
+
+    assert isinstance(mlflow_df.source, SparkDatasetSource)
+
+
+def test_from_pandas_delta_datasource(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.delta")
+    df_spark.write.format("delta").mode("overwrite").save(path)
+
+    delta_datasource = DeltaDatasetSource(path=path)
+    mlflow_df = mlflow.data.from_pandas(df, source=delta_datasource)
+
+    assert isinstance(mlflow_df, PandasDataset)
+    assert mlflow_df.df.equals(df)
+    assert mlflow_df.schema == _infer_schema(df)
+    assert mlflow_df.profile == {
+        "num_rows": len(df),
+        "num_elements": df.size,
+    }
+
+    assert isinstance(mlflow_df.source, DeltaDatasetSource)
+
+
+def test_from_pandas_no_source_specified():
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    mlflow_df = mlflow.data.from_pandas(df)
+
+    assert isinstance(mlflow_df, PandasDataset)
+
+    assert isinstance(mlflow_df.source, CodeDatasetSource)
+    assert "mlflow.source.name" in mlflow_df.source.to_json()
+
+
+def test_to_evaluation_dataset():
+    import numpy as np
+
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    dataset = PandasDataset(
+        df=df,
+        source=source,
+        targets="c",
+        name="testname",
+    )
+    evaluation_dataset = dataset.to_evaluation_dataset()
+
+    assert evaluation_dataset.name is not None
+    assert evaluation_dataset.digest is not None
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert evaluation_dataset.features_data.equals(df.drop("c", axis=1))
+    assert np.array_equal(evaluation_dataset.labels_data, df["c"].to_numpy())
+
+
+def test_df_hashing_with_strings():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+
+    dataset1 = PandasDataset(
+        df=pd.DataFrame([["a", 2, 3], ["a", 2, 3]], columns=["text_column", "b", "c"]),
+        source=source,
+        name="testname",
+    )
+
+    dataset2 = PandasDataset(
+        df=pd.DataFrame([["b", 2, 3], ["b", 2, 3]], columns=["text_column", "b", "c"]),
+        source=source,
+        name="testname",
+    )
+
+    assert dataset1.digest != dataset2.digest
+
+
+def test_df_hashing_with_dicts():
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+
+    df = pd.DataFrame([
+        {"a": [1, 2, 3], "b": {"b": "b", "c": {"c": "c"}}, "c": 3, "d": "d"},
+        {"a": [2, 3], "b": {"b": "b"}, "c": 3, "d": "d"},
+    ])
+    dataset1 = PandasDataset(df=df, source=source, name="testname")
+    dataset2 = PandasDataset(df=df, source=source, name="testname")
+    assert dataset1.digest == dataset2.digest
+
+    evaluation_dataset = dataset1.to_evaluation_dataset()
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert evaluation_dataset.features_data.equals(df)
+    evaluation_dataset2 = dataset2.to_evaluation_dataset()
+    assert evaluation_dataset.hash == evaluation_dataset2.hash

tests/data/test_polars_dataset.py

@@ -0,0 +1,258 @@
+from __future__ import annotations
+
+import json
+import re
+from datetime import date, datetime
+from pathlib import Path
+
+import pandas as pd
+import polars as pl
+import pytest
+
+from mlflow.data.code_dataset_source import CodeDatasetSource
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.filesystem_dataset_source import FileSystemDatasetSource
+from mlflow.data.polars_dataset import PolarsDataset, from_polars, infer_schema
+from mlflow.data.pyfunc_dataset_mixin import PyFuncInputsOutputs
+from mlflow.exceptions import MlflowException
+from mlflow.types.schema import Array, ColSpec, DataType, Object, Property, Schema
+
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+@pytest.fixture(name="source", scope="module")
+def sample_source() -> SampleDatasetSource:
+    source_uri = "test:/my/test/uri"
+    return SampleDatasetSource._resolve(source_uri)
+
+
+def test_infer_schema() -> None:
+    data = [
+        [
+            b"asd",
+            True,
+            datetime(2024, 1, 1, 12, 34, 56, 789),
+            10,
+            10,
+            10,
+            10,
+            10,
+            10,
+            "asd",
+            "😆",
+            "category",
+            "val2",
+            date(2024, 1, 1),
+            10,
+            10,
+            10,
+            [1, 2, 3],
+            [1, 2, 3],
+            {"col1": 1},
+        ]
+    ]
+    schema = {
+        "Binary": pl.Binary,
+        "Boolean": pl.Boolean,
+        "Datetime": pl.Datetime,
+        "Float32": pl.Float32,
+        "Float64": pl.Float64,
+        "Int8": pl.Int8,
+        "Int16": pl.Int16,
+        "Int32": pl.Int32,
+        "Int64": pl.Int64,
+        "String": pl.String,
+        "Utf8": pl.Utf8,
+        "Categorical": pl.Categorical,
+        "Enum": pl.Enum(["val1", "val2"]),
+        "Date": pl.Date,
+        "UInt8": pl.UInt8,
+        "UInt16": pl.UInt16,
+        "UInt32": pl.UInt32,
+        "List": pl.List(pl.Int8),
+        "Array": pl.Array(pl.Int8, 3),
+        "Struct": pl.Struct({"col1": pl.Int8}),
+    }
+    df = pl.DataFrame(data=data, schema=schema)
+
+    assert infer_schema(df) == Schema([
+        ColSpec(name="Binary", type=DataType.binary),
+        ColSpec(name="Boolean", type=DataType.boolean),
+        ColSpec(name="Datetime", type=DataType.datetime),
+        ColSpec(name="Float32", type=DataType.float),
+        ColSpec(name="Float64", type=DataType.double),
+        ColSpec(name="Int8", type=DataType.integer),
+        ColSpec(name="Int16", type=DataType.integer),
+        ColSpec(name="Int32", type=DataType.integer),
+        ColSpec(name="Int64", type=DataType.long),
+        ColSpec(name="String", type=DataType.string),
+        ColSpec(name="Utf8", type=DataType.string),
+        ColSpec(name="Categorical", type=DataType.string),
+        ColSpec(name="Enum", type=DataType.string),
+        ColSpec(name="Date", type=DataType.datetime),
+        ColSpec(name="UInt8", type=DataType.integer),
+        ColSpec(name="UInt16", type=DataType.integer),
+        ColSpec(name="UInt32", type=DataType.long),
+        ColSpec(name="List", type=Array(DataType.integer)),
+        ColSpec(name="Array", type=Array(DataType.integer)),
+        ColSpec(name="Struct", type=Object([Property(name="col1", dtype=DataType.integer)])),
+    ])
+
+
+def test_conversion_to_json(source: SampleDatasetSource) -> None:
+    dataset = PolarsDataset(
+        df=pl.DataFrame([1, 2, 3], schema=["Numbers"]), source=source, name="testname"
+    )
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+
+    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
+    assert Schema.from_json(schema_json) == dataset.schema
+
+
+def test_digest_property_has_expected_value(source: SampleDatasetSource) -> None:
+    dataset = PolarsDataset(df=pl.DataFrame([1, 2, 3], schema=["Numbers"]), source=source)
+    assert dataset.digest == dataset._compute_digest()
+    # Digest value varies across Polars versions due to hash_rows() implementation changes
+    assert re.match(r"^\d+$", dataset.digest)
+
+
+def test_digest_consistent(source: SampleDatasetSource) -> None:
+    dataset1 = PolarsDataset(
+        df=pl.DataFrame({"numbers": [1, 2, 3], "strs": ["a", "b", "c"]}), source=source
+    )
+
+    dataset2 = PolarsDataset(
+        df=pl.DataFrame({"numbers": [2, 3, 1], "strs": ["b", "c", "a"]}), source=source
+    )
+    assert dataset1.digest == dataset2.digest
+
+
+def test_digest_change(source: SampleDatasetSource) -> None:
+    dataset1 = PolarsDataset(
+        df=pl.DataFrame({"numbers": [1, 2, 3], "strs": ["a", "b", "c"]}), source=source
+    )
+
+    dataset2 = PolarsDataset(
+        df=pl.DataFrame({"numbers": [10, 20, 30], "strs": ["aa", "bb", "cc"]}), source=source
+    )
+    assert dataset1.digest != dataset2.digest
+
+
+def test_df_property(source: SampleDatasetSource) -> None:
+    df = pl.DataFrame({"numbers": [1, 2, 3]})
+    dataset = PolarsDataset(df=df, source=source)
+    assert dataset.df.equals(df)
+
+
+def test_targets_none(source: SampleDatasetSource) -> None:
+    df_no_targets = pl.DataFrame({"numbers": [1, 2, 3]})
+    dataset_no_targets = PolarsDataset(df=df_no_targets, source=source)
+    assert dataset_no_targets._targets is None
+
+
+def test_targets_not_none(source: SampleDatasetSource) -> None:
+    df_with_targets = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+    dataset_with_targets = PolarsDataset(df=df_with_targets, source=source, targets="c")
+    assert dataset_with_targets._targets == "c"
+
+
+def test_targets_invalid(source: SampleDatasetSource) -> None:
+    df = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+    with pytest.raises(
+        MlflowException,
+        match="DataFrame does not contain specified targets column: 'd'",
+    ):
+        PolarsDataset(df=df, source=source, targets="d")
+
+
+def test_to_pyfunc_wo_outputs(source: SampleDatasetSource) -> None:
+    df = pl.DataFrame({"numbers": [1, 2, 3]})
+    dataset = PolarsDataset(df=df, source=source)
+
+    input_outputs = dataset.to_pyfunc()
+
+    assert isinstance(input_outputs, PyFuncInputsOutputs)
+    assert len(input_outputs.inputs) == 1
+    assert isinstance(input_outputs.inputs[0], pd.DataFrame)
+    assert input_outputs.inputs[0].equals(pd.DataFrame({"numbers": [1, 2, 3]}))
+
+
+def test_to_pyfunc_with_outputs(source: SampleDatasetSource) -> None:
+    df = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+    dataset = PolarsDataset(df=df, source=source, targets="c")
+
+    input_outputs = dataset.to_pyfunc()
+
+    assert isinstance(input_outputs, PyFuncInputsOutputs)
+    assert len(input_outputs.inputs) == 1
+    assert isinstance(input_outputs.inputs[0], pd.DataFrame)
+    assert input_outputs.inputs[0].equals(pd.DataFrame({"a": [1, 1], "b": [2, 2]}))
+    assert len(input_outputs.outputs) == 1
+    assert isinstance(input_outputs.outputs[0], pd.Series)
+    assert input_outputs.outputs[0].equals(pd.Series([3, 3], name="c"))
+
+
+def test_from_polars_with_targets(tmp_path: Path) -> None:
+    df = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+    path = tmp_path / "temp.csv"
+    df.write_csv(path)
+
+    dataset = from_polars(df, targets="c", source=str(path))
+    input_outputs = dataset.to_pyfunc()
+
+    assert isinstance(input_outputs, PyFuncInputsOutputs)
+    assert len(input_outputs.inputs) == 1
+    assert isinstance(input_outputs.inputs[0], pd.DataFrame)
+    assert input_outputs.inputs[0].equals(pd.DataFrame({"a": [1, 1], "b": [2, 2]}))
+    assert len(input_outputs.outputs) == 1
+    assert isinstance(input_outputs.outputs[0], pd.Series)
+    assert input_outputs.outputs[0].equals(pd.Series([3, 3], name="c"))
+
+
+def test_from_polars_file_system_datasource(tmp_path: Path) -> None:
+    df = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+    path = tmp_path / "temp.csv"
+    df.write_csv(path)
+
+    mlflow_df = from_polars(df, source=str(path))
+
+    assert isinstance(mlflow_df, PolarsDataset)
+    assert mlflow_df.df.equals(df)
+    assert mlflow_df.schema == infer_schema(df)
+    assert mlflow_df.profile == {"num_rows": 2, "num_elements": 6}
+    assert isinstance(mlflow_df.source, FileSystemDatasetSource)
+
+
+def test_from_polars_no_source_specified() -> None:
+    df = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+
+    mlflow_df = from_polars(df)
+
+    assert isinstance(mlflow_df, PolarsDataset)
+    assert isinstance(mlflow_df.source, CodeDatasetSource)
+    assert "mlflow.source.name" in mlflow_df.source.to_json()
+
+
+def test_to_evaluation_dataset(source: SampleDatasetSource) -> None:
+    import numpy as np
+
+    df = pl.DataFrame({"a": [1, 1], "b": [2, 2], "c": [3, 3]})
+    dataset = PolarsDataset(df=df, source=source, targets="c", name="testname")
+    evaluation_dataset = dataset.to_evaluation_dataset()
+
+    assert evaluation_dataset.name is not None
+    assert evaluation_dataset.digest is not None
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert isinstance(evaluation_dataset.features_data, pd.DataFrame)
+    assert evaluation_dataset.features_data.equals(df.drop("c").to_pandas())
+    assert isinstance(evaluation_dataset.labels_data, np.ndarray)
+    assert np.array_equal(evaluation_dataset.labels_data, df["c"].to_numpy())

tests/data/test_spark_dataset.py

@@ -0,0 +1,432 @@
+import json
+import os
+from typing import TYPE_CHECKING, Any
+
+import pandas as pd
+import pytest
+from packaging.version import Version
+
+import mlflow.data
+from mlflow.data.code_dataset_source import CodeDatasetSource
+from mlflow.data.delta_dataset_source import DeltaDatasetSource
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.spark_dataset import SparkDataset
+from mlflow.data.spark_dataset_source import SparkDatasetSource
+from mlflow.exceptions import MlflowException
+from mlflow.types.schema import Schema
+from mlflow.types.utils import _infer_schema
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+
+@pytest.fixture(scope="module")
+def spark_session(tmp_path_factory: pytest.TempPathFactory):
+    import pyspark
+    from pyspark.sql import SparkSession
+
+    pyspark_version = Version(pyspark.__version__)
+    if pyspark_version.major >= 4:
+        delta_package = "io.delta:delta-spark_2.13:4.0.0"
+    else:
+        delta_package = "io.delta:delta-spark_2.12:3.0.0"
+
+    tmp_dir = tmp_path_factory.mktemp("spark_tmp")
+    with (
+        SparkSession.builder
+        .master("local[*]")
+        .config("spark.jars.packages", delta_package)
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog"
+        )
+        .config("spark.sql.warehouse.dir", str(tmp_dir))
+        .getOrCreate()
+    ) as session:
+        yield session
+
+
+@pytest.fixture(autouse=True)
+def drop_tables(spark_session: "SparkSession"):
+    yield
+    for row in spark_session.sql("SHOW TABLES").collect():
+        spark_session.sql(f"DROP TABLE IF EXISTS {row.tableName}")
+
+
+@pytest.fixture
+def df():
+    return pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+
+
+def _assert_dataframes_equal(df1, df2):
+    if df1.schema == df2.schema:
+        diff = df1.exceptAll(df2)
+        assert diff.rdd.isEmpty()
+    else:
+        assert False
+
+
+def _validate_profile_approx_count(parsed_json: dict[str, Any]) -> None:
+    """Validate approx_count in profile data, handling platform/version differences."""
+    # On Windows with certain PySpark versions, Spark datasets may return "unknown" for approx_count
+    # instead of the actual count. We should check that the profile is valid JSON and contains
+    # the expected key, but not assert on the exact value.
+    profile_data = json.loads(parsed_json["profile"])
+    assert "approx_count" in profile_data
+    assert profile_data["approx_count"] in [1, 2, "unknown"]
+
+
+def _check_spark_dataset(dataset, original_df, df_spark, expected_source_type, expected_name=None):
+    assert isinstance(dataset, SparkDataset)
+    _assert_dataframes_equal(dataset.df, df_spark)
+    assert dataset.schema == _infer_schema(original_df)
+    assert isinstance(dataset.profile, dict)
+    approx_count = dataset.profile.get("approx_count")
+    assert isinstance(approx_count, int) or approx_count == "unknown"
+    assert isinstance(dataset.source, expected_source_type)
+    # NB: In real-world scenarios, Spark dataset sources may not match Spark DataFrames precisely.
+    # For example, users may transform Spark DataFrames after loading contents from source files.
+    # To ensure that source loading works properly for the purpose of the test cases in this suite,
+    # we require the source to match the DataFrame and make the following equality assertion
+    _assert_dataframes_equal(dataset.source.load(), df_spark)
+    if expected_name is not None:
+        assert dataset.name == expected_name
+
+
+def test_conversion_to_json_spark_dataset_source(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    source = SparkDatasetSource(path=path)
+
+    dataset = SparkDataset(
+        df=df_spark,
+        source=source,
+        name="testname",
+    )
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    _validate_profile_approx_count(parsed_json)
+
+    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
+    assert Schema.from_json(schema_json) == dataset.schema
+
+
+def test_conversion_to_json_delta_dataset_source(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.format("delta").save(path)
+
+    source = DeltaDatasetSource(path=path)
+
+    dataset = SparkDataset(
+        df=df_spark,
+        source=source,
+        name="testname",
+    )
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    _validate_profile_approx_count(parsed_json)
+
+    schema_json = json.dumps(json.loads(parsed_json["schema"])["mlflow_colspec"])
+    assert Schema.from_json(schema_json) == dataset.schema
+
+
+def test_digest_property_has_expected_value(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    source = SparkDatasetSource(path=path)
+
+    dataset = SparkDataset(
+        df=df_spark,
+        source=source,
+        name="testname",
+    )
+    assert dataset.digest == dataset._compute_digest()
+    # Note that digests are stable within a session, but may not be stable across sessions
+    # Hence we are not checking the digest value here
+
+
+def test_df_property_has_expected_value(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    source = SparkDatasetSource(path=path)
+
+    dataset = SparkDataset(
+        df=df_spark,
+        source=source,
+        name="testname",
+    )
+    assert dataset.df == df_spark
+
+
+def test_targets_property(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    source = SparkDatasetSource(path=path)
+    dataset_no_targets = SparkDataset(
+        df=df_spark,
+        source=source,
+        name="testname",
+    )
+    assert dataset_no_targets.targets is None
+    dataset_with_targets = SparkDataset(
+        df=df_spark,
+        source=source,
+        targets="c",
+        name="testname",
+    )
+    assert dataset_with_targets.targets == "c"
+
+    with pytest.raises(
+        MlflowException,
+        match="The specified Spark dataset does not contain the specified targets column",
+    ):
+        SparkDataset(
+            df=df_spark,
+            source=source,
+            targets="nonexistent",
+            name="testname",
+        )
+
+
+def test_predictions_property(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    source = SparkDatasetSource(path=path)
+    dataset_no_predictions = SparkDataset(
+        df=df_spark,
+        source=source,
+        name="testname",
+    )
+    assert dataset_no_predictions.predictions is None
+    dataset_with_predictions = SparkDataset(
+        df=df_spark,
+        source=source,
+        predictions="b",
+        name="testname",
+    )
+    assert dataset_with_predictions.predictions == "b"
+
+    with pytest.raises(
+        MlflowException,
+        match="The specified Spark dataset does not contain the specified predictions column",
+    ):
+        SparkDataset(
+            df=df_spark,
+            source=source,
+            predictions="nonexistent",
+            name="testname",
+        )
+
+
+def test_from_spark_no_source_specified(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    mlflow_df = mlflow.data.from_spark(df_spark)
+
+    assert isinstance(mlflow_df, SparkDataset)
+
+    assert isinstance(mlflow_df.source, CodeDatasetSource)
+    assert "mlflow.source.name" in mlflow_df.source.to_json()
+
+
+def test_from_spark_with_sql_and_version(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+    with pytest.raises(
+        MlflowException,
+        match="`version` may not be specified when `sql` is specified. `version` may only be"
+        " specified when `table_name` or `path` is specified.",
+    ):
+        mlflow.data.from_spark(df_spark, sql="SELECT * FROM table", version=1)
+
+
+def test_from_spark_path(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    dir_path = str(tmp_path / "df_dir")
+    df_spark.write.parquet(dir_path)
+    assert os.path.isdir(dir_path)
+
+    mlflow_df_from_dir = mlflow.data.from_spark(df_spark, path=dir_path)
+    _check_spark_dataset(mlflow_df_from_dir, df, df_spark, SparkDatasetSource)
+
+    file_path = str(tmp_path / "df.parquet")
+    df_spark.toPandas().to_parquet(file_path)
+    assert not os.path.isdir(file_path)
+
+    mlflow_df_from_file = mlflow.data.from_spark(df_spark, path=file_path)
+    _check_spark_dataset(mlflow_df_from_file, df, df_spark, SparkDatasetSource)
+
+
+def test_from_spark_delta_path(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.delta")
+    df_spark.write.format("delta").save(path)
+
+    mlflow_df = mlflow.data.from_spark(df_spark, path=path)
+
+    _check_spark_dataset(mlflow_df, df, df_spark, DeltaDatasetSource)
+
+
+def test_from_spark_sql(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.createOrReplaceTempView("table")
+
+    mlflow_df = mlflow.data.from_spark(df_spark, sql="SELECT * FROM table")
+
+    _check_spark_dataset(mlflow_df, df, df_spark, SparkDatasetSource)
+
+
+def test_from_spark_table_name(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.createOrReplaceTempView("my_spark_table")
+
+    mlflow_df = mlflow.data.from_spark(df_spark, table_name="my_spark_table")
+
+    _check_spark_dataset(mlflow_df, df, df_spark, SparkDatasetSource)
+
+
+def test_from_spark_table_name_with_version(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.createOrReplaceTempView("my_spark_table")
+
+    with pytest.raises(
+        MlflowException,
+        match="Version '1' was specified, but could not find a Delta table "
+        "with name 'my_spark_table'",
+    ):
+        mlflow.data.from_spark(df_spark, table_name="my_spark_table", version=1)
+
+
+def test_from_spark_delta_table_name(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    # write to delta table
+    df_spark.write.format("delta").mode("overwrite").saveAsTable("my_delta_table")
+
+    mlflow_df = mlflow.data.from_spark(df_spark, table_name="my_delta_table")
+
+    _check_spark_dataset(mlflow_df, df, df_spark, DeltaDatasetSource)
+
+
+def test_from_spark_delta_table_name_and_version(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    # write to delta table
+    df_spark.write.format("delta").mode("overwrite").saveAsTable("my_delta_table")
+
+    mlflow_df = mlflow.data.from_spark(df_spark, table_name="my_delta_table", version=0)
+
+    _check_spark_dataset(mlflow_df, df, df_spark, DeltaDatasetSource)
+
+
+def test_load_delta_with_no_source_info():
+    with pytest.raises(
+        MlflowException,
+        match="Must specify exactly one of `table_name` or `path`.",
+    ):
+        mlflow.data.load_delta()
+
+
+def test_load_delta_with_both_table_name_and_path():
+    with pytest.raises(
+        MlflowException,
+        match="Must specify exactly one of `table_name` or `path`.",
+    ):
+        mlflow.data.load_delta(table_name="my_table", path="my_path")
+
+
+def test_load_delta_path(spark_session, tmp_path, df):
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.delta")
+    df_spark.write.format("delta").mode("overwrite").save(path)
+
+    mlflow_df = mlflow.data.load_delta(path=path)
+
+    _check_spark_dataset(mlflow_df, df, df_spark, DeltaDatasetSource)
+
+
+def test_load_delta_path_with_version(spark_session, tmp_path, df):
+    path = str(tmp_path / "temp.delta")
+
+    df_v0 = pd.DataFrame([[4, 5, 6], [4, 5, 6]], columns=["a", "b", "c"])
+    assert not df_v0.equals(df)
+    df_v0_spark = spark_session.createDataFrame(df_v0)
+    df_v0_spark.write.format("delta").mode("overwrite").save(path)
+
+    # write again to create a new version
+    df_v1_spark = spark_session.createDataFrame(df)
+    df_v1_spark.write.format("delta").mode("overwrite").save(path)
+
+    mlflow_df = mlflow.data.load_delta(path=path, version=1)
+    _check_spark_dataset(mlflow_df, df, df_v1_spark, DeltaDatasetSource)
+
+
+def test_load_delta_table_name(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    # write to delta table
+    df_spark.write.format("delta").mode("overwrite").saveAsTable("my_delta_table")
+
+    mlflow_df = mlflow.data.load_delta(table_name="my_delta_table")
+
+    _check_spark_dataset(mlflow_df, df, df_spark, DeltaDatasetSource, "my_delta_table@v0")
+
+
+def test_load_delta_table_name_with_version(spark_session, df):
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.format("delta").mode("overwrite").saveAsTable("my_delta_table_versioned")
+
+    df2 = pd.DataFrame([[4, 5, 6], [4, 5, 6]], columns=["a", "b", "c"])
+    assert not df2.equals(df)
+    df2_spark = spark_session.createDataFrame(df2)
+    df2_spark.write.format("delta").mode("overwrite").saveAsTable("my_delta_table_versioned")
+
+    mlflow_df = mlflow.data.load_delta(table_name="my_delta_table_versioned", version=1)
+
+    _check_spark_dataset(
+        mlflow_df, df2, df2_spark, DeltaDatasetSource, "my_delta_table_versioned@v1"
+    )
+    pd.testing.assert_frame_equal(mlflow_df.df.toPandas(), df2)
+
+
+def test_to_evaluation_dataset(spark_session, tmp_path, df):
+    import numpy as np
+
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    source = SparkDatasetSource(path=path)
+
+    dataset = SparkDataset(
+        df=df_spark,
+        source=source,
+        targets="c",
+        name="testname",
+        predictions="b",
+    )
+    evaluation_dataset = dataset.to_evaluation_dataset()
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert evaluation_dataset.features_data.equals(df_spark.toPandas()[["a"]])
+    assert np.array_equal(evaluation_dataset.labels_data, df_spark.toPandas()["c"].values)
+    assert np.array_equal(evaluation_dataset.predictions_data, df_spark.toPandas()["b"].values)

tests/data/test_spark_dataset_source.py

@@ -0,0 +1,91 @@
+import json
+
+import pandas as pd
+import pytest
+
+from mlflow.data.dataset_source_registry import get_dataset_source_from_json
+from mlflow.data.spark_dataset_source import SparkDatasetSource
+from mlflow.exceptions import MlflowException
+
+
+@pytest.fixture(scope="module")
+def spark_session():
+    from pyspark.sql import SparkSession
+
+    with (
+        SparkSession.builder
+        .master("local[*]")
+        .config("spark.jars.packages", "io.delta:delta-spark_2.12:3.0.0")
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension")
+        .config(
+            "spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog"
+        )
+        .getOrCreate()
+    ) as session:
+        yield session
+
+
+def test_spark_dataset_source_from_path(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    path = str(tmp_path / "temp.parquet")
+    df_spark.write.parquet(path)
+
+    spark_datasource = SparkDatasetSource(path=path)
+    assert spark_datasource.to_json() == json.dumps({"path": path})
+    loaded_df_spark = spark_datasource.load()
+    assert loaded_df_spark.count() == df_spark.count()
+
+    reloaded_source = get_dataset_source_from_json(
+        spark_datasource.to_json(), source_type=spark_datasource._get_source_type()
+    )
+    assert isinstance(reloaded_source, SparkDatasetSource)
+    assert type(spark_datasource) == type(reloaded_source)
+    assert reloaded_source.to_json() == spark_datasource.to_json()
+
+
+def test_spark_dataset_source_from_table(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.mode("overwrite").saveAsTable("temp", path=tmp_path)
+
+    spark_datasource = SparkDatasetSource(table_name="temp")
+    assert spark_datasource.to_json() == json.dumps({"table_name": "temp"})
+    loaded_df_spark = spark_datasource.load()
+    assert loaded_df_spark.count() == df_spark.count()
+
+    reloaded_source = get_dataset_source_from_json(
+        spark_datasource.to_json(), source_type=spark_datasource._get_source_type()
+    )
+    assert isinstance(reloaded_source, SparkDatasetSource)
+    assert type(spark_datasource) == type(reloaded_source)
+    assert reloaded_source.to_json() == spark_datasource.to_json()
+
+
+def test_spark_dataset_source_from_sql(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.mode("overwrite").saveAsTable("temp_sql", path=tmp_path)
+
+    spark_datasource = SparkDatasetSource(sql="SELECT * FROM temp_sql")
+    assert spark_datasource.to_json() == json.dumps({"sql": "SELECT * FROM temp_sql"})
+    loaded_df_spark = spark_datasource.load()
+    assert loaded_df_spark.count() == df_spark.count()
+
+    reloaded_source = get_dataset_source_from_json(
+        spark_datasource.to_json(), source_type=spark_datasource._get_source_type()
+    )
+    assert isinstance(reloaded_source, SparkDatasetSource)
+    assert type(spark_datasource) == type(reloaded_source)
+    assert reloaded_source.to_json() == spark_datasource.to_json()
+
+
+def test_spark_dataset_source_too_many_inputs(spark_session, tmp_path):
+    df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
+    df_spark = spark_session.createDataFrame(df)
+    df_spark.write.mode("overwrite").saveAsTable("temp", path=tmp_path)
+
+    with pytest.raises(
+        MlflowException, match='Must specify exactly one of "path", "table_name", or "sql"'
+    ):
+        SparkDatasetSource(path=tmp_path, table_name="temp")

tests/data/test_tensorflow_dataset.py

@@ -0,0 +1,395 @@
+import json
+
+import numpy as np
+import pytest
+import tensorflow as tf
+
+import mlflow.data
+from mlflow.data.code_dataset_source import CodeDatasetSource
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.data.pyfunc_dataset_mixin import PyFuncInputsOutputs
+from mlflow.data.schema import TensorDatasetSchema
+from mlflow.data.tensorflow_dataset import TensorFlowDataset
+from mlflow.exceptions import MlflowException
+from mlflow.types.utils import _infer_schema
+
+from tests.resources.data.dataset_source import SampleDatasetSource
+
+
+def test_dataset_construction_validates_features_and_targets():
+    x = np.random.sample((100, 2))
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    tf_tensor = tf.convert_to_tensor(x)
+
+    with pytest.raises(
+        MlflowException,
+        match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*NoneType",
+    ):
+        mlflow.data.from_tensorflow(features=None)
+    with pytest.raises(
+        MlflowException,
+        match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*str",
+    ):
+        mlflow.data.from_tensorflow(features="foo")
+    with pytest.raises(
+        MlflowException,
+        match="features' must be an instance of tf.data.Dataset or a TensorFlow Tensor.*str",
+    ):
+        mlflow.data.from_tensorflow(features="foo", targets=tf_tensor)
+
+    mlflow.data.from_tensorflow(features=tf_tensor, targets=tf_tensor)
+    mlflow.data.from_tensorflow(features=tf_tensor, targets=None)
+    with pytest.raises(
+        MlflowException,
+        match=(
+            "If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
+            " Tensor.*str"
+        ),
+    ):
+        mlflow.data.from_tensorflow(features=tf_tensor, targets="foo")
+    with pytest.raises(
+        MlflowException,
+        match=(
+            "If 'features' is a TensorFlow Tensor, then 'targets' must also be a TensorFlow"
+            " Tensor.*Dataset"
+        ),
+    ):
+        mlflow.data.from_tensorflow(features=tf_tensor, targets=tf_dataset)
+
+    mlflow.data.from_tensorflow(features=tf_dataset, targets=tf_dataset)
+    mlflow.data.from_tensorflow(features=tf_dataset, targets=None)
+    with pytest.raises(
+        MlflowException,
+        match=(
+            "If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
+            " instance of tf.data.Dataset.*str"
+        ),
+    ):
+        mlflow.data.from_tensorflow(features=tf_dataset, targets="foo")
+    with pytest.raises(
+        MlflowException,
+        match=(
+            "If 'features' is an instance of tf.data.Dataset, then 'targets' must also be an"
+            " instance of tf.data.Dataset.*Tensor"
+        ),
+    ):
+        mlflow.data.from_tensorflow(features=tf_dataset, targets=tf_tensor)
+
+
+def test_conversion_to_json():
+    source_uri = "test:/my/test/uri"
+    x = np.random.sample((100, 2))
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+
+    parsed_schema = json.loads(parsed_json["schema"])
+    assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
+
+
+@pytest.mark.parametrize(
+    ("features", "targets"),
+    [
+        (
+            tf.data.Dataset.from_tensors({
+                "a": np.random.sample((100, 2)),
+                "b": np.random.sample((100, 4)),
+            }),
+            tf.data.Dataset.from_tensors({
+                "c": np.random.sample((100, 1)),
+                "d": np.random.sample((100,)),
+            }),
+        ),
+        (
+            tf.data.Dataset.from_tensors((
+                np.random.sample((100, 2)),
+                np.random.sample((100, 4)),
+            )),
+            tf.data.Dataset.from_tensors((
+                np.random.sample((100, 1)),
+                np.random.sample((100,)),
+            )),
+        ),
+        (
+            tf.data.Dataset.from_tensors((
+                np.random.sample((100, 2)),
+                np.random.sample((100, 4)),
+            )),
+            tf.data.Dataset.from_tensors({
+                "c": np.random.sample((100, 1)),
+                "d": np.random.sample((100,)),
+            }),
+        ),
+        (
+            tf.data.Dataset.from_tensors((
+                np.random.sample((100, 2)),
+                np.random.sample((100, 4)),
+            )),
+            None,
+        ),
+    ],
+)
+def test_conversion_to_json_with_multi_tensor_datasets(features, targets):
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=features, targets=targets, source=source, name="testname")
+
+    dataset_json = dataset.to_json()
+    parsed_json = json.loads(dataset_json)
+    assert parsed_json.keys() <= {"name", "digest", "source", "source_type", "schema", "profile"}
+    assert parsed_json["name"] == dataset.name
+    assert parsed_json["digest"] == dataset.digest
+    assert parsed_json["source"] == dataset.source.to_json()
+    assert parsed_json["source_type"] == dataset.source._get_source_type()
+    assert parsed_json["profile"] == json.dumps(dataset.profile)
+
+    parsed_schema = json.loads(parsed_json["schema"])
+    assert TensorDatasetSchema.from_dict(parsed_schema) == dataset.schema
+
+
+def test_schema_and_profile_with_multi_tensor_tuple_datasets():
+    features_dataset = tf.data.Dataset.from_tensors((
+        np.random.sample((100, 2)),
+        np.random.sample((100, 4)),
+    ))
+    targets_dataset = tf.data.Dataset.from_tensors((
+        np.random.sample((100, 1)),
+        np.random.sample((100,)),
+    ))
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(
+        features=features_dataset, targets=targets_dataset, source=source, name="testname"
+    )
+    assert dataset.schema.features == _infer_schema({
+        "0": np.random.sample((100, 2)),
+        "1": np.random.sample((100, 4)),
+    })
+    assert dataset.schema.targets == _infer_schema({
+        "0": np.random.sample((100, 1)),
+        "1": np.random.sample((100,)),
+    })
+    assert dataset.profile == {
+        "features_cardinality": 1,
+        "targets_cardinality": 1,
+    }
+    assert dataset.profile == {
+        "features_cardinality": features_dataset.cardinality().numpy(),
+        "targets_cardinality": targets_dataset.cardinality().numpy(),
+    }
+
+
+def test_schema_and_profile_with_multi_tensor_dict_datasets():
+    features_dataset = tf.data.Dataset.from_tensors({
+        "a": np.random.sample((100, 2)),
+        "b": np.random.sample((100, 4)),
+    })
+    targets_dataset = tf.data.Dataset.from_tensors({
+        "c": np.random.sample((100, 1)),
+        "d": np.random.sample((100,)),
+    })
+    source_uri = "test:/my/test/uri"
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(
+        features=features_dataset, targets=targets_dataset, source=source, name="testname"
+    )
+    assert dataset.schema.features == _infer_schema({
+        "a": np.random.sample((100, 2)),
+        "b": np.random.sample((100, 4)),
+    })
+    assert dataset.schema.targets == _infer_schema({
+        "c": np.random.sample((100, 1)),
+        "d": np.random.sample((100,)),
+    })
+    assert dataset.profile == {
+        "features_cardinality": 1,
+        "targets_cardinality": 1,
+    }
+    assert dataset.profile == {
+        "features_cardinality": features_dataset.cardinality().numpy(),
+        "targets_cardinality": targets_dataset.cardinality().numpy(),
+    }
+
+
+def test_digest_property_has_expected_value():
+    source_uri = "test:/my/test/uri"
+    x = [[1, 2, 3], [4, 5, 6]]
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
+    assert dataset.digest == dataset._compute_digest()
+    assert dataset.digest == "666a9820"
+
+
+def test_data_property_has_expected_value():
+    source_uri = "test:/my/test/uri"
+    x = [[1, 2, 3], [4, 5, 6]]
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
+    assert dataset.data == tf_dataset
+
+
+def test_source_property_has_expected_value():
+    source_uri = "test:/my/test/uri"
+    x = [[1, 2, 3], [4, 5, 6]]
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
+    assert dataset.source == source
+
+
+def test_profile_property_has_expected_value_dataset():
+    source_uri = "test:/my/test/uri"
+    x = [[1, 2, 3], [4, 5, 6]]
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
+    assert dataset.profile == {
+        "features_cardinality": tf_dataset.cardinality().numpy(),
+    }
+
+
+def test_profile_property_has_expected_value_tensors():
+    source_uri = "test:/my/test/uri"
+    x = [[1, 2, 3], [4, 5, 6]]
+    tf_tensor = tf.convert_to_tensor(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_tensor, source=source, name="testname")
+    assert dataset.profile == {
+        "features_cardinality": tf.size(tf_tensor).numpy(),
+    }
+
+
+def test_to_pyfunc():
+    source_uri = "test:/my/test/uri"
+    x = np.random.sample((100, 2))
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(features=tf_dataset, source=source, name="testname")
+    assert isinstance(dataset.to_pyfunc(), PyFuncInputsOutputs)
+
+
+def test_to_evaluation_dataset():
+    source_uri = "test:/my/test/uri"
+    x = np.random.sample((2, 2))
+    y = np.random.sample((2, 1))
+    x_tensors = tf.convert_to_tensor(x)
+    y_tensors = tf.convert_to_tensor(y)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(
+        features=x_tensors, source=source, targets=y_tensors, name="testname"
+    )
+    evaluation_dataset = dataset.to_evaluation_dataset()
+    assert isinstance(evaluation_dataset, EvaluationDataset)
+    assert np.array_equal(evaluation_dataset.features_data, dataset.data.numpy())
+    assert np.array_equal(evaluation_dataset.labels_data, dataset.targets.numpy())
+
+
+def test_to_evaluation_dataset_with_tensorflow_dataset_data():
+    source_uri = "test:/my/test/uri"
+    x = np.random.sample((2, 2))
+    y = np.random.sample((2, 1))
+    x_tf_data = tf.data.Dataset.from_tensors(x)
+    y_tf_data = tf.data.Dataset.from_tensors(y)
+    source = SampleDatasetSource._resolve(source_uri)
+    dataset = TensorFlowDataset(
+        features=x_tf_data, source=source, targets=y_tf_data, name="testname"
+    )
+    with pytest.raises(
+        MlflowException, match="Data must be a Tensor to convert to an EvaluationDataset"
+    ):
+        dataset.to_evaluation_dataset()
+
+
+def test_from_tensorflow_dataset_constructs_expected_dataset():
+    x = np.random.sample((100, 2))
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    mlflow_ds = mlflow.data.from_tensorflow(tf_dataset, source="my_source")
+    assert isinstance(mlflow_ds, TensorFlowDataset)
+    assert mlflow_ds.data == tf_dataset
+    assert mlflow_ds.schema == TensorDatasetSchema(
+        features=_infer_schema(next(tf_dataset.as_numpy_iterator()))
+    )
+    assert mlflow_ds.profile == {
+        "features_cardinality": tf_dataset.cardinality().numpy(),
+    }
+
+
+def test_from_tensorflow_dataset_with_targets_constructs_expected_dataset():
+    x = np.random.sample((100, 2))
+    y = np.random.sample((100, 1))
+    tf_dataset_x = tf.data.Dataset.from_tensors(x)
+    tf_dataset_y = tf.data.Dataset.from_tensors(y)
+    mlflow_ds = mlflow.data.from_tensorflow(tf_dataset_x, source="my_source", targets=tf_dataset_y)
+    assert isinstance(mlflow_ds, TensorFlowDataset)
+    assert mlflow_ds.data == tf_dataset_x
+    assert mlflow_ds.targets == tf_dataset_y
+    assert mlflow_ds.schema == TensorDatasetSchema(
+        features=_infer_schema(next(tf_dataset_x.as_numpy_iterator())),
+        targets=_infer_schema(next(tf_dataset_y.as_numpy_iterator())),
+    )
+    assert mlflow_ds.profile == {
+        "features_cardinality": tf_dataset_x.cardinality().numpy(),
+        "targets_cardinality": tf_dataset_y.cardinality().numpy(),
+    }
+
+
+def test_from_tensorflow_tensor_constructs_expected_dataset():
+    x = np.random.sample((100, 2))
+    tf_tensor = tf.convert_to_tensor(x)
+    mlflow_ds = mlflow.data.from_tensorflow(tf_tensor, source="my_source")
+    assert isinstance(mlflow_ds, TensorFlowDataset)
+    # compare if two tensors are equal using tensorflow utils
+    assert tf.reduce_all(tf.math.equal(mlflow_ds.data, tf_tensor))
+    assert mlflow_ds.schema == TensorDatasetSchema(features=_infer_schema(tf_tensor.numpy()))
+    assert mlflow_ds.profile == {
+        "features_cardinality": tf.size(tf_tensor).numpy(),
+    }
+
+
+def test_from_tensorflow_tensor_with_targets_constructs_expected_dataset():
+    x = np.random.sample((100, 2))
+    y = np.random.sample((100, 1))
+    tf_tensor_x = tf.convert_to_tensor(x)
+    tf_tensor_y = tf.convert_to_tensor(y)
+    mlflow_ds = mlflow.data.from_tensorflow(tf_tensor_x, source="my_source", targets=tf_tensor_y)
+    assert isinstance(mlflow_ds, TensorFlowDataset)
+    assert tf.reduce_all(tf.math.equal(mlflow_ds.data, tf_tensor_x))
+    assert tf.reduce_all(tf.math.equal(mlflow_ds.targets, tf_tensor_y))
+    assert mlflow_ds.schema == TensorDatasetSchema(
+        features=_infer_schema(tf_tensor_x.numpy()),
+        targets=_infer_schema(tf_tensor_y.numpy()),
+    )
+    assert mlflow_ds.profile == {
+        "features_cardinality": tf.size(tf_tensor_x).numpy(),
+        "targets_cardinality": tf.size(tf_tensor_y).numpy(),
+    }
+
+
+def test_from_tensorflow_no_source_specified():
+    x = np.random.sample((100, 2))
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    mlflow_ds = mlflow.data.from_tensorflow(tf_dataset)
+
+    assert isinstance(mlflow_ds, TensorFlowDataset)
+
+    assert isinstance(mlflow_ds.source, CodeDatasetSource)
+    assert "mlflow.source.name" in mlflow_ds.source.to_json()
+
+
+def test_digest_computation_succeeds_with_none_element_in_numpy_iterator():
+    x = np.array([[0, 1], [1, 2]])
+    tf_dataset = tf.data.Dataset.from_tensors(x)
+    tf_dataset.as_numpy_iterator = lambda: [None, x]
+    mlflow_ds = mlflow.data.from_tensorflow(tf_dataset)
+    assert mlflow_ds.digest == "bc8ef018"