wehub-resource-sync
228 lines
8.6 KiB
Python
228 lines
8.6 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
from pandas.testing import assert_frame_equal
|
|
|
|
import ray
|
|
from ray.data.exceptions import UserCodeException
|
|
from ray.data.preprocessors import Concatenator, OneHotEncoder
|
|
|
|
|
|
class TestConcatenator:
|
|
def test_basic(self):
|
|
df = pd.DataFrame(
|
|
{
|
|
"a": [1, 2, 3, 4],
|
|
"b": [5, 6, 7, 8],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["a", "b"], output_column_name="c")
|
|
new_ds = prep.transform(ds)
|
|
for i, row in enumerate(new_ds.take()):
|
|
assert np.array_equal(row["c"], np.array([i + 1, i + 5]))
|
|
|
|
def test_raise_if_missing(self):
|
|
df = pd.DataFrame({"a": [1, 2, 3, 4]})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(
|
|
columns=["a", "b"], output_column_name="c", raise_if_missing=True
|
|
)
|
|
|
|
with pytest.raises(UserCodeException):
|
|
with pytest.raises(ValueError, match="'b'"):
|
|
prep.transform(ds).materialize()
|
|
|
|
def test_exclude_column(self):
|
|
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5], "c": [3, 4, 5, 6]})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["a", "c"])
|
|
new_ds = prep.transform(ds)
|
|
for _, row in enumerate(new_ds.take()):
|
|
assert set(row) == {"concat_out", "b"}
|
|
|
|
def test_include_columns(self):
|
|
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5], "c": [3, 4, 5, 6]})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["a", "b"])
|
|
new_ds = prep.transform(ds)
|
|
for _, row in enumerate(new_ds.take()):
|
|
assert set(row) == {"concat_out", "c"}
|
|
|
|
def test_change_column_order(self):
|
|
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5]})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["b", "a"])
|
|
new_ds = prep.transform(ds)
|
|
expected_df = pd.DataFrame({"concat_out": [[2, 1], [3, 2], [4, 3], [5, 4]]})
|
|
print(new_ds.to_pandas())
|
|
assert_frame_equal(new_ds.to_pandas(), expected_df)
|
|
|
|
def test_strings(self):
|
|
df = pd.DataFrame({"a": ["string", "string2", "string3"]})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["a"], output_column_name="huh")
|
|
new_ds = prep.transform(ds)
|
|
assert "huh" in set(new_ds.schema().names)
|
|
|
|
def test_preserves_order(self):
|
|
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5]})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["a", "b"], output_column_name="c")
|
|
prep = prep.fit(ds)
|
|
|
|
df = pd.DataFrame({"a": [5, 6, 7, 8], "b": [6, 7, 8, 9]})
|
|
concatenated_df = prep.transform_batch(df)
|
|
expected_df = pd.DataFrame({"c": [[5, 6], [6, 7], [7, 8], [8, 9]]})
|
|
assert_frame_equal(concatenated_df, expected_df)
|
|
|
|
other_df = pd.DataFrame({"a": [9, 10, 11, 12], "b": [10, 11, 12, 13]})
|
|
concatenated_other_df = prep.transform_batch(other_df)
|
|
expected_df = pd.DataFrame(
|
|
{
|
|
"c": [
|
|
[9, 10],
|
|
[10, 11],
|
|
[11, 12],
|
|
[12, 13],
|
|
]
|
|
}
|
|
)
|
|
assert_frame_equal(concatenated_other_df, expected_df)
|
|
|
|
@pytest.mark.parametrize("col_b", [[[2, 3], [3, 4], [4, 5], [5, 6]], [2, 3, 4, 5]])
|
|
@pytest.mark.parametrize("flatten", [True, False])
|
|
def test_flatten(self, col_b, flatten):
|
|
col_a = [1, 2, 3, 4]
|
|
col_b = [np.array(v) for v in col_b] if isinstance(col_b[0], list) else col_b
|
|
df = pd.DataFrame({"a": col_a, "b": col_b})
|
|
ds = ray.data.from_pandas(df)
|
|
prep = Concatenator(columns=["a", "b"], flatten=flatten)
|
|
new_ds = prep.transform(ds)
|
|
|
|
for i, row in enumerate(new_ds.take()):
|
|
if flatten or not isinstance(col_b[i], np.ndarray):
|
|
# When flatten=True or when col_b contains simple values
|
|
if isinstance(col_b[i], np.ndarray):
|
|
expected = np.concatenate([np.array([col_a[i]]), col_b[i]])
|
|
else:
|
|
expected = np.array([col_a[i], col_b[i]])
|
|
assert np.array_equal(row["concat_out"], expected)
|
|
else:
|
|
# When flatten=False and col_b contains numpy arrays
|
|
# The output should be a list containing the scalar and the array
|
|
assert len(row["concat_out"]) == 2
|
|
assert row["concat_out"][0] == col_a[i]
|
|
assert np.array_equal(row["concat_out"][1], col_b[i])
|
|
|
|
@pytest.mark.parametrize("flatten", [True, False])
|
|
def test_concatenate_with_onehotencoder(self, flatten):
|
|
df = pd.DataFrame(
|
|
{
|
|
"color": ["red", "green", "blue", "red"],
|
|
"value": [1, 2, 3, 4],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(df)
|
|
# OneHot encode the color column
|
|
encoder = OneHotEncoder(columns=["color"], output_columns=["color_encoded"])
|
|
encoder = encoder.fit(ds)
|
|
encoded_ds = encoder.transform(ds)
|
|
# Concatenate the one-hot encoded column with the value column
|
|
prep = Concatenator(
|
|
columns=["color_encoded", "value"],
|
|
output_column_name="features",
|
|
flatten=flatten,
|
|
)
|
|
new_ds = prep.transform(encoded_ds)
|
|
# Get the expected one-hot vectors
|
|
color_map = {"blue": [1, 0, 0], "green": [0, 1, 0], "red": [0, 0, 1]}
|
|
for i, row in enumerate(new_ds.take()):
|
|
if flatten:
|
|
expected = color_map[df["color"][i]] + [df["value"][i]]
|
|
assert np.array_equal(row["features"], np.array(expected))
|
|
else:
|
|
expected = [np.array(color_map[df["color"][i]]), df["value"][i]]
|
|
assert np.array_equal(row["features"][0], expected[0])
|
|
assert row["features"][1] == expected[1]
|
|
|
|
@pytest.mark.parametrize("flatten", [True, False])
|
|
def test_nested_list_with_dtype(self, flatten: bool):
|
|
# Tests Concatenator with nested lists and dtype: flattens and coerces when flatten=True,
|
|
# raises ValueError when flatten=False.
|
|
output_column = "c"
|
|
df = pd.DataFrame(
|
|
{
|
|
"a": [12.0],
|
|
"b": [[1, 0, 0, 0]],
|
|
}
|
|
)
|
|
prep = Concatenator(
|
|
columns=["a", "b"],
|
|
output_column_name=output_column,
|
|
dtype=np.float32,
|
|
flatten=flatten,
|
|
)
|
|
|
|
if flatten:
|
|
pd_ds = prep._transform_pandas(df)
|
|
expected_pd = pd.DataFrame(
|
|
{output_column: pd.Series([[12.0, 1.0, 0.0, 0.0, 0.0]])}
|
|
)
|
|
assert_frame_equal(pd_ds, expected_pd)
|
|
else:
|
|
# Only for flattened output do we expect the dtype coercion to apply
|
|
with pytest.raises(ValueError):
|
|
pd_ds = prep._transform_pandas(df)
|
|
|
|
def test_concatenator_deserialize_backward_compat(self):
|
|
p1 = Concatenator(columns=["A"], flatten=True)
|
|
delattr(p1, "_flatten")
|
|
data = p1.serialize()
|
|
p2 = Concatenator.deserialize(data)
|
|
assert isinstance(p2, Concatenator)
|
|
assert p2.flatten is False
|
|
|
|
def test_concatenator_serialization(self):
|
|
"""Test Concatenator serialization and deserialization functionality."""
|
|
from ray.data.preprocessor import SerializablePreprocessorBase
|
|
|
|
# Create concatenator
|
|
concatenator = Concatenator(
|
|
columns=["A", "B"],
|
|
output_column_name="combined",
|
|
dtype=np.float32,
|
|
flatten=True,
|
|
)
|
|
|
|
# Serialize using CloudPickle
|
|
serialized = concatenator.serialize()
|
|
|
|
# Verify it's binary CloudPickle format
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
|
|
|
# Deserialize
|
|
deserialized = Concatenator.deserialize(serialized)
|
|
|
|
# Verify type and field values
|
|
assert isinstance(deserialized, Concatenator)
|
|
assert deserialized.columns == ["A", "B"]
|
|
assert deserialized.output_column_name == "combined"
|
|
assert deserialized.dtype == np.float32
|
|
assert deserialized.flatten is True
|
|
|
|
# Verify it works correctly
|
|
df = pd.DataFrame({"A": [[1, 2]], "B": [[3, 4]]})
|
|
result = deserialized.transform_batch(df)
|
|
|
|
# Verify concatenation was applied correctly
|
|
assert "combined" in result.columns
|
|
assert len(result["combined"][0]) == 4
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|