From b1ec0fa79059f5d1619f330f8f9daef773439f25 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Wed, 25 Feb 2026 22:02:18 -0700 Subject: [PATCH] Fix nanmedian for all-NaN groups Closes #494 --- flox/aggregate_flox.py | 4 +++- tests/test_xarray.py | 19 +++++++++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/flox/aggregate_flox.py b/flox/aggregate_flox.py index 2fa7330c..8d597f38 100644 --- a/flox/aggregate_flox.py +++ b/flox/aggregate_flox.py @@ -123,8 +123,10 @@ def quantile_(array, inv_idx, *, q, axis, skipna, group_idx, dtype=None, out=Non # TODO: could support all the interpolations here gamma = np.broadcast_to(virtual_index, idxshape) - lo_ result = _lerp(loval, hival, t=gamma, out=out, dtype=dtype) + tomask = actual_sizes == -1 if not skipna and np.any(nanmask): - result[..., nanmask] = np.nan + tomask |= nanmask + result[..., tomask] = np.nan return result diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 7351786f..798ef847 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -825,3 +825,22 @@ def test_xarray_indexing_array_support(): coords={"labels": ("y", ["a", "a", "b", "b"])}, ) assert is_supported_aggregation(da.variable._data, "sum") + + +def test_nanmedian_gh_494(): + n = 100 + dt = pd.date_range("2020-01-01", periods=n, freq="D") + + data = np.random.randn(n) + data[0:35] = np.nan # groups 0, 1, 2 fully NaN; group 3 partially NaN + data[50:80] = np.nan # groups 5, 6, 7 fully NaN <-- interior all-NaN groups + data[90:] = np.nan # group 9 fully NaN + + da = xr.DataArray(data, dims=["time"], coords={"time": dt}, name="data") + group = xr.DataArray([i // 10 for i in range(n)], dims=["time"], coords={"time": dt}, name="group") + + result_flox = xarray_reduce(da, group, func="nanmedian").compute().to_pandas() + + result_pandas = pd.Series(data).rename("data").groupby(group.values).agg("median") + result_pandas.index.name = "group" + pd.testing.assert_series_equal(result_flox, result_pandas)