Re-Chunking Data

Re-organizing stored data such that it matches the analysis use-case.

Inspiration from: NCAR/rechunk_retro_nwm_v21

Note

• The rechunker documentation contains several examples and a tutorial covering how to re-chunk data. Much of what is here replicates concepts covered in that material. This document uses data that looks like HyTest data (variable names, extent, etc), so may offer a smoother intro to the concepts using familiar data.

• The zarr data standard has a nice tutorial also which covers details of optimizing chunking strategies.

Intro

What is chunking and why should you care?

The idea of data ‘chunks’ is closely aligned with the NetCDF and zarr standards for storing N-dimensional arrays of typed data.

Chunks become more important as the size of the array increases. For very large arrays, it is helpful to organize the memory it occupies into sub-units. These sub-units are the chunks – note that this is not another dimension to the array, but merely a map to how the large array is partitioned into more palatable sized units for manipulation in memory. Array-handling libraries (numpy, xarray, pandas, and others) will handle all of the record-keeping to know which chunk holds a given unit of the array.

A quick side-bar to illustrate two chunking patterns for a simple 2D array. This is a simplified use-case. Consider a square array of integer values. Just for exposition, let’s use a small array 10x10.

illustration

That array can be organized in memory in a few ways… two common options are row-major order, and column-major order:

• Row-Major – A row of data occupies a contiguous block of memory. This implies that cells which are logicall adjacent vertically are not physicall near one another in memory. The ‘distance’ from r0c0 to r0c1 (a one-cell logical move within the row) is short, while the ‘distance’ to r1c0 (a one-cell logical move within the column) is long.

• Column-Major – A column of the array occupies a contiguious block of memory. This implies that cells which are adjacent horizontally are not near one another physically in memory.

In either chunk mapping, r3c5 (for example) still fetches the same value – the array still indexes/addresses in the same way – but the chunking plan determines how nearby an ‘adjacent’ index is.

As the size of the array increases, the chunk pattern becomes more relevant. Suppose your data is chunked by rows, and you need to process a column of data – your process will need to read a lot of data, skipping most of it (possibly going to disk to read more data if the dataset is very large), to get the \(i^{th}\) column value for each row. For this analysis, it would be better if the array could be ‘re-chunked’ from row-major order to column-major order. This would favor column operations.

Array size is important to the chunking plan… it could be that an entire row of data in a large 2D array won’t fit into memory (or it can’t fit into a contiguous block of memory). This would require chunking each row of data, in addition to chunking rows over columns (for a row-major plan) As dimensions are added to the array, this chunk-mapping becomes more complex and it becomes much more relevant to chunk the data to match the analysis.

Pros & Cons

Data that is well-organized to optimize one kind of analysis may not suit another kind of analysis on the same data. Re-chunking is time-consuming, and it produces a separate copy of the dataset, increasing storage requirements. The initial time commitment is a one-time operation so that future analyses can run quickly. The space commitment can be substantial if a complex dataset needs to be organized for many different analyses.

Examining a Small Dataset

Let’s read a sample dataset and examine how it is chunked.

As a test datasaet, we’ve taken a random sampling of 400 stream gages for the month of July, 2000 from the National Water Model Reanalysis Version 2.1. The full dataset is part of the AWS Open Data Program, available via the S3 bucket at

s3://noaa-nwm-retrospective-2-1-zarr-pds/noaa-nwm-retrospective-2-1-zarr-pds/chrtout.zarr

Our subset of that data for use in this tutorial is included in the HyTEST catalog:

%run ../../../environment_set_up/Help_AWS_Credentials.ipynb
## Establish AWS credentials

import xarray as xr 
import intake
url = 'https://raw.githubusercontent.com/hytest-org/hytest/main/dataset_catalog/hytest_intake_catalog.yml'
cat = intake.open_catalog(url)
sampleData = cat['rechunking-tutorial-osn'].to_dask()
sampleData

