Improvements to reindex

Author: dcherian

flox/core.py

@@ -323,10 +323,23 @@ def rechunk_for_blockwise(array, axis, labels):
         return array.rechunk({axis: newchunks})
 
 
-def reindex_(array: np.ndarray, from_, to, fill_value=None, axis: int = -1) -> np.ndarray:
+def reindex_multiple(array, groups, expected_groups, fill_value, promote=False):
+    for ax, (group, expect) in enumerate(zip(groups, expected_groups)):
+        array = reindex_(
+            array, group, expect, fill_value=fill_value, axis=ax - len(groups), promote=promote
+        )
+    return array
+
 
-    assert isinstance(to, pd.Index)
-    assert axis in (0, -1)
+def reindex_(
+    array: np.ndarray, from_, to, fill_value=None, axis: int = -1, promote: bool = False
+) -> np.ndarray:
+
+    if not isinstance(to, pd.Index):
+        if promote:
+            to = pd.Index(to)
+        else:
+            raise ValueError("reindex requires a pandas.Index or promote=True")
 
     if to.ndim > 1:
         raise ValueError(f"Cannot reindex to a multidimensional array: {to}")
@@ -353,15 +366,12 @@ def reindex_(array: np.ndarray, from_, to, fill_value=None, axis: int = -1) -> n
     if any(idx == -1):
         if fill_value is None:
             raise ValueError("Filling is required. fill_value cannot be None.")
-        if axis == 0:
-            loc = (idx == -1, ...)
-        else:
-            loc = (..., idx == -1)
+        indexer[axis] = idx == -1
         # This allows us to match xarray's type promotion rules
         if fill_value is xrdtypes.NA or np.isnan(fill_value):
             new_dtype, fill_value = xrdtypes.maybe_promote(reindexed.dtype)
             reindexed = reindexed.astype(new_dtype, copy=False)
-        reindexed[loc] = fill_value
+        reindexed[tuple(indexer)] = fill_value
     return reindexed
 
 

tests/test_core.py

@@ -487,13 +487,24 @@ def test_groupby_all_nan_blocks(engine):
     assert_equal(actual, expected)
 
 
-def test_reindex():
-    array = np.array([1, 2])
-    groups = np.array(["a", "b"])
-    expected_groups = ["a", "b", "c"]
+@pytest.mark.parametrize("axis", (0, 1, 2, -1))
+def test_reindex(axis):
+    shape = [2, 2, 2]
     fill_value = 0
-    result = reindex_(array, groups, pd.Index(expected_groups), fill_value, axis=-1)
-    assert_equal(result, np.array([1, 2, 0]))
+
+    array = np.broadcast_to(np.array([1, 2]), shape)
+    groups = np.array(["a", "b"])
+    expected_groups = pd.Index(["a", "b", "c"])
+    actual = reindex_(array, groups, expected_groups, fill_value=fill_value, axis=axis)
+
+    if axis < 0:
+        axis = array.ndim + axis
+    result_shape = tuple(len(expected_groups) if ax == axis else s for ax, s in enumerate(shape))
+    slicer = tuple(slice(None, s) for s in shape)
+    expected = np.full(result_shape, fill_value)
+    expected[slicer] = array
+
+    assert_equal(actual, expected)
 
 
 @pytest.mark.xfail