wehub-resource-sync
338 lines
10 KiB
Python
338 lines
10 KiB
Python
from datetime import datetime
|
|
|
|
import pytest
|
|
from pyspark.sql import Row, SparkSession
|
|
from pyspark.sql.types import (
|
|
ArrayType,
|
|
BinaryType,
|
|
BooleanType,
|
|
DateType,
|
|
DoubleType,
|
|
FloatType,
|
|
IntegerType,
|
|
LongType,
|
|
ShortType,
|
|
StringType,
|
|
StructField,
|
|
StructType,
|
|
TimestampType,
|
|
)
|
|
from pyspark.testing import assertDataFrameEqual
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.models.utils import _enforce_schema
|
|
from mlflow.types import ColSpec, DataType, Schema
|
|
from mlflow.types.schema import Array, Object, Property
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def spark():
|
|
with SparkSession.builder.getOrCreate() as spark:
|
|
yield spark
|
|
|
|
|
|
def test_enforce_schema_spark_dataframe(spark):
|
|
spark_df_schema = StructType([
|
|
StructField("smallint", ShortType(), True),
|
|
StructField("int", IntegerType(), True),
|
|
StructField("bigint", LongType(), True),
|
|
StructField("float", FloatType(), True),
|
|
StructField("double", DoubleType(), True),
|
|
StructField("boolean", BooleanType(), True),
|
|
StructField("date", DateType(), True),
|
|
StructField("timestamp", TimestampType(), True),
|
|
StructField("string", StringType(), True),
|
|
StructField("binary", BinaryType(), True),
|
|
])
|
|
|
|
data = [
|
|
(
|
|
1, # smallint
|
|
2, # int
|
|
1234567890123456789, # bigint
|
|
1.23, # float
|
|
3.456789, # double
|
|
True, # boolean
|
|
datetime(2020, 1, 1), # date
|
|
datetime.now(), # timestamp
|
|
"example string", # string
|
|
bytearray("example binary", "utf-8"), # binary
|
|
)
|
|
]
|
|
|
|
input_schema = Schema([
|
|
ColSpec(DataType.integer, "smallint"),
|
|
ColSpec(DataType.integer, "int"),
|
|
ColSpec(DataType.long, "bigint"),
|
|
ColSpec(DataType.float, "float"),
|
|
ColSpec(DataType.double, "double"),
|
|
ColSpec(DataType.boolean, "boolean"),
|
|
ColSpec(DataType.datetime, "date"),
|
|
ColSpec(DataType.datetime, "timestamp"),
|
|
ColSpec(DataType.string, "string"),
|
|
ColSpec(DataType.binary, "binary"),
|
|
])
|
|
|
|
input_df = spark.createDataFrame(data, spark_df_schema)
|
|
result = _enforce_schema(input_df, input_schema)
|
|
assertDataFrameEqual(input_df, result)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("spark_df_schema", "data", "input_schema"),
|
|
[
|
|
(
|
|
StructType([StructField("query", ArrayType(StringType()), True)]),
|
|
[(["sentence_1", "sentence_2"],)],
|
|
Schema([ColSpec(Array(DataType.string), name="query")]),
|
|
),
|
|
(
|
|
StructType([
|
|
StructField(
|
|
"teststruct",
|
|
StructType([
|
|
StructField("smallint", ShortType(), True),
|
|
StructField("int", IntegerType(), True),
|
|
StructField("bigint", LongType(), True),
|
|
StructField("float", FloatType(), True),
|
|
StructField("double", DoubleType(), True),
|
|
StructField("boolean", BooleanType(), True),
|
|
StructField("date", DateType(), True),
|
|
StructField("timestamp", TimestampType(), True),
|
|
StructField("string", StringType(), True),
|
|
StructField("binary", BinaryType(), True),
|
|
]),
|
|
True,
|
|
)
|
|
]),
|
|
[
|
|
Row(
|
|
teststruct=Row(
|
|
smallint=100,
|
|
int=1000,
|
|
bigint=10000000000,
|
|
float=10.5,
|
|
double=20.5,
|
|
boolean=True,
|
|
date=datetime(2020, 1, 1),
|
|
timestamp=datetime.now(),
|
|
string="example",
|
|
binary=b"binary_data",
|
|
)
|
|
),
|
|
Row(
|
|
teststruct=Row(
|
|
smallint=200,
|
|
int=2000,
|
|
bigint=20000000000,
|
|
float=20.5,
|
|
double=30.5,
|
|
boolean=False,
|
|
date=datetime(2020, 1, 1),
|
|
timestamp=datetime.now(),
|
|
string="sample",
|
|
binary=b"sample_data",
|
|
)
|
|
),
|
|
Row(
|
|
teststruct=Row(
|
|
smallint=300,
|
|
int=3000,
|
|
bigint=30000000000,
|
|
float=30.5,
|
|
double=40.5,
|
|
boolean=True,
|
|
date=datetime(2020, 1, 1),
|
|
timestamp=datetime.now(),
|
|
string="data",
|
|
binary=b"data_binary",
|
|
)
|
|
),
|
|
],
|
|
Schema([
|
|
ColSpec(
|
|
Object([
|
|
Property("smallint", DataType.integer),
|
|
Property("int", DataType.integer),
|
|
Property("bigint", DataType.long),
|
|
Property("float", DataType.float),
|
|
Property("double", DataType.double),
|
|
Property("boolean", DataType.boolean),
|
|
Property("date", DataType.datetime),
|
|
Property("timestamp", DataType.datetime),
|
|
Property("string", DataType.string),
|
|
Property("binary", DataType.binary),
|
|
]),
|
|
"teststruct",
|
|
)
|
|
]),
|
|
),
|
|
(
|
|
StructType([
|
|
StructField(
|
|
"array",
|
|
ArrayType(
|
|
StructType([
|
|
StructField("name", StringType(), True),
|
|
StructField("age", DoubleType(), True),
|
|
])
|
|
),
|
|
True,
|
|
)
|
|
]),
|
|
[
|
|
(
|
|
[
|
|
Row(name="Alice", age=30.0),
|
|
Row(name="Bob", age=25.0),
|
|
Row(name="Catherine", age=35.0),
|
|
],
|
|
)
|
|
],
|
|
Schema([
|
|
ColSpec(
|
|
Array(
|
|
Object([
|
|
Property("name", DataType.string),
|
|
Property("age", DataType.double),
|
|
])
|
|
),
|
|
name="array",
|
|
),
|
|
]),
|
|
),
|
|
(
|
|
StructType([StructField("nested_list", ArrayType(ArrayType(IntegerType())), True)]),
|
|
[
|
|
([[1, 2, 3], [4, 5, 6], [7, 8, 9]],),
|
|
([[10, 11], [12, 13, 14]],),
|
|
],
|
|
Schema([ColSpec(Array(Array(DataType.integer)), name="nested_list")]),
|
|
),
|
|
],
|
|
)
|
|
def test_enforce_schema_spark_dataframe_complex(spark_df_schema, data, input_schema, spark):
|
|
input_df = spark.createDataFrame(data, spark_df_schema)
|
|
result = _enforce_schema(input_df, input_schema)
|
|
assertDataFrameEqual(input_df, result)
|
|
|
|
|
|
def test_enforce_schema_spark_dataframe_missing_col(spark):
|
|
spark_df_schema = StructType([
|
|
StructField("smallint", ShortType(), True),
|
|
StructField("int", IntegerType(), True),
|
|
])
|
|
|
|
data = [
|
|
(
|
|
1, # smallint
|
|
2, # int
|
|
)
|
|
]
|
|
|
|
input_schema = Schema([
|
|
ColSpec(DataType.integer, "smallint"),
|
|
ColSpec(DataType.integer, "int"),
|
|
ColSpec(DataType.long, "bigint"),
|
|
])
|
|
|
|
df = spark.createDataFrame(data, spark_df_schema)
|
|
with pytest.raises(MlflowException, match="Model is missing inputs"):
|
|
_enforce_schema(df, input_schema)
|
|
|
|
|
|
def test_enforce_schema_spark_dataframe_incompatible_type(spark):
|
|
spark_df_schema = StructType([
|
|
StructField("a", ShortType(), True),
|
|
StructField("b", DoubleType(), True),
|
|
])
|
|
|
|
data = [
|
|
(
|
|
1, # a
|
|
2.3, # b
|
|
)
|
|
]
|
|
|
|
input_schema = Schema([
|
|
ColSpec(DataType.integer, "a"),
|
|
ColSpec(DataType.integer, "b"),
|
|
])
|
|
|
|
df = spark.createDataFrame(data, spark_df_schema)
|
|
with pytest.raises(MlflowException, match="Incompatible input types"):
|
|
_enforce_schema(df, input_schema)
|
|
|
|
|
|
def test_enforce_schema_spark_dataframe_incompatible_type_complex(spark):
|
|
spark_df_schema = StructType([
|
|
StructField(
|
|
"teststruct",
|
|
StructType([
|
|
StructField("int", IntegerType(), True),
|
|
StructField("double", DoubleType(), True),
|
|
]),
|
|
)
|
|
])
|
|
|
|
data = [
|
|
Row(
|
|
teststruct=Row(
|
|
int=1000,
|
|
double=20.5,
|
|
)
|
|
)
|
|
]
|
|
|
|
input_schema = Schema([
|
|
ColSpec(
|
|
Object([
|
|
Property("int", DataType.integer),
|
|
Property("double", DataType.string),
|
|
])
|
|
)
|
|
])
|
|
|
|
df = spark.createDataFrame(data, spark_df_schema)
|
|
with pytest.raises(MlflowException, match="Failed to enforce schema"):
|
|
_enforce_schema(df, input_schema)
|
|
|
|
|
|
def test_enforce_schema_spark_dataframe_extra_col(spark):
|
|
spark_df_schema = StructType([
|
|
StructField("a", ShortType(), True),
|
|
StructField("b", DoubleType(), True),
|
|
])
|
|
|
|
data = [
|
|
(
|
|
1, # a
|
|
2.3, # b
|
|
)
|
|
]
|
|
|
|
input_schema = Schema([ColSpec(DataType.integer, "a")])
|
|
|
|
df = spark.createDataFrame(data, spark_df_schema)
|
|
result = _enforce_schema(df, input_schema)
|
|
expected_result = df.drop("b")
|
|
assertDataFrameEqual(result, expected_result)
|
|
|
|
|
|
def test_enforce_schema_spark_dataframe_no_schema(spark):
|
|
data = [
|
|
(
|
|
1, # a
|
|
2.3, # b
|
|
)
|
|
]
|
|
|
|
input_schema = Schema([
|
|
ColSpec(DataType.integer, "a"),
|
|
ColSpec(DataType.double, "b"),
|
|
])
|
|
|
|
df = spark.createDataFrame(data, ["a", "b"])
|
|
with pytest.raises(MlflowException, match="Incompatible input types"):
|
|
_enforce_schema(df, input_schema)
|