Rechunker

https://rechunker.readthedocs.io/en/latest/
Admin
Rechunker is a Python package for efficient and scalable manipulation of chunk structures in array...
Tokens:6,008
Snippets:71
Trust Score:9.6
Update:5 months ago
Show doc for...
Tokens:
Raw
### Install Rechunker

Source: https://rechunker.readthedocs.io/en/latest/index.html

Use pip to install the rechunker package.

```bash
>>> pip install rechunker

```

--------------------------------

### Import Libraries and Load Example Data

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Imports necessary libraries (xarray, zarr, dask) and loads a sample dataset. It then chunks the dataset for demonstration purposes and prepares it for writing to Zarr.

```python
import xarray as xr

xr.set_options(display_style="text")
import zarr
import dask.array as dsa

ds = xr.tutorial.open_dataset("air_temperature")
# create initial chunk structure
ds = ds.chunk({"time": 100})
ds.air.encoding = {}  # helps when writing to zarr
ds
```

--------------------------------

### Execute Rechunk Plan with Python Executor

Source: https://rechunker.readthedocs.io/en/latest/executors.html

Use the PythonExecutor for simple, loop-based execution, which is helpful for debugging rechunking plans. Ensure optional dependencies for executors are installed.

```python
>>> import zarr
>>> from rechunker import rechunk
>>> source = zarr.ones((4, 4), chunks=(2, 2), store="source.zarr")
>>> intermediate = "intermediate.zarr"
>>> target = "target.zarr"
>>> rechunked = rechunk(
...    source, target_chunks=(4, 1), target_store=target,
...    max_mem=256000000, temp_store=intermediate,
...    executor="python"
... )
>>> rechunked.execute()

```

--------------------------------

### Rechunk Function Example

Source: https://rechunker.readthedocs.io/en/latest/api.html

Demonstrates how to use the rechunk function to rechunk a Zarr array. Ensure to call execute() on the returned Rechunked object and manually delete the intermediate store after completion.

```python
>>> source = zarr.ones((4, 4), chunks=(2, 2), store="source.zarr")
>>> intermediate = "intermediate.zarr"
>>> target = "target.zarr"
>>> rechunked = rechunk(source, target_chunks=(4, 1), target_store=target,
...                     max_mem=256000,
...                     temp_store=intermediate)
>>> rechunked
<Rechunked>
* Source      : <zarr.core.Array (4, 4) float64>
* Intermediate: dask.array<from-zarr, ... >
* Target      : <zarr.core.Array (4, 4) float64>
>>> rechunked.execute()
<zarr.core.Array (4, 4) float64>
```

--------------------------------

### Define Target Chunks for Group Rechunking

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Specify target chunk sizes for each array within a group. Use `None` to indicate arrays that should not be rechunked. This setup is for rechunking an entire group.

```python
target_chunks = {
    "air": {"time": 2920, "lat": 25, "lon": 1},
    "time": None,  # don't rechunk this array
    "lon": None,
    "lat": None,
}
max_mem = "1MB"

target_store = "group_rechunked.zarr"
temp_store = "group_rechunked-tmp.zarr"

# need to remove the existing stores or it won't work
!rm -rf group_rechunked.zarr group_rechunked-tmp.zarr
array_plan = rechunk(
    source_group, target_chunks, max_mem, target_store, temp_store=temp_store
)
array_plan
```

--------------------------------

### Get Zarr Array Chunks

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Retrieves the chunking scheme of a Zarr array. This information is crucial for understanding how the data is partitioned and for planning rechunking operations.

```python
source_array.chunks
```

--------------------------------

### Rechunk Group with Dimension-Based Chunks

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

For groups where dimensions should have equivalent chunks across variables, provide a simpler dictionary mapping dimension names to chunk sizes. This example rechunks directly from an Xarray Dataset.

```python
# extend the dataset with some more variables
ds_complex = ds
ds_complex["air_slice"] = ds.air.isel(lat=10)
ds_complex["air_timeseries"] = ds.air.isel(lat=10, lon=10)
ds_complex

target_chunks = {"time": 2920, "lat": 25, "lon": 1}
max_mem = "1MB"

target_store = "group_complex_rechunked.zarr"
temp_store = "group_complex_rechunked-tmp.zarr"

# need to remove the existing stores or it won't work
!rm -rf group_complex_rechunked.zarr group_complex_rechunked-tmp.zarr

# rechunk directly from dataset this time
array_plan = rechunk(
    ds_complex, target_chunks, max_mem, target_store, temp_store=temp_store
)
array_plan
```

