mlflow--mlflow/tests/data/test_numpy_dataset.py

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()