perf(load): short-circuit is_chunked_array for numpy arrays#11354
Open
FBumann wants to merge 3 commits into
Open
perf(load): short-circuit is_chunked_array for numpy arrays#11354FBumann wants to merge 3 commits into
FBumann wants to merge 3 commits into
Conversation
For datasets with many variables, Dataset.load() called is_chunked_array
once per variable in its dict comprehension, then again per variable via
Variable.load() -> to_duck_array(). The function itself called
is_duck_array twice (once directly, once via is_duck_dask_array).
Add a numpy fast-path and consolidate the duck-array check to one call.
For non-numpy inputs the behavior is unchanged: any duck-array with a
dask graph or a `chunks` attribute is still reported as chunked.
Measured on isel(...).load() of a 400-scalar-var dataset
(asv_bench/benchmarks/indexing.py::Indexing.time_indexing_basic_ds_large):
base: 0.524 ms / call (best of 5x50, GC off)
branch: 0.335 ms / call ~1.56x
Profile attribution previously showed ~25% of the load wall time inside
the is_chunked_array dispatch chain; that portion is now near-free.
Closes #2 on the fork.
Co-authored-by: Claude <noreply@anthropic.com>
for more information, see https://pre-commit.ci
6 tasks
The previous `isinstance(x, np.ndarray)` short-circuit incorrectly
returned False for ndarray subclasses with a `chunks` attribute (e.g.
DummyChunkedArray in test_parallelcompat.py, or any third-party chunked
array implementation that subclasses ndarray), breaking chunked-array
detection on those types.
Narrow the fast path to `isinstance + not hasattr("chunks")` so plain
ndarrays and non-chunked subclasses (MaskedArray, np.matrix) still skip
the duck-array dispatch, while subclasses that advertise chunks fall
through to the full check.
Co-authored-by: Claude <noreply@anthropic.com>
Illviljan
requested changes
May 26, 2026
Contributor
There was a problem hiding this comment.
I think this can be done with is_duck_array, it has the numpy short-circuit already.
Triggering isinstance twice for eager (and lazy) arrays seems wasteful too.
xarray/xarray/namedarray/utils.py
Lines 78 to 91 in d022da5
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Dataset.load()/DataArray.load()walk every variable and callis_chunked_array(...)on each — first inside the dict comprehension inDataset.load(xarray/core/dataset.py:563), then again per variable viaVariable.load -> to_duck_array.is_chunked_arrayitself calledis_duck_arraytwice (once directly, once viais_duck_dask_array).For the dominant case — a
numpy.ndarray— none of that work is needed. This PR:np.ndarrayfast path returningFalseimmediately.is_duck_array(x)check feeding both the dask-collection branch (now viais_dask_collectiondirectly) and thehasattr(x, "chunks")fallback.Behavior is unchanged on non-numpy inputs: any duck-array with a dask graph or a
chunksattribute is still reported as chunked.Where this fires
is_chunked_arrayis called from ~25 places across xarray — every site that branches "dask vs. numpy". On numpy-backed data, all of them now skip the dispatch chain. Notable categories:ds.load(),da.load(),compute(),xr.load_dataset/dataarray/datatree,persist(),.values,.to_numpy(),.to_dataframe(),.to_pandas(), plotting, repr previews.coding/times.py:1037,1249,coding/strings.py:182,221,coding/common.py:104.apply_ufunc(computation/apply_ufunc.py:739,763,875),corr/cov/polyval/polyfit,interp,interpolate_na, reductions.Dataset.chunk(),interp, vectorized indexer dispatch (isel/sel),Variable._shuffle,contains_only_chunked_or_numpy.groupers.py:121,242,427,groupby.py:351,557,679.backends/common.py:400(ArrayWriter.add).Not affected: arithmetic on lazy/dask objects (stays lazy); arithmetic on numpy-backed objects (already pure numpy, never reached
is_chunked_array).Benchmark numbers
asv_bench/benchmarks/indexing.py::Indexing.time_indexing_basic_ds_large(added in #9003 for this exact concern), best of 5×50 iterations, GC off:mainScaling check on synthetic
isel().load()workloads:Profiler attribution previously showed ~25% of
load()wall time inside theis_chunked_arraydispatch chain; that portion is now near-free.Checklist
chunksattribute or a dask graph is still reported as chunkedpytest xarray/tests/test_variable.py— 544 passed, 69 skipped, 8 xfailed, 3 xpasseddoc/whats-new.rstentry under Internal ChangesFollow-up
A complementary PR — skipping the entire
Variable.loaddispatch on numpy data — will follow.AI Disclosure
Tools: Claude (Claude Code)
[This is Claude Code on behalf of Felix Bumann]