Allow specifying output dtype

Closes pydata/xarray#6902

Author: dcherian

flox/core.py

@@ -1364,6 +1364,7 @@ def groupby_reduce(
     isbin: bool = False,
     axis=None,
     fill_value=None,
+    dtype=None,
     min_count: int | None = None,
     split_out: int = 1,
     method: str = "map-reduce",
@@ -1566,8 +1567,13 @@ def groupby_reduce(
         # overwrite than when min_count is set
         fill_value = np.nan
 
+    if dtype is not None and not isinstance(dtype, np.dtype):
+        dtype = np.dtype(dtype)
+
     kwargs = dict(axis=axis, fill_value=fill_value, engine=engine)
-    agg = _initialize_aggregation(func, array.dtype, fill_value, min_count, finalize_kwargs)
+    agg = _initialize_aggregation(
+        func, array.dtype if dtype is None else dtype, fill_value, min_count, finalize_kwargs
+    )
 
     if not has_dask:
         results = _reduce_blockwise(

flox/xarray.py

@@ -60,6 +60,7 @@ def xarray_reduce(
     dim: Hashable = None,
     split_out: int = 1,
     fill_value=None,
+    dtype=None,
     method: str = "map-reduce",
     engine: str = "flox",
     keep_attrs: bool | None = True,
@@ -95,6 +96,8 @@ def xarray_reduce(
     fill_value
         Value used for missing groups in the output i.e. when one of the labels
         in ``expected_groups`` is not actually present in ``by``.
+    dtype: str
+        DType for the output (DataArray only).
     method : {"map-reduce", "blockwise", "cohorts", "split-reduce"}, optional
         Strategy for reduction of dask arrays only:
           * ``"map-reduce"``:
@@ -341,7 +344,9 @@ def wrapper(array, *by, func, skipna, **kwargs):
         exclude_dims=set(dim),
         output_core_dims=[group_names],
         dask="allowed",
-        dask_gufunc_kwargs=dict(output_sizes=group_sizes),
+        dask_gufunc_kwargs=dict(
+            output_sizes=group_sizes, output_dtypes=[dtype] if dtype is not None else None
+        ),
         keep_attrs=keep_attrs,
         kwargs={
             "func": func,
@@ -357,6 +362,7 @@ def wrapper(array, *by, func, skipna, **kwargs):
             "expected_groups": tuple(expected_groups),
             "isbin": isbin,
             "finalize_kwargs": finalize_kwargs,
+            "dtype": dtype,
         },
     )
 

tests/test_core.py

@@ -1009,3 +1009,14 @@ def grouped_median(group_idx, array, *, axis=-1, size=None, fill_value=None, dty
         method="blockwise",
     )
     assert_equal(expected, actual)
+
+
+@pytest.mark.parametrize("func", ALL_FUNCS)
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_dtype(func, dtype, engine):
+    if "arg" in func or func in ["any", "all"]:
+        pytest.skip()
+    arr = np.ones((4, 12), dtype=dtype)
+    labels = np.array(["a", "a", "c", "c", "c", "b", "b", "c", "c", "b", "b", "f"])
+    actual, _ = groupby_reduce(arr, labels, func=func, dtype=np.float64)
+    assert actual.dtype == np.dtype("float64")

tests/test_xarray.py

@@ -485,3 +485,20 @@ def test_mixed_grouping(chunk):
         fill_value=0,
     )
     assert (r.sel(v1=[3, 4, 5]) == 0).all().data
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_dtype(dtype, engine):
+    arr = xr.DataArray(
+        data=np.ones((4, 12), dtype=dtype),
+        dims=("x", "t"),
+        coords={
+            "labels": ("t", np.array(["a", "a", "c", "c", "c", "b", "b", "c", "c", "b", "b", "f"]))
+        },
+    )
+    actual = xarray_reduce(arr, "labels", func="mean", dtype=np.float64)
+    assert actual.dtype == np.dtype("float64")
+
+    actual = xarray_reduce(arr.chunk({"x": 1}), arr.labels, func="mean", dtype=np.float64)
+    assert actual.dtype == np.dtype("float64")
+    assert actual.compute().dtype == np.dtype("float64")