--------------------------------

### Initialize GCSFileSystem for Cloud Storage

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Sets up a GCSFileSystem object to interact with Google Cloud Storage. It uses requester pays for accessing the specified bucket.

```python
import gcsfs

# a zarr group lives here
url = "gs://pangeo-cmems-duacs"
gcs = gcsfs.GCSFileSystem(requester_pays=True)
source_store = gcs.get_mapper(url)
```

--------------------------------

### Initialize Dask Gateway Cluster

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Sets up a Dask Gateway cluster with a specified number of workers. This is useful for scaling up processing in a cloud environment.

```python
from dask_gateway import GatewayCluster

cluster = GatewayCluster()
cluster.scale(20)
cluster
```

--------------------------------

### Open Zarr Store with Consolidated Metadata

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Opens a Zarr store using xarray. It attempts to use consolidated metadata for faster opening. If this fails, it falls back to non-consolidated metadata, which is slower. Consider consolidating metadata or explicitly setting the `consolidated` parameter.

```python
xr.open_zarr("group_complex_rechunked.zarr")
```

--------------------------------

### Initialize Dask Distributed Cluster

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Sets up a local Dask cluster and client, which rechunker will utilize for distributed execution if available.

```python
from dask.distributed import Client, LocalCluster, progress

cluster = LocalCluster()
client = Client(cluster)
```

--------------------------------

### Create and Push Release Tag

Source: https://rechunker.readthedocs.io/en/latest/contributing.html

Use these Git commands to create an empty commit for a new release, tag it, and push the tag to the upstream repository.

```git
git commit --allow-empty -m "RLS: v0.2.0"
git tag -a 0.2.0 -m "RLS: 0.2.0"
git push upstream master --follow-tags

```

--------------------------------

### Open Zarr Group and Print Tree

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Opens the 'air_temperature.zarr' file as a Zarr group and prints its hierarchical structure.

```python
source_group = zarr.open("air_temperature.zarr")
print(source_group.tree())
```

--------------------------------

### Write Data to Zarr and Clean Up

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Removes any existing Zarr data and writes the current dataset to a Zarr store named 'air_temperature.zarr'.

```python
! rm -rf *.zarr # clean up any existing temporary data
ds.to_zarr("air_temperature.zarr")
```

--------------------------------

### Basic Rechunking with Rechunker and Dask

Source: https://rechunker.readthedocs.io/en/latest/index.html

Demonstrates how to rechunk a Zarr array using rechunker. This involves specifying source, target, intermediate storage, and memory limits. The rechunk operation is then executed.

```python
>>> import zarr
>>> from rechunker import rechunk
>>> source = zarr.ones((4, 4), chunks=(2, 2), store="source.zarr")
>>> intermediate = "intermediate.zarr"
>>> target = "target.zarr"
>>> rechunked = rechunk(source, target_chunks=(4, 1), target_store=target,
...                     max_mem=256000,
...                     temp_store=intermediate)
>>> rechunked
<Rechunked>
* Source      : <zarr.core.Array (4, 4) float64>
* Intermediate: dask.array<from-zarr, ... >
* Target      : <zarr.core.Array (4, 4) float64>
>>> rechunked.execute()
<zarr.core.Array (4, 4) float64>

```

--------------------------------

### Execute Rechunking with Progress Bar

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Executes the rechunking process and displays a progress bar to monitor its completion, useful for long-running operations.

```python
from dask.diagnostics import ProgressBar

with ProgressBar():
    array_plan.execute()
```

--------------------------------

### Execute Rechunking and Check Chunks

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Executes the rechunking plan created previously and verifies the resulting chunk structure of the rechunked array.

```python
result = array_plan.execute()
result.chunks
```

--------------------------------

### Initialize Dask Distributed Client

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Connects a Dask distributed client to an existing Dask Gateway cluster. This client is used to manage and interact with the distributed computation.

```python
from dask.distributed import Client

client = Client(cluster)
client
```

--------------------------------

### Open Rechunked Array with Zarr

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Opens the rechunked Zarr array using the Zarr library, allowing direct access to the rechunked data on disk.

```python
target_array = zarr.open("air_rechunked.zarr")
target_array
```

--------------------------------

### Execute Rechunking Plan

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Executes the previously defined rechunking plan. This operation processes the data according to the plan and writes the rechunked data to the target store.