/home/conda/global/9764ad0fe205ee3ccb00b4474c3e6708e73fb3a628843d562645d7d1e21ccb8d-20231221-205006-476455-52-pangeo/lib/python3.11/site-packages/xarray/core/utils.py:494: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  warnings.warn(

<xarray.Dataset>
Dimensions:     (feature_id: 400, time: 744)
Coordinates:
    elevation   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
  * feature_id  (feature_id) int32 3109 189899 239166 ... 947070134 1010003783
    gage_id     (feature_id) |S15 dask.array<chunksize=(400,), meta=np.ndarray>
    latitude    (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    longitude   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    order       (feature_id) int32 dask.array<chunksize=(400,), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-31T23:00:00
Data variables:
    streamflow  (time, feature_id) float64 dask.array<chunksize=(256, 16), meta=np.ndarray>
    velocity    (time, feature_id) float64 dask.array<chunksize=(256, 16), meta=np.ndarray>
Attributes:
    TITLE:                OUTPUT FROM WRF-Hydro v5.2.0-beta2
    code_version:         v5.2.0-beta2
    featureType:          timeSeries
    model_configuration:  retrospective
    proj4:                +proj=lcc +units=m +a=6370000.0 +b=6370000.0 +lat_1...

xarray.Dataset

• Dimensions:
  • feature_id: 400
  • time: 744
• Coordinates: (7)
  • elevation
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature elevation

    units :

    meters

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • feature_id
    (feature_id)
    int32
    3109 189899 ... 1010003783

    cf_role :

    timeseries_id

    comment :

    NHDPlusv2 ComIDs within CONUS, arbitrary Reach IDs outside of CONUS

    long_name :

    Reach ID

    array([      3109,     189899,     239166, ...,  945010401,  947070134,
           1010003783], dtype=int32)

  • gage_id
    (feature_id)
    |S15
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    NHD Gage Event ID from SOURCE_FEA field in Gages feature class

    Array Chunk Bytes 5.86 kiB 5.86 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type |S15 numpy.ndarray

  • latitude
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature latitude

    standard_name :

    latitude

    units :

    degrees_north

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • longitude
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature longitude

    standard_name :

    longitude

    units :

    degrees_east

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • order
    (feature_id)
    int32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    stream order

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type int32 numpy.ndarray

  • time
    (time)
    datetime64[ns]
    2000-07-01 ... 2000-07-31T23:00:00

    long_name :

    valid output time

    standard_name :

    time

    array(['2000-07-01T00:00:00.000000000', '2000-07-01T01:00:00.000000000',
           '2000-07-01T02:00:00.000000000', ..., '2000-07-31T21:00:00.000000000',
           '2000-07-31T22:00:00.000000000', '2000-07-31T23:00:00.000000000'],
          dtype='datetime64[ns]')

• Data variables: (2)
  • streamflow
    (time, feature_id)
    float64
    dask.array<chunksize=(256, 16), meta=np.ndarray>

    grid_mapping :

    crs

    long_name :

    River Flow

    units :

    m3 s-1

    Array Chunk Bytes 2.27 MiB 32.00 kiB Shape (744, 400) (256, 16) Dask graph 75 chunks in 2 graph layers Data type float64 numpy.ndarray

  • velocity
    (time, feature_id)
    float64
    dask.array<chunksize=(256, 16), meta=np.ndarray>

    grid_mapping :

    crs

    long_name :

    River Velocity

    units :

    m s-1

    Array Chunk Bytes 2.27 MiB 32.00 kiB Shape (744, 400) (256, 16) Dask graph 75 chunks in 2 graph layers Data type float64 numpy.ndarray

• Indexes: (2)
  • feature_id
    PandasIndex

    PandasIndex(Int64Index([      3109,     189899,     239166,     334022,     397590,
                    418429,     481940,     591268,     682279,     734307,
                ...
                  41027830,   41028190,   41039416,   41044540,  166743866,
                 932040162,  940080052,  945010401,  947070134, 1010003783],
               dtype='int64', name='feature_id', length=400))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-07-01 00:00:00', '2000-07-01 01:00:00',
                   '2000-07-01 02:00:00', '2000-07-01 03:00:00',
                   '2000-07-01 04:00:00', '2000-07-01 05:00:00',
                   '2000-07-01 06:00:00', '2000-07-01 07:00:00',
                   '2000-07-01 08:00:00', '2000-07-01 09:00:00',
                   ...
                   '2000-07-31 14:00:00', '2000-07-31 15:00:00',
                   '2000-07-31 16:00:00', '2000-07-31 17:00:00',
                   '2000-07-31 18:00:00', '2000-07-31 19:00:00',
                   '2000-07-31 20:00:00', '2000-07-31 21:00:00',
                   '2000-07-31 22:00:00', '2000-07-31 23:00:00'],
                  dtype='datetime64[ns]', name='time', length=744, freq=None))

