Check method only for dask reductions. (#241)

Author: dcherian

flox/core.py

@@ -1875,12 +1875,6 @@ def groupby_reduce(
         axis_ = np.core.numeric.normalize_axis_tuple(axis, array.ndim)  # type: ignore
     nax = len(axis_)
 
-    if method in ["blockwise", "cohorts"] and nax != by_.ndim:
-        raise NotImplementedError(
-            "Must reduce along all dimensions of `by` when method != 'map-reduce'."
-            f"Received method={method!r}"
-        )
-
     # TODO: make sure expected_groups is unique
     if nax == 1 and by_.ndim > 1 and expected_groups is None:
         if not any_by_dask:
@@ -1949,6 +1943,12 @@ def groupby_reduce(
                 f"\n\n Received: {func}"
             )
 
+        if method in ["blockwise", "cohorts"] and nax != by_.ndim:
+            raise NotImplementedError(
+                "Must reduce along all dimensions of `by` when method != 'map-reduce'."
+                f"Received method={method!r}"
+            )
+
         # TODO: just do this in dask_groupby_agg
         # we always need some fill_value (see above) so choose the default if needed
         if kwargs["fill_value"] is None:

tests/test_core.py

@@ -1347,3 +1347,37 @@ def test_expected_index_conversion_passthrough_range_index(sort):
         expected_groups=(index,), isbin=(False,), sort=(sort,)
     )
     assert actual[0] is index
+
+
+def test_method_check_numpy():
+    bins = [-2, -1, 0, 1, 2]
+    field = np.ones((5, 3))
+    by = np.array([[-1.5, -1.5, 0.5, 1.5, 1.5] * 3]).reshape(5, 3)
+    actual, _ = groupby_reduce(
+        field,
+        by,
+        expected_groups=pd.IntervalIndex.from_breaks(bins),
+        func="count",
+        method="cohorts",
+        fill_value=np.nan,
+    )
+    expected = np.array([6, np.nan, 3, 6])
+    assert_equal(actual, expected)
+
+    actual, _ = groupby_reduce(
+        field,
+        by,
+        expected_groups=pd.IntervalIndex.from_breaks(bins),
+        func="count",
+        fill_value=np.nan,
+        method="cohorts",
+        axis=0,
+    )
+    expected = np.array(
+        [
+            [2.0, np.nan, 1.0, 2.0],
+            [2.0, np.nan, 1.0, 2.0],
+            [2.0, np.nan, 1.0, 2.0],
+        ]
+    )
+    assert_equal(actual, expected)