wehub-resource-sync
1678 lines
60 KiB
Python
1678 lines
60 KiB
Python
from typing import Any, Dict
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
|
from ray.data.exceptions import UserCodeException
|
|
from ray.data.preprocessor import (
|
|
PreprocessorNotFittedException,
|
|
SerializablePreprocessorBase as SerializablePreprocessor,
|
|
)
|
|
from ray.data.preprocessors import (
|
|
Categorizer,
|
|
LabelEncoder,
|
|
MultiHotEncoder,
|
|
OneHotEncoder,
|
|
OrdinalEncoder,
|
|
)
|
|
|
|
|
|
# Helper functions for parameterized OrdinalEncoder tests
|
|
def _create_pandas_stats(unique_values: Dict[str, list]) -> Dict[str, Dict[Any, int]]:
|
|
"""Create stats in pandas dict format: {value: index}."""
|
|
return {
|
|
f"unique_values({col})": {v: i for i, v in enumerate(sorted(values))}
|
|
for col, values in unique_values.items()
|
|
}
|
|
|
|
|
|
def _create_arrow_stats(
|
|
unique_values: Dict[str, list],
|
|
) -> Dict[str, tuple]:
|
|
"""Create stats in Arrow tuple format: (keys_array, values_array)."""
|
|
result = {}
|
|
for col, values in unique_values.items():
|
|
sorted_values = sorted(values)
|
|
keys_array = pa.array(sorted_values)
|
|
values_array = pa.array(range(len(sorted_values)), type=pa.int64())
|
|
result[f"unique_values({col})"] = (keys_array, values_array)
|
|
return result
|
|
|
|
|
|
def _stats_to_dict(stats_value) -> Dict[Any, int]:
|
|
"""Convert stats to dict format regardless of whether it's Arrow or pandas format."""
|
|
if isinstance(stats_value, dict):
|
|
return stats_value
|
|
elif isinstance(stats_value, tuple):
|
|
# Arrow format: (keys_array, values_array)
|
|
keys_array, values_array = stats_value
|
|
return {k.as_py(): v.as_py() for k, v in zip(keys_array, values_array)}
|
|
else:
|
|
raise ValueError(f"Unknown stats format: {type(stats_value)}")
|
|
|
|
|
|
def _assert_stats_equal(actual_stats: Dict, expected_stats: Dict):
|
|
"""Assert that stats are equal, regardless of Arrow or pandas format."""
|
|
for key, expected_value in expected_stats.items():
|
|
assert key in actual_stats, f"Missing key: {key}"
|
|
actual_value = _stats_to_dict(actual_stats[key])
|
|
assert (
|
|
actual_value == expected_value
|
|
), f"Stats mismatch for {key}: expected {expected_value}, got {actual_value}"
|
|
|
|
|
|
def test_ordinal_encoder_strings():
|
|
"""Test the OrdinalEncoder for strings."""
|
|
|
|
input_dataframe = pd.DataFrame({"sex": ["male"] * 2000 + ["female"]})
|
|
|
|
ds = ray.data.from_pandas(input_dataframe)
|
|
encoder = OrdinalEncoder(columns=["sex"])
|
|
encoded_ds = encoder.fit_transform(ds)
|
|
encoded_ds_pd = encoded_ds.to_pandas()
|
|
|
|
# Check if the "sex" column exists and is correctly encoded as integers
|
|
assert (
|
|
"sex" in encoded_ds_pd.columns
|
|
), "The 'sex' column is missing in the encoded DataFrame"
|
|
assert pd.api.types.is_integer_dtype(
|
|
encoded_ds_pd["sex"].dtype
|
|
), "The 'sex' column is not encoded as integers"
|
|
|
|
# Verify that the encoding worked as expected.
|
|
# We expect "male" to be encoded as 0 and "female" as 1
|
|
unique_values = encoded_ds_pd["sex"].unique()
|
|
assert set(unique_values) == {
|
|
0,
|
|
1,
|
|
}, f"Unexpected unique values in 'sex' column: {unique_values}"
|
|
expected_encoding = {"male": 1, "female": 0}
|
|
for original, encoded in zip(input_dataframe["sex"], encoded_ds_pd["sex"]):
|
|
assert (
|
|
encoded == expected_encoding[original]
|
|
), f"Expected {original} to be encoded as {expected_encoding[original]}, but got {encoded}" # noqa: E501
|
|
|
|
|
|
def test_ordinal_encoder_arrow_transform():
|
|
"""Test the OrdinalEncoder _transform_arrow method."""
|
|
# Create test data
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
col_c = [1, 10, 5, 10]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
|
|
encoder = OrdinalEncoder(["B", "C"])
|
|
|
|
# B: sorted unique = [cold, hot, warm] -> indices [0, 1, 2]
|
|
# C: sorted unique = [1, 5, 10] -> indices [0, 1, 2]
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"], "C": [1, 5, 10]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Create Arrow table for transformation
|
|
table = pa.Table.from_pandas(in_df)
|
|
|
|
# Transform using Arrow
|
|
result_table = encoder._transform_arrow(table)
|
|
|
|
# Verify result is an Arrow table
|
|
assert isinstance(result_table, pa.Table)
|
|
|
|
# Convert to pandas for easier comparison
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Expected encoding: sorted unique values get indices 0, 1, 2, ...
|
|
# B: cold=0, hot=1, warm=2
|
|
# C: 1=0, 5=1, 10=2
|
|
expected_col_b = [2, 0, 1, 0] # warm=2, cold=0, hot=1, cold=0
|
|
expected_col_c = [0, 2, 1, 2] # 1=0, 10=2, 5=1, 10=2
|
|
|
|
assert result_df["A"].tolist() == col_a, "Column A should be unchanged"
|
|
assert (
|
|
result_df["B"].tolist() == expected_col_b
|
|
), f"Column B mismatch: {result_df['B'].tolist()}"
|
|
assert (
|
|
result_df["C"].tolist() == expected_col_c
|
|
), f"Column C mismatch: {result_df['C'].tolist()}"
|
|
|
|
|
|
def test_ordinal_encoder_arrow_transform_append_mode():
|
|
"""Test the OrdinalEncoder _transform_arrow method in append mode."""
|
|
col_a = ["red", "green", "blue"]
|
|
col_b = ["warm", "cold", "hot"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
|
|
|
encoder = OrdinalEncoder(["B"], output_columns=["B_encoded"])
|
|
|
|
# B: sorted unique = [cold, hot, warm] -> indices [0, 1, 2]
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Original columns should be unchanged
|
|
assert result_df["A"].tolist() == col_a
|
|
assert result_df["B"].tolist() == col_b
|
|
|
|
# New column should have encoded values
|
|
# B: cold=0, hot=1, warm=2
|
|
expected_b_encoded = [2, 0, 1] # warm=2, cold=0, hot=1
|
|
assert result_df["B_encoded"].tolist() == expected_b_encoded
|
|
|
|
|
|
def test_ordinal_encoder_arrow_transform_unknown_values():
|
|
"""Test the OrdinalEncoder _transform_arrow method with unknown values."""
|
|
encoder = OrdinalEncoder(["B"])
|
|
|
|
# Fit encoder with only "warm" and "cold" (not "unknown")
|
|
# B: sorted unique = [cold, warm] -> indices [0, 1]
|
|
fit_df = pd.DataFrame({"B": ["cold", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Transform data with an unknown value
|
|
test_df = pd.DataFrame({"B": ["warm", "cold", "unknown"]})
|
|
table = pa.Table.from_pandas(test_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# warm=1, cold=0, unknown should be null
|
|
# pc.index_in returns null for values not found
|
|
assert result_df["B"].tolist()[0] == 1 # warm
|
|
assert result_df["B"].tolist()[1] == 0 # cold
|
|
assert pd.isna(result_df["B"].tolist()[2]) # unknown -> null
|
|
|
|
|
|
# =============================================================================
|
|
# Parameterized tests for OrdinalEncoder (testing both pandas and arrow paths)
|
|
# =============================================================================
|
|
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "arrow"])
|
|
def test_ordinal_encoder_transform_scalars(batch_format):
|
|
"""Test OrdinalEncoder transformation for scalar values with both pandas and arrow."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
col_c = [1, 10, 5, 10]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
|
|
encoder = OrdinalEncoder(["B", "C"])
|
|
|
|
# B: sorted unique = [cold, hot, warm] -> indices [0, 1, 2]
|
|
# C: sorted unique = [1, 5, 10] -> indices [0, 1, 2]
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"], "C": [1, 5, 10]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For pandas batch_format test, convert Arrow stats to pandas format
|
|
if batch_format == "pandas":
|
|
unique_values = {"B": ["cold", "hot", "warm"], "C": [1, 5, 10]}
|
|
encoder.stats_ = _create_pandas_stats(unique_values)
|
|
|
|
# Transform using the appropriate method
|
|
if batch_format == "pandas":
|
|
result_df = encoder._transform_pandas(in_df.copy())
|
|
else:
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Expected encoding: sorted unique values get indices 0, 1, 2, ...
|
|
# B: cold=0, hot=1, warm=2
|
|
# C: 1=0, 5=1, 10=2
|
|
expected_col_b = [2, 0, 1, 0] # warm=2, cold=0, hot=1, cold=0
|
|
expected_col_c = [0, 2, 1, 2] # 1=0, 10=2, 5=1, 10=2
|
|
|
|
assert result_df["A"].tolist() == col_a, "Column A should be unchanged"
|
|
assert (
|
|
result_df["B"].tolist() == expected_col_b
|
|
), f"Column B mismatch: {result_df['B'].tolist()}"
|
|
assert (
|
|
result_df["C"].tolist() == expected_col_c
|
|
), f"Column C mismatch: {result_df['C'].tolist()}"
|
|
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "arrow"])
|
|
def test_ordinal_encoder_transform_append_mode(batch_format):
|
|
"""Test OrdinalEncoder append mode with both pandas and arrow."""
|
|
col_a = ["red", "green", "blue"]
|
|
col_b = ["warm", "cold", "hot"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
|
|
|
encoder = OrdinalEncoder(["B"], output_columns=["B_encoded"])
|
|
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For pandas batch_format test, convert Arrow stats to pandas format
|
|
if batch_format == "pandas":
|
|
unique_values = {"B": ["cold", "hot", "warm"]}
|
|
encoder.stats_ = _create_pandas_stats(unique_values)
|
|
|
|
# Transform using the appropriate method
|
|
if batch_format == "pandas":
|
|
result_df = encoder._transform_pandas(in_df.copy())
|
|
else:
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Original columns should be unchanged
|
|
assert result_df["A"].tolist() == col_a
|
|
assert result_df["B"].tolist() == col_b
|
|
|
|
# New column should have encoded values
|
|
# B: cold=0, hot=1, warm=2
|
|
expected_b_encoded = [2, 0, 1] # warm=2, cold=0, hot=1
|
|
assert result_df["B_encoded"].tolist() == expected_b_encoded
|
|
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "arrow"])
|
|
def test_ordinal_encoder_transform_unknown_values(batch_format):
|
|
"""Test OrdinalEncoder with unknown values using both pandas and arrow."""
|
|
encoder = OrdinalEncoder(["B"])
|
|
|
|
# Fit encoder with only "warm" and "cold" (not "unknown")
|
|
fit_df = pd.DataFrame({"B": ["cold", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For pandas batch_format test, convert Arrow stats to pandas format
|
|
if batch_format == "pandas":
|
|
unique_values = {"B": ["cold", "warm"]}
|
|
encoder.stats_ = _create_pandas_stats(unique_values)
|
|
|
|
# Transform data with an unknown value
|
|
test_df = pd.DataFrame({"B": ["warm", "cold", "unknown"]})
|
|
|
|
if batch_format == "pandas":
|
|
result_df = encoder._transform_pandas(test_df.copy())
|
|
else:
|
|
table = pa.Table.from_pandas(test_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# warm=1, cold=0, unknown should be null/None
|
|
assert result_df["B"].tolist()[0] == 1 # warm
|
|
assert result_df["B"].tolist()[1] == 0 # cold
|
|
assert pd.isna(result_df["B"].tolist()[2]) # unknown -> null
|
|
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "arrow"])
|
|
def test_ordinal_encoder_transform_multiple_columns(batch_format):
|
|
"""Test OrdinalEncoder with multiple columns using both pandas and arrow."""
|
|
in_df = pd.DataFrame(
|
|
{
|
|
"color": ["red", "blue", "green", "red"],
|
|
"size": ["small", "large", "medium", "small"],
|
|
"count": [1, 3, 2, 1],
|
|
}
|
|
)
|
|
|
|
encoder = OrdinalEncoder(["color", "size", "count"])
|
|
|
|
fit_df = pd.DataFrame(
|
|
{
|
|
"color": ["blue", "green", "red"],
|
|
"size": ["large", "medium", "small"],
|
|
"count": [1, 2, 3],
|
|
}
|
|
)
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For pandas batch_format test, convert Arrow stats to pandas format
|
|
if batch_format == "pandas":
|
|
unique_values = {
|
|
"color": ["blue", "green", "red"],
|
|
"size": ["large", "medium", "small"],
|
|
"count": [1, 2, 3],
|
|
}
|
|
encoder.stats_ = _create_pandas_stats(unique_values)
|
|
|
|
if batch_format == "pandas":
|
|
result_df = encoder._transform_pandas(in_df.copy())
|
|
else:
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Verify encodings
|
|
# color: blue=0, green=1, red=2 -> [2, 0, 1, 2]
|
|
# size: large=0, medium=1, small=2 -> [2, 0, 1, 2]
|
|
# count: 1=0, 2=1, 3=2 -> [0, 2, 1, 0]
|
|
assert result_df["color"].tolist() == [2, 0, 1, 2]
|
|
assert result_df["size"].tolist() == [2, 0, 1, 2]
|
|
assert result_df["count"].tolist() == [0, 2, 1, 0]
|
|
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "arrow"])
|
|
def test_ordinal_encoder_transform_integers(batch_format):
|
|
"""Test OrdinalEncoder with integer columns using both pandas and arrow."""
|
|
in_df = pd.DataFrame({"values": [100, 50, 200, 50, 100]})
|
|
|
|
encoder = OrdinalEncoder(["values"])
|
|
|
|
fit_df = pd.DataFrame({"values": [50, 100, 200]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For pandas batch_format test, convert Arrow stats to pandas format
|
|
if batch_format == "pandas":
|
|
unique_values = {"values": [50, 100, 200]}
|
|
encoder.stats_ = _create_pandas_stats(unique_values)
|
|
|
|
if batch_format == "pandas":
|
|
result_df = encoder._transform_pandas(in_df.copy())
|
|
else:
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# 50=0, 100=1, 200=2 -> [1, 0, 2, 0, 1]
|
|
assert result_df["values"].tolist() == [1, 0, 2, 0, 1]
|
|
|
|
|
|
def test_ordinal_encoder_list_fallback_to_pandas():
|
|
"""Test that Arrow transform falls back to pandas for list columns."""
|
|
# This test verifies the fallback behavior when Arrow encounters list columns
|
|
col_d = [["warm", "cold"], ["hot"], ["warm", "hot", "cold"]]
|
|
in_df = pd.DataFrame({"D": col_d})
|
|
|
|
encoder = OrdinalEncoder(["D"], encode_lists=True)
|
|
|
|
# Fit encoder on data with list values containing all unique elements
|
|
fit_df = pd.DataFrame({"D": [["cold", "hot", "warm"]]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For list columns with fallback, we need pandas-format stats
|
|
# (Arrow transform will fall back to pandas for list columns)
|
|
encoder.stats_ = {"unique_values(D)": {"cold": 0, "hot": 1, "warm": 2}}
|
|
|
|
# Create Arrow table with list column
|
|
table = pa.Table.from_pandas(in_df)
|
|
|
|
# Verify column is detected as list type
|
|
assert pa.types.is_list(table.schema.field("D").type)
|
|
|
|
# Transform should fall back to pandas and work correctly
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Verify encoding: cold=0, hot=1, warm=2
|
|
expected = [[2, 0], [1], [2, 1, 0]]
|
|
result_lists = [list(arr) for arr in result_df["D"]]
|
|
assert result_lists == expected
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for vectorized Arrow encoding
|
|
# =============================================================================
|
|
|
|
|
|
def test_ordinal_encoder_encode_column_vectorized():
|
|
"""Test _encode_column_vectorized method directly."""
|
|
encoder = OrdinalEncoder(["col"])
|
|
|
|
fit_df = pd.DataFrame({"col": ["a", "b", "c"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Create a chunked array to encode
|
|
column = pa.chunked_array([["b", "a", "c", "a", "b"]])
|
|
|
|
result = encoder._encode_column_vectorized(column, "col")
|
|
|
|
# a=0, b=1, c=2
|
|
assert result.to_pylist() == [1, 0, 2, 0, 1]
|
|
|
|
|
|
def test_ordinal_encoder_encode_column_with_unknown_values():
|
|
"""Test encoding handles unknown values correctly."""
|
|
encoder = OrdinalEncoder(["col"])
|
|
|
|
# Fit encoder with only "a" and "b" (not "c")
|
|
fit_df = pd.DataFrame({"col": ["a", "b"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Column with unknown value "c"
|
|
column = pa.chunked_array([["a", "b", "c"]])
|
|
|
|
result = encoder._encode_column_vectorized(column, "col")
|
|
assert result.to_pylist()[0] == 0 # a
|
|
assert result.to_pylist()[1] == 1 # b
|
|
assert result.to_pylist()[2] is None # c (unknown)
|
|
|
|
|
|
def test_ordinal_encoder_vectorized_multiple_columns():
|
|
"""Test vectorized encoding works correctly with multiple columns."""
|
|
col_a = ["x", "y"] * 50
|
|
col_b = [1, 2, 3] * 34
|
|
col_b = col_b[:100]
|
|
|
|
in_df = pd.DataFrame({"A": col_a, "B": col_b})
|
|
|
|
encoder = OrdinalEncoder(["A", "B"])
|
|
|
|
fit_df = pd.DataFrame({"A": ["x", "y", "x"], "B": [1, 2, 3]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Verify both columns are encoded correctly
|
|
expected_a = [{"x": 0, "y": 1}[v] for v in col_a]
|
|
expected_b = [{1: 0, 2: 1, 3: 2}[v] for v in col_b]
|
|
|
|
assert result_df["A"].tolist() == expected_a
|
|
assert result_df["B"].tolist() == expected_b
|
|
|
|
|
|
def test_ordinal_encoder():
|
|
"""Tests basic OrdinalEncoder functionality."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
col_c = [1, 10, 5, 10]
|
|
col_d = [["warm"], [], ["hot", "warm", "cold"], ["cold", "cold"]]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c, "D": col_d})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = OrdinalEncoder(["B", "C", "D"])
|
|
|
|
# Transform with unfitted preprocessor.
|
|
with pytest.raises(PreprocessorNotFittedException):
|
|
encoder.transform(ds)
|
|
|
|
# Fit data.
|
|
encoder.fit(ds)
|
|
# Stats may be in Arrow tuple format or pandas dict format depending on
|
|
# preferred_batch_format. Use helper to verify regardless of format.
|
|
_assert_stats_equal(
|
|
encoder.stats_,
|
|
{
|
|
"unique_values(B)": {"cold": 0, "hot": 1, "warm": 2},
|
|
"unique_values(C)": {1: 0, 5: 1, 10: 2},
|
|
"unique_values(D)": {"cold": 0, "hot": 1, "warm": 2},
|
|
},
|
|
)
|
|
|
|
# Transform data.
|
|
transformed = encoder.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
processed_col_a = col_a
|
|
processed_col_b = [2, 0, 1, 0]
|
|
processed_col_c = [0, 2, 1, 2]
|
|
processed_col_d = [[2], [], [1, 2, 0], [0, 0]]
|
|
expected_df = ArrowBlockAccessor(
|
|
pa.Table.from_pydict(
|
|
{
|
|
"A": processed_col_a,
|
|
"B": processed_col_b,
|
|
"C": processed_col_c,
|
|
"D": processed_col_d,
|
|
}
|
|
)
|
|
).to_pandas()
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df)
|
|
|
|
# Transform batch.
|
|
pred_col_a = ["blue", "yellow", None]
|
|
pred_col_b = ["cold", "warm", "other"]
|
|
pred_col_c = [10, 1, 20]
|
|
pred_col_d = [["cold", "warm"], [], ["other", "cold"]]
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c, "D": pred_col_d}
|
|
)
|
|
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
|
|
pred_processed_col_a = pred_col_a
|
|
pred_processed_col_b = [0, 2, None]
|
|
pred_processed_col_c = [2, 0, None]
|
|
pred_processed_col_d = [[0, 2], [], [None, 0]]
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_processed_col_a,
|
|
"B": pred_processed_col_b,
|
|
"C": pred_processed_col_c,
|
|
"D": pred_processed_col_d,
|
|
}
|
|
)
|
|
|
|
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df)
|
|
|
|
# append mode
|
|
with pytest.raises(ValueError):
|
|
OrdinalEncoder(columns=["B", "C", "D"], output_columns=["B_encoded"])
|
|
|
|
encoder = OrdinalEncoder(
|
|
columns=["B", "C", "D"], output_columns=["B_encoded", "C_encoded", "D_encoded"]
|
|
)
|
|
encoder.fit(ds)
|
|
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c, "D": pred_col_d}
|
|
)
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_col_a,
|
|
"B": pred_col_b,
|
|
"C": pred_col_c,
|
|
"D": pred_col_d,
|
|
"B_encoded": pred_processed_col_b,
|
|
"C_encoded": pred_processed_col_c,
|
|
"D_encoded": pred_processed_col_d,
|
|
}
|
|
)
|
|
|
|
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
|
|
|
# Test null behavior.
|
|
null_col = [1, None]
|
|
nonnull_col = [1, 1]
|
|
null_df = pd.DataFrame.from_dict({"A": null_col})
|
|
null_ds = ray.data.from_pandas(null_df)
|
|
nonnull_df = pd.DataFrame.from_dict({"A": nonnull_col})
|
|
nonnull_ds = ray.data.from_pandas(nonnull_df)
|
|
null_encoder = OrdinalEncoder(["A"])
|
|
|
|
# Verify fit fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.fit(null_ds)
|
|
null_encoder.fit(nonnull_ds)
|
|
|
|
# Verify transform fails for null values.
|
|
with pytest.raises((UserCodeException, ValueError)):
|
|
null_encoder.transform(null_ds).materialize()
|
|
null_encoder.transform(nonnull_ds)
|
|
|
|
# Verify transform_batch fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.transform_batch(null_df)
|
|
null_encoder.transform_batch(nonnull_df)
|
|
|
|
|
|
def test_ordinal_encoder_no_encode_list():
|
|
"""Tests OrdinalEncoder with encode_lists=False."""
|
|
in_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": ["red", "green", "blue", "red"],
|
|
"B": ["warm", "cold", "hot", "cold"],
|
|
"C": [1, 10, 5, 10],
|
|
"D": [["warm"], [], ["hot", "warm", "cold"], ["cold", "cold"]],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = OrdinalEncoder(["B", "C", "D"], encode_lists=False)
|
|
|
|
# Transform with unfitted preprocessor.
|
|
with pytest.raises(PreprocessorNotFittedException):
|
|
encoder.transform(ds)
|
|
|
|
# Fit data.
|
|
encoder.fit(ds)
|
|
# Stats may be in Arrow tuple format or pandas dict format
|
|
assert _stats_to_dict(encoder.stats_["unique_values(B)"]) == {
|
|
"cold": 0,
|
|
"hot": 1,
|
|
"warm": 2,
|
|
}
|
|
assert _stats_to_dict(encoder.stats_["unique_values(C)"]) == {1: 0, 5: 1, 10: 2}
|
|
hash_dict = _stats_to_dict(encoder.stats_["unique_values(C)"])
|
|
assert len(set(hash_dict.keys())) == len(set(hash_dict.values())) == len(hash_dict)
|
|
assert max(hash_dict.values()) == len(hash_dict) - 1
|
|
|
|
# Transform data.
|
|
transformed = encoder.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
arrow_in_df = ArrowBlockAccessor(pa.Table.from_pandas(in_df)).to_pandas()
|
|
assert out_df["A"].equals(arrow_in_df["A"])
|
|
assert out_df["B"].equals(
|
|
ArrowBlockAccessor(pa.table({"B": [2, 0, 1, 0]})).to_pandas()["B"]
|
|
)
|
|
assert out_df["C"].equals(
|
|
ArrowBlockAccessor(pa.table({"C": [0, 2, 1, 2]})).to_pandas()["C"]
|
|
)
|
|
assert set(out_df["D"].to_list()) == {3, 0, 2, 1}
|
|
|
|
# Transform batch.
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": ["blue", "yellow", None],
|
|
"B": ["cold", "warm", "other"],
|
|
"C": [10, 1, 20],
|
|
"D": [["cold", "cold"], [], ["other", "cold"]],
|
|
}
|
|
)
|
|
|
|
pred_out_df: pd.DataFrame = encoder.transform_batch(pred_in_df)
|
|
assert pred_out_df["A"].equals(pred_in_df["A"])
|
|
assert pred_out_df["B"].equals(pd.Series([0, 2, None]))
|
|
assert pred_out_df["C"].equals(pd.Series([2, 0, None]))
|
|
assert pd.isnull(pred_out_df["D"].iloc[-1]), "Expected last value to be null"
|
|
assert (
|
|
len(pred_out_df["D"].iloc[:-1].dropna().drop_duplicates())
|
|
== len(pred_out_df) - 1
|
|
), "All values excluding last one must be unique and non-null"
|
|
|
|
|
|
def _assert_one_hot_equal(actual_series, expected_values):
|
|
"""Assert one-hot encoded columns are equal, handling both list and numpy array types."""
|
|
assert len(actual_series) == len(expected_values)
|
|
for actual, expected in zip(actual_series, expected_values):
|
|
assert list(actual) == list(expected)
|
|
|
|
|
|
def _assert_list_column_equal(actual_series, expected_series):
|
|
"""Assert list columns are equal, handling Arrow round-trip type changes."""
|
|
assert len(actual_series) == len(expected_series)
|
|
for actual, expected in zip(actual_series, expected_series):
|
|
assert list(actual) == list(expected)
|
|
|
|
|
|
# =============================================================================
|
|
# Tests for OneHotEncoder Arrow transform
|
|
# =============================================================================
|
|
|
|
|
|
def test_one_hot_encoder_arrow_transform():
|
|
"""Test the OneHotEncoder _transform_arrow method."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
col_c = [1, 10, 5, 10]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
|
|
encoder = OneHotEncoder(["B", "C"])
|
|
|
|
# B: cold=0, hot=1, warm=2
|
|
# C: 1=0, 5=1, 10=2
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"], "C": [1, 5, 10]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Create Arrow table for transformation
|
|
table = pa.Table.from_pandas(in_df)
|
|
|
|
# Transform using Arrow
|
|
result_table = encoder._transform_arrow(table)
|
|
|
|
# Verify result is an Arrow table
|
|
assert isinstance(result_table, pa.Table)
|
|
|
|
# Convert to pandas for easier comparison
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Expected one-hot encoding:
|
|
# B: warm=[0,0,1], cold=[1,0,0], hot=[0,1,0]
|
|
# C: 1=[1,0,0], 10=[0,0,1], 5=[0,1,0]
|
|
expected_col_b = [[0, 0, 1], [1, 0, 0], [0, 1, 0], [1, 0, 0]]
|
|
expected_col_c = [[1, 0, 0], [0, 0, 1], [0, 1, 0], [0, 0, 1]]
|
|
|
|
assert result_df["A"].tolist() == col_a, "Column A should be unchanged"
|
|
_assert_one_hot_equal(result_df["B"], expected_col_b)
|
|
_assert_one_hot_equal(result_df["C"], expected_col_c)
|
|
|
|
|
|
def test_one_hot_encoder_arrow_transform_append_mode():
|
|
"""Test the OneHotEncoder _transform_arrow method in append mode."""
|
|
col_a = ["red", "green", "blue"]
|
|
col_b = ["warm", "cold", "hot"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
|
|
|
encoder = OneHotEncoder(["B"], output_columns=["B_encoded"])
|
|
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Original columns should be unchanged
|
|
assert result_df["A"].tolist() == col_a
|
|
assert result_df["B"].tolist() == col_b
|
|
|
|
# New column should have one-hot encoded values
|
|
expected_b_encoded = [[0, 0, 1], [1, 0, 0], [0, 1, 0]]
|
|
_assert_one_hot_equal(result_df["B_encoded"], expected_b_encoded)
|
|
|
|
|
|
def test_one_hot_encoder_arrow_transform_unknown_values():
|
|
"""Test the OneHotEncoder _transform_arrow method with unknown values."""
|
|
encoder = OneHotEncoder(["B"])
|
|
|
|
# Fit encoder with only "warm" and "cold" (not "unknown")
|
|
fit_df = pd.DataFrame({"B": ["cold", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Transform data with an unknown value
|
|
test_df = pd.DataFrame({"B": ["warm", "cold", "unknown"]})
|
|
table = pa.Table.from_pandas(test_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# warm=[0,1], cold=[1,0], unknown=[0,0] (all zeros for unknown)
|
|
_assert_one_hot_equal(result_df["B"], [[0, 1], [1, 0], [0, 0]])
|
|
|
|
|
|
def test_one_hot_encoder_list_fallback_to_pandas():
|
|
"""Test that Arrow transform falls back to pandas for list columns."""
|
|
col_d = [["warm", "cold"], ["hot"], ["warm", "hot", "cold"]]
|
|
in_df = pd.DataFrame({"D": col_d})
|
|
|
|
encoder = OneHotEncoder(["D"])
|
|
|
|
# Fit encoder on data with list values (lists are treated as categories)
|
|
fit_df = pd.DataFrame({"D": [["cold"], ["hot"], ["warm"]]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# For list columns with fallback, we need pandas-format stats
|
|
# (Arrow transform will fall back to pandas for list columns)
|
|
encoder.stats_ = {"unique_values(D)": {("cold",): 0, ("hot",): 1, ("warm",): 2}}
|
|
|
|
# Create Arrow table with list column
|
|
table = pa.Table.from_pandas(in_df)
|
|
|
|
# Verify column is detected as list type
|
|
assert pa.types.is_list(table.schema.field("D").type)
|
|
|
|
# Transform should fall back to pandas and work correctly
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Verify one-hot encoding for list columns (handled by pandas fallback)
|
|
# Each list element maps to one-hot vector
|
|
assert len(result_df["D"]) == 3
|
|
|
|
|
|
def test_one_hot_encoder_multi_chunk_column():
|
|
"""Test OneHotEncoder with multi-chunk ChunkedArray input.
|
|
|
|
This test ensures that _encode_column_one_hot correctly handles ChunkedArrays
|
|
with multiple chunks, which can occur with partitioned or concatenated data.
|
|
The implementation uses zero_copy_only=False when calling to_numpy() for
|
|
compatibility across PyArrow versions.
|
|
"""
|
|
encoder = OneHotEncoder(["col"])
|
|
|
|
fit_df = pd.DataFrame({"col": ["a", "b", "c"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Create a table with a multi-chunk column (simulates partitioned/concatenated data)
|
|
chunk1 = pa.array(["a", "b", "c"])
|
|
chunk2 = pa.array(["b", "a", "c"])
|
|
chunk3 = pa.array(["c", "c", "a"])
|
|
multi_chunk_column = pa.chunked_array([chunk1, chunk2, chunk3])
|
|
|
|
# Verify we have multiple chunks in the input
|
|
assert multi_chunk_column.num_chunks == 3
|
|
|
|
# Create table with the multi-chunk column
|
|
table = pa.table({"col": multi_chunk_column})
|
|
|
|
# Transform using Arrow path - this exercises _encode_column_one_hot
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Verify correct one-hot encoding
|
|
# a=[1,0,0], b=[0,1,0], c=[0,0,1]
|
|
expected = [
|
|
[1, 0, 0], # a
|
|
[0, 1, 0], # b
|
|
[0, 0, 1], # c
|
|
[0, 1, 0], # b
|
|
[1, 0, 0], # a
|
|
[0, 0, 1], # c
|
|
[0, 0, 1], # c
|
|
[0, 0, 1], # c
|
|
[1, 0, 0], # a
|
|
]
|
|
_assert_one_hot_equal(result_df["col"], expected)
|
|
|
|
|
|
@pytest.mark.parametrize("batch_format", ["pandas", "arrow"])
|
|
def test_one_hot_encoder_transform_scalars(batch_format):
|
|
"""Test OneHotEncoder transformation for scalar values with both pandas and arrow."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
|
|
|
encoder = OneHotEncoder(["B"])
|
|
|
|
fit_df = pd.DataFrame({"B": ["cold", "hot", "warm"]})
|
|
encoder.fit(ray.data.from_pandas(fit_df))
|
|
|
|
# Transform using the appropriate method
|
|
if batch_format == "pandas":
|
|
result_df = encoder._transform_pandas(in_df.copy())
|
|
else:
|
|
table = pa.Table.from_pandas(in_df)
|
|
result_table = encoder._transform_arrow(table)
|
|
result_df = result_table.to_pandas()
|
|
|
|
# Expected one-hot encoding: cold=[1,0,0], hot=[0,1,0], warm=[0,0,1]
|
|
expected_col_b = [[0, 0, 1], [1, 0, 0], [0, 1, 0], [1, 0, 0]]
|
|
|
|
assert result_df["A"].tolist() == col_a, "Column A should be unchanged"
|
|
_assert_one_hot_equal(result_df["B"], expected_col_b)
|
|
|
|
|
|
def test_one_hot_encoder():
|
|
"""Tests basic OneHotEncoder functionality."""
|
|
in_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": ["red", "green", "blue", "red"],
|
|
"B": ["warm", "cold", "hot", "cold"],
|
|
"C": [1, 10, 5, 10],
|
|
"D": [["warm"], [], ["hot", "warm", "cold"], ["cold", "cold"]],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = OneHotEncoder(["B", "C", "D"])
|
|
|
|
# Transform with unfitted preprocessor.
|
|
with pytest.raises(PreprocessorNotFittedException):
|
|
encoder.transform(ds)
|
|
|
|
# Fit data.
|
|
encoder.fit(ds)
|
|
|
|
assert encoder.stats_["unique_values(B)"] == {
|
|
"cold": 0,
|
|
"hot": 1,
|
|
"warm": 2,
|
|
}
|
|
assert encoder.stats_["unique_values(C)"] == {1: 0, 5: 1, 10: 2}
|
|
hash_dict = encoder.stats_["unique_values(D)"]
|
|
assert len(set(hash_dict.keys())) == len(set(hash_dict.values())) == len(hash_dict)
|
|
assert max(hash_dict.values()) == len(hash_dict) - 1
|
|
|
|
# Transform data.
|
|
transformed = encoder.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
arrow_in_df = ArrowBlockAccessor(pa.Table.from_pandas(in_df)).to_pandas()
|
|
assert out_df["A"].equals(arrow_in_df["A"])
|
|
_assert_one_hot_equal(out_df["B"], [[0, 0, 1], [1, 0, 0], [0, 1, 0], [1, 0, 0]])
|
|
_assert_one_hot_equal(out_df["C"], [[1, 0, 0], [0, 0, 1], [0, 1, 0], [0, 0, 1]])
|
|
assert {tuple(row) for row in out_df["D"]} == {
|
|
tuple(row)
|
|
for row in pd.Series([[0, 0, 0, 1], [1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0]])
|
|
}
|
|
|
|
# Transform batch.
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": ["blue", "yellow", None],
|
|
"B": ["cold", "warm", "other"],
|
|
"C": [10, 1, 20],
|
|
"D": [["cold", "cold"], [], ["other", "cold"]],
|
|
}
|
|
)
|
|
|
|
pred_out_df: pd.DataFrame = encoder.transform_batch(pred_in_df.copy())
|
|
|
|
assert pred_out_df["A"].equals(pred_in_df["A"])
|
|
_assert_one_hot_equal(pred_out_df["B"], [[1, 0, 0], [0, 0, 1], [0, 0, 0]])
|
|
_assert_one_hot_equal(pred_out_df["C"], [[0, 0, 1], [1, 0, 0], [0, 0, 0]])
|
|
assert list(pred_out_df["D"].iloc[-1]) == [0, 0, 0, 0]
|
|
assert (
|
|
len(
|
|
{
|
|
i
|
|
for row in pred_out_df["D"].iloc[:-1]
|
|
for i, val in enumerate(row)
|
|
if val == 1
|
|
}
|
|
)
|
|
== 2
|
|
)
|
|
|
|
# append mode
|
|
with pytest.raises(ValueError):
|
|
OneHotEncoder(columns=["B", "C", "D"], output_columns=["B_encoded"])
|
|
|
|
encoder = OneHotEncoder(
|
|
columns=["B", "C", "D"],
|
|
output_columns=["B_onehot_encoded", "C_onehot_encoded", "D_onehot_encoded"],
|
|
)
|
|
encoder.fit(ds)
|
|
pred_out_append_df: pd.DataFrame = encoder.transform_batch(pred_in_df.copy())
|
|
assert pred_out_append_df["A"].equals(pred_in_df["A"])
|
|
assert pred_out_append_df["B"].equals(pred_in_df["B"])
|
|
assert pred_out_append_df["C"].equals(pred_in_df["C"])
|
|
# List column D may have type changes after Arrow round-trip
|
|
_assert_list_column_equal(pred_out_append_df["D"], pred_in_df["D"])
|
|
_assert_one_hot_equal(
|
|
pred_out_append_df["B_onehot_encoded"], pred_out_df["B"].tolist()
|
|
)
|
|
_assert_one_hot_equal(
|
|
pred_out_append_df["C_onehot_encoded"], pred_out_df["C"].tolist()
|
|
)
|
|
_assert_one_hot_equal(
|
|
pred_out_append_df["D_onehot_encoded"], pred_out_df["D"].tolist()
|
|
)
|
|
|
|
# Test null behavior.
|
|
null_col = [1, None]
|
|
nonnull_col = [1, 1]
|
|
null_df = pd.DataFrame.from_dict({"A": null_col})
|
|
null_ds = ray.data.from_pandas(null_df)
|
|
nonnull_df = pd.DataFrame.from_dict({"A": nonnull_col})
|
|
nonnull_ds = ray.data.from_pandas(nonnull_df)
|
|
null_encoder = OneHotEncoder(["A"])
|
|
|
|
# Verify fit fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.fit(null_ds)
|
|
null_encoder.fit(nonnull_ds)
|
|
|
|
# Verify transform fails for null values.
|
|
with pytest.raises((UserCodeException, ValueError)):
|
|
null_encoder.transform(null_ds).materialize()
|
|
null_encoder.transform(nonnull_ds)
|
|
|
|
# Verify transform_batch fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.transform_batch(null_df)
|
|
null_encoder.transform_batch(nonnull_df)
|
|
|
|
|
|
def test_one_hot_encoder_with_max_categories():
|
|
"""Tests basic OneHotEncoder functionality with limit."""
|
|
col_a = ["red", "green", "blue", "red", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold", "hot"]
|
|
col_c = [1, 10, 5, 10, 10]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = OneHotEncoder(["B", "C"], max_categories={"B": 2})
|
|
|
|
ds_out = encoder.fit_transform(ds)
|
|
df_out = ds_out.to_pandas()
|
|
assert len(ds_out.to_pandas().columns) == 3
|
|
|
|
expected_df = ArrowBlockAccessor(
|
|
pa.table(
|
|
{
|
|
"A": pa.array(col_a),
|
|
"B": pa.array(
|
|
[[0, 0], [1, 0], [0, 1], [1, 0], [0, 1]],
|
|
type=pa.list_(pa.uint8(), 2),
|
|
),
|
|
"C": pa.array(
|
|
[[1, 0, 0], [0, 0, 1], [0, 1, 0], [0, 0, 1], [0, 0, 1]],
|
|
type=pa.list_(pa.uint8(), 3),
|
|
),
|
|
}
|
|
)
|
|
).to_pandas()
|
|
pd.testing.assert_frame_equal(df_out, expected_df, check_like=True)
|
|
|
|
|
|
def test_one_hot_encoder_max_categories_global_sum_across_partitions():
|
|
"""Tests that max_categories sums counts across partitions before picking top-k.
|
|
|
|
Edge case: a category that is NOT in the per-partition top-k of ANY single
|
|
partition can still become a global top-k category once counts are summed.
|
|
|
|
Setup: 2 partitions with column "X" having these value counts:
|
|
Partition 1: {A: 5, B: 4, C: 3} → per-partition top-2: {A, B}
|
|
Partition 2: {D: 5, E: 4, C: 3} → per-partition top-2: {D, E}
|
|
|
|
Global counts: {A: 5, B: 4, C: 6, D: 5, E: 4}
|
|
Global top-2: {C, A} or {C, D} (C=6 is highest, then A=5 and D=5 tie)
|
|
|
|
If top-k were applied per-partition first, C would be excluded from both
|
|
partitions' top-2, and the union {A, B, D, E} would never include C.
|
|
"""
|
|
part1 = pd.DataFrame({"X": ["A"] * 5 + ["B"] * 4 + ["C"] * 3})
|
|
part2 = pd.DataFrame({"X": ["D"] * 5 + ["E"] * 4 + ["C"] * 3})
|
|
ds = ray.data.from_pandas([part1, part2])
|
|
|
|
encoder = OneHotEncoder(["X"], max_categories={"X": 2})
|
|
encoder.fit(ds)
|
|
|
|
stats = encoder.stats_
|
|
encoded_categories = set(stats["unique_values(X)"].keys())
|
|
assert len(encoded_categories) == 2, (
|
|
f"Expected 2 categories from global top-k, got {len(encoded_categories)}: "
|
|
f"{encoded_categories}"
|
|
)
|
|
# C must be included since it has the highest global count (6).
|
|
assert (
|
|
"C" in encoded_categories
|
|
), f"Expected C (global count=6) to be in top-2, got {encoded_categories}"
|
|
# The second category should be one of A or D (both have count=5).
|
|
remaining = encoded_categories - {"C"}
|
|
assert remaining <= {
|
|
"A",
|
|
"D",
|
|
}, f"Expected second category to be A or D (count=5), got {remaining}"
|
|
|
|
|
|
def test_multi_hot_encoder():
|
|
"""Tests basic MultiHotEncoder functionality."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
col_c = [1, 10, 5, 10]
|
|
col_d = [["warm"], [], ["hot", "warm", "cold"], ["cold", "cold"]]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c, "D": col_d})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = MultiHotEncoder(["B", "C", "D"])
|
|
|
|
with pytest.raises(PreprocessorNotFittedException):
|
|
encoder.transform(ds)
|
|
|
|
# Fit data.
|
|
encoder.fit(ds)
|
|
|
|
assert encoder.stats_ == {
|
|
"unique_values(B)": {"cold": 0, "hot": 1, "warm": 2},
|
|
"unique_values(C)": {1: 0, 5: 1, 10: 2},
|
|
"unique_values(D)": {"cold": 0, "hot": 1, "warm": 2},
|
|
}
|
|
|
|
# Transform data.
|
|
transformed = encoder.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
processed_col_a = col_a
|
|
processed_col_b = [[0, 0, 1], [1, 0, 0], [0, 1, 0], [1, 0, 0]]
|
|
processed_col_c = [[1, 0, 0], [0, 0, 1], [0, 1, 0], [0, 0, 1]]
|
|
processed_col_d = [[0, 0, 1], [0, 0, 0], [1, 1, 1], [2, 0, 0]]
|
|
expected_df = ArrowBlockAccessor(
|
|
pa.Table.from_pydict(
|
|
{
|
|
"A": processed_col_a,
|
|
"B": processed_col_b,
|
|
"C": processed_col_c,
|
|
"D": processed_col_d,
|
|
}
|
|
)
|
|
).to_pandas()
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df)
|
|
|
|
# Transform batch.
|
|
pred_col_a = ["blue", "yellow", None]
|
|
pred_col_b = ["cold", "warm", "other"]
|
|
pred_col_c = [10, 1, 20]
|
|
pred_col_d = [["cold", "warm"], [], ["other", "cold"]]
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c, "D": pred_col_d}
|
|
)
|
|
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
print(pred_out_df.to_string())
|
|
|
|
pred_processed_col_a = ["blue", "yellow", None]
|
|
pred_processed_col_b = [[1, 0, 0], [0, 0, 1], [0, 0, 0]]
|
|
pred_processed_col_c = [[0, 0, 1], [1, 0, 0], [0, 0, 0]]
|
|
pred_processed_col_d = [[1, 0, 1], [0, 0, 0], [1, 0, 0]]
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_processed_col_a,
|
|
"B": pred_processed_col_b,
|
|
"C": pred_processed_col_c,
|
|
"D": pred_processed_col_d,
|
|
}
|
|
)
|
|
|
|
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df)
|
|
|
|
# append mode
|
|
with pytest.raises(ValueError):
|
|
MultiHotEncoder(columns=["B", "C", "D"], output_columns=["B_encoded"])
|
|
|
|
encoder = MultiHotEncoder(
|
|
columns=["B", "C", "D"],
|
|
output_columns=[
|
|
"B_multihot_encoded",
|
|
"C_multihot_encoded",
|
|
"D_multihot_encoded",
|
|
],
|
|
)
|
|
encoder.fit(ds)
|
|
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c, "D": pred_col_d}
|
|
)
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_col_a,
|
|
"B": pred_col_b,
|
|
"C": pred_col_c,
|
|
"D": pred_col_d,
|
|
"B_multihot_encoded": pred_processed_col_b,
|
|
"C_multihot_encoded": pred_processed_col_c,
|
|
"D_multihot_encoded": pred_processed_col_d,
|
|
}
|
|
)
|
|
|
|
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df)
|
|
|
|
# Test null behavior.
|
|
null_col = [1, None]
|
|
nonnull_col = [1, 1]
|
|
null_df = pd.DataFrame.from_dict({"A": null_col})
|
|
null_ds = ray.data.from_pandas(null_df)
|
|
nonnull_df = pd.DataFrame.from_dict({"A": nonnull_col})
|
|
nonnull_ds = ray.data.from_pandas(nonnull_df)
|
|
null_encoder = MultiHotEncoder(["A"])
|
|
|
|
# Verify fit fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.fit(null_ds)
|
|
null_encoder.fit(nonnull_ds)
|
|
|
|
# Verify transform fails for null values.
|
|
with pytest.raises((UserCodeException, ValueError)):
|
|
null_encoder.transform(null_ds).materialize()
|
|
null_encoder.transform(nonnull_ds)
|
|
|
|
# Verify transform_batch fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.transform_batch(null_df)
|
|
null_encoder.transform_batch(nonnull_df)
|
|
|
|
|
|
def test_multi_hot_encoder_with_max_categories():
|
|
"""Tests basic MultiHotEncoder functionality with limit."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold"]
|
|
col_c = [1, 10, 5, 10]
|
|
col_d = [["warm"], [], ["hot", "warm", "cold"], ["cold", "cold"]]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c, "D": col_d})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = MultiHotEncoder(["B", "C", "D"], max_categories={"B": 2})
|
|
|
|
ds_out = encoder.fit_transform(ds)
|
|
assert len(ds_out.to_pandas()["B"].iloc[0]) == 2
|
|
assert len(ds_out.to_pandas()["C"].iloc[0]) == 3
|
|
assert len(ds_out.to_pandas()["D"].iloc[0]) == 3
|
|
|
|
|
|
def test_label_encoder():
|
|
"""Tests basic LabelEncoder functionality."""
|
|
col_a = ["red", "green", "blue", "red"]
|
|
col_b = ["warm", "cold", "cold", "hot"]
|
|
col_c = [1, 2, 3, 4]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
encoder = LabelEncoder("A")
|
|
|
|
# Transform with unfitted preprocessor.
|
|
with pytest.raises(PreprocessorNotFittedException):
|
|
encoder.transform(ds)
|
|
|
|
# Fit data.
|
|
encoder.fit(ds)
|
|
|
|
assert encoder.stats_ == {"unique_values(A)": {"blue": 0, "green": 1, "red": 2}}
|
|
|
|
# Transform data.
|
|
transformed = encoder.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
processed_col_a = [2, 1, 0, 2]
|
|
processed_col_b = col_b
|
|
processed_col_c = col_c
|
|
expected_df = ArrowBlockAccessor(
|
|
pa.Table.from_pydict(
|
|
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
|
)
|
|
).to_pandas()
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
# append mode
|
|
append_encoder = LabelEncoder("A", output_column="A_encoded")
|
|
append_encoder.fit(ds)
|
|
append_transformed = append_encoder.transform(ds)
|
|
out_df = append_transformed.to_pandas()
|
|
|
|
expected_df = ArrowBlockAccessor(
|
|
pa.Table.from_pydict(
|
|
{"A": col_a, "B": col_b, "C": col_c, "A_encoded": processed_col_a}
|
|
)
|
|
).to_pandas()
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
# Inverse transform data.
|
|
inverse_transformed = encoder.inverse_transform(transformed)
|
|
inverse_df = inverse_transformed.to_pandas()
|
|
|
|
arrow_in_df = ArrowBlockAccessor(pa.Table.from_pandas(in_df)).to_pandas()
|
|
pd.testing.assert_frame_equal(inverse_df, arrow_in_df, check_like=True)
|
|
|
|
inverse_append_transformed = append_encoder.inverse_transform(append_transformed)
|
|
inverse_append_df = inverse_append_transformed.to_pandas()
|
|
expected_df = ArrowBlockAccessor(
|
|
pa.Table.from_pydict(
|
|
{"A": col_a, "B": col_b, "C": col_c, "A_encoded": processed_col_a}
|
|
)
|
|
).to_pandas()
|
|
pd.testing.assert_frame_equal(inverse_append_df, expected_df, check_like=True)
|
|
|
|
# Inverse transform without fitting.
|
|
new_encoder = LabelEncoder("A")
|
|
|
|
with pytest.raises(RuntimeError):
|
|
new_encoder.inverse_transform(ds)
|
|
|
|
# Inverse transform on fitted preprocessor that hasn't transformed anything.
|
|
new_encoder.fit(ds)
|
|
inv_non_fitted = new_encoder.inverse_transform(transformed)
|
|
inv_non_fitted_df = inv_non_fitted.to_pandas()
|
|
|
|
assert inv_non_fitted_df.equals(arrow_in_df)
|
|
|
|
# Transform batch.
|
|
pred_col_a = ["blue", "red", "yellow"]
|
|
pred_col_b = ["cold", "unknown", None]
|
|
pred_col_c = [10, 20, None]
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
|
)
|
|
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
|
|
pred_processed_col_a = [0, 2, None]
|
|
pred_processed_col_b = pred_col_b
|
|
pred_processed_col_c = pred_col_c
|
|
pred_expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": pred_processed_col_a,
|
|
"B": pred_processed_col_b,
|
|
"C": pred_processed_col_c,
|
|
}
|
|
)
|
|
assert pred_out_df.equals(pred_expected_df)
|
|
|
|
# Test null behavior.
|
|
null_col = [1, None]
|
|
nonnull_col = [1, 1]
|
|
null_df = pd.DataFrame.from_dict({"A": null_col})
|
|
null_ds = ray.data.from_pandas(null_df)
|
|
nonnull_df = pd.DataFrame.from_dict({"A": nonnull_col})
|
|
nonnull_ds = ray.data.from_pandas(nonnull_df)
|
|
null_encoder = LabelEncoder("A")
|
|
|
|
# Verify fit fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.fit(null_ds)
|
|
null_encoder.fit(nonnull_ds)
|
|
|
|
# Verify transform fails for null values.
|
|
with pytest.raises((UserCodeException, ValueError)):
|
|
null_encoder.transform(null_ds).materialize()
|
|
null_encoder.transform(nonnull_ds)
|
|
|
|
# Verify transform_batch fails for null values.
|
|
with pytest.raises(ValueError):
|
|
null_encoder.transform_batch(null_df)
|
|
null_encoder.transform_batch(nonnull_df)
|
|
|
|
|
|
@pytest.mark.parametrize("predefined_dtypes", [True, False])
|
|
def test_categorizer(predefined_dtypes):
|
|
"""Tests basic Categorizer functionality."""
|
|
col_a = ["red", "green", "blue", "red", "red"]
|
|
col_b = ["warm", "cold", "hot", "cold", None]
|
|
col_c = [1, 10, 5, 10, 1]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
columns = ["B", "C"]
|
|
if predefined_dtypes:
|
|
expected_dtypes = {
|
|
"B": pd.CategoricalDtype(["cold", "hot", "warm"], ordered=True),
|
|
"C": pd.CategoricalDtype([1, 5, 10]),
|
|
}
|
|
dtypes = {"B": pd.CategoricalDtype(["cold", "hot", "warm"], ordered=True)}
|
|
else:
|
|
expected_dtypes = {
|
|
"B": pd.CategoricalDtype(["cold", "hot", "warm"]),
|
|
"C": pd.CategoricalDtype([1, 5, 10]),
|
|
}
|
|
columns = ["B", "C"]
|
|
dtypes = None
|
|
|
|
encoder = Categorizer(columns, dtypes)
|
|
|
|
# Transform with unfitted preprocessor.
|
|
with pytest.raises(PreprocessorNotFittedException):
|
|
encoder.transform(ds)
|
|
|
|
# Fit data.
|
|
encoder.fit(ds)
|
|
assert encoder.stats_ == expected_dtypes
|
|
|
|
# Transform data.
|
|
transformed = encoder.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
arrow_passthrough_dtype = (
|
|
ArrowBlockAccessor(pa.Table.from_pandas(in_df[["A"]])).to_pandas().dtypes["A"]
|
|
)
|
|
assert out_df.dtypes["A"] == arrow_passthrough_dtype
|
|
assert out_df.dtypes["B"] == expected_dtypes["B"]
|
|
assert out_df.dtypes["C"] == expected_dtypes["C"]
|
|
|
|
# Transform batch.
|
|
pred_col_a = ["blue", "yellow", None]
|
|
pred_col_b = ["cold", "warm", "other"]
|
|
pred_col_c = [10, 1, 20]
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
|
)
|
|
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
|
|
assert pred_out_df.dtypes["A"] == np.object_
|
|
assert pred_out_df.dtypes["B"] == expected_dtypes["B"]
|
|
assert pred_out_df.dtypes["C"] == expected_dtypes["C"]
|
|
|
|
# append mode
|
|
with pytest.raises(ValueError):
|
|
Categorizer(columns=["B", "C"], output_columns=["B_categorized"])
|
|
|
|
encoder = Categorizer(
|
|
columns=["B", "C"],
|
|
output_columns=["B_categorized", "C_categorized"],
|
|
dtypes=dtypes,
|
|
)
|
|
encoder.fit(ds)
|
|
|
|
pred_in_df = pd.DataFrame.from_dict(
|
|
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
|
)
|
|
pred_out_df = encoder.transform_batch(pred_in_df)
|
|
|
|
assert pred_out_df.dtypes["A"] == np.object_
|
|
assert pred_out_df.dtypes["B"] == np.object_
|
|
assert pred_out_df.dtypes["C"] == np.int64
|
|
assert pred_out_df.dtypes["B_categorized"] == expected_dtypes["B"]
|
|
assert pred_out_df.dtypes["C_categorized"] == expected_dtypes["C"]
|
|
|
|
|
|
class TestEncoderSerialization:
|
|
"""Test basic serialization/deserialization functionality for all encoder preprocessors."""
|
|
|
|
def setup_method(self):
|
|
"""Set up test data for encoders."""
|
|
# Data for categorical encoders
|
|
self.categorical_df = pd.DataFrame(
|
|
{
|
|
"category": ["A", "B", "C", "A", "B", "C", "A"],
|
|
"grade": ["high", "medium", "low", "high", "medium", "low", "high"],
|
|
"region": ["north", "south", "east", "west", "north", "south", "east"],
|
|
}
|
|
)
|
|
|
|
# Data for multi-hot encoder (with lists)
|
|
self.multihot_df = pd.DataFrame(
|
|
{
|
|
"tags": [
|
|
["red", "car"],
|
|
["blue", "bike"],
|
|
["red", "truck"],
|
|
["green", "car"],
|
|
],
|
|
"features": [
|
|
["fast", "loud"],
|
|
["quiet"],
|
|
["fast", "heavy"],
|
|
["quiet", "light"],
|
|
],
|
|
}
|
|
)
|
|
|
|
# Data for label encoder
|
|
self.label_df = pd.DataFrame(
|
|
{
|
|
"target": ["cat", "dog", "bird", "cat", "dog", "bird"],
|
|
"other": [1, 2, 3, 4, 5, 6],
|
|
}
|
|
)
|
|
|
|
def test_ordinal_encoder_serialization(self):
|
|
"""Test OrdinalEncoder save/load functionality."""
|
|
# Create and fit encoder
|
|
encoder = OrdinalEncoder(columns=["category", "grade"])
|
|
dataset = ray.data.from_pandas(self.categorical_df)
|
|
fitted_encoder = encoder.fit(dataset)
|
|
|
|
# Test CloudPickle serialization (primary format)
|
|
serialized = fitted_encoder.serialize()
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessor.MAGIC_CLOUDPICKLE)
|
|
|
|
# Test deserialization
|
|
deserialized = SerializablePreprocessor.deserialize(serialized)
|
|
assert isinstance(deserialized, OrdinalEncoder)
|
|
assert deserialized._fitted
|
|
assert deserialized.columns == ["category", "grade"]
|
|
assert deserialized.encode_lists is True # default value
|
|
|
|
# Test functional equivalence
|
|
test_df = pd.DataFrame({"category": ["A", "B"], "grade": ["high", "low"]})
|
|
|
|
original_result = fitted_encoder.transform_batch(test_df.copy())
|
|
deserialized_result = deserialized.transform_batch(test_df.copy())
|
|
|
|
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
|
|
|
def test_onehot_encoder_serialization(self):
|
|
"""Test OneHotEncoder save/load functionality."""
|
|
# Create and fit encoder
|
|
encoder = OneHotEncoder(columns=["category"], max_categories={"category": 3})
|
|
dataset = ray.data.from_pandas(self.categorical_df)
|
|
fitted_encoder = encoder.fit(dataset)
|
|
|
|
# Test CloudPickle serialization (primary format)
|
|
serialized = fitted_encoder.serialize()
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessor.MAGIC_CLOUDPICKLE)
|
|
|
|
# Test deserialization
|
|
deserialized = SerializablePreprocessor.deserialize(serialized)
|
|
assert isinstance(deserialized, OneHotEncoder)
|
|
assert deserialized._fitted
|
|
assert deserialized.columns == ["category"]
|
|
assert deserialized.max_categories == {"category": 3}
|
|
|
|
# Test functional equivalence
|
|
test_df = pd.DataFrame({"category": ["A", "B", "C"]})
|
|
|
|
original_result = fitted_encoder.transform_batch(test_df.copy())
|
|
deserialized_result = deserialized.transform_batch(test_df.copy())
|
|
|
|
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
|
|
|
def test_multihot_encoder_serialization(self):
|
|
"""Test MultiHotEncoder save/load functionality."""
|
|
# Create and fit encoder
|
|
encoder = MultiHotEncoder(columns=["tags"], max_categories={"tags": 5})
|
|
dataset = ray.data.from_pandas(self.multihot_df)
|
|
fitted_encoder = encoder.fit(dataset)
|
|
|
|
# Test CloudPickle serialization (primary format)
|
|
serialized = fitted_encoder.serialize()
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessor.MAGIC_CLOUDPICKLE)
|
|
|
|
# Test deserialization
|
|
deserialized = SerializablePreprocessor.deserialize(serialized)
|
|
assert isinstance(deserialized, MultiHotEncoder)
|
|
assert deserialized._fitted
|
|
assert deserialized.columns == ["tags"]
|
|
assert deserialized.max_categories == {"tags": 5}
|
|
|
|
# Test functional equivalence
|
|
test_df = pd.DataFrame({"tags": [["red", "car"], ["blue", "bike"]]})
|
|
|
|
original_result = fitted_encoder.transform_batch(test_df.copy())
|
|
deserialized_result = deserialized.transform_batch(test_df.copy())
|
|
|
|
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
|
|
|
def test_label_encoder_serialization(self):
|
|
"""Test LabelEncoder save/load functionality."""
|
|
# Create and fit encoder
|
|
encoder = LabelEncoder(label_column="target")
|
|
dataset = ray.data.from_pandas(self.label_df)
|
|
fitted_encoder = encoder.fit(dataset)
|
|
|
|
# Test CloudPickle serialization (primary format)
|
|
serialized = fitted_encoder.serialize()
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessor.MAGIC_CLOUDPICKLE)
|
|
|
|
# Test deserialization
|
|
deserialized = SerializablePreprocessor.deserialize(serialized)
|
|
assert isinstance(deserialized, LabelEncoder)
|
|
assert deserialized._fitted
|
|
assert deserialized.label_column == "target"
|
|
assert deserialized.output_column == "target" # default
|
|
|
|
# Test functional equivalence
|
|
test_df = pd.DataFrame({"target": ["cat", "dog", "bird"]})
|
|
|
|
original_result = fitted_encoder.transform_batch(test_df.copy())
|
|
deserialized_result = deserialized.transform_batch(test_df.copy())
|
|
|
|
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
|
|
|
def test_categorizer_serialization(self):
|
|
"""Test Categorizer save/load functionality."""
|
|
# Create categorizer with predefined dtypes
|
|
sex_dtype = pd.CategoricalDtype(categories=["male", "female"], ordered=False)
|
|
grade_dtype = pd.CategoricalDtype(
|
|
categories=["high", "medium", "low"], ordered=True
|
|
)
|
|
|
|
categorizer = Categorizer(
|
|
columns=["category", "grade"],
|
|
dtypes={"category": sex_dtype, "grade": grade_dtype},
|
|
)
|
|
|
|
# Test CloudPickle serialization (primary format, even without fitting)
|
|
serialized = categorizer.serialize()
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessor.MAGIC_CLOUDPICKLE)
|
|
|
|
# Test deserialization
|
|
deserialized = SerializablePreprocessor.deserialize(serialized)
|
|
assert isinstance(deserialized, Categorizer)
|
|
assert deserialized.columns == ["category", "grade"]
|
|
|
|
# Test dtypes preservation
|
|
assert len(deserialized.dtypes) == 2
|
|
assert isinstance(deserialized.dtypes["category"], pd.CategoricalDtype)
|
|
assert isinstance(deserialized.dtypes["grade"], pd.CategoricalDtype)
|
|
|
|
# Check category preservation
|
|
assert list(deserialized.dtypes["category"].categories) == ["male", "female"]
|
|
assert deserialized.dtypes["category"].ordered is False
|
|
|
|
assert list(deserialized.dtypes["grade"].categories) == [
|
|
"high",
|
|
"medium",
|
|
"low",
|
|
]
|
|
assert deserialized.dtypes["grade"].ordered is True
|
|
|
|
def test_categorizer_fitted_serialization(self):
|
|
"""Test Categorizer save/load functionality after fitting."""
|
|
# Create and fit categorizer (without predefined dtypes)
|
|
categorizer = Categorizer(columns=["category", "grade"])
|
|
dataset = ray.data.from_pandas(self.categorical_df)
|
|
fitted_categorizer = categorizer.fit(dataset)
|
|
|
|
# Test CloudPickle serialization (primary format)
|
|
serialized = fitted_categorizer.serialize()
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessor.MAGIC_CLOUDPICKLE)
|
|
|
|
# Test deserialization
|
|
deserialized = SerializablePreprocessor.deserialize(serialized)
|
|
assert isinstance(deserialized, Categorizer)
|
|
assert deserialized._fitted
|
|
assert deserialized.columns == ["category", "grade"]
|
|
|
|
# Test functional equivalence
|
|
test_df = pd.DataFrame({"category": ["A", "B"], "grade": ["high", "low"]})
|
|
|
|
original_result = fitted_categorizer.transform_batch(test_df.copy())
|
|
deserialized_result = deserialized.transform_batch(test_df.copy())
|
|
|
|
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
|
|
|
def test_encoder_serialization_formats(self):
|
|
"""Test that encoders work with different serialization formats."""
|
|
encoder = OrdinalEncoder(columns=["category"])
|
|
dataset = ray.data.from_pandas(self.categorical_df)
|
|
fitted_encoder = encoder.fit(dataset)
|
|
|
|
# Test CloudPickle format (default)
|
|
cloudpickle_serialized = fitted_encoder.serialize()
|
|
assert isinstance(cloudpickle_serialized, bytes)
|
|
|
|
# Test Pickle format (legacy)
|
|
pickle_serialized = fitted_encoder.serialize()
|
|
assert isinstance(pickle_serialized, bytes)
|
|
|
|
# Both should deserialize to equivalent objects
|
|
cloudpickle_deserialized = SerializablePreprocessor.deserialize(
|
|
cloudpickle_serialized
|
|
)
|
|
pickle_deserialized = SerializablePreprocessor.deserialize(pickle_serialized)
|
|
|
|
# Test functional equivalence
|
|
test_df = pd.DataFrame({"category": ["A", "B"]})
|
|
|
|
cloudpickle_result = cloudpickle_deserialized.transform_batch(test_df.copy())
|
|
pickle_result = pickle_deserialized.transform_batch(test_df.copy())
|
|
|
|
pd.testing.assert_frame_equal(cloudpickle_result, pickle_result)
|
|
|
|
def test_encoder_error_handling(self):
|
|
"""Test error handling for encoder serialization."""
|
|
# Test unknown preprocessor type
|
|
import cloudpickle
|
|
|
|
unknown_data = {
|
|
"type": "NonExistentEncoder",
|
|
"version": 1,
|
|
"fields": {"columns": ["test"]},
|
|
"stats": {},
|
|
"stats_type": "default",
|
|
}
|
|
|
|
fake_serialized = (
|
|
SerializablePreprocessor.MAGIC_CLOUDPICKLE + cloudpickle.dumps(unknown_data)
|
|
)
|
|
|
|
from ray.data.preprocessors.version_support import UnknownPreprocessorError
|
|
|
|
with pytest.raises(UnknownPreprocessorError) as exc_info:
|
|
SerializablePreprocessor.deserialize(fake_serialized)
|
|
|
|
assert exc_info.value.preprocessor_type == "NonExistentEncoder"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|