Skip to content

Commit

Permalink
Update V2 codec pipeline to use concrete classes (#2244)
Browse files Browse the repository at this point in the history
The previous implementation used the codec config, rather than the codec
itself.
  • Loading branch information
TomAugspurger authored Sep 25, 2024
1 parent 4cbb17e commit fafd0bf
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 6 deletions.
12 changes: 6 additions & 6 deletions src/zarr/codecs/_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,16 +8,18 @@

from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec
from zarr.core.buffer import Buffer, NDBuffer, default_buffer_prototype
from zarr.core.common import JSON, to_thread
from zarr.core.common import to_thread
from zarr.registry import get_ndbuffer_class

if TYPE_CHECKING:
import numcodecs.abc

from zarr.core.array_spec import ArraySpec


@dataclass(frozen=True)
class V2Compressor(ArrayBytesCodec):
compressor: dict[str, JSON] | None
compressor: numcodecs.abc.Codec | None

is_fixed_size = False

Expand All @@ -27,9 +29,8 @@ async def _decode_single(
chunk_spec: ArraySpec,
) -> NDBuffer:
if self.compressor is not None:
compressor = numcodecs.get_codec(self.compressor)
chunk_numpy_array = ensure_ndarray(
await to_thread(compressor.decode, chunk_bytes.as_array_like())
await to_thread(self.compressor.decode, chunk_bytes.as_array_like())
)
else:
chunk_numpy_array = ensure_ndarray(chunk_bytes.as_array_like())
Expand All @@ -47,14 +48,13 @@ async def _encode_single(
) -> Buffer | None:
chunk_numpy_array = chunk_array.as_numpy_array()
if self.compressor is not None:
compressor = numcodecs.get_codec(self.compressor)
if (
not chunk_numpy_array.flags.c_contiguous
and not chunk_numpy_array.flags.f_contiguous
):
chunk_numpy_array = chunk_numpy_array.copy(order="A")
encoded_chunk_bytes = ensure_bytes(
await to_thread(compressor.encode, chunk_numpy_array)
await to_thread(self.compressor.encode, chunk_numpy_array)
)
else:
encoded_chunk_bytes = ensure_bytes(chunk_numpy_array)
Expand Down
2 changes: 2 additions & 0 deletions src/zarr/core/metadata/v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,8 @@ def _json_convert(
return o.str
else:
return o.descr
if isinstance(o, numcodecs.abc.Codec):
return o.get_config()
if np.isscalar(o):
out: Any
if hasattr(o, "dtype") and o.dtype.kind == "M" and hasattr(o, "view"):
Expand Down
20 changes: 20 additions & 0 deletions tests/v3/test_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@

import numpy as np
import pytest
from numcodecs import Delta
from numcodecs.blosc import Blosc

import zarr
from zarr import Array
from zarr.store import MemoryStore, StorePath

Expand All @@ -26,3 +29,20 @@ def test_simple(store: StorePath) -> None:

a[:, :] = data
assert np.array_equal(data, a[:, :])


def test_codec_pipeline() -> None:
# https://github.com/zarr-developers/zarr-python/issues/2243
store = MemoryStore(mode="w")
array = zarr.create(
store=store,
shape=(1,),
dtype="i4",
zarr_format=2,
filters=[Delta(dtype="i4").get_config()],
compressor=Blosc().get_config(),
)
array[:] = 1
result = array[:]
expected = np.ones(1)
np.testing.assert_array_equal(result, expected)

0 comments on commit fafd0bf

Please sign in to comment.