• Attributes: (5)

  TITLE :

  OUTPUT FROM WRF-Hydro v5.2.0-beta2

  code_version :

  v5.2.0-beta2

  featureType :

  timeSeries

  model_configuration :

  retrospective

  proj4 :

  +proj=lcc +units=m +a=6370000.0 +b=6370000.0 +lat_1=30.0 +lat_2=60.0 +lat_0=40.0 +lon_0=-97.0 +x_0=0 +y_0=0 +k_0=1.0 +nadgrids=@

The critical items to notice in this output are highlighted here:

<xarray.Dataset>

Dimensions:     (feature_id: 400, time: 744)  <-- NOTE: Two dimensions

                   +--- most coordinates are tied to feature_id dimension 
                   | 
Coordinates:       V
    elevation   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
  * feature_id  (feature_id) int32 3109 189899 239166 ... 947070134 1010003783
    gage_id     (feature_id) |S15 dask.array<chunksize=(400,), meta=np.ndarray>
    latitude    (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    longitude   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    order       (feature_id) int32 dask.array<chunksize=(400,), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-31T23:00:00

Data variables:
    streamflow  (time, feature_id) float64 dask.array<chunksize=(256, 16), meta=np.ndarray>
    velocity    (time, feature_id) float64 dask.array<chunksize=(256, 16), meta=np.ndarray>
                 ^^^^  ^^^^^^^^^^
                 the data variables are addressed by both dimensions; this is 2D data.

This dataset is a ‘stack’ of two 2D arrays. They are named ‘streamflow’ and ‘velocity’. The indices into each of those 2D arrays are time on one axis, and feature_id on the other. The feature id is bound to a number of other coordinates, so you can relate/refer to a given feature by its elevation, gage_id, latitude, longitude, or stream order.

Note the chunksize highlighted in green. This says that the data is stored in blocks mapping to 256 adjacent time-steps for 16 adjacent features. (NOTE: The original data is not chunked this way; we’ve deliberately fiddled with the chunk configuration for this tutorial)

A time-series analysis (i.e. sampling all time-step values for a single feature_id) would require multiple chunks to be fetched.

# Fetch all the time values for a specific feature_id
sampleData['streamflow'].sel(feature_id=1343034)

<xarray.DataArray 'streamflow' (time: 744)>
dask.array<getitem, shape=(744,), dtype=float64, chunksize=(256,), chunktype=numpy.ndarray>
Coordinates:
    elevation   float32 dask.array<chunksize=(), meta=np.ndarray>
    feature_id  int32 1343034
    gage_id     |S15 dask.array<chunksize=(), meta=np.ndarray>
    latitude    float32 dask.array<chunksize=(), meta=np.ndarray>
    longitude   float32 dask.array<chunksize=(), meta=np.ndarray>
    order       int32 dask.array<chunksize=(), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-31T23:00:00
Attributes:
    grid_mapping:  crs
    long_name:     River Flow
    units:         m3 s-1

xarray.DataArray

'streamflow'

• time: 744

• dask.array<chunksize=(256,), meta=np.ndarray>

  Array Chunk Bytes 5.81 kiB 2.00 kiB Shape (744,) (256,) Dask graph 3 chunks in 3 graph layers Data type float64 numpy.ndarray

• Coordinates: (7)
  • elevation
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature elevation

    units :

    meters

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • feature_id
    ()
    int32
    1343034

    cf_role :

    timeseries_id

    comment :

    NHDPlusv2 ComIDs within CONUS, arbitrary Reach IDs outside of CONUS

    long_name :

    Reach ID

    array(1343034, dtype=int32)

  • gage_id
    ()
    |S15
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    NHD Gage Event ID from SOURCE_FEA field in Gages feature class

    Array Chunk Bytes 15 B 15 B Shape () () Dask graph 1 chunks in 3 graph layers Data type |S15 numpy.ndarray

  • latitude
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature latitude

    standard_name :

    latitude

    units :

    degrees_north

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • longitude
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature longitude

    standard_name :

    longitude

    units :

    degrees_east

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • order
    ()
    int32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    stream order

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type int32 numpy.ndarray

  • time
    (time)
    datetime64[ns]
    2000-07-01 ... 2000-07-31T23:00:00

    long_name :

    valid output time

    standard_name :

    time

    array(['2000-07-01T00:00:00.000000000', '2000-07-01T01:00:00.000000000',
           '2000-07-01T02:00:00.000000000', ..., '2000-07-31T21:00:00.000000000',
           '2000-07-31T22:00:00.000000000', '2000-07-31T23:00:00.000000000'],
          dtype='datetime64[ns]')

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-07-01 00:00:00', '2000-07-01 01:00:00',
                   '2000-07-01 02:00:00', '2000-07-01 03:00:00',
                   '2000-07-01 04:00:00', '2000-07-01 05:00:00',
                   '2000-07-01 06:00:00', '2000-07-01 07:00:00',
                   '2000-07-01 08:00:00', '2000-07-01 09:00:00',
                   ...
                   '2000-07-31 14:00:00', '2000-07-31 15:00:00',
                   '2000-07-31 16:00:00', '2000-07-31 17:00:00',
                   '2000-07-31 18:00:00', '2000-07-31 19:00:00',
                   '2000-07-31 20:00:00', '2000-07-31 21:00:00',
                   '2000-07-31 22:00:00', '2000-07-31 23:00:00'],
                  dtype='datetime64[ns]', name='time', length=744, freq=None))

