Backport PR #31290: BUG: Handle IntegerArray in pd.cut (#31374)

Co-authored-by: Tom Augspurger <[email protected]>

Author: 2 people

pandas/core/reshape/tile.py

@@ -14,7 +14,9 @@
     is_datetime64_dtype,
     is_datetime64tz_dtype,
     is_datetime_or_timedelta_dtype,
+    is_extension_array_dtype,
     is_integer,
+    is_integer_dtype,
     is_list_like,
     is_scalar,
     is_timedelta64_dtype,
@@ -205,6 +207,12 @@ def cut(
     x = _preprocess_for_cut(x)
     x, dtype = _coerce_to_type(x)
 
+    # To support cut(IntegerArray), we convert to object dtype with NaN
+    # Will properly support in the future.
+    # https://github.com/pandas-dev/pandas/pull/31290
+    if is_extension_array_dtype(x.dtype) and is_integer_dtype(x.dtype):
+        x = x.to_numpy(dtype=object, na_value=np.nan)
+
     if not np.iterable(bins):
         if is_scalar(bins) and bins < 1:
             raise ValueError("`bins` should be a positive integer.")

pandas/tests/arrays/test_integer.py

@@ -1059,6 +1059,19 @@ def test_value_counts_na():
     tm.assert_series_equal(result, expected)
 
 
+@pytest.mark.parametrize("bins", [3, [0, 5, 15]])
+@pytest.mark.parametrize("right", [True, False])
+@pytest.mark.parametrize("include_lowest", [True, False])
+def test_cut(bins, right, include_lowest):
+    a = np.random.randint(0, 10, size=50).astype(object)
+    a[::2] = np.nan
+    result = pd.cut(
+        pd.array(a, dtype="Int64"), bins, right=right, include_lowest=include_lowest
+    )
+    expected = pd.cut(a, bins, right=right, include_lowest=include_lowest)
+    tm.assert_categorical_equal(result, expected)
+
+
 # TODO(jreback) - these need testing / are broken
 
 # shift