chore: import upstream snapshot with attribution

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

@@ -0,0 +1,312 @@
+import pandas as pd
+import pytest
+
+import ray
+from ray.data._internal.util import rows_same
+from ray.data.preprocessors import CustomKBinsDiscretizer, UniformKBinsDiscretizer
+
+
+@pytest.mark.parametrize("bins", (3, {"A": 4, "B": 3}))
+@pytest.mark.parametrize(
+    "dtypes",
+    (
+        None,
+        {"A": int, "B": int},
+        {"A": int, "B": pd.CategoricalDtype(["cat1", "cat2", "cat3"], ordered=True)},
+    ),
+)
+@pytest.mark.parametrize("right", (True, False))
+@pytest.mark.parametrize("include_lowest", (True, False))
+def test_uniform_kbins_discretizer(
+    bins,
+    dtypes,
+    right,
+    include_lowest,
+):
+    """Tests basic UniformKBinsDiscretizer functionality."""
+
+    col_a = [0.2, 1.4, 2.5, 6.2, 9.7, 2.1]
+    col_b = col_a.copy()
+    col_c = col_a.copy()
+    in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
+    ds = ray.data.from_pandas(in_df).repartition(2)
+
+    discretizer = UniformKBinsDiscretizer(
+        ["A", "B"], bins=bins, dtypes=dtypes, right=right, include_lowest=include_lowest
+    )
+
+    transformed = discretizer.fit_transform(ds)
+    out_df = transformed.to_pandas()
+
+    if isinstance(bins, dict):
+        bins_A = bins["A"]
+        bins_B = bins["B"]
+    else:
+        bins_A = bins_B = bins
+
+    labels_A = False
+    ordered_A = True
+    labels_B = False
+    ordered_B = True
+    if isinstance(dtypes, dict):
+        if isinstance(dtypes.get("A"), pd.CategoricalDtype):
+            labels_A = dtypes.get("A").categories
+            ordered_A = dtypes.get("A").ordered
+        if isinstance(dtypes.get("B"), pd.CategoricalDtype):
+            labels_B = dtypes.get("B").categories
+            ordered_B = dtypes.get("B").ordered
+
+    # Create expected dataframe with transformed columns
+    expected_df = in_df.copy()
+    expected_df["A"] = pd.cut(
+        in_df["A"],
+        bins_A,
+        labels=labels_A,
+        ordered=ordered_A,
+        right=right,
+        include_lowest=include_lowest,
+    )
+    expected_df["B"] = pd.cut(
+        in_df["B"],
+        bins_B,
+        labels=labels_B,
+        ordered=ordered_B,
+        right=right,
+        include_lowest=include_lowest,
+    )
+
+    # Use rows_same to compare regardless of row ordering
+    assert rows_same(out_df, expected_df)
+
+    # append mode
+    expected_message = "The length of columns and output_columns must match."
+    with pytest.raises(ValueError, match=expected_message):
+        UniformKBinsDiscretizer(["A", "B"], bins=bins, output_columns=["A_discretized"])
+
+    discretizer = UniformKBinsDiscretizer(
+        ["A", "B"],
+        bins=bins,
+        dtypes=dtypes,
+        right=right,
+        include_lowest=include_lowest,
+        output_columns=["A_discretized", "B_discretized"],
+    )
+
+    transformed = discretizer.fit_transform(ds)
+    out_df = transformed.to_pandas()
+
+    # Create expected dataframe with appended columns
+    expected_df = in_df.copy()
+    expected_df["A_discretized"] = pd.cut(
+        in_df["A"],
+        bins_A,
+        labels=labels_A,
+        ordered=ordered_A,
+        right=right,
+        include_lowest=include_lowest,
+    )
+    expected_df["B_discretized"] = pd.cut(
+        in_df["B"],
+        bins_B,
+        labels=labels_B,
+        ordered=ordered_B,
+        right=right,
+        include_lowest=include_lowest,
+    )
+
+    # Use rows_same to compare regardless of row ordering
+    assert rows_same(out_df, expected_df)
+
+
+@pytest.mark.parametrize(
+    "bins", ([3, 4, 6, 9], {"A": [3, 4, 6, 8, 9], "B": [3, 4, 6, 9]})
+)
+@pytest.mark.parametrize(
+    "dtypes",
+    (
+        None,
+        {"A": int, "B": int},
+        {"A": int, "B": pd.CategoricalDtype(["cat1", "cat2", "cat3"], ordered=True)},
+    ),
+)
+@pytest.mark.parametrize("right", (True, False))
+@pytest.mark.parametrize("include_lowest", (True, False))
+def test_custom_kbins_discretizer(
+    bins,
+    dtypes,
+    right,
+    include_lowest,
+):
+    """Tests basic CustomKBinsDiscretizer functionality."""
+
+    col_a = [0.2, 1.4, 2.5, 6.2, 9.7, 2.1]
+    col_b = col_a.copy()
+    col_c = col_a.copy()
+    in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
+    ds = ray.data.from_pandas(in_df).repartition(2)
+
+    discretizer = CustomKBinsDiscretizer(
+        ["A", "B"], bins=bins, dtypes=dtypes, right=right, include_lowest=include_lowest
+    )
+
+    transformed = discretizer.transform(ds)
+    out_df = transformed.to_pandas()
+
+    if isinstance(bins, dict):
+        bins_A = bins["A"]
+        bins_B = bins["B"]
+    else:
+        bins_A = bins_B = bins
+
+    labels_A = False
+    ordered_A = True
+    labels_B = False
+    ordered_B = True
+    if isinstance(dtypes, dict):
+        if isinstance(dtypes.get("A"), pd.CategoricalDtype):
+            labels_A = dtypes.get("A").categories
+            ordered_A = dtypes.get("A").ordered
+        if isinstance(dtypes.get("B"), pd.CategoricalDtype):
+            labels_B = dtypes.get("B").categories
+            ordered_B = dtypes.get("B").ordered
+
+    # Create expected dataframe with transformed columns
+    expected_df = in_df.copy()
+    expected_df["A"] = pd.cut(
+        in_df["A"],
+        bins_A,
+        labels=labels_A,
+        ordered=ordered_A,
+        right=right,
+        include_lowest=include_lowest,
+    )
+    expected_df["B"] = pd.cut(
+        in_df["B"],
+        bins_B,
+        labels=labels_B,
+        ordered=ordered_B,
+        right=right,
+        include_lowest=include_lowest,
+    )
+
+    # Use rows_same to compare regardless of row ordering
+    assert rows_same(out_df, expected_df)
+
+    # append mode
+    expected_message = "The length of columns and output_columns must match."
+    with pytest.raises(ValueError, match=expected_message):
+        CustomKBinsDiscretizer(["A", "B"], bins=bins, output_columns=["A_discretized"])
+
+    discretizer = CustomKBinsDiscretizer(
+        ["A", "B"],
+        bins=bins,
+        dtypes=dtypes,
+        right=right,
+        include_lowest=include_lowest,
+        output_columns=["A_discretized", "B_discretized"],
+    )
+
+    transformed = discretizer.fit_transform(ds)
+    out_df = transformed.to_pandas()
+
+    # Create expected dataframe with appended columns
+    expected_df = in_df.copy()
+    expected_df["A_discretized"] = pd.cut(
+        in_df["A"],
+        bins_A,
+        labels=labels_A,
+        ordered=ordered_A,
+        right=right,
+        include_lowest=include_lowest,
+    )
+    expected_df["B_discretized"] = pd.cut(
+        in_df["B"],
+        bins_B,
+        labels=labels_B,
+        ordered=ordered_B,
+        right=right,
+        include_lowest=include_lowest,
+    )
+
+    # Use rows_same to compare regardless of row ordering
+    assert rows_same(out_df, expected_df)
+
+
+def test_custom_kbins_discretizer_serialization():
+    """Test CustomKBinsDiscretizer serialization and deserialization functionality."""
+    from ray.data.preprocessor import SerializablePreprocessorBase
+
+    # Create discretizer
+    discretizer = CustomKBinsDiscretizer(
+        columns=["A"], bins={"A": [0, 1, 2, 3]}, right=True
+    )
+
+    # Serialize using CloudPickle
+    serialized = discretizer.serialize()
+
+    # Verify it's binary CloudPickle format
+    assert isinstance(serialized, bytes)
+    assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
+
+    # Deserialize
+    deserialized = CustomKBinsDiscretizer.deserialize(serialized)
+
+    # Verify type and field values
+    assert isinstance(deserialized, CustomKBinsDiscretizer)
+    assert deserialized.columns == ["A"]
+    assert deserialized.bins == {"A": [0, 1, 2, 3]}
+    assert deserialized.right is True
+
+    # Verify it works correctly
+    df = pd.DataFrame({"A": [0.5, 1.5, 2.5]})
+    result = deserialized.transform_batch(df)
+
+    # Verify discretization was applied correctly
+    assert "A" in result.columns
+    assert len(result) == 3
+
+
+def test_uniform_kbins_discretizer_serialization():
+    """Test UniformKBinsDiscretizer serialization and deserialization functionality."""
+    import ray
+    from ray.data.preprocessor import SerializablePreprocessorBase
+
+    # Create and fit discretizer
+    discretizer = UniformKBinsDiscretizer(columns=["A"], bins=3)
+    df = pd.DataFrame({"A": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]})
+    ds = ray.data.from_pandas(df)
+    fitted_discretizer = discretizer.fit(ds)
+
+    # Serialize using CloudPickle
+    serialized = fitted_discretizer.serialize()
+
+    # Verify it's binary CloudPickle format
+    assert isinstance(serialized, bytes)
+    assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
+
+    # Deserialize
+    deserialized = UniformKBinsDiscretizer.deserialize(serialized)
+
+    # Verify type and field values
+    assert isinstance(deserialized, UniformKBinsDiscretizer)
+    assert deserialized._fitted
+    assert deserialized.columns == ["A"]
+    assert deserialized.bins == 3
+
+    # Verify stats are preserved: bin edges for 3 bins = 4 edge values
+    assert "A" in deserialized.stats_
+    assert len(deserialized.stats_["A"]) == 4
+
+    # Verify it works correctly
+    test_df = pd.DataFrame({"A": [1.5, 3.5, 5.5]})
+    result = deserialized.transform_batch(test_df)
+
+    # Verify discretization was applied correctly
+    assert "A" in result.columns
+    assert len(result) == 3
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))