chore: import upstream snapshot with attribution

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

@@ -0,0 +1,570 @@
+import re
+import warnings
+from typing import Dict, Union
+from unittest.mock import patch
+
+import numpy as np
+import pandas as pd
+import pyarrow
+import pytest
+
+import ray
+from ray.data.aggregate import Mean
+from ray.data.constants import MAX_REPR_LENGTH
+from ray.data.preprocessor import Preprocessor
+from ray.data.preprocessors import (
+    Categorizer,
+    Chain,
+    Concatenator,
+    CountVectorizer,
+    FeatureHasher,
+    HashingVectorizer,
+    LabelEncoder,
+    MaxAbsScaler,
+    MinMaxScaler,
+    MultiHotEncoder,
+    Normalizer,
+    OneHotEncoder,
+    OrdinalEncoder,
+    PowerTransformer,
+    RobustScaler,
+    SimpleImputer,
+    StandardScaler,
+    Tokenizer,
+    TorchVisionPreprocessor,
+)
+from ray.data.util.data_batch_conversion import BatchFormat
+
+
+@pytest.fixture
+def create_dummy_preprocessors():
+    class DummyPreprocessorWithNothing(Preprocessor):
+        _is_fittable = False
+
+    class DummyPreprocessorWithPandas(DummyPreprocessorWithNothing):
+        def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
+            return df
+
+    class DummyPreprocessorWithNumpy(DummyPreprocessorWithNothing):
+        batch_format = "numpy"
+
+        def _transform_numpy(
+            self, np_data: Union[np.ndarray, Dict[str, np.ndarray]]
+        ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
+            return np_data
+
+    class DummyPreprocessorWithPandasAndNumpy(DummyPreprocessorWithNothing):
+        def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
+            return df
+
+        def _transform_numpy(
+            self, np_data: Union[np.ndarray, Dict[str, np.ndarray]]
+        ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
+            return np_data
+
+    class DummyPreprocessorWithPandasAndNumpyPreferred(DummyPreprocessorWithNothing):
+        def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
+            return df
+
+        def _transform_numpy(
+            self, np_data: Union[np.ndarray, Dict[str, np.ndarray]]
+        ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
+            return np_data
+
+        def preferred_batch_format(cls) -> BatchFormat:
+            return BatchFormat.NUMPY
+
+    yield (
+        DummyPreprocessorWithNothing(),
+        DummyPreprocessorWithPandas(),
+        DummyPreprocessorWithNumpy(),
+        DummyPreprocessorWithPandasAndNumpy(),
+        DummyPreprocessorWithPandasAndNumpyPreferred(),
+    )
+
+
+@pytest.mark.parametrize(
+    "preprocessor",
+    [
+        Categorizer(columns=["X"]),
+        CountVectorizer(columns=["X"]),
+        Chain(StandardScaler(columns=["X"]), MinMaxScaler(columns=["X"])),
+        FeatureHasher(columns=["X"], num_features=1, output_column="X_transformed"),
+        HashingVectorizer(columns=["X"], num_features=1),
+        LabelEncoder(label_column="X"),
+        MaxAbsScaler(columns=["X"]),
+        MinMaxScaler(columns=["X"]),
+        MultiHotEncoder(columns=["X"]),
+        Normalizer(columns=["X"]),
+        OneHotEncoder(columns=["X"]),
+        OrdinalEncoder(columns=["X"]),
+        PowerTransformer(columns=["X"], power=1),
+        RobustScaler(columns=["X"]),
+        SimpleImputer(columns=["X"]),
+        StandardScaler(columns=["X"]),
+        Concatenator(columns=["X"]),
+        Tokenizer(columns=["X"]),
+    ],
+)
+def test_repr(preprocessor):
+    representation = repr(preprocessor)
+
+    assert len(representation) < MAX_REPR_LENGTH
+    pattern = re.compile(f"^{preprocessor.__class__.__name__}\\((.*)\\)$")
+    assert pattern.match(representation)
+
+
+def test_fitted_preprocessor_without_stats():
+    """Tests that Preprocessors can be fitted without needing to set self.stats_."""
+
+    class FittablePreprocessor(Preprocessor):
+        def _fit(self, ds):
+            return self
+
+    preprocessor = FittablePreprocessor()
+    ds = ray.data.from_items([1])
+    _ = preprocessor.fit(ds)
+    assert preprocessor.fit_status() == Preprocessor.FitStatus.FITTED
+
+
+def test_fitted_preprocessor_with_stats():
+    """Tests that Preprocessors can be fitted by setting an attribute that ends
+    with _."""
+
+    class FittablePreprocessor(Preprocessor):
+        ...
+
+    preprocessor = FittablePreprocessor()
+    preprocessor.stats_ = True
+    assert preprocessor.fit_status() == Preprocessor.FitStatus.FITTED
+
+
+@patch.object(warnings, "warn")
+def test_fit_twice(mocked_warn):
+    """Tests that a warning msg should be printed."""
+    col_a = [-1, 0, 1]
+    col_b = [1, 3, 5]
+    col_c = [1, 1, None]
+    in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
+    ds = ray.data.from_pandas(in_df)
+
+    scaler = MinMaxScaler(["B", "C"])
+
+    # Fit data.
+    scaler.fit(ds)
+    assert scaler.stats_ == {"min(B)": 1, "max(B)": 5, "min(C)": 1, "max(C)": 1}
+
+    ds = ds.map_batches(lambda x: {k: v * 2 for k, v in x.items()})
+    # Fit again
+    scaler.fit(ds)
+    # Assert that the fitted state is corresponding to the second ds.
+    assert scaler.stats_ == {"min(B)": 2, "max(B)": 10, "min(C)": 2, "max(C)": 2}
+    msg = (
+        "`fit` has already been called on the preprocessor (or at least one "
+        "contained preprocessors if this is a chain). "
+        "All previously fitted state will be overwritten!"
+    )
+    mocked_warn.assert_called_once_with(msg)
+
+
+def test_fit_twice_clears_stale_stats():
+    """Tests that fit() clears stale stats when stat keys are data-dependent.
+
+    When a preprocessor's stat keys depend on the data (e.g., auto-detected columns),
+    calling fit() again on a different dataset should not retain stale stats from
+    the previous fit. This ensures that fit(A).fit(B) is equivalent to fit(B).
+    """
+
+    class DataDependentPreprocessor(Preprocessor):
+        """A preprocessor whose stat keys depend on the data columns present."""
+
+        _is_fittable = True
+
+        def _fit(self, ds):
+            # Dynamically detect columns from the dataset schema
+            schema = ds.schema()
+            column_names = list(schema.names)
+            self.stat_computation_plan.add_aggregator(
+                aggregator_fn=Mean,
+                columns=column_names,
+            )
+            return self
+
+        def _transform_pandas(self, df):
+            return df
+
+    # Dataset A has columns: "a", "b"
+    dataset_a = ray.data.from_items(
+        [
+            {"a": 1.0, "b": 10.0},
+            {"a": 2.0, "b": 20.0},
+            {"a": 3.0, "b": 30.0},
+        ]
+    )
+
+    # Dataset B has columns: "b", "c" (note: "a" is missing, "c" is new)
+    dataset_b = ray.data.from_items(
+        [
+            {"b": 100.0, "c": 1000.0},
+            {"b": 200.0, "c": 2000.0},
+            {"b": 300.0, "c": 3000.0},
+        ]
+    )
+
+    preprocessor = DataDependentPreprocessor()
+
+    # First fit on dataset A
+    preprocessor.fit(dataset_a)
+    assert preprocessor.stats_ == {"mean(a)": 2.0, "mean(b)": 20.0}
+
+    # Second fit on dataset B - stale stats should be cleared
+    preprocessor.fit(dataset_b)
+
+    # Verify stale stat "mean(a)" is NOT present
+    # Verify stats are correct after refit, and stale stats are cleared.
+    expected_stats = {"mean(b)": 200.0, "mean(c)": 2000.0}
+    assert preprocessor.stats_ == expected_stats, (
+        f"Stats after refit are incorrect. "
+        f"Expected: {expected_stats}, Got: {preprocessor.stats_}"
+    )
+
+
+def test_transform_all_configs():
+    batch_size = 2
+    num_cpus = 2
+    concurrency = 2
+    memory = 1024
+
+    class DummyPreprocessor(Preprocessor):
+        _is_fittable = False
+
+        def _get_transform_config(self):
+            return {"batch_size": batch_size}
+
+        def _transform_numpy(self, data):
+            assert ray.get_runtime_context().get_assigned_resources()["CPU"] == num_cpus
+            assert (
+                ray.get_runtime_context().get_assigned_resources()["memory"] == memory
+            )
+            # Read(10 rows) → Limit(5) → Transform(batch_size=2)
+            assert (
+                len(data["value"]) <= batch_size
+            )  # The last batch is size 1, and limit pushdown resulted in the transform occurring for fewer rows.
+            return data
+
+        def _transform_pandas(self, data):
+            raise RuntimeError(
+                "Pandas transform should not be called with numpy batch format."
+            )
+
+        def _determine_transform_to_use(self):
+            return "numpy"
+
+    prep = DummyPreprocessor()
+    ds = ray.data.from_pandas(pd.DataFrame({"value": list(range(10))}))
+    ds = prep.transform(
+        ds,
+        num_cpus=num_cpus,
+        memory=memory,
+        concurrency=concurrency,
+    )
+    assert [x["value"] for x in ds.take(5)] == [0, 1, 2, 3, 4]
+
+
+@pytest.mark.parametrize("dataset_format", ["simple", "pandas", "arrow"])
+def test_transform_all_formats(create_dummy_preprocessors, dataset_format):
+    (
+        with_nothing,
+        with_pandas,
+        with_numpy,
+        with_pandas_and_numpy,
+        with_pandas_and_numpy_preferred,
+    ) = create_dummy_preprocessors
+
+    if dataset_format == "simple":
+        ds = ray.data.range(10)
+    elif dataset_format == "pandas":
+        df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
+        ds = ray.data.from_pandas(df)
+    elif dataset_format == "arrow":
+        df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
+        ds = ray.data.from_arrow(pyarrow.Table.from_pandas(df))
+    else:
+        raise ValueError(f"Untested dataset_format configuration: {dataset_format}.")
+
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform(ds)
+
+    patcher = patch.object(ray.data.dataset.Dataset, "map_batches")
+
+    with patcher as mock_map_batches:
+        with_pandas.transform(ds)
+        mock_map_batches.assert_called_once_with(
+            with_pandas._transform_pandas,
+            batch_format=BatchFormat.PANDAS,
+            zero_copy_batch=True,
+        )
+
+    with patcher as mock_map_batches:
+        with_numpy.transform(ds)
+        mock_map_batches.assert_called_once_with(
+            with_numpy._transform_numpy,
+            batch_format=BatchFormat.NUMPY,
+            zero_copy_batch=True,
+        )
+
+    # Pandas preferred by default.
+    with patcher as mock_map_batches:
+        with_pandas_and_numpy.transform(ds)
+    mock_map_batches.assert_called_once_with(
+        with_pandas_and_numpy._transform_pandas,
+        batch_format=BatchFormat.PANDAS,
+        zero_copy_batch=True,
+    )
+
+    with patcher as mock_map_batches:
+        with_pandas_and_numpy_preferred.transform(ds)
+    mock_map_batches.assert_called_once_with(
+        with_pandas_and_numpy_preferred._transform_numpy,
+        batch_format=BatchFormat.NUMPY,
+        zero_copy_batch=True,
+    )
+
+
+def test_numpy_pandas_support_transform_batch_wrong_format(create_dummy_preprocessors):
+    # Case 1: simple dataset. No support
+    (
+        with_nothing,
+        with_pandas,
+        with_numpy,
+        with_pandas_and_numpy,
+        with_pandas_and_numpy_preferred,
+    ) = create_dummy_preprocessors
+
+    batch = [1, 2, 3]
+    with pytest.raises(ValueError):
+        with_nothing.transform_batch(batch)
+
+    with pytest.raises(ValueError):
+        with_pandas.transform_batch(batch)
+
+    with pytest.raises(ValueError):
+        with_numpy.transform_batch(batch)
+
+    with pytest.raises(ValueError):
+        with_pandas_and_numpy.transform_batch(batch)
+
+    with pytest.raises(ValueError):
+        with_pandas_and_numpy_preferred.transform_batch(batch)
+
+
+def test_numpy_pandas_support_transform_batch_pandas(create_dummy_preprocessors):
+    # Case 2: pandas dataset
+    (
+        with_nothing,
+        with_pandas,
+        with_numpy,
+        with_pandas_and_numpy,
+        with_pandas_and_numpy_preferred,
+    ) = create_dummy_preprocessors
+
+    df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
+    df_single_column = pd.DataFrame([1, 2, 3], columns=["A"])
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(df)
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(df_single_column)
+
+    assert isinstance(with_pandas.transform_batch(df), pd.DataFrame)
+    assert isinstance(with_pandas.transform_batch(df_single_column), pd.DataFrame)
+
+    assert isinstance(with_numpy.transform_batch(df), (np.ndarray, dict))
+    # We can get pd.DataFrame after returning numpy data from UDF
+    assert isinstance(with_numpy.transform_batch(df_single_column), (np.ndarray, dict))
+
+    assert isinstance(with_pandas_and_numpy.transform_batch(df), pd.DataFrame)
+    assert isinstance(
+        with_pandas_and_numpy.transform_batch(df_single_column), pd.DataFrame
+    )
+
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(df), (np.ndarray, dict)
+    )
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(df_single_column),
+        (np.ndarray, dict),
+    )
+
+
+def test_numpy_pandas_support_transform_batch_arrow(create_dummy_preprocessors):
+    # Case 3: arrow dataset
+    (
+        with_nothing,
+        with_pandas,
+        with_numpy,
+        with_pandas_and_numpy,
+        with_pandas_and_numpy_preferred,
+    ) = create_dummy_preprocessors
+
+    df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
+    df_single_column = pd.DataFrame([1, 2, 3], columns=["A"])
+
+    table = pyarrow.Table.from_pandas(df)
+    table_single_column = pyarrow.Table.from_pandas(df_single_column)
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(table)
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(table_single_column)
+
+    assert isinstance(with_pandas.transform_batch(table), pd.DataFrame)
+    assert isinstance(with_pandas.transform_batch(table_single_column), pd.DataFrame)
+
+    assert isinstance(with_numpy.transform_batch(table), (np.ndarray, dict))
+    # We can get pyarrow.Table after returning numpy data from UDF
+    assert isinstance(
+        with_numpy.transform_batch(table_single_column), (np.ndarray, dict)
+    )
+
+    assert isinstance(with_pandas_and_numpy.transform_batch(table), pd.DataFrame)
+    assert isinstance(
+        with_pandas_and_numpy.transform_batch(table_single_column), pd.DataFrame
+    )
+
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(table), (np.ndarray, dict)
+    )
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(table_single_column),
+        (np.ndarray, dict),
+    )
+
+
+def test_numpy_pandas_support_transform_batch_tensor(create_dummy_preprocessors):
+    # Case 4: tensor dataset created by from numpy data directly
+    (
+        with_nothing,
+        with_pandas,
+        with_numpy,
+        with_pandas_and_numpy,
+        with_pandas_and_numpy_preferred,
+    ) = create_dummy_preprocessors
+    np_data = np.arange(12).reshape(3, 2, 2)
+    np_single_column = {"A": np.arange(12).reshape(3, 2, 2)}
+    np_multi_column = {
+        "A": np.arange(12).reshape(3, 2, 2),
+        "B": np.arange(12, 24).reshape(3, 2, 2),
+    }
+
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(np_data)
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(np_single_column)
+    with pytest.raises(NotImplementedError):
+        with_nothing.transform_batch(np_multi_column)
+
+    assert isinstance(with_pandas.transform_batch(np_data), pd.DataFrame)
+    assert isinstance(with_pandas.transform_batch(np_single_column), pd.DataFrame)
+    assert isinstance(with_pandas.transform_batch(np_multi_column), pd.DataFrame)
+
+    assert isinstance(with_numpy.transform_batch(np_data), np.ndarray)
+    assert isinstance(with_numpy.transform_batch(np_single_column), dict)
+    assert isinstance(with_numpy.transform_batch(np_multi_column), dict)
+
+    assert isinstance(with_pandas_and_numpy.transform_batch(np_data), pd.DataFrame)
+    assert isinstance(
+        with_pandas_and_numpy.transform_batch(np_single_column), pd.DataFrame
+    )
+    assert isinstance(
+        with_pandas_and_numpy.transform_batch(np_multi_column), pd.DataFrame
+    )
+
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(np_data), np.ndarray
+    )
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(np_single_column), dict
+    )
+    assert isinstance(
+        with_pandas_and_numpy_preferred.transform_batch(np_multi_column), dict
+    )
+
+
+def test_get_input_output_columns():
+    """Tests get_input_columns() and get_output_columns() methods."""
+    # Test with preprocessors that have columns attribute
+    scaler = StandardScaler(columns=["A", "B"])
+    assert scaler.get_input_columns() == ["A", "B"]
+    assert scaler.get_output_columns() == ["A", "B"]
+
+    # Test with output_columns specified
+    scaler_with_output = StandardScaler(
+        columns=["A", "B"], output_columns=["A_scaled", "B_scaled"]
+    )
+    assert scaler_with_output.get_input_columns() == ["A", "B"]
+    assert scaler_with_output.get_output_columns() == ["A_scaled", "B_scaled"]
+
+    # Test with encoders
+    encoder = OneHotEncoder(columns=["X", "Y"])
+    assert encoder.get_input_columns() == ["X", "Y"]
+    assert encoder.get_output_columns() == ["X", "Y"]
+
+    encoder_with_output = OneHotEncoder(
+        columns=["X", "Y"], output_columns=["X_encoded", "Y_encoded"]
+    )
+    assert encoder_with_output.get_input_columns() == ["X", "Y"]
+    assert encoder_with_output.get_output_columns() == ["X_encoded", "Y_encoded"]
+
+    # Test LabelEncoder without output_column (in-place transformation)
+    label_encoder = LabelEncoder(label_column="target")
+    assert label_encoder.get_input_columns() == ["target"]
+    assert label_encoder.get_output_columns() == ["target"]
+
+    # Test LabelEncoder with output_column (append mode)
+    label_encoder = LabelEncoder(label_column="target", output_column="target_encoded")
+    assert label_encoder.get_input_columns() == ["target"]
+    assert label_encoder.get_output_columns() == ["target_encoded"]
+
+    # Test Concatenator (uses output_column_name instead of output_columns)
+    concatenator = Concatenator(columns=["A", "B"])
+    assert concatenator.get_input_columns() == ["A", "B"]
+    assert concatenator.get_output_columns() == ["concat_out"]
+
+    concatenator_with_output = Concatenator(
+        columns=["A", "B"], output_column_name="AB_concat"
+    )
+    assert concatenator_with_output.get_input_columns() == ["A", "B"]
+    assert concatenator_with_output.get_output_columns() == ["AB_concat"]
+
+    # Test FeatureHasher (uses output_column instead of output_columns)
+    feature_hasher = FeatureHasher(
+        columns=["token1", "token2"], num_features=8, output_column="hashed"
+    )
+    assert feature_hasher.get_input_columns() == ["token1", "token2"]
+    assert feature_hasher.get_output_columns() == ["hashed"]
+
+    # Test TorchVisionPreprocessor (uses _columns and _output_columns)
+    torch_preprocessor = TorchVisionPreprocessor(
+        columns=["image"], transform=lambda x: x
+    )
+    assert torch_preprocessor.get_input_columns() == ["image"]
+    assert torch_preprocessor.get_output_columns() == ["image"]
+
+    torch_preprocessor_with_output = TorchVisionPreprocessor(
+        columns=["image"], transform=lambda x: x, output_columns=["image_transformed"]
+    )
+    assert torch_preprocessor_with_output.get_input_columns() == ["image"]
+    assert torch_preprocessor_with_output.get_output_columns() == ["image_transformed"]
+
+    # Test with preprocessor without columns attribute
+    class CustomPreprocessor(Preprocessor):
+        _is_fittable = False
+
+    custom = CustomPreprocessor()
+    assert custom.get_input_columns() == []
+    assert custom.get_output_columns() == []
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))