wehub-resource-sync
259 lines
9.3 KiB
Python
259 lines
9.3 KiB
Python
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())
|