Better handling of RangeIndex for sparse reindexing

Author: dcherian

flox/core.py

@@ -2327,6 +2327,8 @@ def _factorize_multiple(
     if any_by_dask:
         import dask.array
 
+        from . import dask_array_ops  # noqa
+
         # unifying chunks will make sure all arrays in `by` are dask arrays
         # with compatible chunks, even if there was originally a numpy array
         inds = tuple(range(by[0].ndim))

tests/test_core.py

@@ -2085,12 +2085,14 @@ def test_datetime_timedelta_first_last(engine, func) -> None:
 
 @requires_dask
 @requires_sparse
-def test_reindex_sparse():
+@pytest.mark.xdist_group(name="sparse-group")
+@pytest.mark.parametrize("size", [2**62 - 1, 11])
+def test_reindex_sparse(size):
     import sparse
 
     array = dask.array.ones((2, 12), chunks=(-1, 3))
     func = "sum"
-    expected_groups = pd.Index(np.arange(11))
+    expected_groups = pd.RangeIndex(size)
     by = dask.array.from_array(np.repeat(np.arange(6) * 2, 2), chunks=(3,))
     dense = np.zeros((2, 11))
     dense[..., np.arange(6) * 2] = 2
@@ -2115,9 +2117,10 @@ def mocked_reindex(*args, **kwargs):
         actual, *_ = groupby_reduce(
             array, by, func=func, reindex=reindex, expected_groups=expected_groups, fill_value=0
         )
-        assert_equal(actual, expected)
-        # once during graph construction, 10 times afterward
-        assert mocked_func.call_count > 1
+        if size == 11:
+            assert_equal(actual, expected)
+            # once during graph construction, 10 times afterward
+            assert mocked_func.call_count > 1
 
 
 def test_sparse_errors():