• Attributes: (3)

  grid_mapping :

  crs

  long_name :

  River Flow

  units :

  m3 s-1

This data has 744 time-steps available, chunked into chunks of 256 values each. Three chunks are needed to hold this time-series for one feature. Not too bad, but not good either.

On the other hand, an analysis which samples all locations for a single point in time would need to fetch multiple chunks also.

# Fetch all the gage values for a single day
sampleData['streamflow'].sel(time='07-01-2000')

<xarray.DataArray 'streamflow' (time: 24, feature_id: 400)>
dask.array<getitem, shape=(24, 400), dtype=float64, chunksize=(24, 16), chunktype=numpy.ndarray>
Coordinates:
    elevation   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
  * feature_id  (feature_id) int32 3109 189899 239166 ... 947070134 1010003783
    gage_id     (feature_id) |S15 dask.array<chunksize=(400,), meta=np.ndarray>
    latitude    (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    longitude   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    order       (feature_id) int32 dask.array<chunksize=(400,), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-01T23:00:00
Attributes:
    grid_mapping:  crs
    long_name:     River Flow
    units:         m3 s-1

xarray.DataArray

'streamflow'

• time: 24
• feature_id: 400

• dask.array<chunksize=(24, 16), meta=np.ndarray>

  Array Chunk Bytes 75.00 kiB 3.00 kiB Shape (24, 400) (24, 16) Dask graph 25 chunks in 3 graph layers Data type float64 numpy.ndarray

• Coordinates: (7)
  • elevation
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature elevation

    units :

    meters

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • feature_id
    (feature_id)
    int32
    3109 189899 ... 1010003783

    cf_role :

    timeseries_id

    comment :

    NHDPlusv2 ComIDs within CONUS, arbitrary Reach IDs outside of CONUS

    long_name :

    Reach ID

    array([      3109,     189899,     239166, ...,  945010401,  947070134,
           1010003783], dtype=int32)

  • gage_id
    (feature_id)
    |S15
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    NHD Gage Event ID from SOURCE_FEA field in Gages feature class

    Array Chunk Bytes 5.86 kiB 5.86 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type |S15 numpy.ndarray

  • latitude
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature latitude

    standard_name :

    latitude

    units :

    degrees_north

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • longitude
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature longitude

    standard_name :

    longitude

    units :

    degrees_east

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • order
    (feature_id)
    int32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    stream order

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type int32 numpy.ndarray

  • time
    (time)
    datetime64[ns]
    2000-07-01 ... 2000-07-01T23:00:00

    long_name :

    valid output time

    standard_name :

    time

    array(['2000-07-01T00:00:00.000000000', '2000-07-01T01:00:00.000000000',
           '2000-07-01T02:00:00.000000000', '2000-07-01T03:00:00.000000000',
           '2000-07-01T04:00:00.000000000', '2000-07-01T05:00:00.000000000',
           '2000-07-01T06:00:00.000000000', '2000-07-01T07:00:00.000000000',
           '2000-07-01T08:00:00.000000000', '2000-07-01T09:00:00.000000000',
           '2000-07-01T10:00:00.000000000', '2000-07-01T11:00:00.000000000',
           '2000-07-01T12:00:00.000000000', '2000-07-01T13:00:00.000000000',
           '2000-07-01T14:00:00.000000000', '2000-07-01T15:00:00.000000000',
           '2000-07-01T16:00:00.000000000', '2000-07-01T17:00:00.000000000',
           '2000-07-01T18:00:00.000000000', '2000-07-01T19:00:00.000000000',
           '2000-07-01T20:00:00.000000000', '2000-07-01T21:00:00.000000000',
           '2000-07-01T22:00:00.000000000', '2000-07-01T23:00:00.000000000'],
          dtype='datetime64[ns]')

• Indexes: (2)
  • feature_id
    PandasIndex

    PandasIndex(Int64Index([      3109,     189899,     239166,     334022,     397590,
                    418429,     481940,     591268,     682279,     734307,
                ...
                  41027830,   41028190,   41039416,   41044540,  166743866,
                 932040162,  940080052,  945010401,  947070134, 1010003783],
               dtype='int64', name='feature_id', length=400))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-07-01 00:00:00', '2000-07-01 01:00:00',
                   '2000-07-01 02:00:00', '2000-07-01 03:00:00',
                   '2000-07-01 04:00:00', '2000-07-01 05:00:00',
                   '2000-07-01 06:00:00', '2000-07-01 07:00:00',
                   '2000-07-01 08:00:00', '2000-07-01 09:00:00',
                   '2000-07-01 10:00:00', '2000-07-01 11:00:00',
                   '2000-07-01 12:00:00', '2000-07-01 13:00:00',
                   '2000-07-01 14:00:00', '2000-07-01 15:00:00',
                   '2000-07-01 16:00:00', '2000-07-01 17:00:00',
                   '2000-07-01 18:00:00', '2000-07-01 19:00:00',
                   '2000-07-01 20:00:00', '2000-07-01 21:00:00',
                   '2000-07-01 22:00:00', '2000-07-01 23:00:00'],
                  dtype='datetime64[ns]', name='time', freq=None))

