docs

Author: jhamman

changes/3534.feature.md

@@ -0,0 +1 @@
+Adds support for `RectilinearChunkGrid`, enabling arrays with variable chunk sizes along each dimension in Zarr v3. Users can now specify irregular chunking patterns using nested sequences: `chunks=[[10, 20, 30], [25, 25, 25, 25]]` creates an array with 3 chunks of sizes 10, 20, and 30 along the first dimension, and 4 chunks of size 25 along the second dimension. This feature is useful for data with non-uniform structure or when aligning chunks with existing data partitions. Note that `RectilinearChunkGrid` is only supported in Zarr format 3 and cannot be used with sharding or when creating arrays from existing data via `from_array()`.

docs/user-guide/arrays.md

@@ -566,6 +566,124 @@ In this example a shard shape of (1000, 1000) and a chunk shape of (100, 100) is
 This means that `10*10` chunks are stored in each shard, and there are `10*10` shards in total.
 Without the `shards` argument, there would be 10,000 chunks stored as individual files.
 
+## Variable Chunking (Zarr v3)
+
+In addition to regular chunking where all chunks have the same size, Zarr v3 supports
+**variable chunking** (also called rectilinear chunking), where chunks can have different
+sizes along each dimension. This is useful when your data has non-uniform structure or
+when you need to align chunks with existing data partitions.
+
+### Basic usage
+
+To create an array with variable chunking, provide a nested sequence to the `chunks`
+parameter instead of a regular tuple:
+
+```python exec="true" session="arrays" source="above" result="ansi"
+# Create an array with variable chunk sizes
+z = zarr.create_array(
+    store='data/example-21.zarr',
+    shape=(60, 100),
+    chunks=[[10, 20, 30], [25, 25, 25, 25]],  # Variable chunks
+    dtype='float32',
+    zarr_format=3
+)
+print(z)
+print(f"Chunk grid type: {type(z.metadata.chunk_grid).__name__}")
+```
+
+In this example, the first dimension is divided into 3 chunks with sizes 10, 20, and 30
+(totaling 60), and the second dimension is divided into 4 chunks of size 25 (totaling 100).
+
+### Reading and writing
+
+Arrays with variable chunking support the same read/write operations as regular arrays:
+
+```python exec="true" session="arrays" source="above" result="ansi"
+# Write data
+data = np.arange(60 * 100, dtype='float32').reshape(60, 100)
+z[:] = data
+
+# Read data back
+result = z[:]
+print(f"Data matches: {np.all(result == data)}")
+print(f"Slice [10:30, 50:75]: {z[10:30, 50:75].shape}")
+```
+
+### Accessing chunk information
+
+With variable chunking, the standard `.chunks` property is not available since chunks
+have different sizes. Instead, access chunk information through the chunk grid:
+
+```python exec="true" session="arrays" source="above" result="ansi"
+from zarr.core.chunk_grids import RectilinearChunkGrid
+
+# Access the chunk grid
+chunk_grid = z.metadata.chunk_grid
+print(f"Chunk grid type: {type(chunk_grid).__name__}")
+
+# Get chunk shapes for each dimension
+if isinstance(chunk_grid, RectilinearChunkGrid):
+    print(f"Dimension 0 chunk sizes: {chunk_grid.chunk_shapes[0]}")
+    print(f"Dimension 1 chunk sizes: {chunk_grid.chunk_shapes[1]}")
+    print(f"Total number of chunks: {chunk_grid.get_nchunks((60, 100))}")
+```
+
+### Use cases
+
+Variable chunking is particularly useful for:
+
+1. **Irregular time series**: When your data has non-uniform time intervals, you can
+   create chunks that align with your sampling periods.
+
+2. **Aligning with partitions**: When you need to match chunk boundaries with existing
+   data partitions or structural boundaries in your data.
+
+3. **Optimizing access patterns**: When certain regions of your array are accessed more
+   frequently, you can use smaller chunks there for finer-grained access.
+
+### Example: Time series with irregular intervals
+
+```python exec="true" session="arrays" source="above" result="ansi"
+# Daily measurements for one year, chunked by month
+# Each chunk corresponds to one month (varying from 28-31 days)
+z_timeseries = zarr.create_array(
+    store='data/example-22.zarr',
+    shape=(365, 100),  # 365 days, 100 measurements per day
+    chunks=[[31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31], [100]],  # Days per month
+    dtype='float64',
+    zarr_format=3
+)
+print(f"Created array with shape {z_timeseries.shape}")
+print(f"Chunk shapes: {z_timeseries.metadata.chunk_grid.chunk_shapes}")
+print(f"Number of chunks: {len(z_timeseries.metadata.chunk_grid.chunk_shapes[0])} months")
+```
+
+### Limitations
+
+Variable chunking has some important limitations:
+
+1. **Zarr v3 only**: This feature is only available when using `zarr_format=3`.
+   Attempting to use variable chunks with `zarr_format=2` will raise an error.
+
+2. **Not compatible with sharding**: You cannot use variable chunking together with
+   the sharding feature. Arrays must use either variable chunking or sharding, but not both.
+
+3. **Not compatible with `from_array()`**: Variable chunking cannot be used when creating
+   arrays from existing data using [`zarr.from_array`][]. This is because the function needs
+   to partition the input data, which requires regular chunk sizes.
+
+4. **No `.chunks` property**: For arrays with variable chunking, accessing the `.chunks`
+   property will raise a `NotImplementedError`. Use `.metadata.chunk_grid.chunk_shapes`
+   instead.
+
+```python exec="true" session="arrays" source="above" result="ansi"
+# This will raise an error
+try:
+    _ = z.chunks
+except NotImplementedError as e:
+    print(f"Error: {e}")
+```
+
 ## Missing features in 3.0
 
 The following features have not been ported to 3.0 yet.

