Skip to content

Commit 33ee8a1

Browse files
d-v-bclaude
andcommitted
perf: vectorize is_regular_1d for ndarray inputs
Per-element Python iteration over a 100K-element int64 array dominated create_array runtime on the (10**8,) chunks=(1000,) regression case (~6 ms in is_regular_nd, downstream of ChunksTuple normalization). Dispatch on np.ndarray and use a single vectorized comparison instead. End-to-end create_array on the regression case: ~7.9 ms -> ~0.6 ms. Promote numpy to a runtime import (was TYPE_CHECKING-only) for the isinstance dispatch. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 12a325a commit 33ee8a1

1 file changed

Lines changed: 4 additions & 2 deletions

File tree

  • src/zarr/core/metadata

src/zarr/core/metadata/v3.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from dataclasses import dataclass, field, replace
66
from typing import TYPE_CHECKING, Any, Final, Literal, NotRequired, TypeGuard, cast
77

8+
import numpy as np
89
from typing_extensions import TypedDict
910

1011
from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
@@ -39,8 +40,6 @@
3940
if TYPE_CHECKING:
4041
from typing import Self
4142

42-
import numpy as np
43-
4443
from zarr.core.buffer import Buffer, BufferPrototype
4544
from zarr.core.chunk_grids import ChunksTuple
4645
from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
@@ -385,6 +384,9 @@ def is_regular_1d(
385384
if len(dim_chunks) <= 1:
386385
return True
387386
first = dim_chunks[0]
387+
if isinstance(dim_chunks, np.ndarray):
388+
# Vectorized comparison avoids per-element Python iteration over int64 arrays.
389+
return bool((dim_chunks[1:-1] == first).all() and dim_chunks[-1] <= first)
388390
for c in dim_chunks[1:-1]:
389391
if c != first:
390392
return False

0 commit comments

Comments
 (0)