• Attributes: (3)

  grid_mapping :

  crs

  long_name :

  River Flow

  units :

  m3 s-1

This dataset has 400 features, broken into chunks of 16 data values each. Many more chunks to fetch. This is much worse: the I/O engine needs to find and retrieve 25 chunks vs 3 in the previous example.

If we were going to do either of those analyses on a very large dataset with this pattern, we’d want to re-chunk the data to optimize for our read pattern.

Re-Chunking the Sample Data

This is a trivial example, due to the small size of the dataset – It all fits in memory easily. But it is worth doing, as concepts will apply when we take this to the full-sized data.

First thing we need is a chunk plan to describe the chunk layout we want. This can be generated using various methods. For this dataset, it’s easy enough to write it manually:

# Numbers are *size* of the chunk. 
chunk_plan = {
    'streamflow': {'time': 744, 'feature_id': 1}, # all time records in one chunk for each feature_id
    'velocity': {'time': 744, 'feature_id': 1},
    'elevation': (400,),
    'gage_id': (400,),
    'latitude': (400,),
    'longitude': (400,),    
    'order': (400,),    
    'time': (744,),
    'feature_id': (400,)
}

With this plan, we can ask rechunker to re-write the data using the prescribed chunking pattern.

import rechunker
outfile = r"/tmp/outfile.zarr"
result = rechunker.rechunk(
    sampleData,
    chunk_plan,
    "2GB",                #<--- Max Memory
    outfile ,
    temp_store="/tmp/scratch.zarr" 
)
_ = result.execute() # Note that we must specifically direct rechunk to calculate.
# without the call to execute(), the zarr dataset will be empty, and result will hold only
# a 'task graph' outlining the calculation steps.

Note that rechunker.rechunk does not overwrite any data. If it sees that /tmp/outfile.zarr or /tmp/scratch.zarr already exist, it will balk and likely raise an exception. Be sure that these locations do not exist.

The rechunker also writes a minimalist data group. Meaning that variable metadata is not consolidated. This is not a required step, but it will really spead up future workflows when the data is read back in.

import zarr
_ = zarr.consolidate_metadata(outfile)

Results

Let’s read in the resulting re-chunked dataset to see how it looks:

reChunkedData = xr.open_zarr(outfile)
reChunkedData

