[v3] Inner chunk size validation behavior for ShardingCodec when downstream of TransposeCodec

Zarr version

3.0.0a0

Numcodecs version

0.13.0

Python Version

3.12.4

Operating System

Linux

Installation

using pip into conda env

Description

I expect a ShardingCodec downstream of a TransposeCodec to consume the transposed array. As a result, I would expect the inner chunk size would have to be "transposed" in the same way that the array was transposed.

However, transposing the inner chunk size is caught as invalid in some cases (erroneously, in my view).
Not transposing the inner chunk size seems to pass validation, but causes a different error.

If the size of shards/chunks along different dimensions do not share a common factor, there is no way (currently) to save a transposed and sharded array.

The code below reproduces the error.

Steps to reproduce

# data is an array with shape [15,32]
data = np.arange(15*32, dtype=np.single).reshape(15, 32)

# first codec (array -> array) transpose the array of size [15,32] to size [32,15]
transpose = zarr.codecs.TransposeCodec(order=[1, 0])

# second codec (array -> bytes) is sharding 
# I expect it to operate on the [32,15] array output by the TransposeCodec
# as a result, its chunk shape should evenly divide [32,15], for example [16,5] should work
sharding = zarr.codecs.ShardingCodec(chunk_shape=(16, 5), codecs=[zarr.codecs.BytesCodec()], index_codecs=[zarr.codecs.BytesCodec()])
codecs=[transpose, sharding]

store = zarr.store.LocalStore('<path>/test.zarr', mode='w')
z = zarr.array(data, path='transposed_sharded', chunk_shape=(15, 32), codecs=codecs, store=store)
# ValueError: The array's `chunk_shape` needs to be divisible by the shard's inner `chunk_shape`.

If instead, we don't transpose the sharding codec's chunk_shape, it seems to pass validation
but crashes later with a ZeroDivisionError

data = np.arange(15*32, dtype=np.single).reshape(15, 32)
transpose = zarr.codecs.TransposeCodec(order=[1,0])
sharding = zarr.codecs.ShardingCodec(chunk_shape=(5, 16), codecs=[zarr.codecs.BytesCodec()], index_codecs=[zarr.codecs.BytesCodec()])
codecs=[transpose, sharding]

store = zarr.store.LocalStore('<path>/test.zarr', mode='w')
z = zarr.array(data, path='transposed_sharded_uhoh', chunk_shape=(15, 32), codecs=codecs, store=store)
# ZeroDivisionError: integer modulo by zero

Additional output

No response

watching

I am guessing this happens because we validate the sharding codec against the static shape of the zarr array, not against the shape of the array that codec will receive in the data encoding process, which might depend on any number of array -> array codecs. So to properly validate the sharding codecs, we need to be able to statically resolve the shape of its input, based on all previous codecs. Right now I think we only have the transpose codec to worry about, so this shouldn't be too hard to fix.

edit: we would need to change this routine:

zarr-python/src/zarr/core/metadata.py

Lines 226 to 228 in 726fdfb

    
           for codec in self.codecs: 
        
               codec.validate(shape=self.shape, dtype=self.data_type, chunk_grid=self.chunk_grid)

The replacement should probably be a stand-alone function that takes an input shape, dtype, and a list of codecs and internally tracks the shape changes through the chain of codecs.

There is some related discussion to this here: https://ossci.zulipchat.com/#narrow/stream/423692-Zarr/topic/Transpose.20codec.20interpretation.

🤖 AI text below 🤖

On zarr 3.4.0 and on main at 8dc50b8, the check no longer ends in an error. An array it accepts writes and reads back different data, and an array it refuses is valid:

import numpy as np
import zarr
from zarr.codecs import BytesCodec, ShardingCodec, TransposeCodec


def round_trip(shape, chunks):
    store = zarr.storage.MemoryStore()
    array = zarr.create(
        store=store,
        shape=shape,
        chunks=chunks,
        dtype="float32",
        zarr_format=3,
        codecs=[
            TransposeCodec(order=(1, 0)),
            ShardingCodec(chunk_shape=(4, 3), codecs=[BytesCodec()]),
        ],
    )
    data = np.arange(np.prod(shape), dtype="float32").reshape(shape)
    array[...] = data
    return np.array_equal(zarr.open_array(store, mode="r")[...], data)


print(round_trip((8, 12), (4, 6)))  # False
print(round_trip((12, 8), (6, 4)))
# ValueError: Chunk edge length 6 in dimension 0 is not divisible by the shard's inner chunk size 4.

The first array's chunks reach the sharding codec transposed, as (6, 4), which (4, 3) does not divide. The check compares (4, 3) with the array's (4, 6), so it passes.
The second array's chunks reach it as (4, 6), which (4, 3) divides. The check compares with (6, 4), so it refuses.

The check is ShardingCodec.validate. It is called with the array's chunk grid, at v3.py L540 and L571.

A sharding codec inside another's codecs is never validated at all. That has the same root cause and is filed as #4437.

Navigation Menu

Repository navigation

[v3] Inner chunk size validation behavior for `ShardingCodec` when downstream of `TransposeCodec` #2050

Zarr version

Numcodecs version

Python Version

Operating System

Installation

Description

Steps to reproduce

Additional output

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	for codec in self.codecs:
	codec.validate(shape=self.shape, dtype=self.data_type, chunk_grid=self.chunk_grid)

[v3] Inner chunk size validation behavior for ShardingCodec when downstream of TransposeCodec #2050

Description

Zarr version

Numcodecs version

Python Version

Operating System

Installation

Description

Steps to reproduce

Additional output

Activity

dchen116 commented on Aug 13, 2024

d-v-b commented on Sep 12, 2024

LDeakin commented on Sep 12, 2024

d-v-b commented on Sep 27, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

[v3] Inner chunk size validation behavior for `ShardingCodec` when downstream of `TransposeCodec` #2050