```python
result = r.execute()
result
```

--------------------------------

### Execute Group Rechunking Plan

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

After defining the rechunking plan, execute it to perform the actual data rechunking. This operation writes the rechunked data to the specified target store.

```python
array_plan.execute()
```

--------------------------------

### Open Rechunked Array with Dask

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Opens the rechunked Zarr array using Dask, enabling lazy loading and further Dask-based operations on the rechunked data.

```python
target_array_dask = dsa.from_zarr("air_rechunked.zarr")
target_array_dask
```

--------------------------------

### Create Rechunking Plan

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Defines a rechunking plan for a Zarr array. It specifies the target chunk sizes, maximum memory per chunk, and the target and temporary storage locations.

```python
max_mem = "1GB"
target_chunks = (8901, 72, 72)
# you must have write access to this location
store_tmp = gcs.get_mapper("pangeo-scratch/rabernat/rechunker_demo/temp.zarr")
store_target = gcs.get_mapper("pangeo-scratch/rabernat/rechunker_demo/target.zarr")
r = rechunk(source_array, target_chunks, max_mem, store_target, temp_store=store_tmp)
r
```

--------------------------------

### BeamExecutor

Source: https://rechunker.readthedocs.io/en/latest/api.html

An execution engine based on Apache Beam. Supports copying between any arrays that implement `__getitem__` and `__setitem__` for tuples of `slice` objects. Array must also be serializable by Beam (i.e., with pickle). Execution plans for BeamExecutor are beam.PTransform objects.

```APIDOC
## rechunker.executors.beam.BeamExecutor()

### Description
An execution engine based on Apache Beam.
Supports copying between any arrays that implement `__getitem__` and `__setitem__` for tuples of `slice` objects. Array must also be serializable by Beam (i.e., with pickle).
Execution plans for BeamExecutor are beam.PTransform objects.

### Methods
#### execute_plan(plan, **kwargs)
Execute a plan.

#### prepare_plan(specs)
Convert copy specifications into a plan.
```

--------------------------------

### Access and Inspect Source Array

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Retrieves the 'air' array from the Zarr group and displays its information, including data type, shape, chunk shape, and compression details.

```python
source_array = source_group["air"]
source_array.info
```

--------------------------------

### Open Rechunked Group with Xarray

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Once the group has been rechunked, it can be opened using Xarray. This allows for further analysis and manipulation of the rechunked data.

```python
xr.open_zarr("group_rechunked.zarr")
```

--------------------------------

### Create Xarray DataArray from Zarr Store

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Creates an Xarray DataArray from a Zarr store. This is typically done after rechunking to load the rechunked data into an Xarray object for further analysis.

```python
dsa.from_zarr(result)
```

--------------------------------

### Open Consolidated Zarr Group

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Opens a Zarr group from a GCS store using consolidated metadata. This provides access to the Zarr arrays within the group.

```python
group = zarr.open_consolidated(source_store, mode="r")
source_array = group["sla"]
source_array
```

--------------------------------

### Examine Data Chunk Structure

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Displays the current chunk structure of the 'air' data variable using Dask's array representation.

```python
ds.air.data
```

--------------------------------

### Rechunk Array with Dictionary Specification

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Rechunks the source array using a dictionary to specify target chunk sizes for each dimension. It first removes existing stores to ensure a clean operation.

```python
target_chunks_dict = {"time": 2920, "lat": 25, "lon": 1}

# need to remove the existing stores or it won't work
!rm -rf air_rechunked.zarr air_rechunked-tmp.zarr
array_plan = rechunk(
    source_array, target_chunks_dict, max_mem, target_store, temp_store=temp_store
)
array_plan
```

--------------------------------

### Rechunk Array with Tuple Specification

Source: https://rechunker.readthedocs.io/en/latest/tutorial.html

Rechunks the source array to have contiguous time chunks and specified spatial chunking. It uses a small `max_mem` to force intermediate dataset creation and specifies target and temporary storage locations.

```python
from rechunker import rechunk

target_chunks = (2920, 25, 1)
max_mem = "1MB"

target_store = "air_rechunked.zarr"
temp_store = "air_rechunked-tmp.zarr"

array_plan = rechunk(
    source_array, target_chunks, max_mem, target_store, temp_store=temp_store
)
array_plan
```

--------------------------------

### rechunk() Function

Source: https://rechunker.readthedocs.io/en/latest/api.html