<xarray.Dataset>
Dimensions:     (feature_id: 400, time: 744)
Coordinates:
    elevation   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
  * feature_id  (feature_id) int32 3109 189899 239166 ... 947070134 1010003783
    gage_id     (feature_id) |S15 dask.array<chunksize=(400,), meta=np.ndarray>
    latitude    (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    longitude   (feature_id) float32 dask.array<chunksize=(400,), meta=np.ndarray>
    order       (feature_id) int32 dask.array<chunksize=(400,), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-31T23:00:00
Data variables:
    streamflow  (time, feature_id) float64 dask.array<chunksize=(744, 1), meta=np.ndarray>
    velocity    (time, feature_id) float64 dask.array<chunksize=(744, 1), meta=np.ndarray>
Attributes:
    TITLE:                OUTPUT FROM WRF-Hydro v5.2.0-beta2
    code_version:         v5.2.0-beta2
    featureType:          timeSeries
    model_configuration:  retrospective
    proj4:                +proj=lcc +units=m +a=6370000.0 +b=6370000.0 +lat_1...

xarray.Dataset

• Dimensions:
  • feature_id: 400
  • time: 744
• Coordinates: (7)
  • elevation
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature elevation

    units :

    meters

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • feature_id
    (feature_id)
    int32
    3109 189899 ... 1010003783

    cf_role :

    timeseries_id

    comment :

    NHDPlusv2 ComIDs within CONUS, arbitrary Reach IDs outside of CONUS

    long_name :

    Reach ID

    array([      3109,     189899,     239166, ...,  945010401,  947070134,
           1010003783], dtype=int32)

  • gage_id
    (feature_id)
    |S15
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    NHD Gage Event ID from SOURCE_FEA field in Gages feature class

    Array Chunk Bytes 5.86 kiB 5.86 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type |S15 numpy.ndarray

  • latitude
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature latitude

    standard_name :

    latitude

    units :

    degrees_north

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • longitude
    (feature_id)
    float32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    feature longitude

    standard_name :

    longitude

    units :

    degrees_east

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type float32 numpy.ndarray

  • order
    (feature_id)
    int32
    dask.array<chunksize=(400,), meta=np.ndarray>

    long_name :

    stream order

    Array Chunk Bytes 1.56 kiB 1.56 kiB Shape (400,) (400,) Dask graph 1 chunks in 2 graph layers Data type int32 numpy.ndarray

  • time
    (time)
    datetime64[ns]
    2000-07-01 ... 2000-07-31T23:00:00

    long_name :

    valid output time

    standard_name :

    time

    array(['2000-07-01T00:00:00.000000000', '2000-07-01T01:00:00.000000000',
           '2000-07-01T02:00:00.000000000', ..., '2000-07-31T21:00:00.000000000',
           '2000-07-31T22:00:00.000000000', '2000-07-31T23:00:00.000000000'],
          dtype='datetime64[ns]')

• Data variables: (2)
  • streamflow
    (time, feature_id)
    float64
    dask.array<chunksize=(744, 1), meta=np.ndarray>

    grid_mapping :

    crs

    long_name :

    River Flow

    units :

    m3 s-1

    Array Chunk Bytes 2.27 MiB 5.81 kiB Shape (744, 400) (744, 1) Dask graph 400 chunks in 2 graph layers Data type float64 numpy.ndarray

  • velocity
    (time, feature_id)
    float64
    dask.array<chunksize=(744, 1), meta=np.ndarray>

    grid_mapping :

    crs

    long_name :

    River Velocity

    units :

    m s-1

    Array Chunk Bytes 2.27 MiB 5.81 kiB Shape (744, 400) (744, 1) Dask graph 400 chunks in 2 graph layers Data type float64 numpy.ndarray

• Indexes: (2)
  • feature_id
    PandasIndex

    PandasIndex(Int64Index([      3109,     189899,     239166,     334022,     397590,
                    418429,     481940,     591268,     682279,     734307,
                ...
                  41027830,   41028190,   41039416,   41044540,  166743866,
                 932040162,  940080052,  945010401,  947070134, 1010003783],
               dtype='int64', name='feature_id', length=400))

  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-07-01 00:00:00', '2000-07-01 01:00:00',
                   '2000-07-01 02:00:00', '2000-07-01 03:00:00',
                   '2000-07-01 04:00:00', '2000-07-01 05:00:00',
                   '2000-07-01 06:00:00', '2000-07-01 07:00:00',
                   '2000-07-01 08:00:00', '2000-07-01 09:00:00',
                   ...
                   '2000-07-31 14:00:00', '2000-07-31 15:00:00',
                   '2000-07-31 16:00:00', '2000-07-31 17:00:00',
                   '2000-07-31 18:00:00', '2000-07-31 19:00:00',
                   '2000-07-31 20:00:00', '2000-07-31 21:00:00',
                   '2000-07-31 22:00:00', '2000-07-31 23:00:00'],
                  dtype='datetime64[ns]', name='time', length=744, freq=None))

