diff --git a/flox/core.py b/flox/core.py index 4179bc236..728d8b186 100644 --- a/flox/core.py +++ b/flox/core.py @@ -1465,9 +1465,11 @@ def _assert_by_is_aligned(shape, by): for idx, b in enumerate(by): if not all(j in [i, 1] for i, j in zip(shape[-b.ndim :], b.shape)): raise ValueError( - "`array` and `by` arrays must be aligned " - "i.e. array.shape[-by.ndim :] == by.shape. " - "for every array in `by`." + "`array` and `by` arrays must be 'aligned' " + "so that such that by_ is broadcastable to array.shape[-by.ndim:] " + "for every array `by_` in `by`. " + "Either array.shape[-by_.ndim :] == by_.shape or the only differences " + "should be size-1 dimensions in by_." f"Received array of shape {shape} but " f"array {idx} in `by` has shape {b.shape}." ) diff --git a/flox/xarray.py b/flox/xarray.py index 35d99b6fa..2a2b27bdc 100644 --- a/flox/xarray.py +++ b/flox/xarray.py @@ -265,6 +265,14 @@ def xarray_reduce( # broadcast to make sure grouper dimensions are present in the array. exclude_dims = tuple(d for d in ds.dims if d not in grouper_dims and d not in dim_tuple) + + try: + xr.align(ds, *by_da, join="exact") + except ValueError as e: + raise ValueError( + "Object being grouped must be exactly aligned with every array in `by`." + ) from e + ds_broad = xr.broadcast(ds, *by_da, exclude=exclude_dims)[0] if any(d not in grouper_dims and d not in obj.dims for d in dim_tuple): diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 0bee41c15..42e6ee17f 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -499,6 +499,12 @@ def test_mixed_grouping(chunk): assert (r.sel(v1=[3, 4, 5]) == 0).all().data +def test_alignment_error(): + da = xr.DataArray(np.arange(10), dims="x", coords={"x": np.arange(10)}) + with pytest.raises(ValueError): + xarray_reduce(da, da.x.sel(x=slice(5)), func="count") + + @pytest.mark.parametrize("add_nan", [True, False]) @pytest.mark.parametrize("dtype_out", [np.float64, "float64", np.dtype("float64")]) @pytest.mark.parametrize("dtype", [np.float32, np.float64])