pydata / pydata/xarray

Should performance be equivalent when opening with chunks or re-chunking a dataset?

Open
#3,486 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

topic-performance
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

I was wondering if the chunking behaviour would be expected to be equivalent under two different use cases:

(1) When opening a dataset using the chunks option;
(2) When re-chunking an existing dataset using Dataset.chunk method.

I'm interested in performance for slicing across different dimensions. In my case the performance is quite different, please see the example below:

Open dataset with one single chunk along station dimension (fast for slicing one time)
In [1]: import xarray as xr

In [2]: dset = xr.open_dataset( 
    ...: "/source/wavespectra/tests/sample_files/spec20170101T00_spec.nc", 
    ...: chunks={"station": None}
    ...: )

In [3]: dset
Out[3]: 
<xarray.Dataset>
Dimensions:       (direction: 24, frequency: 25, station: 14048, time: 249)
Coordinates:
  * time          (time) datetime64[ns] 2017-01-01 ... 2017-02-01
  * station       (station) float64 1.0 2.0 3.0 ... 1.405e+04 1.405e+04
  * frequency     (frequency) float32 0.04118 0.045298003 ... 0.40561208
  * direction     (direction) float32 90.0 75.0 60.0 45.0 ... 135.0 120.0 105.0
Data variables:
    longitude     (time, station) float32 dask.array<chunksize=(249, 14048), meta=np.ndarray>
    latitude      (time, station) float32 dask.array<chunksize=(249, 14048), meta=np.ndarray>
    efth          (time, station, frequency, direction) float32 dask.array<chunksize=(249, 14048, 25, 24), meta=np.ndarray>

In [4]: %time lats = dset.latitude.isel(time=0).values
CPU times: user 171 ms, sys: 49.2 ms, total: 220 ms
Wall time: 219 ms
Open dataset with many size=1 chunks along station dimension (fast for slicing one station, slow for slicing one time)
In [5]: dset = xr.open_dataset( 
    ...: "/source/wavespectra/tests/sample_files/spec20170101T00_spec.nc", 
    ...: chunks={"station": 1}
    ...: )

In [6]: dset
Out[6]: 
<xarray.Dataset>
Dimensions:       (direction: 24, frequency: 25, station: 14048, time: 249)
Coordinates:
  * time          (time) datetime64[ns] 2017-01-01 ... 2017-02-01
  * station       (station) float64 1.0 2.0 3.0 ... 1.405e+04 1.405e+04
  * frequency     (frequency) float32 0.04118 0.045298003 ... 0.40561208
  * direction     (direction) float32 90.0 75.0 60.0 45.0 ... 135.0 120.0 105.0
Data variables:
    longitude     (time, station) float32 dask.array<chunksize=(249, 1), meta=np.ndarray>
    latitude      (time, station) float32 dask.array<chunksize=(249, 1), meta=np.ndarray>
    efth          (time, station, frequency, direction) float32 dask.array<chunksize=(249, 1, 25, 24), meta=np.ndarray>

In [7]: %time lats = dset.latitude.isel(time=0).values
CPU times: user 13.1 s, sys: 1.94 s, total: 15 s
Wall time: 11.1 s
Try rechunk station into one single chunk (still slow to slice one time)
In [8]: dset = dset.chunk({"station": None})

In [8]: dset
Out[8]: 
<xarray.Dataset>
Dimensions:       (direction: 24, frequency: 25, station: 14048, time: 249)
Coordinates:
  * time          (time) datetime64[ns] 2017-01-01 ... 2017-02-01
  * station       (station) float64 1.0 2.0 3.0 ... 1.405e+04 1.405e+04
  * frequency     (frequency) float32 0.04118 0.045298003 ... 0.40561208
  * direction     (direction) float32 90.0 75.0 60.0 45.0 ... 135.0 120.0 105.0
Data variables:
    longitude     (time, station) float32 dask.array<chunksize=(249, 14048), meta=np.ndarray>
    latitude      (time, station) float32 dask.array<chunksize=(249, 14048), meta=np.ndarray>
    efth          (time, station, frequency, direction) float32 dask.array<chunksize=(249, 14048, 25, 24), meta=np.ndarray>

In [9]: %time lats = dset.latitude.isel(time=0).values
CPU times: user 9.06 s, sys: 1.13 s, total: 10.2 s
Wall time: 7.7 s

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with xarray.open_dataset and Dataset.chunk using the sample file at wavespectra/tests/sample_files/spec20170101T00_spec.nc. Reproduce the two chunking paths and compare the isel(time=0) slice timings and resulting chunks. Done means explaining or correcting the performance difference, with a regression test or benchmark covering both paths.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.