• Attributes: (5)

  TITLE :

  OUTPUT FROM WRF-Hydro v5.2.0-beta2

  code_version :

  v5.2.0-beta2

  featureType :

  timeSeries

  model_configuration :

  retrospective

  proj4 :

  +proj=lcc +units=m +a=6370000.0 +b=6370000.0 +lat_1=30.0 +lat_2=60.0 +lat_0=40.0 +lon_0=-97.0 +x_0=0 +y_0=0 +k_0=1.0 +nadgrids=@

Note here that for both streamflow and velocity, the chunksize in the time dimension is 744 (the total number of time steps). Analyses which favor fetching all time-step values for a given facility_id will prefer this chunking strategy.

Comparison

Before Re-Chunking:

sampleData['streamflow'].sel(feature_id=1343034)
# Note: three chunks needed to service a single feature_id

<xarray.DataArray 'streamflow' (time: 744)>
dask.array<getitem, shape=(744,), dtype=float64, chunksize=(256,), chunktype=numpy.ndarray>
Coordinates:
    elevation   float32 dask.array<chunksize=(), meta=np.ndarray>
    feature_id  int32 1343034
    gage_id     |S15 dask.array<chunksize=(), meta=np.ndarray>
    latitude    float32 dask.array<chunksize=(), meta=np.ndarray>
    longitude   float32 dask.array<chunksize=(), meta=np.ndarray>
    order       int32 dask.array<chunksize=(), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-31T23:00:00
Attributes:
    grid_mapping:  crs
    long_name:     River Flow
    units:         m3 s-1

xarray.DataArray

'streamflow'

• time: 744

• dask.array<chunksize=(256,), meta=np.ndarray>

  Array Chunk Bytes 5.81 kiB 2.00 kiB Shape (744,) (256,) Dask graph 3 chunks in 3 graph layers Data type float64 numpy.ndarray

• Coordinates: (7)
  • elevation
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature elevation

    units :

    meters

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • feature_id
    ()
    int32
    1343034

    cf_role :

    timeseries_id

    comment :

    NHDPlusv2 ComIDs within CONUS, arbitrary Reach IDs outside of CONUS

    long_name :

    Reach ID

    array(1343034, dtype=int32)

  • gage_id
    ()
    |S15
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    NHD Gage Event ID from SOURCE_FEA field in Gages feature class

    Array Chunk Bytes 15 B 15 B Shape () () Dask graph 1 chunks in 3 graph layers Data type |S15 numpy.ndarray

  • latitude
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature latitude

    standard_name :

    latitude

    units :

    degrees_north

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • longitude
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature longitude

    standard_name :

    longitude

    units :

    degrees_east

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • order
    ()
    int32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    stream order

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type int32 numpy.ndarray

  • time
    (time)
    datetime64[ns]
    2000-07-01 ... 2000-07-31T23:00:00

    long_name :

    valid output time

    standard_name :

    time

    array(['2000-07-01T00:00:00.000000000', '2000-07-01T01:00:00.000000000',
           '2000-07-01T02:00:00.000000000', ..., '2000-07-31T21:00:00.000000000',
           '2000-07-31T22:00:00.000000000', '2000-07-31T23:00:00.000000000'],
          dtype='datetime64[ns]')

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-07-01 00:00:00', '2000-07-01 01:00:00',
                   '2000-07-01 02:00:00', '2000-07-01 03:00:00',
                   '2000-07-01 04:00:00', '2000-07-01 05:00:00',
                   '2000-07-01 06:00:00', '2000-07-01 07:00:00',
                   '2000-07-01 08:00:00', '2000-07-01 09:00:00',
                   ...
                   '2000-07-31 14:00:00', '2000-07-31 15:00:00',
                   '2000-07-31 16:00:00', '2000-07-31 17:00:00',
                   '2000-07-31 18:00:00', '2000-07-31 19:00:00',
                   '2000-07-31 20:00:00', '2000-07-31 21:00:00',
                   '2000-07-31 22:00:00', '2000-07-31 23:00:00'],
                  dtype='datetime64[ns]', name='time', length=744, freq=None))

