Fix datetime binning

Author: dcherian

flox/core.py

@@ -400,10 +400,12 @@ def factorize_(by: tuple, axis, expected_groups: tuple[pd.Index, ...] = None):
             # this makes the reindexing logic simpler.
             if expect is None:
                 raise ValueError("Please pass bin edges in expected_groups.")
-            # idx = np.digitize(groupvar.ravel(), expect) - 1
-            idx = pd.cut(groupvar.ravel(), bins=expect, labels=False).codes.copy()
-            # same sentinel value as factorize
+            # TODO: fix for binning
             found_groups.append(expect)
+            # 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()]])
+            idx = pd.cut(groupvar.ravel(), bins=expect).codes.copy()
         else:
             idx, groups = pd.factorize(groupvar.ravel())
             found_groups.append(np.array(groups))
@@ -1246,13 +1248,15 @@ def _assert_by_is_aligned(shape, by):
 
 
 def _convert_expected_groups_to_index(expected_groups, isbin: bool) -> pd.Index | None:
-    if isinstance(expected_groups, pd.Index):
+    if isinstance(expected_groups, pd.IntervalIndex) or (
+        isinstance(expected_groups, pd.Index) and not isbin
+    ):
         return expected_groups
     if isbin:
         return pd.IntervalIndex.from_arrays(expected_groups[:-1], expected_groups[1:])
     elif expected_groups is not None:
         return pd.Index(expected_groups)
-    return None
+    return expected_groups
 
 
 def groupby_reduce(

tests/test_core.py

@@ -4,7 +4,9 @@
 from numpy_groupies.aggregate_numpy import aggregate
 
 from flox.core import (
+    _convert_expected_groups_to_index,
     _get_optimal_chunks_for_groups,
+    factorize_,
     find_group_cohorts,
     groupby_reduce,
     rechunk_for_cohorts,
@@ -758,3 +760,18 @@ def test_empty_bins(func, engine):
     )
     expected = np.array([1.0, 1.0, np.nan])
     assert_equal(actual, expected)
+
+
+def test_datetime_binning():
+    time_bins = pd.date_range(start="2010-08-01", end="2010-08-15", freq="24H")
+    by = pd.date_range("2010-08-01", "2010-08-15", freq="15min")
+
+    actual = _convert_expected_groups_to_index(time_bins, isbin=True)
+    expected = pd.IntervalIndex.from_arrays(time_bins[:-1], time_bins[1:])
+    assert_equal(actual, expected)
+
+    ret = factorize_((by.to_numpy(),), axis=0, expected_groups=(actual,))
+    group_idx = ret[0]
+    expected = pd.cut(by, time_bins).codes.copy()
+    expected[0] = 14  # factorize doesn't return -1 for nans
+    assert_equal(group_idx, expected)