Skip to content

Commit 01f4445

Browse files
d-v-bclaude
andcommitted
fix: update tests for simplified pipeline, add n_workers parameter
- Remove test_read_plan.py and test_write_plan.py (tested removed layout abstraction) - Fix test_evolve_from_array_spec to check _sync_transform instead of layout - Replace test_simple_layout_decode_skips_indexer with test_sync_transform_encode_decode_roundtrip - Add n_workers parameter to read_sync/write_sync for thread-pool parallelism Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 5fa3fdf commit 01f4445

4 files changed

Lines changed: 13 additions & 630 deletions

File tree

src/zarr/core/codec_pipeline.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -938,6 +938,7 @@ def read_sync(
938938
batch_info: Iterable[tuple[ByteGetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
939939
out: NDBuffer,
940940
drop_axes: tuple[int, ...] = (),
941+
n_workers: int = 0,
941942
) -> tuple[GetResult, ...]:
942943
"""Synchronous read: fetch -> decode -> scatter, per chunk."""
943944
assert self._sync_transform is not None
@@ -978,6 +979,7 @@ def write_sync(
978979
batch_info: Iterable[tuple[ByteSetter, ArraySpec, SelectorTuple, SelectorTuple, bool]],
979980
value: NDBuffer,
980981
drop_axes: tuple[int, ...] = (),
982+
n_workers: int = 0,
981983
) -> None:
982984
"""Synchronous write: merge -> encode -> store, per chunk."""
983985
assert self._sync_transform is not None

tests/test_phased_codec_pipeline.py

Lines changed: 11 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -62,13 +62,13 @@ def test_construction(codecs: tuple[Any, ...]) -> None:
6262

6363

6464
def test_evolve_from_array_spec() -> None:
65-
"""evolve_from_array_spec creates a ChunkLayout."""
65+
"""evolve_from_array_spec creates a sync transform."""
6666
from zarr.core.array_spec import ArrayConfig, ArraySpec
6767
from zarr.core.buffer import default_buffer_prototype
6868
from zarr.core.dtype import get_data_type_from_native_dtype
6969

7070
pipeline = PhasedCodecPipeline.from_codecs((BytesCodec(),))
71-
assert pipeline.layout is None
71+
assert pipeline._sync_transform is None
7272

7373
zdtype = get_data_type_from_native_dtype(np.dtype("float64"))
7474
spec = ArraySpec(
@@ -79,7 +79,7 @@ def test_evolve_from_array_spec() -> None:
7979
prototype=default_buffer_prototype(),
8080
)
8181
evolved = pipeline.evolve_from_array_spec(spec)
82-
assert evolved.layout is not None
82+
assert evolved._sync_transform is not None
8383

8484

8585
@pytest.mark.parametrize(
@@ -294,8 +294,8 @@ async def test_sync_write_async_read_roundtrip() -> None:
294294
)
295295

296296

297-
def test_simple_layout_decode_skips_indexer() -> None:
298-
"""Non-sharded decode should not create BasicIndexer or ChunkGrid."""
297+
def test_sync_transform_encode_decode_roundtrip() -> None:
298+
"""Sync transform can encode and decode a chunk."""
299299
from zarr.core.array_spec import ArrayConfig, ArraySpec
300300
from zarr.core.buffer import default_buffer_prototype
301301
from zarr.core.dtype import Float64
@@ -311,22 +311,17 @@ def test_simple_layout_decode_skips_indexer() -> None:
311311
config=ArrayConfig(order="C", write_empty_chunks=True),
312312
)
313313
pipeline = pipeline.evolve_from_array_spec(spec)
314+
assert pipeline._sync_transform is not None
314315

315-
# Encode some data
316+
# Encode
316317
proto = default_buffer_prototype()
317318
data = proto.nd_buffer.from_numpy_array(np.arange(100, dtype="float64"))
318-
assert pipeline.layout is not None
319-
encoded = pipeline.layout.inner_transform.encode_chunk(data)
319+
encoded = pipeline._sync_transform.encode_chunk(data)
320320
assert encoded is not None
321321

322-
# Decode via decode_chunks_from_index using leaf_transform
323-
from zarr.core.codec_pipeline import ShardIndex, decode_chunks_from_index
324-
325-
index = pipeline.layout.resolve_index(None, "c/0")
326-
raw_chunks = {(0,): encoded}
327-
result = decode_chunks_from_index(raw_chunks, index, spec)
328-
assert result is not None
329-
np.testing.assert_array_equal(result.as_numpy_array(), np.arange(100, dtype="float64"))
322+
# Decode
323+
decoded = pipeline._sync_transform.decode_chunk(encoded)
324+
np.testing.assert_array_equal(decoded.as_numpy_array(), np.arange(100, dtype="float64"))
330325

331326

332327
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)