wehub-resource-sync
689 lines
24 KiB
Python
689 lines
24 KiB
Python
import base64
|
|
import datetime
|
|
import json
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from google.protobuf.text_format import Parse as ParseTextIntoProto
|
|
|
|
from mlflow.entities import Experiment, Metric
|
|
from mlflow.entities.model_registry import ModelVersion, RegisteredModel
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.model_registry_pb2 import RegisteredModel as ProtoRegisteredModel
|
|
from mlflow.protos.service_pb2 import Experiment as ProtoExperiment
|
|
from mlflow.protos.service_pb2 import Metric as ProtoMetric
|
|
from mlflow.types import ColSpec, DataType, Schema, TensorSpec
|
|
from mlflow.types.schema import Array, Map, Object, Property
|
|
from mlflow.types.utils import _infer_schema
|
|
from mlflow.utils.proto_json_utils import (
|
|
MlflowFailedTypeConversion,
|
|
_CustomJsonEncoder,
|
|
cast_df_types_according_to_schema,
|
|
dataframe_from_parsed_json,
|
|
dataframe_from_raw_json,
|
|
message_to_json,
|
|
parse_dict,
|
|
parse_tf_serving_input,
|
|
)
|
|
|
|
from tests.protos.test_message_pb2 import SampleMessage
|
|
|
|
|
|
def test_message_to_json():
|
|
json_out = message_to_json(Experiment("123", "name", "arty", "active").to_proto())
|
|
assert json.loads(json_out) == {
|
|
"experiment_id": "123",
|
|
"name": "name",
|
|
"artifact_location": "arty",
|
|
"lifecycle_stage": "active",
|
|
"workspace": "default",
|
|
}
|
|
|
|
original_proto_message = RegisteredModel(
|
|
name="model_1",
|
|
creation_timestamp=111,
|
|
last_updated_timestamp=222,
|
|
description="Test model",
|
|
latest_versions=[
|
|
ModelVersion(
|
|
name="mv-1",
|
|
version="1",
|
|
creation_timestamp=333,
|
|
last_updated_timestamp=444,
|
|
description="v 1",
|
|
user_id="u1",
|
|
current_stage="Production",
|
|
source="A/B",
|
|
run_id="9245c6ce1e2d475b82af84b0d36b52f4",
|
|
status="READY",
|
|
status_message=None,
|
|
),
|
|
ModelVersion(
|
|
name="mv-2",
|
|
version="2",
|
|
creation_timestamp=555,
|
|
last_updated_timestamp=666,
|
|
description="v 2",
|
|
user_id="u2",
|
|
current_stage="Staging",
|
|
source="A/C",
|
|
run_id="123",
|
|
status="READY",
|
|
status_message=None,
|
|
),
|
|
],
|
|
).to_proto()
|
|
json_out = message_to_json(original_proto_message)
|
|
json_dict = json.loads(json_out)
|
|
assert json_dict == {
|
|
"name": "model_1",
|
|
"creation_timestamp": 111,
|
|
"last_updated_timestamp": 222,
|
|
"description": "Test model",
|
|
"latest_versions": [
|
|
{
|
|
"name": "mv-1",
|
|
"version": "1",
|
|
"creation_timestamp": 333,
|
|
"last_updated_timestamp": 444,
|
|
"current_stage": "Production",
|
|
"description": "v 1",
|
|
"user_id": "u1",
|
|
"source": "A/B",
|
|
"run_id": "9245c6ce1e2d475b82af84b0d36b52f4",
|
|
"status": "READY",
|
|
},
|
|
{
|
|
"name": "mv-2",
|
|
"version": "2",
|
|
"creation_timestamp": 555,
|
|
"last_updated_timestamp": 666,
|
|
"current_stage": "Staging",
|
|
"description": "v 2",
|
|
"user_id": "u2",
|
|
"source": "A/C",
|
|
"run_id": "123",
|
|
"status": "READY",
|
|
},
|
|
],
|
|
}
|
|
new_proto_message = ProtoRegisteredModel()
|
|
parse_dict(json_dict, new_proto_message)
|
|
assert original_proto_message == new_proto_message
|
|
|
|
test_message = ParseTextIntoProto(
|
|
"""
|
|
field_int32: 11
|
|
field_int64: 12
|
|
field_uint32: 13
|
|
field_uint64: 14
|
|
field_sint32: 15
|
|
field_sint64: 16
|
|
field_fixed32: 17
|
|
field_fixed64: 18
|
|
field_sfixed32: 19
|
|
field_sfixed64: 20
|
|
field_bool: true
|
|
field_string: "Im a string"
|
|
field_with_default1: 111
|
|
field_repeated_int64: [1, 2, 3]
|
|
field_enum: ENUM_VALUE1
|
|
field_inner_message {
|
|
field_inner_int64: 101
|
|
field_inner_repeated_int64: [102, 103]
|
|
}
|
|
field_inner_message {
|
|
field_inner_int64: 104
|
|
field_inner_repeated_int64: [105, 106]
|
|
}
|
|
oneof1: 207
|
|
[mlflow.ExtensionMessage.field_extended_int64]: 100
|
|
field_map1: [{key: 51 value: "52"}, {key: 53 value: "54"}]
|
|
field_map2: [{key: "61" value: 62}, {key: "63" value: 64}]
|
|
field_map3: [{key: 561 value: 562}, {key: 563 value: 564}]
|
|
field_map4: [{key: 71
|
|
value: {field_inner_int64: 72
|
|
field_inner_repeated_int64: [81, 82]
|
|
field_inner_string: "str1"}},
|
|
{key: 73
|
|
value: {field_inner_int64: 74
|
|
field_inner_repeated_int64: 83
|
|
field_inner_string: "str2"}}]
|
|
""",
|
|
SampleMessage(),
|
|
)
|
|
json_out = message_to_json(test_message)
|
|
json_dict = json.loads(json_out)
|
|
assert json_dict == {
|
|
"field_int32": 11,
|
|
"field_int64": 12,
|
|
"field_uint32": 13,
|
|
"field_uint64": 14,
|
|
"field_sint32": 15,
|
|
"field_sint64": 16,
|
|
"field_fixed32": 17,
|
|
"field_fixed64": 18,
|
|
"field_sfixed32": 19,
|
|
"field_sfixed64": 20,
|
|
"field_bool": True,
|
|
"field_string": "Im a string",
|
|
"field_with_default1": 111,
|
|
"field_repeated_int64": [1, 2, 3],
|
|
"field_enum": "ENUM_VALUE1",
|
|
"field_inner_message": [
|
|
{"field_inner_int64": 101, "field_inner_repeated_int64": [102, 103]},
|
|
{"field_inner_int64": 104, "field_inner_repeated_int64": [105, 106]},
|
|
],
|
|
"oneof1": 207,
|
|
# JSON doesn't support non-string keys, so the int keys will be converted to strings.
|
|
"field_map1": {"51": "52", "53": "54"},
|
|
"field_map2": {"63": 64, "61": 62},
|
|
"field_map3": {"561": 562, "563": 564},
|
|
"field_map4": {
|
|
"73": {
|
|
"field_inner_int64": 74,
|
|
"field_inner_repeated_int64": [83],
|
|
"field_inner_string": "str2",
|
|
},
|
|
"71": {
|
|
"field_inner_int64": 72,
|
|
"field_inner_repeated_int64": [81, 82],
|
|
"field_inner_string": "str1",
|
|
},
|
|
},
|
|
"[mlflow.ExtensionMessage.field_extended_int64]": "100",
|
|
}
|
|
new_test_message = SampleMessage()
|
|
parse_dict(json_dict, new_test_message)
|
|
assert new_test_message == test_message
|
|
|
|
|
|
def test_parse_dict():
|
|
in_json = {"experiment_id": "123", "name": "name", "unknown": "field"}
|
|
message = ProtoExperiment()
|
|
parse_dict(in_json, message)
|
|
experiment = Experiment.from_proto(message)
|
|
assert experiment.experiment_id == "123"
|
|
assert experiment.name == "name"
|
|
assert experiment.artifact_location == ""
|
|
|
|
|
|
def test_parse_dict_int_as_string_backcompat():
|
|
in_json = {"timestamp": "123"}
|
|
message = ProtoMetric()
|
|
parse_dict(in_json, message)
|
|
experiment = Metric.from_proto(message)
|
|
assert experiment.timestamp == 123
|
|
|
|
|
|
def assert_result(result, expected_result):
|
|
assert result.keys() == expected_result.keys()
|
|
for key in result:
|
|
assert (result[key] == expected_result[key]).all()
|
|
assert result[key].dtype == expected_result[key].dtype
|
|
|
|
|
|
def test_parse_tf_serving_dictionary():
|
|
# instances are correctly aggregated to dict of input name -> tensor
|
|
tfserving_input = {
|
|
"instances": [
|
|
{"a": "s1", "b": 1.1, "c": [1, 2, 3]},
|
|
{"a": "s2", "b": 2.2, "c": [4, 5, 6]},
|
|
{"a": "s3", "b": 3.3, "c": [7, 8, 9]},
|
|
]
|
|
}
|
|
# Without Schema
|
|
result = parse_tf_serving_input(tfserving_input)
|
|
expected_result_no_schema = {
|
|
"a": np.array(["s1", "s2", "s3"]),
|
|
"b": np.array([1.1, 2.2, 3.3]),
|
|
"c": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]),
|
|
}
|
|
assert_result(result, expected_result_no_schema)
|
|
|
|
# With schema
|
|
schema = Schema([
|
|
TensorSpec(np.dtype("str"), [-1], "a"),
|
|
TensorSpec(np.dtype("float32"), [-1], "b"),
|
|
TensorSpec(np.dtype("int32"), [-1], "c"),
|
|
])
|
|
df_schema = Schema([ColSpec("string", "a"), ColSpec("float", "b"), ColSpec("integer", "c")])
|
|
result = parse_tf_serving_input(tfserving_input, schema)
|
|
expected_result_schema = {
|
|
"a": np.array(["s1", "s2", "s3"], dtype=np.dtype("str")),
|
|
"b": np.array([1.1, 2.2, 3.3], dtype="float32"),
|
|
"c": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype="int32"),
|
|
}
|
|
assert_result(result, expected_result_schema)
|
|
# With df Schema
|
|
result = parse_tf_serving_input(tfserving_input, df_schema)
|
|
assert_result(result, expected_result_schema)
|
|
# With df Schema containing array
|
|
new_schema = _infer_schema(tfserving_input["instances"])
|
|
result = parse_tf_serving_input(tfserving_input, new_schema)
|
|
expected_result = {
|
|
"a": np.array(["s1", "s2", "s3"]),
|
|
"b": np.array([1.1, 2.2, 3.3], dtype="float64"),
|
|
"c": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype="int64"),
|
|
}
|
|
assert_result(result, expected_result)
|
|
|
|
# input provided as a dict
|
|
tfserving_input = {
|
|
"inputs": {
|
|
"a": ["s1", "s2", "s3"],
|
|
"b": [1.1, 2.2, 3.3],
|
|
"c": [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
|
|
}
|
|
}
|
|
# Without Schema
|
|
result = parse_tf_serving_input(tfserving_input)
|
|
assert_result(result, expected_result_no_schema)
|
|
|
|
# With Schema
|
|
result = parse_tf_serving_input(tfserving_input, schema)
|
|
assert_result(result, expected_result_schema)
|
|
|
|
# With df Schema
|
|
result = parse_tf_serving_input(tfserving_input, df_schema)
|
|
assert_result(result, expected_result_schema)
|
|
|
|
|
|
def test_parse_tf_serving_arbitrary_input_dictionary():
|
|
# input provided as a columnar dict with an arbitrary shape for each input, specifically a
|
|
# different 0th dimension.
|
|
tfserving_input_arbitrary = {
|
|
"inputs": {
|
|
"a": [["s1", "s2", "s3"], ["s4", "s5", "s6"]], # [2, 3]
|
|
"b": [1.1, 2.2, 3.3], # [3, ]
|
|
"c": [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]], # [4, 3]
|
|
}
|
|
}
|
|
|
|
schema = Schema([
|
|
TensorSpec(np.dtype("str"), [-1, 3], "a"),
|
|
TensorSpec(np.dtype("float32"), [-1], "b"),
|
|
TensorSpec(np.dtype("int32"), [-1, 4], "c"),
|
|
])
|
|
df_schema = Schema([ColSpec("string", "a"), ColSpec("float", "b"), ColSpec("integer", "c")])
|
|
|
|
expected_result_no_schema_arbitrary = {
|
|
"a": np.array([["s1", "s2", "s3"], ["s4", "s5", "s6"]]),
|
|
"b": np.array([1.1, 2.2, 3.3]),
|
|
"c": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]]),
|
|
}
|
|
expected_result_schema_arbitrary = {
|
|
"a": np.array([["s1", "s2", "s3"], ["s4", "s5", "s6"]], dtype=np.dtype("str")),
|
|
"b": np.array([1.1, 2.2, 3.3], dtype="float32"),
|
|
"c": np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]], dtype="int32"),
|
|
}
|
|
|
|
# Without Schema
|
|
result = parse_tf_serving_input(tfserving_input_arbitrary)
|
|
assert_result(result, expected_result_no_schema_arbitrary)
|
|
|
|
# With Schema
|
|
result = parse_tf_serving_input(tfserving_input_arbitrary, schema)
|
|
assert_result(result, expected_result_schema_arbitrary)
|
|
|
|
# With df Schema
|
|
result = parse_tf_serving_input(tfserving_input_arbitrary, df_schema)
|
|
assert_result(result, expected_result_schema_arbitrary)
|
|
|
|
|
|
def test_parse_tf_serving_single_array():
|
|
def assert_result(result, expected_result):
|
|
assert (result == expected_result).all()
|
|
|
|
# values for each column are properly converted to a tensor
|
|
arr = [
|
|
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
|
|
[[3, 2, 1], [6, 5, 4], [9, 8, 7]],
|
|
]
|
|
tfserving_instances = {"instances": arr}
|
|
tfserving_inputs = {"inputs": arr}
|
|
|
|
# Without schema
|
|
instance_result = parse_tf_serving_input(tfserving_instances)
|
|
assert instance_result.shape == (2, 3, 3)
|
|
assert_result(instance_result, np.array(arr, dtype="int64"))
|
|
|
|
input_result = parse_tf_serving_input(tfserving_inputs)
|
|
assert input_result.shape == (2, 3, 3)
|
|
assert_result(input_result, np.array(arr, dtype="int64"))
|
|
|
|
# Unnamed schema
|
|
schema = Schema([TensorSpec(np.dtype("float32"), [-1])])
|
|
instance_result = parse_tf_serving_input(tfserving_instances, schema)
|
|
assert_result(instance_result, np.array(arr, dtype="float32"))
|
|
|
|
input_result = parse_tf_serving_input(tfserving_inputs, schema)
|
|
assert_result(input_result, np.array(arr, dtype="float32"))
|
|
|
|
# named schema
|
|
schema = Schema([TensorSpec(np.dtype("float32"), [-1], "a")])
|
|
instance_result = parse_tf_serving_input(tfserving_instances, schema)
|
|
assert isinstance(instance_result, dict)
|
|
assert len(instance_result.keys()) == 1
|
|
assert "a" in instance_result
|
|
assert_result(instance_result["a"], np.array(arr, dtype="float32"))
|
|
|
|
input_result = parse_tf_serving_input(tfserving_inputs, schema)
|
|
assert isinstance(input_result, dict)
|
|
assert len(input_result.keys()) == 1
|
|
assert "a" in input_result
|
|
assert_result(input_result["a"], np.array(arr, dtype="float32"))
|
|
|
|
|
|
def test_parse_tf_serving_raises_expected_errors():
|
|
# input is bad if a column value is missing for a row/instance
|
|
tfserving_instances = {
|
|
"instances": [
|
|
{"a": "s1", "b": 1},
|
|
{"a": "s2", "b": 2, "c": [4, 5, 6]},
|
|
{"a": "s3", "b": 3, "c": [7, 8, 9]},
|
|
]
|
|
}
|
|
with pytest.raises(
|
|
MlflowException, match="The length of values for each input/column name are not the same"
|
|
):
|
|
parse_tf_serving_input(tfserving_instances)
|
|
|
|
# cannot specify both instance and inputs
|
|
tfserving_input = {
|
|
"instances": [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
|
|
"inputs": {"a": ["s1", "s2", "s3"], "b": [1, 2, 3], "c": [[1, 2, 3], [4, 5, 6], [7, 8, 9]]},
|
|
}
|
|
match = 'Invalid input. One of "instances" and "inputs" must be specified'
|
|
with pytest.raises(MlflowException, match=match):
|
|
parse_tf_serving_input(tfserving_input)
|
|
|
|
# cannot specify signature name
|
|
tfserving_input = {
|
|
"signature_name": "hello",
|
|
"inputs": {"a": ["s1", "s2", "s3"], "b": [1, 2, 3], "c": [[1, 2, 3], [4, 5, 6], [7, 8, 9]]},
|
|
}
|
|
match = '"signature_name" parameter is currently not supported'
|
|
with pytest.raises(MlflowException, match=match):
|
|
parse_tf_serving_input(tfserving_input)
|
|
|
|
|
|
def test_dataframe_from_json():
|
|
source = pd.DataFrame(
|
|
{
|
|
"boolean": [True, False, True],
|
|
"string": ["a", "b", "c"],
|
|
"float": np.array([1.2, 2.3, 3.4], dtype=np.float32),
|
|
"double": np.array([1.2, 2.3, 3.4], dtype=np.float64),
|
|
"integer": np.array([3, 4, 5], dtype=np.int32),
|
|
"long": np.array([3, 4, 5], dtype=np.int64),
|
|
"binary": [bytes([1, 2, 3]), bytes([4, 5]), bytes([6])],
|
|
"date_string": ["2018-02-03", "1996-03-02", "2021-03-05"],
|
|
},
|
|
columns=[
|
|
"boolean",
|
|
"string",
|
|
"float",
|
|
"double",
|
|
"integer",
|
|
"long",
|
|
"binary",
|
|
"date_string",
|
|
],
|
|
)
|
|
|
|
jsonable_df = pd.DataFrame(source, copy=True)
|
|
jsonable_df["binary"] = jsonable_df["binary"].map(base64.b64encode)
|
|
schema = Schema([
|
|
ColSpec("boolean", "boolean"),
|
|
ColSpec("string", "string"),
|
|
ColSpec("float", "float"),
|
|
ColSpec("double", "double"),
|
|
ColSpec("integer", "integer"),
|
|
ColSpec("long", "long"),
|
|
ColSpec("binary", "binary"),
|
|
ColSpec("string", "date_string"),
|
|
])
|
|
parsed = dataframe_from_raw_json(
|
|
jsonable_df.to_json(orient="split"), pandas_orient="split", schema=schema
|
|
)
|
|
pd.testing.assert_frame_equal(parsed, source)
|
|
parsed = dataframe_from_raw_json(
|
|
jsonable_df.to_json(orient="records"), pandas_orient="records", schema=schema
|
|
)
|
|
pd.testing.assert_frame_equal(parsed, source)
|
|
# try parsing with tensor schema
|
|
tensor_schema = Schema([
|
|
TensorSpec(np.dtype("bool"), [-1], "boolean"),
|
|
TensorSpec(np.dtype("str"), [-1], "string"),
|
|
TensorSpec(np.dtype("float32"), [-1], "float"),
|
|
TensorSpec(np.dtype("float64"), [-1], "double"),
|
|
TensorSpec(np.dtype("int32"), [-1], "integer"),
|
|
TensorSpec(np.dtype("int64"), [-1], "long"),
|
|
TensorSpec(np.dtype(bytes), [-1], "binary"),
|
|
])
|
|
parsed = dataframe_from_raw_json(
|
|
jsonable_df.to_json(orient="split"), pandas_orient="split", schema=tensor_schema
|
|
)
|
|
|
|
# NB: tensor schema does not automatically decode base64 encoded bytes.
|
|
pd.testing.assert_frame_equal(parsed, jsonable_df)
|
|
parsed = dataframe_from_raw_json(
|
|
jsonable_df.to_json(orient="records"), pandas_orient="records", schema=tensor_schema
|
|
)
|
|
|
|
# NB: tensor schema does not automatically decode base64 encoded bytes.
|
|
pd.testing.assert_frame_equal(parsed, jsonable_df)
|
|
|
|
# Test parse with TensorSchema with a single tensor
|
|
tensor_schema = Schema([TensorSpec(np.dtype("float32"), [-1, 3])])
|
|
source = pd.DataFrame(
|
|
{
|
|
"a": np.array([1, 2, 3], dtype=np.float32),
|
|
"b": np.array([4.1, 5.2, 6.3], dtype=np.float32),
|
|
"c": np.array([7, 8, 9], dtype=np.float32),
|
|
},
|
|
columns=["a", "b", "c"],
|
|
)
|
|
pd.testing.assert_frame_equal(
|
|
source,
|
|
dataframe_from_raw_json(
|
|
source.to_json(orient="split"), pandas_orient="split", schema=tensor_schema
|
|
),
|
|
)
|
|
pd.testing.assert_frame_equal(
|
|
source,
|
|
dataframe_from_raw_json(
|
|
source.to_json(orient="records"), pandas_orient="records", schema=tensor_schema
|
|
),
|
|
)
|
|
|
|
schema = Schema([ColSpec("datetime", "datetime")])
|
|
parsed = dataframe_from_raw_json(
|
|
"""
|
|
[
|
|
{"datetime": "2022-01-01T00:00:00"},
|
|
{"datetime": "2022-01-02T03:04:05"}
|
|
]
|
|
""",
|
|
pandas_orient="records",
|
|
schema=schema,
|
|
)
|
|
expected = pd.DataFrame(
|
|
{
|
|
"datetime": pd.to_datetime([
|
|
"2022-01-01T00:00:00",
|
|
"2022-01-02T03:04:05",
|
|
])
|
|
},
|
|
)
|
|
pd.testing.assert_frame_equal(parsed, expected)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("dt", "expected"),
|
|
[
|
|
(datetime.datetime(2022, 1, 1), '"2022-01-01T00:00:00"'),
|
|
(datetime.datetime(2022, 1, 2, 3, 4, 5), '"2022-01-02T03:04:05"'),
|
|
(datetime.date(2022, 1, 1), '"2022-01-01"'),
|
|
(datetime.time(0, 0, 0), '"00:00:00"'),
|
|
(pd.Timestamp(2022, 1, 1), '"2022-01-01T00:00:00"'),
|
|
],
|
|
)
|
|
def test_datetime_encoder(dt, expected):
|
|
assert json.dumps(dt, cls=_CustomJsonEncoder) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("dataframe", "schema", "expected"),
|
|
[
|
|
(
|
|
pd.DataFrame(columns=["foo"], data=[1, 2, 3]),
|
|
Schema([TensorSpec(np.dtype("float64"), [-1], "foo")]),
|
|
np.dtype("float64"),
|
|
),
|
|
(
|
|
pd.DataFrame(columns=["foo"], data=[[[1, 2, 3]], [[4, 5, 6]]]),
|
|
Schema([TensorSpec(np.dtype("float64"), [-1, 1], "foo")]),
|
|
np.dtype("object"),
|
|
),
|
|
(
|
|
pd.DataFrame(index=[1, 2, 3], columns=["foo"], data=[1, 2, 3]),
|
|
Schema([TensorSpec(np.dtype("float64"), [-1], "foo")]),
|
|
np.dtype("float64"),
|
|
),
|
|
(
|
|
pd.DataFrame(columns=["foo"], data=[1, 2, 3]),
|
|
Schema([ColSpec("double", "foo")]),
|
|
np.dtype("float64"),
|
|
),
|
|
],
|
|
)
|
|
def test_cast_df_types_according_to_schema_success(dataframe, schema, expected):
|
|
casted_pdf = cast_df_types_according_to_schema(dataframe, schema)
|
|
assert casted_pdf["foo"].dtype == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("dataframe", "schema", "error_message"),
|
|
[
|
|
(
|
|
pd.DataFrame(columns=["foo"], data=[1, 2, 3]),
|
|
Schema([ColSpec("binary", "foo")]),
|
|
r"TypeError\('encoding without a string argument'\)",
|
|
),
|
|
(
|
|
pd.DataFrame(columns=["foo"], data=["a", "b", "c"]),
|
|
Schema([ColSpec("double", "foo")]),
|
|
r'ValueError\("could not convert string to float: \'a\'"\)',
|
|
),
|
|
],
|
|
)
|
|
def test_cast_df_types_according_to_schema_error_message(dataframe, schema, error_message):
|
|
with pytest.raises(MlflowFailedTypeConversion, match=error_message):
|
|
cast_df_types_according_to_schema(dataframe, schema)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("data", "schema", "instances_data"),
|
|
[
|
|
({"query": "sentence"}, Schema([ColSpec(DataType.string, name="query")]), None),
|
|
(
|
|
{"query": ["sentence_1", "sentence_2"]},
|
|
Schema([ColSpec(Array(DataType.string), name="query")]),
|
|
None,
|
|
),
|
|
(
|
|
{"query": ["sentence_1", "sentence_2"], "table": "some_table"},
|
|
Schema([
|
|
ColSpec(Array(DataType.string), name="query"),
|
|
ColSpec(DataType.string, name="table"),
|
|
]),
|
|
None,
|
|
),
|
|
(
|
|
{"query": [{"name": "value", "age": 10}, {"name": "value"}], "table": ["some_table"]},
|
|
Schema([
|
|
ColSpec(
|
|
Array(
|
|
Object([
|
|
Property("name", DataType.string),
|
|
Property("age", DataType.long, required=False),
|
|
])
|
|
),
|
|
name="query",
|
|
),
|
|
ColSpec(Array(DataType.string), name="table"),
|
|
]),
|
|
None,
|
|
),
|
|
(
|
|
[{"query": "sentence"}, {"query": "sentence"}],
|
|
Schema([ColSpec(DataType.string, name="query")]),
|
|
{"query": ["sentence", "sentence"]},
|
|
),
|
|
(
|
|
[
|
|
{"query": ["sentence_1", "sentence_2"], "table": "some_table"},
|
|
{"query": ["sentence_1", "sentence_2"]},
|
|
],
|
|
Schema([
|
|
ColSpec(Array(DataType.string), name="query"),
|
|
ColSpec(DataType.string, name="table", required=False),
|
|
]),
|
|
{
|
|
"query": [["sentence_1", "sentence_2"], ["sentence_1", "sentence_2"]],
|
|
"table": ["some_table"],
|
|
},
|
|
),
|
|
(
|
|
[
|
|
{"query": {"a": "sentence_1", "b": "sentence_2"}, "table": "some_table"},
|
|
{"query": {"a": "sentence_1"}, "table": "some_table"},
|
|
],
|
|
Schema([
|
|
ColSpec(
|
|
Object([
|
|
Property("a", DataType.string),
|
|
Property("b", DataType.string, required=False),
|
|
]),
|
|
name="query",
|
|
),
|
|
ColSpec(DataType.string, name="table"),
|
|
]),
|
|
{
|
|
"query": [{"a": "sentence_1", "b": "sentence_2"}, {"a": "sentence_1"}],
|
|
"table": ["some_table", "some_table"],
|
|
},
|
|
),
|
|
(
|
|
{
|
|
"query": [{"name": "value", "age": "10"}, {"name": "value"}],
|
|
"table": {"k": "some_table"},
|
|
"data": {"k1": ["a", "b"], "k2": ["c"]},
|
|
},
|
|
Schema([
|
|
ColSpec(
|
|
Array(Map(value_type=DataType.string)),
|
|
name="query",
|
|
),
|
|
ColSpec(Map(value_type=DataType.string), name="table"),
|
|
ColSpec(Map(value_type=Array(DataType.string)), name="data"),
|
|
]),
|
|
None,
|
|
),
|
|
],
|
|
)
|
|
def test_parse_tf_serving_input_for_dictionaries_and_lists_and_maps(data, schema, instances_data):
|
|
np.testing.assert_equal(parse_tf_serving_input({"inputs": data}, schema), data)
|
|
if instances_data is None:
|
|
np.testing.assert_equal(parse_tf_serving_input({"instances": data}, schema), data)
|
|
else:
|
|
np.testing.assert_equal(parse_tf_serving_input({"instances": data}, schema), instances_data)
|
|
df = pd.DataFrame(data) if isinstance(data, list) else pd.DataFrame([data])
|
|
df_split = df.to_dict(orient="split")
|
|
pd.testing.assert_frame_equal(dataframe_from_parsed_json(df_split, "split", schema), df)
|
|
df_records = df.to_dict(orient="records")
|
|
pd.testing.assert_frame_equal(dataframe_from_parsed_json(df_records, "records", schema), df)
|