chore: import upstream snapshot with attribution

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"