Expand groupby_reduce property tests (#385)

* Expand groupby_reduce property tests

* Add back var, std

* cast quantile result

* Revert "Add back var, std"

This reverts commit 805b8d3.

* pin numpy in benchmark env

* Add benchmarks as test

Author: dcherian

ci/benchmark.yml

@@ -6,7 +6,7 @@ dependencies:
   - build
   - cachey
   - dask-core
-  - numpy>=1.22
+  - numpy<2
   - mamba
   - pip
   - python=3.10

flox/aggregate_flox.py

@@ -160,7 +160,7 @@ def _np_grouped_op(
 
 def _nan_grouped_op(group_idx, array, func, fillna, *args, **kwargs):
     if fillna in [dtypes.INF, dtypes.NINF]:
-        fillna = dtypes._get_fill_value(kwargs.get("dtype", array.dtype), fillna)
+        fillna = dtypes._get_fill_value(kwargs.get("dtype", None) or array.dtype, fillna)
     result = func(group_idx, np.where(isnull(array), fillna, array), *args, **kwargs)
     # np.nanmax([np.nan, np.nan]) = np.nan
     # To recover this behaviour, we need to search for the fillna value

tests/test_asv.py

@@ -0,0 +1,24 @@
+# Run asv benchmarks as tests
+
+import pytest
+
+pytest.importorskip("dask")
+
+from asv_bench.benchmarks import reduce
+
+
+@pytest.mark.parametrize(
+    "problem", [reduce.ChunkReduce1D, reduce.ChunkReduce2D, reduce.ChunkReduce2DAllAxes]
+)
+def test_reduce(problem) -> None:
+    testcase = problem()
+    testcase.setup()
+    for args in zip(*testcase.time_reduce.params):
+        testcase.time_reduce(*args)
+
+
+def test_reduce_bare() -> None:
+    testcase = reduce.ChunkReduce1D()
+    testcase.setup()
+    for args in zip(*testcase.time_reduce_bare.params):
+        testcase.time_reduce_bare(*args)

tests/test_properties.py

@@ -61,7 +61,11 @@ def not_overflowing_array(array: np.ndarray[Any, Any]) -> bool:
     return result
 
 
-@given(data=st.data(), array=numeric_arrays, func=func_st)
+@given(
+    data=st.data(),
+    array=st.one_of(numeric_arrays, chunked_arrays(arrays=numeric_arrays)),
+    func=func_st,
+)
 def test_groupby_reduce(data, array, func: str) -> None:
     # overflow behaviour differs between bincount and sum (for example)
     assume(not_overflowing_array(array))
@@ -93,7 +97,13 @@ def test_groupby_reduce(data, array, func: str) -> None:
 
         # numpy-groupies always does the calculation in float64
         if (
-            ("var" in func or "std" in func or "sum" in func or "mean" in func)
+            (
+                "var" in func
+                or "std" in func
+                or "sum" in func
+                or "mean" in func
+                or "quantile" in func
+            )
             and array.dtype.kind == "f"
             and array.dtype.itemsize != 8
         ):