Repository: modin-project/modin

Issue:
REFACTOR: Add query compiler implementation of pd.cut
porting https://github.com/ponder-org/modin-public/commit/582cbba6b1becd9fd4d0b15a2f7d6aad97c78f68

note: if it's easy, can add an implementation for dask/ray/unidist as well


Discussion:

Resolving patch:
diff --git a/modin/core/storage_formats/base/query_compiler.py b/modin/core/storage_formats/base/query_compiler.py
--- a/modin/core/storage_formats/base/query_compiler.py
+++ b/modin/core/storage_formats/base/query_compiler.py
@@ -3788,6 +3788,38 @@ def repeat(self, repeats):
         """
         return SeriesDefault.register(pandas.Series.repeat)(self, repeats=repeats)
 
+    @doc_utils.add_refer_to("cut")
+    def cut(
+        self,
+        bins,
+        **kwargs,
+    ):
+        """
+        Bin values into discrete intervals.
+
+        Parameters
+        ----------
+        bins : int, array of ints, or IntervalIndex
+            The criteria to bin by.
+        **kwargs : dict
+            The keyword arguments to pass through.
+
+        Returns
+        -------
+        BaseQueryCompiler or np.ndarray or list[np.ndarray]
+            Returns the result of pd.cut.
+        """
+
+        def squeeze_and_cut(df, *args, **kwargs):
+            # We need this function to ensure we squeeze our internal
+            # representation (a dataframe) to a Series.
+            series = df.squeeze(axis=1)
+            return pandas.cut(series, *args, **kwargs)
+
+        # We use `default_to_pandas` here since the type and number of
+        # results can change depending on the input arguments.
+        return self.default_to_pandas(squeeze_and_cut, bins, **kwargs)
+
     # Indexing
 
     index = property(_get_axis(0), _set_axis(0))
diff --git a/modin/pandas/__init__.py b/modin/pandas/__init__.py
--- a/modin/pandas/__init__.py
+++ b/modin/pandas/__init__.py
@@ -26,7 +26,6 @@
     warnings.simplefilter("ignore")
     from pandas import (
         eval,
-        cut,
         factorize,
         test,
         date_range,
@@ -247,6 +246,7 @@ def init_remote_ray(partition):
     wide_to_long,
     to_timedelta,
     pivot_table,
+    cut,
 )
 
 from .plotting import Plotting as plotting
diff --git a/modin/pandas/general.py b/modin/pandas/general.py
--- a/modin/pandas/general.py
+++ b/modin/pandas/general.py
@@ -290,6 +290,60 @@ def qcut(
     return x._qcut(q, **kwargs)
 
 
+@_inherit_docstrings(pandas.cut, apilink="pandas.cut")
+@enable_logging
+def cut(
+    x,
+    bins,
+    right: bool = True,
+    labels=None,
+    retbins: bool = False,
+    precision: int = 3,
+    include_lowest: bool = False,
+    duplicates: str = "raise",
+    ordered: bool = True,
+):
+    if isinstance(x, DataFrame):
+        raise ValueError("Input array must be 1 dimensional")
+    if not isinstance(x, Series):
+        ErrorMessage.default_to_pandas(
+            reason=f"pd.cut is not supported on objects of type {type(x)}"
+        )
+        import pandas
+
+        return pandas.cut(
+            x,
+            bins,
+            right=right,
+            labels=labels,
+            retbins=retbins,
+            precision=precision,
+            include_lowest=include_lowest,
+            duplicates=duplicates,
+            ordered=ordered,
+        )
+
+    def _wrap_in_series_object(qc_result):
+        if isinstance(qc_result, type(x._query_compiler)):
+            return Series(query_compiler=qc_result)
+        if isinstance(qc_result, (tuple, list)):
+            return tuple([_wrap_in_series_object(result) for result in qc_result])
+        return qc_result
+
+    return _wrap_in_series_object(
+        x._query_compiler.cut(
+            bins,
+            right=right,
+            labels=labels,
+            retbins=retbins,
+            precision=precision,
+            include_lowest=include_lowest,
+            duplicates=duplicates,
+            ordered=ordered,
+        )
+    )
+
+
 @_inherit_docstrings(pandas.unique, apilink="pandas.unique")
 @enable_logging
 def unique(values):  # noqa: PR01, RT01, D200


Test patch:
diff --git a/modin/pandas/test/test_general.py b/modin/pandas/test/test_general.py
--- a/modin/pandas/test/test_general.py
+++ b/modin/pandas/test/test_general.py
@@ -758,6 +758,85 @@ def test_qcut(retbins):
     df_equals(modin_result, pandas_result)
 
 
+@pytest.mark.parametrize(
+    "bins, labels",
+    [
+        pytest.param(
+            [-int(1e18), -1000, 0, 1000, 2000, int(1e18)],
+            [
+                "-inf_to_-1000",
+                "-1000_to_0",
+                "0_to_1000",
+                "1000_to_2000",
+                "2000_to_inf",
+            ],
+            id="bin_list_spanning_entire_range_with_custom_labels",
+        ),
+        pytest.param(
+            [-int(1e18), -1000, 0, 1000, 2000, int(1e18)],
+            None,
+            id="bin_list_spanning_entire_range_with_default_labels",
+        ),
+        pytest.param(
+            [-1000, 0, 1000, 2000], None, id="bin_list_not_spanning_entire_range"
+        ),
+        pytest.param(
+            10,
+            [f"custom_label{i}" for i in range(9)],
+            id="int_bin_10_with_custom_labels",
+        ),
+        pytest.param(1, None, id="int_bin_1_with_default_labels"),
+        pytest.param(-1, None, id="int_bin_-1_with_default_labels"),
+        pytest.param(111, None, id="int_bin_111_with_default_labels"),
+    ],
+)
+@pytest.mark.parametrize("retbins", bool_arg_values, ids=bool_arg_keys)
+def test_cut(retbins, bins, labels):
+    # Would use `eval_general` here, but `eval_general` expects the operation
+    # to be supported by Modin, and so errors out when we give the defaulting
+    # to pandas UserWarning. We could get around this by using
+    # @pytest.mark.filterwarnings("ignore"), but then `eval_general` fails because
+    # sometimes the return type of pd.cut is an np.ndarray, and `eval_general` does
+    # not know how to handle that.
+    try:
+        pd_result = pandas.cut(
+            pandas.Series(range(1000)), retbins=retbins, bins=bins, labels=labels
+        )
+    except Exception as pd_e:
+        with pytest.raises(Exception) as md_e:
+            with warns_that_defaulting_to_pandas():
+                md_result = pd.cut(
+                    pd.Series(range(1000)), retbins=retbins, bins=bins, labels=labels
+                )
+        assert isinstance(
+            md_e.value, type(pd_e)
+        ), f"Got Modin Exception type {type(md_e.value)}, but pandas Exception type {type(pd_e)} was expected"
+    else:
+        with warns_that_defaulting_to_pandas():
+            md_result = pd.cut(
+                pd.Series(range(1000)), retbins=retbins, bins=bins, labels=labels
+            )
+        if not isinstance(pd_result, tuple):
+            df_equals(md_result, pd_result)
+        else:
+            assert isinstance(
+                md_result, tuple
+            ), "Modin returned single value, but pandas returned tuple of values"
+            for pd_res, md_res in zip(pd_result, md_result):
+                if isinstance(pd_res, pandas.Series):
+                    df_equals(pd_res, md_res)
+                else:
+                    np.testing.assert_array_equal(pd_res, md_res)
+
+
+def test_cut_fallback():
+    # Test case for falling back to pandas for cut.
+    pandas_result = pandas.cut(range(5), 4)
+    with warns_that_defaulting_to_pandas():
+        modin_result = pd.cut(range(5), 4)
+    df_equals(modin_result, pandas_result)
+
+
 @pytest.mark.parametrize(
     "data", [test_data_values[0], []], ids=["test_data_values[0]", "[]"]
 )