Correctly factorize values outside bin edges

Author: dcherian

flox/core.py

@@ -431,6 +431,7 @@ def factorize_(
             # pd.cut with bins = IntervalIndex[datetime64] doesn't work...
             if groupvar.dtype.kind == "M":
                 expect = np.concatenate([expect.left.to_numpy(), [expect.right[-1].to_numpy()]])
+            # code is -1 for values outside the bounds of all intervals
             idx = pd.cut(groupvar.ravel(), bins=expect).codes.copy()
         else:
             if expect is not None and reindex:
@@ -455,9 +456,12 @@ def factorize_(
     grp_shape = tuple(len(grp) for grp in found_groups)
     ngroups = np.prod(grp_shape)
     if len(by) > 1:
-        group_idx = np.ravel_multi_index(factorized, grp_shape, mode="wrap").reshape(by[0].shape)
-        nan_by_mask = reduce(np.logical_or, [isnull(b) for b in by])
+        group_idx = np.ravel_multi_index(factorized, grp_shape, mode="wrap")
+        # NaNs; as well as values outside the bins are coded by -1
+        # Restore these after the raveling
+        nan_by_mask = reduce(np.logical_or, [(f == -1) for f in factorized])
         group_idx[nan_by_mask] = -1
+        group_idx = group_idx.reshape(by[0].shape)
     else:
         group_idx = factorized[0]
 

tests/test_core.py

@@ -888,3 +888,20 @@ def test_group_by_datetime(engine, method):
     )
     expected = np.broadcast_to(expected, (2, 3, expected.shape[-1]))
     assert_equal(expected, actual)
+
+
+def test_factorize_values_outside_bins():
+
+    vals = factorize_(
+        (np.arange(10).reshape(5, 2), np.arange(10).reshape(5, 2)),
+        axis=(0, 1),
+        expected_groups=(
+            pd.IntervalIndex.from_breaks(np.arange(2, 8, 1)),
+            pd.IntervalIndex.from_breaks(np.arange(2, 8, 1)),
+        ),
+        reindex=True,
+        fastpath=True,
+    )
+    actual = vals[0]
+    expected = np.array([[-1, -1], [-1, 0], [6, 12], [18, 24], [-1, -1]])
+    assert_equal(expected, actual)