Skip to content

Fix dask rolling bottleneck dtype metadata#11449

Merged
dcherian merged 4 commits into
pydata:mainfrom
mrocklin:codex-fix-rolling-bottleneck-dtype
Jul 13, 2026
Merged

Fix dask rolling bottleneck dtype metadata#11449
dcherian merged 4 commits into
pydata:mainfrom
mrocklin:codex-fix-rolling-bottleneck-dtype

Conversation

@mrocklin

Copy link
Copy Markdown
Contributor

Summary

  • Fix dask-backed bottleneck rolling reductions declaring metadata dtypes that
    differ from the matching numpy-backed bottleneck result.
  • In particular, boolean rolling reductions now declare float64 metadata
    instead of object, matching the actual bottleneck result dtype.
  • Preserve centered rolling behavior by using the original bool dtype through
    object padding, while continuing to infer centered integer metadata from the
    padded dtype.
  • Add bool and integer coverage comparing numpy-backed and dask-backed rolling
    values, result dtype, dask _meta, computed block dtype, and full computed
    dtype.

Background

The same rolling reduction can currently report different dtypes depending on
the backend:

import xarray as xr, numpy as np, dask.array as dda

n = (
    xr.DataArray(np.ones((100, 4)), dims=("t", "a")) > 0
).rolling(t=72, min_periods=72).sum()

d = (
    xr.DataArray(dda.ones((100, 4), chunks=(100, 4)), dims=("t", "a")) > 0
).rolling(t=72, min_periods=72).sum()

print(n.dtype, d.dtype)
print(d.data._meta.dtype, d.compute().dtype)

On xarray 2025.12 and current main, the numpy-backed result is float64, while
the dask-backed result declares object metadata. With legacy dask, this is
partly masked because computed blocks are still float64 under the declared
object meta. Stricter dask backends can faithfully materialize the declared
object dtype, surfacing the inconsistency.

Diagnosis

xarray/compat/dask_array_ops.py::dask_rolling_wrapper previously declared the
dask map_overlap dtype using:

dtype = dtypes.maybe_promote(a.dtype)[0]

For bool input, maybe_promote(bool) returns object. That promotion is
appropriate for generic missing-value operations such as padding, filling, and
where, but it is too broad for bottleneck moving-window reductions.

Bottleneck accepts bool arrays directly and returns float64, which already
has room for the NaN values introduced by min_periods. The dask wrapper was
therefore declaring metadata from generic NA-promotion rules instead of from the
actual moving function result type.

The centered path has one extra wrinkle: xarray pads before calling the dask
rolling wrapper. For centered bool rolling, padding turns the intermediate dask
array into object, so dtype inference needs to use the original bool dtype.
For centered integer rolling, the padded dtype is already the behavior that
matches the numpy-backed path, so that should be preserved.

Implementation

The dask bottleneck wrapper now derives metadata for numeric and boolean inputs
from the moving function's actual NumPy result dtype. Non-numeric inputs keep
the previous maybe_promote fallback.

_bottleneck_reduce now passes an explicit input_dtype into the dask wrapper:
bool inputs use the original rolling object dtype, while other inputs use the
padded dtype. This fixes centered bool rolling without changing centered
integer behavior.

This change intentionally does not modify dtypes.maybe_promote. The
maybe_promote(bool) -> object rule may be worth revisiting separately, but
that function is shared by broader NA-handling paths and should not be changed
as part of this targeted bottleneck metadata fix.

Tests

  • Added rolling coverage for bool, int8, and int64 inputs.
  • Covered bottleneck reductions across centered and non-centered rolling.
  • Asserted numpy-backed and dask-backed results agree on values and dtype.
  • Asserted dask _meta, a computed block, and the full computed result all
    agree on dtype.
  • Skipped only centered bool std and median, where the numpy-backed
    bottleneck path itself currently errors.

Checklist

  • Tests added
  • User visible changes (including notable bug fixes) are documented in whats-new.rst

AI Disclosure

  • This PR contains AI-generated content.
    • I have tested any AI-generated content in my PR.

    • I take responsibility for any AI-generated content in my PR.

      Tools: Codex

Derive dask map_overlap metadata from the bottleneck moving function result for numeric rolling inputs instead of using the generic NA-promotion dtype. Preserve the padded dtype for centered integer rolling while carrying the original bool dtype through centered padding so bool reductions no longer declare object metadata.

Add bool and integer rolling coverage comparing numpy and dask dtype/value behavior, including dask meta and computed block dtypes.

Co-Authored-By: OpenAI Codex <codex@openai.com>
Comment thread xarray/compat/dask_array_ops.py
@dcherian dcherian enabled auto-merge (squash) July 13, 2026 15:57
@mrocklin

Copy link
Copy Markdown
Contributor Author

Thanks @dcherian !

Remove a C++-style comment that was introduced on the PR branch and caused xarray imports to fail before CI reached tests.

Co-Authored-By: OpenAI Codex <codex@openai.com>
@dcherian dcherian merged commit 2f6713f into pydata:main Jul 13, 2026
43 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants