chore: import upstream snapshot with attribution

python/ray/data/tests/preprocessors/test_chain.py

@@ -0,0 +1,243 @@
+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__]))