def test_spark_schema_inference(pandas_df_with_all_types):
import pyspark
from pyspark.sql.types import _parse_datatype_string, StructField, StructType
schema = _infer_schema(pandas_df_with_all_types)
assert schema == Schema([ColSpec(x, x) for x in pandas_df_with_all_types.columns])
spark_session = pyspark.sql.SparkSession(pyspark.SparkContext.getOrCreate())
spark_schema = StructType(
[StructField(t.name, _parse_datatype_string(t.name), True)
for t in schema.column_types()])
sparkdf = spark_session.createDataFrame(pandas_df_with_all_types, schema=spark_schema)
schema = _infer_schema(sparkdf)
assert schema == Schema([ColSpec(x, x) for x in pandas_df_with_all_types.columns])
def test_spark_type_mapping(pandas_df_with_all_types):
import pyspark
from pyspark.sql.types import BooleanType, IntegerType, LongType, FloatType, DoubleType, \
StringType, BinaryType
from pyspark.sql.types import StructField, StructType
assert isinstance(DataType.boolean.to_spark(), BooleanType)
assert isinstance(DataType.integer.to_spark(), IntegerType)
assert isinstance(DataType.long.to_spark(), LongType)
assert isinstance(DataType.float.to_spark(), FloatType)
assert isinstance(DataType.double.to_spark(), DoubleType)
assert isinstance(DataType.string.to_spark(), StringType)
assert isinstance(DataType.binary.to_spark(), BinaryType)
schema = _infer_schema(pandas_df_with_all_types)
expected_spark_schema = StructType(
[StructField(t.name, t.to_spark(), True)
for t in schema.column_types()])
actual_spark_schema = schema.as_spark_schema()
assert expected_spark_schema.jsonValue() == actual_spark_schema.jsonValue()
spark_session = pyspark.sql.SparkSession(pyspark.SparkContext.getOrCreate())
sparkdf = spark_session.createDataFrame(pandas_df_with_all_types,
schema=actual_spark_schema)
schema2 = _infer_schema(sparkdf)
assert schema == schema2
# test unnamed columns
schema = Schema([ColSpec(col.type) for col in schema.columns])
expected_spark_schema = StructType(
[StructField(str(i), t.to_spark(), True)
for i, t in enumerate(schema.column_types())])
actual_spark_schema = schema.as_spark_schema()
assert expected_spark_schema.jsonValue() == actual_spark_schema.jsonValue()
# test single unnamed column is mapped to just a single spark type
schema = Schema([ColSpec(DataType.integer)])
spark_type = schema.as_spark_schema()
assert isinstance(spark_type, IntegerType)
def test_model_log():
with TempDir(chdr=True) as tmp:
experiment_id = kiwi.create_experiment("test")
sig = ModelSignature(inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")]))
input_example = {"x": 1, "y": 2}
with kiwi.start_run(experiment_id=experiment_id) as r:
Model.log("some/path", TestFlavor,
signature=sig,
input_example=input_example)
local_path = _download_artifact_from_uri("runs:/{}/some/path".format(r.info.run_id),
output_path=tmp.path(""))
loaded_model = Model.load(os.path.join(local_path, "MLmodel"))
assert loaded_model.run_id == r.info.run_id
assert loaded_model.artifact_path == "some/path"
assert loaded_model.flavors == {
"flavor1": {"a": 1, "b": 2},
"flavor2": {"x": 1, "y": 2},
}
assert loaded_model.signature == sig
path = os.path.join(local_path, loaded_model.saved_input_example_info["artifact_path"])
x = _dataframe_from_json(path)
assert x.to_dict(orient="records")[0] == input_example
def test_model_save_load():
m = Model(artifact_path="some/path",
run_id="123",
flavors={
"flavor1": {"a": 1, "b": 2},
"flavor2": {"x": 1, "y": 2},
},
signature=ModelSignature(
inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")])),
saved_input_example_info={"x": 1, "y": 2})
assert m.get_input_schema() == m.signature.inputs
assert m.get_output_schema() == m.signature.outputs
x = Model(artifact_path="some/other/path", run_id="1234")
assert x.get_input_schema() is None
assert x.get_output_schema() is None
n = Model(artifact_path="some/path",
run_id="123",
flavors={
"flavor1": {"a": 1, "b": 2},
"flavor2": {"x": 1, "y": 2},
},
signature=ModelSignature(
inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]),
outputs=Schema([ColSpec(name=None, type="double")])),
saved_input_example_info={"x": 1, "y": 2})
n.utc_time_created = m.utc_time_created
assert m == n
n.signature = None
assert m != n
with TempDir() as tmp:
m.save(tmp.path("model"))
o = Model.load(tmp.path("model"))
assert m == o
assert m.to_json() == o.to_json()
assert m.to_yaml() == o.to_yaml()
def _infer_schema(data: Any) -> Schema:
"""
Infer an MLflow schema from a dataset.
This method captures the column names and data types from the user data. The signature
represents model input and output as data frames with (optionally) named columns and data
type specified as one of types defined in :py:class:`DataType`. This method will raise
an exception if the user data contains incompatible types or is not passed in one of the
supported formats (containers).
The input should be one of these:
- pandas.DataFrame or pandas.Series
- dictionary of { name -> numpy.ndarray}
- numpy.ndarray
- pyspark.sql.DataFrame
The element types should be mappable to one of :py:class:`mlflow.models.signature.DataType`.
NOTE: Multidimensional (>2d) arrays (aka tensors) are not supported at this time.
:param data: Dataset to infer from.
:return: Schema
"""
if isinstance(data, dict):
res = []
for col in data.keys():
ary = data[col]
if not isinstance(ary, np.ndarray):
raise TypeError(
"Data in the dictionary must be of type numpy.ndarray")
dims = len(ary.shape)
if dims == 1:
res.append(ColSpec(type=_infer_numpy_array(ary), name=col))
else:
raise TensorsNotSupportedException(
"Data in the dictionary must be 1-dimensional, "
"got shape {}".format(ary.shape))
return Schema(res)
elif isinstance(data, pd.Series):
return Schema([ColSpec(type=_infer_numpy_array(data.values))])
elif isinstance(data, pd.DataFrame):
return Schema([
ColSpec(type=_infer_numpy_array(data[col].values), name=col)
for col in data.columns
])
elif isinstance(data, np.ndarray):
if len(data.shape) > 2:
raise TensorsNotSupportedException(
"Attempting to infer schema from numpy array with "
"shape {}".format(data.shape))
if data.dtype == np.object:
data = pd.DataFrame(data).infer_objects()
return Schema([
ColSpec(type=_infer_numpy_array(data[col].values))
for col in data.columns
])
if len(data.shape) == 1:
return Schema([ColSpec(type=_infer_numpy_dtype(data.dtype))])
elif len(data.shape) == 2:
return Schema([ColSpec(type=_infer_numpy_dtype(data.dtype))] *
data.shape[1])
elif _is_spark_df(data):
return Schema([
ColSpec(type=_infer_spark_type(field.dataType), name=field.name)
for field in data.schema.fields
])
raise TypeError(
"Expected one of (pandas.DataFrame, numpy array, "
"dictionary of (name -> numpy.ndarray), pyspark.sql.DataFrame) "
"but got '{}'".format(type(data)))
def test_schema_inference_on_numpy_array(pandas_df_with_all_types):
# drop int and float as we lose type size information when storing as objects and defaults are
# 64b.
pandas_df_with_all_types = pandas_df_with_all_types.drop(columns=["integer", "float"])
schema = _infer_schema(pandas_df_with_all_types.values)
assert schema == Schema([ColSpec(x) for x in pandas_df_with_all_types.columns])
# test objects
schema = _infer_schema(np.array(["a"], dtype=np.object))
assert schema == Schema([ColSpec(DataType.string)])
schema = _infer_schema(np.array([bytes([1])], dtype=np.object))
assert schema == Schema([ColSpec(DataType.binary)])
schema = _infer_schema(np.array([bytearray([1]), None], dtype=np.object))
assert schema == Schema([ColSpec(DataType.binary)])
schema = _infer_schema(np.array([True, None], dtype=np.object))
assert schema == Schema([ColSpec(DataType.boolean)])
schema = _infer_schema(np.array([1.1, None], dtype=np.object))
assert schema == Schema([ColSpec(DataType.double)])
# test bytes
schema = _infer_schema(np.array([bytes([1])], dtype=np.bytes_))
assert schema == Schema([ColSpec(DataType.binary)])
schema = _infer_schema(np.array([bytearray([1])], dtype=np.bytes_))
assert schema == Schema([ColSpec(DataType.binary)])
# test string
schema = _infer_schema(np.array(["a"], dtype=np.str))
assert schema == Schema([ColSpec(DataType.string)])
# test boolean
schema = _infer_schema(np.array([True], dtype=np.bool))
assert schema == Schema([ColSpec(DataType.boolean)])
# test ints
for t in [np.uint8, np.uint16, np.int8, np.int16, np.int32]:
schema = _infer_schema(np.array([1, 2, 3], dtype=t))
assert schema == Schema([ColSpec("integer")])
# test longs
for t in [np.uint32, np.int64]:
schema = _infer_schema(np.array([1, 2, 3], dtype=t))
assert schema == Schema([ColSpec("long")])
# unsigned long is unsupported
with pytest.raises(MlflowException):
_infer_schema(np.array([1, 2, 3], dtype=np.uint64))
# test floats
for t in [np.float16, np.float32]:
schema = _infer_schema(np.array([1.1, 2.2, 3.3], dtype=t))
assert schema == Schema([ColSpec("float")])
# test doubles
schema = _infer_schema(np.array([1.1, 2.2, 3.3], dtype=np.float64))
assert schema == Schema([ColSpec("double")])
# unsupported
if hasattr(np, "float128"):
with pytest.raises(MlflowException):
_infer_schema(np.array([1, 2, 3], dtype=np.float128))
def test_schema_inference_on_dataframe(pandas_df_with_all_types):
schema = _infer_schema(pandas_df_with_all_types)
assert schema == Schema([ColSpec(x, x) for x in pandas_df_with_all_types.columns])
def test_col_spec():
a1 = ColSpec("string", "a")
a2 = ColSpec(DataType.string, "a")
a3 = ColSpec(DataType.integer, "a")
assert a1 != a3
b1 = ColSpec(DataType.string, "b")
assert b1 != a1
assert a1 == a2
with pytest.raises(MlflowException) as ex:
ColSpec("unsupported")
assert "Unsupported type 'unsupported'" in ex.value.message
a4 = ColSpec(**a1.to_dict())
assert a4 == a1
assert ColSpec(**json.loads(json.dumps(a1.to_dict()))) == a1
a5 = ColSpec("string")
a6 = ColSpec("string", None)
assert a5 == a6
assert ColSpec(**json.loads(json.dumps(a5.to_dict()))) == a5