The main function to rechunk Zarr Arrays, Groups, Dask Arrays, or Xarray Datasets. It takes source and target specifications, memory limits, and an optional executor to perform the rechunking operation.

```APIDOC
## rechunk()

### Description
Rechunk a Zarr Array or Group, a Dask Array, or an Xarray Dataset.

### Method
`rechunk()`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters
- **source** (zarr.Array, zarr.Group, dask.array.Array, or xarray.Dataset) - Required - Named dimensions in the Zarr arrays will be parsed according to the Xarray Zarr Encoding Specification.
- **target_chunks** (tuple, dict, or None) - Required - The desired chunks of the array after rechunking. For a single array source, `target_chunks` can be either a tuple (e.g. `(20, 5, 3)`) or a dictionary (e.g. `{'time': 20, 'lat': 5, 'lon': 3}`). Dictionary syntax requires the dimension names be present in the Zarr Array attributes. A value of `None` means that the array will be copied with no change to its chunk structure. For a group of arrays, a dict is required where keys correspond to array names and values are `target_chunks` arguments for the array. Arrays not present in the `target_chunks` dict will be ignored.
- **max_mem** (str or int) - Required - The amount of memory (in bytes) that workers are allowed to use. A string (e.g. `100MB`) can also be used.
- **target_store** (str, MutableMapping, or zarr.Store object) - Required - The location in which to store the final, rechunked result. Will be passed directly to `zarr.creation.create()`.
- **target_options** (Dict, optional) - Optional - Additional keyword arguments used to control array storage. If the source is `xarray.Dataset`, then these options will be used to encode variables in the same manner as the `encoding` parameter in `xarray.Dataset.to_zarr()`. Otherwise, these options will be passed to `zarr.creation.create()`. For a single array source, this should be a single dict. For a group of arrays, a nested dict is required.
- **temp_store** (str, MutableMapping, or zarr.Store object, optional) - Optional - Location of temporary store for intermediate data. Can be deleted once rechunking is complete.
- **temp_options** (Dict, optional) - Optional - Options with same semantics as `target_options` for `temp_store` rather than `target_store`. Defaults to `target_options` and has no effect when source is of type xarray.Dataset.
- **executor** (str or rechunker.types.Executor) - Optional - Implementation of the execution engine for copying between zarr arrays. Defaults to 'dask'. Other options include 'beam', 'prefect', 'python', 'pywren'.
- **array_name** (str, optional) - Optional - Required when rechunking an array if any of the targets is a group.

### Returns
- **rechunked** (`Rechunked` object) - A delayed rechunked result representing the rechunking plan.

### Request Example
```python
>>> import zarr
>>> from rechunker import rechunk
>>> source = zarr.ones((4, 4), chunks=(2, 2), store="source.zarr")
>>> intermediate = "intermediate.zarr"
>>> target = "target.zarr"
>>> rechunked = rechunk(source, target_chunks=(4, 1), target_store=target,
...                     max_mem=256000,
...                     temp_store=intermediate)
>>> rechunked
<Rechunked>
* Source      : <zarr.core.Array (4, 4) float64>
* Intermediate: dask.array<from-zarr, ... >
* Target      : <zarr.core.Array (4, 4) float64>
>>> rechunked.execute()
<zarr.core.Array (4, 4) float64>
```

### Response
#### Success Response (200)
- **Rechunked object** - A `Rechunked` object that holds the rechunking plan.

#### Response Example
```json
{
  "example": "<Rechunked>"
}
```

### Notes
- You must call `execute()` on the `Rechunked` object to perform the rechunking.
- You must manually delete the intermediate store when `execute` is finished.
```

--------------------------------

### Rechunked Object

Source: https://rechunker.readthedocs.io/en/latest/api.html

Represents a delayed rechunking operation. It holds the plan and can be executed to perform the actual rechunking.

```APIDOC
## Rechunked Object

### Description
The `Rechunked` object is returned by the `rechunk()` function. It represents the rechunking plan and allows for delayed execution.

### Class Definition
`class rechunker.Rechunked(_executor_, _plan_, _source_, _intermediate_, _target_)

### Attributes
- **plan**: Returns the executor-specific scheduling plan.

### Methods
- **execute(**kwargs**)**: Executes the rechunking operation.

### Example Usage
```python
>>> # Assuming 'rechunked' is a Rechunked object obtained from rechunk()
>>> rechunked.execute()
```
```
Tokens:
Raw