Skip to content

A2: Reduce is_chunked_array / is_dask_collection call overhead in load() #2

Description

@FBumann

Problem

`Dataset.load()` and `isel().load()` call `is_chunked_array` / `is_dask_collection` once per variable, with full duck-array protocol dispatch each time. Profile on a 400-extra-scalar-var dataset shows:

```
14.53% to_duck_array xarray/namedarray/pycompat.py
13.87% is_dask_collection xarray/namedarray/utils.py
7.47% is_duck_array xarray/namedarray/utils.py
5.37% is_chunked_array xarray/namedarray/pycompat.py
5.32% is_dask_collection dask/base.py
4.20% is_duck_dask_array xarray/namedarray/utils.py
```

805 calls per `isel().load()` for a dataset with ~400 variables. Pure dispatch overhead — none of these checks touch user data.

Proposed fix

Options:

  1. Cache per-variable — store the duck-array-type classification on `Variable` (or in `pycompat._get_cached_duck_array_module`'s cache layer)
  2. Hoist — check once at the `Dataset.load()` level whether any variable is chunked; if not, skip the per-variable check entirely

Option 2 is the lower-risk change with similar payoff for the common "all-numpy dataset" case.

Expected outcome

2–3× faster `.load()` and `.isel().load()` on datasets with many small/scalar variables. No effect on chunked workloads.

Benchmark

`asv_bench/benchmarks/indexing.py::Indexing.time_indexing_basic_ds_large` already exercises this path (the 400-extra-vars setup). Confirm regression numbers there.

Acceptance

  • Existing indexing + dask tests pass
  • asv shows speedup on `time_indexing_basic_ds_large`
  • No regression on chunked-array path

[This is Claude Code on behalf of Felix Bumann]

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions