BUG: Handle IntegerArray in pd.cut

xref pandas-dev#30944.
I think this doesn't close it, since only the pd.cut compoment
is fixed.

Author: TomAugspurger

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,
@@ -209,16 +211,28 @@ def cut(
         if is_scalar(bins) and bins < 1:
             raise ValueError("`bins` should be a positive integer.")
 
-        try:  # for array-like
-            sz = x.size
+        # TODO: Support arbitrary Extension Arrays. We need
+        # For now, we're only attempting to support IntegerArray.
+        # See the note on _bins_to_cuts about what is needed.
+        is_nullable_integer = is_extension_array_dtype(x.dtype) and is_integer_dtype(
+            x.dtype
+        )
+        try:
+            if is_extension_array_dtype(x) and is_integer_dtype(x):
+                sz = len(x)
+            else:
+                sz = x.size
         except AttributeError:
             x = np.asarray(x)
             sz = x.size
 
         if sz == 0:
             raise ValueError("Cannot cut empty array")
 
-        rng = (nanops.nanmin(x), nanops.nanmax(x))
+        if is_nullable_integer:
+            rng = x._reduce("min"), x._reduce("max")
+        else:
+            rng = (nanops.nanmin(x), nanops.nanmax(x))
         mn, mx = [mi + 0.0 for mi in rng]
 
         if np.isinf(mn) or np.isinf(mx):
@@ -383,10 +397,26 @@ def _bins_to_cuts(
             bins = unique_bins
 
     side = "left" if right else "right"
-    ids = ensure_int64(bins.searchsorted(x, side=side))
+    is_nullable_integer = is_extension_array_dtype(x.dtype) and is_integer_dtype(
+        x.dtype
+    )
+
+    if is_nullable_integer:
+        # TODO: Support other extension types somehow. We don't currently
+        # We *could* use factorize here, but that does more that we need.
+        # We just need some integer representation, and the NA values needn't
+        # even be marked specially.
+        x_int = x._ndarray_values
+        ids = ensure_int64(bins.searchsorted(x_int, side=side))
+    else:
+        ids = ensure_int64(bins.searchsorted(x, side=side))
 
     if include_lowest:
-        ids[x == bins[0]] = 1
+        mask = x == bins[0]
+        if is_nullable_integer:
+            # when x is integer
+            mask = mask.to_numpy(na_value=False, dtype=bool)
+        ids[mask] = 1
 
     na_mask = isna(x) | (ids == len(bins)) | (ids == 0)
     has_nas = na_mask.any()

pandas/tests/arrays/test_integer.py

@@ -1061,6 +1061,20 @@ 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(float)
+    a[::2] = np.nan
+    tm.assert_categorical_equal(
+        pd.cut(
+            pd.array(a, dtype="Int64"), bins, right=right, include_lowest=include_lowest
+        ),
+        pd.cut(a, bins, right=right, include_lowest=include_lowest),
+    )
+
+
 # TODO(jreback) - these need testing / are broken
 
 # shift