Add tests

Author: dcherian

tests/test_core.py

@@ -1066,7 +1066,8 @@ def test_cohorts_map_reduce_consistent_dtypes(method, dtype, labels_dtype):
 @pytest.mark.parametrize("func", ALL_FUNCS)
 @pytest.mark.parametrize("axis", (-1, None))
 @pytest.mark.parametrize("method", ["blockwise", "cohorts", "map-reduce"])
-def test_cohorts_nd_by(func, method, axis, engine):
+@pytest.mark.parametrize("by_is_dask", [True, False])
+def test_cohorts_nd_by(by_is_dask, func, method, axis, engine):
     if (
         ("arg" in func and (axis is None or engine in ["flox", "numbagg"]))
         or (method != "blockwise" and func in BLOCKWISE_FUNCS)
@@ -1080,10 +1081,12 @@ def test_cohorts_nd_by(func, method, axis, engine):
     o2 = dask.array.ones((2, 3), chunks=-1)
 
     array = dask.array.block([[o, 2 * o], [3 * o2, 4 * o2]])
-    by = array.compute().astype(np.int64)
+    by = array.astype(np.int64)
     by[0, 1] = 30
     by[2, 1] = 40
     by[0, 4] = 31
+    if not by_is_dask:
+        by = by.compute()
     array = np.broadcast_to(array, (2, 3) + array.shape)
 
     if func in ["any", "all"]:
@@ -1099,10 +1102,19 @@ def test_cohorts_nd_by(func, method, axis, engine):
     assert_equal(groups, sorted_groups)
     assert_equal(actual, expected)
 
-    actual, groups = groupby_reduce(array, by, sort=False, **kwargs)
-    assert_equal(groups, np.array([1, 30, 2, 31, 3, 4, 40], dtype=np.int64))
-    reindexed = reindex_(actual, groups, pd.Index(sorted_groups))
-    assert_equal(reindexed, expected)
+    if isinstance(by, dask.array.Array):
+        cache.clear()
+        actual_cohorts = find_group_cohorts(by, array.chunks[-by.ndim :])
+        expected_cohorts = find_group_cohorts(by.compute(), array.chunks[-by.ndim :])
+        assert actual_cohorts == expected_cohorts
+        # assert cache.nbytes
+
+    if not isinstance(by, dask.array.Array):
+        # Always sorting groups with cohorts and dask array
+        actual, groups = groupby_reduce(array, by, sort=False, **kwargs)
+        assert_equal(groups, np.array([1, 30, 2, 31, 3, 4, 40], dtype=np.int64))
+        reindexed = reindex_(actual, groups, pd.Index(sorted_groups))
+        assert_equal(reindexed, expected)
 
 
 @pytest.mark.parametrize("func", ["sum", "count"])