wehub-resource-sync
244 lines
7.4 KiB
Python
244 lines
7.4 KiB
Python
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data.preprocessor import Preprocessor
|
|
from ray.data.preprocessors import Chain, LabelEncoder, SimpleImputer, StandardScaler
|
|
from ray.data.util.data_batch_conversion import BatchFormat
|
|
|
|
|
|
def test_chain():
|
|
"""Tests basic Chain functionality."""
|
|
col_a = [-1, -1, 1, 1]
|
|
col_b = [1, 1, 1, None]
|
|
col_c = ["sunday", "monday", "tuesday", "tuesday"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
imputer = SimpleImputer(["B"])
|
|
scaler = StandardScaler(["A", "B"])
|
|
encoder = LabelEncoder("C")
|
|
chain = Chain(scaler, imputer, encoder)
|
|
|
|
# Fit data.
|
|
chain.fit(ds)
|
|
# Transform data.
|
|
transformed = chain.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
assert imputer.stats_ == {
|
|
"mean(B)": 0.0,
|
|
}
|
|
assert scaler.stats_ == {
|
|
"mean(A)": 0.0,
|
|
"mean(B)": 1.0,
|
|
"std(A)": 1.0,
|
|
"std(B)": 0.0,
|
|
}
|
|
assert encoder.stats_ == {
|
|
"unique_values(C)": {"monday": 0, "sunday": 1, "tuesday": 2}
|
|
}
|
|
|
|
processed_col_a = [-1.0, -1.0, 1.0, 1.0]
|
|
processed_col_b = [0.0, 0.0, 0.0, 0.0]
|
|
processed_col_c = [1, 0, 2, 2]
|
|
expected_df = pd.DataFrame.from_dict(
|
|
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
|
).astype(out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
# Transform batch.
|
|
pred_col_a = [1, 2, None]
|
|
pred_col_b = [0, None, 2]
|
|
pred_col_c = ["monday", "tuesday", "wednesday"]
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
|
)
|
|
|
|
pred_out_df = chain.transform_batch(pred_in_df)
|
|
|
|
pred_processed_col_a = [1, 2, None]
|
|
pred_processed_col_b = [-1.0, 0.0, 1.0]
|
|
pred_processed_col_c = [0, 2, None]
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_processed_col_a,
|
|
"B": pred_processed_col_b,
|
|
"C": pred_processed_col_c,
|
|
}
|
|
).astype(pred_out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
|
|
|
|
|
def test_nested_chain_state():
|
|
col_a = [-1, -1, 1, 1]
|
|
col_b = [1, 1, 1, None]
|
|
col_c = ["sunday", "monday", "tuesday", "tuesday"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
def create_chain():
|
|
imputer = SimpleImputer(["B"])
|
|
scaler = StandardScaler(["A", "B"])
|
|
encoder = LabelEncoder("C")
|
|
return Chain(Chain(scaler, imputer), encoder)
|
|
|
|
chain = create_chain()
|
|
assert chain.fit_status() == Preprocessor.FitStatus.NOT_FITTED
|
|
|
|
chain = create_chain()
|
|
chain.preprocessors[1].fit(ds)
|
|
assert chain.fit_status() == Preprocessor.FitStatus.PARTIALLY_FITTED
|
|
|
|
chain = create_chain()
|
|
chain.preprocessors[0].fit(ds)
|
|
assert chain.fit_status() == Preprocessor.FitStatus.PARTIALLY_FITTED
|
|
|
|
chain.preprocessors[1].fit(ds)
|
|
assert chain.fit_status() == Preprocessor.FitStatus.FITTED
|
|
|
|
chain = create_chain()
|
|
chain.fit(ds)
|
|
assert chain.fit_status() == Preprocessor.FitStatus.FITTED
|
|
|
|
|
|
def test_nested_chain():
|
|
"""Tests Chain-inside-Chain functionality."""
|
|
col_a = [-1, -1, 1, 1]
|
|
col_b = [1, 1, 1, None]
|
|
col_c = ["sunday", "monday", "tuesday", "tuesday"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
imputer = SimpleImputer(["B"])
|
|
scaler = StandardScaler(["A", "B"])
|
|
encoder = LabelEncoder("C")
|
|
chain = Chain(Chain(scaler, imputer), encoder)
|
|
|
|
# Fit data.
|
|
chain.fit(ds)
|
|
# Transform data.
|
|
transformed = chain.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
assert imputer.stats_ == {
|
|
"mean(B)": 0.0,
|
|
}
|
|
assert scaler.stats_ == {
|
|
"mean(A)": 0.0,
|
|
"mean(B)": 1.0,
|
|
"std(A)": 1.0,
|
|
"std(B)": 0.0,
|
|
}
|
|
assert encoder.stats_ == {
|
|
"unique_values(C)": {"monday": 0, "sunday": 1, "tuesday": 2}
|
|
}
|
|
|
|
processed_col_a = [-1.0, -1.0, 1.0, 1.0]
|
|
processed_col_b = [0.0, 0.0, 0.0, 0.0]
|
|
processed_col_c = [1, 0, 2, 2]
|
|
expected_df = pd.DataFrame.from_dict(
|
|
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
|
).astype(out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
# Transform batch.
|
|
pred_col_a = [1, 2, None]
|
|
pred_col_b = [0, None, 2]
|
|
pred_col_c = ["monday", "tuesday", "wednesday"]
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
|
)
|
|
|
|
pred_out_df = chain.transform_batch(pred_in_df)
|
|
|
|
pred_processed_col_a = [1, 2, None]
|
|
pred_processed_col_b = [-1.0, 0.0, 1.0]
|
|
pred_processed_col_c = [0, 2, None]
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_processed_col_a,
|
|
"B": pred_processed_col_b,
|
|
"C": pred_processed_col_c,
|
|
}
|
|
).astype(pred_out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
|
|
|
|
|
class PreprocessorWithoutTransform(Preprocessor):
|
|
pass
|
|
|
|
|
|
def test_determine_transform_to_use():
|
|
# Test that _determine_transform_to_use doesn't throw any exceptions
|
|
# and selects the transform function of the underlying preprocessor
|
|
# while dealing with the nested Chain case.
|
|
|
|
# Check that error is propagated correctly
|
|
with pytest.raises(NotImplementedError):
|
|
chain = Chain(PreprocessorWithoutTransform())
|
|
chain._determine_transform_to_use()
|
|
|
|
# Should have no errors from here on
|
|
preprocessor = SimpleImputer(["A"])
|
|
chain1 = Chain(preprocessor)
|
|
format1 = chain1._determine_transform_to_use()
|
|
assert format1 == BatchFormat.PANDAS
|
|
|
|
chain2 = Chain(chain1)
|
|
format2 = chain2._determine_transform_to_use()
|
|
|
|
assert format1 == format2
|
|
|
|
|
|
def test_chain_serialization():
|
|
"""Test Chain serialization and deserialization functionality."""
|
|
import ray
|
|
from ray.data.preprocessor import SerializablePreprocessorBase
|
|
from ray.data.preprocessors import Normalizer, StandardScaler
|
|
|
|
# Create and fit chain
|
|
scaler = StandardScaler(columns=["A"])
|
|
normalizer = Normalizer(columns=["A"])
|
|
chain = Chain(scaler, normalizer)
|
|
|
|
df = pd.DataFrame({"A": [1.0, 2.0, 3.0]})
|
|
ds = ray.data.from_pandas(df)
|
|
fitted_chain = chain.fit(ds)
|
|
|
|
# Serialize using CloudPickle
|
|
serialized = fitted_chain.serialize()
|
|
|
|
# Verify it's binary CloudPickle format
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
|
|
|
# Deserialize
|
|
deserialized = Chain.deserialize(serialized)
|
|
|
|
# Verify type and field values
|
|
assert isinstance(deserialized, Chain)
|
|
assert len(deserialized._preprocessors) == 2
|
|
assert isinstance(deserialized._preprocessors[0], StandardScaler)
|
|
assert isinstance(deserialized._preprocessors[1], Normalizer)
|
|
# Verify the StandardScaler is fitted (Normalizer is stateless)
|
|
assert deserialized._preprocessors[0]._fitted
|
|
|
|
# Verify it works correctly
|
|
test_df = pd.DataFrame({"A": [1.5, 2.5]})
|
|
result = deserialized.transform_batch(test_df)
|
|
|
|
# Result should have been transformed by both preprocessors
|
|
assert "A" in result.columns
|
|
assert len(result) == 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|