docs/user-guide/extending.md

@@ -85,4 +85,6 @@ classes by implementing the interface defined in [`zarr.abc.buffer.BufferPrototy
 
 ## Other extensions
 
-In the future, Zarr will support writing custom custom data types and chunk grids.
+Zarr now includes built-in support for `RectilinearChunkGrid` (variable chunking), which allows arrays to have different chunk sizes along each dimension. See the [Variable Chunking](arrays.md#variable-chunking-zarr-v3) section in the Arrays guide for more information.
+
+In the future, Zarr will support writing fully custom chunk grids and custom data types.

src/zarr/core/array.py

@@ -4338,6 +4338,10 @@ async def from_array(
         - tuple[int, ...]: A tuple of integers representing the chunk shape.
 
         If not specified, defaults to "keep" if data is a zarr Array, otherwise "auto".
+
+        .. note::
+            Variable chunking (RectilinearChunkGrid) is not supported when creating arrays from
+            existing data. Use regular chunking (uniform chunk sizes) instead.
     shards : tuple[int, ...], optional
         Shard shape of the array.
         Following values are supported:

tests/test_chunk_grids/test_rectilinear_integration.py

@@ -0,0 +1,163 @@
+"""Integration tests for RectilinearChunkGrid with array creation."""
+
+from typing import Literal
+
+import numpy as np
+import pytest
+
+import zarr
+from zarr.core.chunk_grids import RectilinearChunkGrid
+from zarr.storage import MemoryStore
+
+
+@pytest.mark.parametrize("zarr_format", [3])
+async def test_create_array_with_nested_chunks(zarr_format: Literal[2, 3]) -> None:
+    """
+    Test creating an array with nested chunk specification (RectilinearChunkGrid).
+    This is an end-to-end test for the feature.
+    """
+    store = MemoryStore()
+    arr = await zarr.api.asynchronous.create_array(
+        store=store,
+        shape=(60, 100),
+        chunks=[[10, 20, 30], [25, 25, 25, 25]],
+        dtype="i4",
+        zarr_format=zarr_format,
+    )
+
+    # Verify metadata has RectilinearChunkGrid
+    assert isinstance(arr.metadata.chunk_grid, RectilinearChunkGrid)
+    assert arr.metadata.chunk_grid.chunk_shapes == ((10, 20, 30), (25, 25, 25, 25))
+
+    # Verify array is functional - can write and read data
+    data = np.arange(60 * 100, dtype="i4").reshape(60, 100)
+    await arr.setitem(slice(None), data)
+    result = await arr.getitem(slice(None))
+    np.testing.assert_array_equal(result, data)
+
+
+async def test_create_array_nested_chunks_read_write() -> None:
+    """
+    Test that arrays with RectilinearChunkGrid support standard read/write operations.
+    """
+    store = MemoryStore()
+    arr = await zarr.api.asynchronous.create_array(
+        store=store,
+        shape=(30, 40),
+        chunks=[[10, 10, 10], [10, 10, 10, 10]],
+        dtype="f4",
+        zarr_format=3,
+    )
+
+    # Write data to different chunks
+    arr_data = np.random.random((30, 40)).astype("f4")
+    await arr.setitem(slice(None), arr_data)
+
+    # Read full array
+    result = await arr.getitem(slice(None))
+    np.testing.assert_array_almost_equal(np.asarray(result), arr_data)
+
+    # Read partial slices
+    partial = await arr.getitem((slice(5, 25), slice(10, 30)))
+    np.testing.assert_array_almost_equal(np.asarray(partial), arr_data[5:25, 10:30])
+
+
+async def test_rectilinear_chunk_grid_roundtrip() -> None:
+    """
+    Test that RectilinearChunkGrid persists correctly through save/load.
+    """
+    store = MemoryStore()
+
+    # Create array with nested chunks
+    arr1 = await zarr.api.asynchronous.create_array(
+        store=store,
+        name="test_array",
+        shape=(60, 80),
+        chunks=[[10, 20, 30], [20, 20, 20, 20]],
+        dtype="u1",
+        zarr_format=3,
+    )
+
+    # Write some data
+    data = np.arange(60 * 80, dtype="u1").reshape(60, 80)
+    await arr1.setitem(slice(None), data)
+
+    # Re-open the array
+    arr2 = await zarr.api.asynchronous.open_array(store=store, path="test_array")
+
+    # Verify chunk_grid is preserved
+    assert isinstance(arr2.metadata.chunk_grid, RectilinearChunkGrid)
+    assert arr2.metadata.chunk_grid.chunk_shapes == ((10, 20, 30), (20, 20, 20, 20))
+
+    # Verify data is preserved
+    result = await arr2.getitem(slice(None))
+    np.testing.assert_array_equal(result, data)
+
+
+async def test_from_array_rejects_nested_chunks() -> None:
+    """
+    Test that from_array rejects nested chunks (RectilinearChunkGrid) with has_data=True.
+    """
+    store = MemoryStore()
+    data = np.arange(30 * 40, dtype="i4").reshape(30, 40)
+
+    # Should raise error because RectilinearChunkGrid is not compatible with has_data=True
+    with pytest.raises(
+        ValueError,
+        match="Cannot use RectilinearChunkGrid.*when creating array from data",
+    ):
+        await zarr.api.asynchronous.from_array(
+            store=store,
+            data=data,
+            chunks=[[10, 10, 10], [10, 10, 10, 10]],  # type: ignore[arg-type]
+            zarr_format=3,
+        )
+
+
+async def test_nested_chunks_with_different_sizes() -> None:
+    """
+    Test RectilinearChunkGrid with highly irregular chunk sizes.
+    """
+    store = MemoryStore()
+    arr = await zarr.api.asynchronous.create_array(
+        store=store,
+        shape=(100, 100),
+        chunks=[[5, 10, 15, 20, 50], [100]],  # Very irregular first dim, uniform second
+        dtype="i2",
+        zarr_format=3,
+    )
+
+    assert isinstance(arr.metadata.chunk_grid, RectilinearChunkGrid)
+    assert arr.metadata.chunk_grid.chunk_shapes == ((5, 10, 15, 20, 50), (100,))
+
+    # Verify writes work correctly
+    data = np.arange(100 * 100, dtype="i2").reshape(100, 100)
+    await arr.setitem(slice(None), data)
+    result = await arr.getitem(slice(None))
+    np.testing.assert_array_equal(result, data)
+
+
+async def test_rectilinear_chunk_grid_nchunks_not_supported() -> None:
+    """
+    Test that nchunks property raises NotImplementedError for RectilinearChunkGrid.
+
+    Note: The chunks property (and thus nchunks) is only defined for RegularChunkGrid.
+    For RectilinearChunkGrid, use chunk_grid.get_nchunks() instead.
+    """
+    store = MemoryStore()
+    arr = await zarr.api.asynchronous.create_array(
+        store=store,
+        shape=(60, 100),
+        chunks=[[10, 20, 30], [25, 25, 25, 25]],
+        dtype="u1",
+        zarr_format=3,
+    )
+
+    # The chunks property is not defined for RectilinearChunkGrid
+    with pytest.raises(
+        NotImplementedError, match="only defined for arrays using.*RegularChunkGrid"
+    ):
+        _ = arr.nchunks
+
+    # But we can get nchunks from the chunk_grid directly
+    assert arr.metadata.chunk_grid.get_nchunks((60, 100)) == 12