# Rechunking

Rechunking lets us re-distribute how datasets are split between variables and chunks across a Beam PCollection.

To get started we’ll recreate our dummy data from the data model tutorial:

```
import apache_beam as beam
import numpy as np
import xarray_beam as xbeam
import xarray
def create_records():
for offset in [0, 4]:
key = xbeam.Key({'x': offset, 'y': 0})
data = 2 * offset + np.arange(8).reshape(4, 2)
chunk = xarray.Dataset({
'foo': (('x', 'y'), data),
'bar': (('x', 'y'), 100 + data),
})
yield key, chunk
inputs = list(create_records())
```

## Choosing chunks

Chunking can be essential for some operations. Some operations are very hard or impossible to perform with certain chunking schemes. For example, to make a plot all the data needs to come toether on a single machine. Other calculations such as calculating a median are *possible* to perform on distributed data, but require tricky algorithms and/or approximation.

More broadly, chunking can have critical performance implications, similar to those for Xarray and Dask. As a rule of thumb, chunk sizes of 10-100 MB work well. The optimal chunk size is a balance among a number of considerations, adapted here from Dask docs:

Chunks should be small enough to fit comfortably into memory on a single machine. As an upper limit, chunks over roughly 2 GB in size will not fit into the protocol buffers Beam uses to pass data between workers.

There should be enough chunks for Beam runners (like Cloud Dataflow) to elastically shard work over many workers.

Chunks should be large enough to amortize the overhead of networking and the Python interpreter, which starts to become noticeable for arrays with fewer than 1 million elements.

The `nbytes`

attribute on both NumPy arrays and `xarray.Dataset`

objects is a good easy way to figure out how larger chunks are.

## Adjusting variables

The simplest transformation is splitting (or consoldating) different *variables* in a Dataset with `SplitVariables()`

and `ConsolidateVariables()`

, e.g.,

```
inputs | xbeam.SplitVariables()
```

```
[(Key(offsets={'x': 0, 'y': 0}, vars={'foo'}),
<xarray.Dataset>
Dimensions: (x: 4, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 0 1 2 3 4 5 6 7),
(Key(offsets={'x': 0, 'y': 0}, vars={'bar'}),
<xarray.Dataset>
Dimensions: (x: 4, y: 2)
Dimensions without coordinates: x, y
Data variables:
bar (x, y) int64 100 101 102 103 104 105 106 107),
(Key(offsets={'x': 4, 'y': 0}, vars={'foo'}),
<xarray.Dataset>
Dimensions: (x: 4, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 8 9 10 11 12 13 14 15),
(Key(offsets={'x': 4, 'y': 0}, vars={'bar'}),
<xarray.Dataset>
Dimensions: (x: 4, y: 2)
Dimensions without coordinates: x, y
Data variables:
bar (x, y) int64 108 109 110 111 112 113 114 115)]
```

Tip

Instead of a separate transform for splitting variables, you can also set `split_vars=True`

in `DatasetToChunks`

.

## Adjusting chunks

You can also adjust *chunks* in a dataset to distribute arrays of different sizes. Here you have two choices of API:

The lower level

`SplitChunks`

and`ConsolidateChunks`

. These transformations apply a single splitting (with indexing) or consolidation (with xarray.concat) function to array elements.The high level

`Rechunk`

, which uses a pipeline of multiple split/consolidate steps (as needed) to efficiently rechunk a dataset.

### Low level rechunking

For minor adjustments (e.g., mostly along a single dimension), the more explicit `SplitChunks()`

and `ConsolidateChunks()`

are good options. They take a dict of *desired* chunk sizes as a parameter, which can also be `-1`

to indicate “no chunking” along a dimension:

```
inputs | xbeam.ConsolidateChunks({'x': -1})
```

```
[(Key(offsets={'x': 0, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 8, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
bar (x, y) int64 100 101 102 103 104 105 ... 110 111 112 113 114 115)]
```

Note that because these transformations only split *or* consolidate, they cannot necessary fully rechunk a dataset in a single step if the new chunk sizes are not multiples of old chunks (with consolidate) or do not even divide the old chunks (with split), e.g.,

```
inputs | xbeam.SplitChunks({'x': 5}) # notice that the first two chunks are still separate!
```

```
[(Key(offsets={'x': 0, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 4, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 0 1 2 3 4 5 6 7
bar (x, y) int64 100 101 102 103 104 105 106 107),
(Key(offsets={'x': 4, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 1, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 8 9
bar (x, y) int64 108 109),
(Key(offsets={'x': 5, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 3, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 10 11 12 13 14 15
bar (x, y) int64 110 111 112 113 114 115)]
```

For such uneven cases, you’ll need to use split followed by consolidate:

```
inputs | xbeam.SplitChunks({'x': 5}) | xbeam.ConsolidateChunks({'x': 5})
```

```
[(Key(offsets={'x': 0, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 5, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 0 1 2 3 4 5 6 7 8 9
bar (x, y) int64 100 101 102 103 104 105 106 107 108 109),
(Key(offsets={'x': 5, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 3, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 10 11 12 13 14 15
bar (x, y) int64 110 111 112 113 114 115)]
```

### High level rechunking

Alternatively, the high-level `Rechunk()`

method applies multiple split and consolidate steps based on the Rechunker algorithm:

```
inputs | xbeam.Rechunk(dim_sizes={'x': 6}, source_chunks={'x': 3}, target_chunks={'x': 5}, itemsize=8)
```

```
[(Key(offsets={'x': 0, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 5, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 0 1 2 3 4 5 6 7 8 9
bar (x, y) int64 100 101 102 103 104 105 106 107 108 109),
(Key(offsets={'x': 5, 'y': 0}, vars=None),
<xarray.Dataset>
Dimensions: (x: 3, y: 2)
Dimensions without coordinates: x, y
Data variables:
foo (x, y) int64 10 11 12 13 14 15
bar (x, y) int64 110 111 112 113 114 115)]
```

`Rechunk`

requires specifying a few more parameters, but based on that information **it can be much more efficient for more complex rechunking tasks**, particular in cases where data needs to be distributed into a very different shape (e.g., distributing a matrix across rows vs. columns).

The naive “splitting” approach in such cases may divide datasets into extremely small tasks corresponding to individual array elements, which adds a huge amount of overhead.