• Attributes: (3)

  grid_mapping :

  crs

  long_name :

  River Flow

  units :

  m3 s-1

After re-chunking:

reChunkedData['streamflow'].sel(feature_id=1343034) 
# All data for the specified feature_id is in a single chunk

<xarray.DataArray 'streamflow' (time: 744)>
dask.array<getitem, shape=(744,), dtype=float64, chunksize=(744,), chunktype=numpy.ndarray>
Coordinates:
    elevation   float32 dask.array<chunksize=(), meta=np.ndarray>
    feature_id  int32 1343034
    gage_id     |S15 dask.array<chunksize=(), meta=np.ndarray>
    latitude    float32 dask.array<chunksize=(), meta=np.ndarray>
    longitude   float32 dask.array<chunksize=(), meta=np.ndarray>
    order       int32 dask.array<chunksize=(), meta=np.ndarray>
  * time        (time) datetime64[ns] 2000-07-01 ... 2000-07-31T23:00:00
Attributes:
    grid_mapping:  crs
    long_name:     River Flow
    units:         m3 s-1

xarray.DataArray

'streamflow'

• time: 744

• dask.array<chunksize=(744,), meta=np.ndarray>

  Array Chunk Bytes 5.81 kiB 5.81 kiB Shape (744,) (744,) Dask graph 1 chunks in 3 graph layers Data type float64 numpy.ndarray

• Coordinates: (7)
  • elevation
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature elevation

    units :

    meters

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • feature_id
    ()
    int32
    1343034

    cf_role :

    timeseries_id

    comment :

    NHDPlusv2 ComIDs within CONUS, arbitrary Reach IDs outside of CONUS

    long_name :

    Reach ID

    array(1343034, dtype=int32)

  • gage_id
    ()
    |S15
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    NHD Gage Event ID from SOURCE_FEA field in Gages feature class

    Array Chunk Bytes 15 B 15 B Shape () () Dask graph 1 chunks in 3 graph layers Data type |S15 numpy.ndarray

  • latitude
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature latitude

    standard_name :

    latitude

    units :

    degrees_north

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • longitude
    ()
    float32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    feature longitude

    standard_name :

    longitude

    units :

    degrees_east

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type float32 numpy.ndarray

  • order
    ()
    int32
    dask.array<chunksize=(), meta=np.ndarray>

    long_name :

    stream order

    Array Chunk Bytes 4 B 4 B Shape () () Dask graph 1 chunks in 3 graph layers Data type int32 numpy.ndarray

  • time
    (time)
    datetime64[ns]
    2000-07-01 ... 2000-07-31T23:00:00

    long_name :

    valid output time

    standard_name :

    time

    array(['2000-07-01T00:00:00.000000000', '2000-07-01T01:00:00.000000000',
           '2000-07-01T02:00:00.000000000', ..., '2000-07-31T21:00:00.000000000',
           '2000-07-31T22:00:00.000000000', '2000-07-31T23:00:00.000000000'],
          dtype='datetime64[ns]')

• Indexes: (1)
  • time
    PandasIndex

    PandasIndex(DatetimeIndex(['2000-07-01 00:00:00', '2000-07-01 01:00:00',
                   '2000-07-01 02:00:00', '2000-07-01 03:00:00',
                   '2000-07-01 04:00:00', '2000-07-01 05:00:00',
                   '2000-07-01 06:00:00', '2000-07-01 07:00:00',
                   '2000-07-01 08:00:00', '2000-07-01 09:00:00',
                   ...
                   '2000-07-31 14:00:00', '2000-07-31 15:00:00',
                   '2000-07-31 16:00:00', '2000-07-31 17:00:00',
                   '2000-07-31 18:00:00', '2000-07-31 19:00:00',
                   '2000-07-31 20:00:00', '2000-07-31 21:00:00',
                   '2000-07-31 22:00:00', '2000-07-31 23:00:00'],
                  dtype='datetime64[ns]', name='time', length=744, freq=None))

• Attributes: (3)

  grid_mapping :

  crs

  long_name :

  River Flow

  units :

  m3 s-1

Cleaning Up

import shutil
if os.path.exists(outfile):
    print(f"removing {outfile}")
    shutil.rmtree(outfile)
if os.path.exists(r"/tmp/scratch.zarr"):
    print("removing scratch space")
    shutil.rmtree(r"/tmp/scratch.zarr")

removing /tmp/outfile.zarr
removing scratch space