chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,685 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import math
|
||||
import numbers
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
||||
from ray.data._internal.util import (
|
||||
_check_import,
|
||||
_is_local_scheme,
|
||||
iterate_with_retry,
|
||||
)
|
||||
from ray.data.block import Block, BlockMetadata
|
||||
from ray.data.datasource.datasource import Datasource, ReadTask
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fsspec.spec import AbstractFileSystem
|
||||
from pyarrow import fs as pyarrow_fs
|
||||
from zarr import Array as ZarrArray
|
||||
from zarr.hierarchy import Group as ZarrGroup
|
||||
|
||||
from ray.data.context import DataContext
|
||||
|
||||
ZarrRoot = ZarrGroup | ZarrArray
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ZarrArrayMeta:
|
||||
"""``shape``/``chunks``/``dtype`` for a single Zarr v2 array."""
|
||||
|
||||
shape: tuple[int, ...]
|
||||
chunks: tuple[int, ...]
|
||||
dtype: str
|
||||
|
||||
@classmethod
|
||||
def from_zarr_array(cls, arr: "ZarrArray") -> ZarrArrayMeta:
|
||||
return cls(
|
||||
shape=tuple(int(s) for s in arr.shape),
|
||||
chunks=tuple(int(c) for c in arr.chunks),
|
||||
dtype=str(arr.dtype),
|
||||
)
|
||||
|
||||
@property
|
||||
def rank(self) -> int:
|
||||
return len(self.shape)
|
||||
|
||||
@property
|
||||
def itemsize(self) -> int:
|
||||
return np.dtype(self.dtype).itemsize
|
||||
|
||||
def effective_chunks(
|
||||
self,
|
||||
array_name: str,
|
||||
user_chunk_shape: tuple[int, ...] | dict[str, tuple[int, ...]] | None,
|
||||
) -> tuple[int, ...]:
|
||||
"""Resolve the user's ``chunk_shapes`` override(s) against this array's chunks.
|
||||
|
||||
A single sequence overrides the leading axes (trailing axes keep the
|
||||
native chunks), so one ``chunk_shapes=[16]`` applies across arrays of
|
||||
different ranks. A dict maps array path → that array's override prefix;
|
||||
arrays absent from it keep native chunks. ``None`` keeps native chunks;
|
||||
an override longer than the array's rank raises ``ValueError``.
|
||||
"""
|
||||
if user_chunk_shape is None:
|
||||
return self.chunks
|
||||
|
||||
if isinstance(user_chunk_shape, dict):
|
||||
user_chunk_shape = user_chunk_shape.get(array_name)
|
||||
if user_chunk_shape is None:
|
||||
return self.chunks
|
||||
if len(user_chunk_shape) > self.rank:
|
||||
raise ValueError(
|
||||
f"chunk_shapes override for array {array_name!r} has "
|
||||
f"{len(user_chunk_shape)} axes but array of shape "
|
||||
f"{self.shape!r} has rank {self.rank}. Each chunk_shapes "
|
||||
f"override may not be longer than its target array's rank."
|
||||
)
|
||||
return user_chunk_shape + self.chunks[len(user_chunk_shape) :]
|
||||
|
||||
def grid_shape(self, chunks: tuple[int, ...]) -> tuple[int, ...]:
|
||||
"""Number of chunks along each axis under the given chunk shape."""
|
||||
return tuple(math.ceil(s / c) for s, c in zip(self.shape, chunks))
|
||||
|
||||
def chunk_slices(
|
||||
self, chunk_index: tuple[int, ...], chunks: tuple[int, ...]
|
||||
) -> tuple[tuple[int, int], ...]:
|
||||
"""Per-axis ``(start, stop)`` for ``array[chunk_index]`` under ``chunks``.
|
||||
|
||||
Trailing-edge chunks are clamped to ``shape[i]``, so they may be
|
||||
shorter than ``chunks[i]``. No padding is applied.
|
||||
"""
|
||||
return tuple(
|
||||
(i * c, min((i + 1) * c, s))
|
||||
for i, c, s in zip(chunk_index, chunks, self.shape)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Chunk reading
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _read_chunk(
|
||||
root: ZarrRoot,
|
||||
array_name: str,
|
||||
chunk_slices: tuple[tuple[int, int], ...],
|
||||
retry_match: Optional[List[str]] = None,
|
||||
) -> np.ndarray:
|
||||
"""Read ``array[chunk_slices]`` as an ndarray.
|
||||
|
||||
The underlying filesystem's own retry policy still applies underneath.
|
||||
"""
|
||||
|
||||
def _read() -> np.ndarray:
|
||||
indexer = tuple(slice(s, e) for s, e in chunk_slices)
|
||||
arr = root if array_name == "" else root[array_name]
|
||||
return np.asarray(arr[indexer])
|
||||
|
||||
if not retry_match:
|
||||
return _read()
|
||||
# TODO(Artur): This would be more elegant with a general retry helper for non-iterables.
|
||||
return next(
|
||||
iterate_with_retry(
|
||||
lambda: [_read()], description="read a Zarr chunk", match=retry_match
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _ChunkRange:
|
||||
"""A contiguous slice ``[flat_start, flat_stop)`` of an array's chunk grid.
|
||||
|
||||
The flat indices address the row-major flattening of the chunk grid; the
|
||||
read fn unravels each to an N-D ``chunk_index`` lazily on the worker. Keeping
|
||||
a range (not a materialized per-chunk list) makes read-task planning
|
||||
O(parallelism) rather than O(total chunks) -- important for stores with very
|
||||
many chunks.
|
||||
"""
|
||||
|
||||
array_name: str
|
||||
meta: ZarrArrayMeta
|
||||
chunks: tuple[int, ...]
|
||||
grid: tuple[int, ...]
|
||||
flat_start: int
|
||||
flat_stop: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _AlignedChunkDescriptor:
|
||||
"""One wide row: a global axis-0 range ``[t_start, t_stop)`` across the
|
||||
aligned arrays. With ``overlap > 0`` the row's data extends to
|
||||
``t_stop_data = min(t_stop + overlap, shape[0])`` (lookahead so windows
|
||||
starting in this row reach their tail without crossing a row boundary).
|
||||
"""
|
||||
|
||||
chunk_index: int
|
||||
t_start: int
|
||||
t_stop: int
|
||||
t_stop_data: int
|
||||
|
||||
|
||||
def _create_read_fn(
|
||||
chunk_range: _ChunkRange,
|
||||
root: ZarrRoot,
|
||||
per_task_row_limit: Optional[int],
|
||||
retry_match: Optional[List[str]],
|
||||
) -> Callable[[], Iterable[Block]]:
|
||||
"""Build a callable that materializes one block for a chunk-grid range.
|
||||
|
||||
This is the case where arrays are not aligned. Chunks are enumerated lazily
|
||||
(on the worker) from ``chunk_range``. ``per_task_row_limit`` caps how many
|
||||
chunks this task reads so a downstream ``limit`` reads only what it needs
|
||||
(``None`` reads the whole range).
|
||||
"""
|
||||
cr = chunk_range
|
||||
stop = cr.flat_stop
|
||||
if per_task_row_limit is not None:
|
||||
stop = min(stop, cr.flat_start + per_task_row_limit)
|
||||
|
||||
def read_fn() -> Iterable[Block]:
|
||||
builder = DelegatingBlockBuilder()
|
||||
for flat_index in range(cr.flat_start, stop):
|
||||
chunk_index = tuple(int(i) for i in np.unravel_index(flat_index, cr.grid))
|
||||
chunk_slices = cr.meta.chunk_slices(chunk_index, cr.chunks)
|
||||
builder.add(
|
||||
{
|
||||
"array": cr.array_name,
|
||||
"chunk_index": chunk_index,
|
||||
"chunk_slices": chunk_slices,
|
||||
"chunk": _read_chunk(
|
||||
root, cr.array_name, chunk_slices, retry_match
|
||||
),
|
||||
}
|
||||
)
|
||||
yield builder.build()
|
||||
|
||||
return read_fn
|
||||
|
||||
|
||||
def _create_aligned_read_fn(
|
||||
batch: list[_AlignedChunkDescriptor],
|
||||
aligned_array_names: list[str],
|
||||
root: ZarrRoot,
|
||||
per_task_row_limit: Optional[int],
|
||||
retry_match: Optional[List[str]],
|
||||
) -> Callable[[], Iterable[Block]]:
|
||||
"""Build a callable for aligned (wide-row) reads.
|
||||
|
||||
Each output row carries ``t_start``, ``t_stop``, and one column per
|
||||
aligned array holding that array's ``[t_start:t_stop, ...]`` slice at
|
||||
its natural shape (edge rows may be shorter). All arrays in one row
|
||||
share the same axis-0 range.
|
||||
|
||||
This is the case where arrays are aligned on axis 0. ``per_task_row_limit``
|
||||
caps how many rows this task reads (``None`` reads the whole batch).
|
||||
"""
|
||||
batch = batch[:per_task_row_limit]
|
||||
|
||||
def read_fn() -> Iterable[Block]:
|
||||
builder = DelegatingBlockBuilder()
|
||||
for d in batch:
|
||||
row: dict[str, Any] = {"t_start": d.t_start, "t_stop": d.t_stop}
|
||||
for name in aligned_array_names:
|
||||
row[name] = _read_chunk(
|
||||
root, name, ((d.t_start, d.t_stop_data),), retry_match
|
||||
)
|
||||
builder.add(row)
|
||||
yield builder.build()
|
||||
|
||||
return read_fn
|
||||
|
||||
|
||||
def _is_positive_int(x) -> bool:
|
||||
"""True for a positive integer, including NumPy integers; False for bool."""
|
||||
return not isinstance(x, bool) and isinstance(x, numbers.Integral) and int(x) > 0
|
||||
|
||||
|
||||
def _validate_chunk_shapes_dict(chunk_shapes: dict) -> dict[str, tuple[int, ...]]:
|
||||
"""Normalize chunk_shapes keys to store paths and validate their values."""
|
||||
from zarr.util import normalize_storage_path
|
||||
|
||||
normalized: dict[str, tuple[int, ...]] = {}
|
||||
for k, v in chunk_shapes.items():
|
||||
if (
|
||||
not isinstance(v, (tuple, list))
|
||||
or not v
|
||||
or not all(_is_positive_int(x) for x in v)
|
||||
):
|
||||
raise ValueError(
|
||||
f"chunk_shapes[{k!r}] must be a non-empty sequence of positive "
|
||||
f"integers (list or tuple), got {v!r}"
|
||||
)
|
||||
normalized[normalize_storage_path(k)] = tuple(int(x) for x in v)
|
||||
return normalized
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Datasource
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class ZarrV2Datasource(Datasource):
|
||||
"""Reads one or more Zarr v2 arrays into a Ray Data ``Dataset``.
|
||||
|
||||
Emits long-form rows (one per chunk per array) or, with
|
||||
``align_axis_0=True``, wide rows (one per axis-0 chunk, one column per
|
||||
array). See :func:`ray.data.read_zarr` for the row schemas and full API.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path: str,
|
||||
filesystem: pyarrow_fs.FileSystem | AbstractFileSystem | None = None,
|
||||
chunk_shapes: dict[str, list] | list | None = None,
|
||||
array_paths: list[str] | None = None,
|
||||
allow_full_metadata_scan: bool = False,
|
||||
align_axis_0: bool = False,
|
||||
overlap: int = 0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
_check_import(self, module="zarr", package="zarr")
|
||||
import zarr
|
||||
|
||||
_check_import(self, module="fsspec", package="fsspec")
|
||||
from fsspec.spec import AbstractFileSystem
|
||||
|
||||
if int(zarr.__version__.split(".")[0]) >= 3:
|
||||
raise ImportError(
|
||||
f"read_zarr supports zarr-python 2.x (Zarr v2 stores), but found "
|
||||
f"zarr=={zarr.__version__}. Install a compatible version with "
|
||||
f"`pip install 'zarr<3'`."
|
||||
)
|
||||
|
||||
self.allow_full_metadata_scan = allow_full_metadata_scan
|
||||
self.paths = [str(path)]
|
||||
# ``local://`` stores live on the driver's local disk, so pin reads to
|
||||
# the driver node (workers on other nodes can't see those files).
|
||||
self._supports_distributed_reads = not _is_local_scheme(self.paths)
|
||||
|
||||
# Resolve filesystem + store path. The order of precedence:
|
||||
# 1. Explicit ``filesystem=`` always wins.
|
||||
# 2. ``.zip`` URL/path: auto-wrap with fsspec's ZipFileSystem.
|
||||
# 3. Otherwise delegate to Ray Data's standard URL to filesystem
|
||||
# helper (the same one every other ``read_*`` API uses).
|
||||
# "store path" is the path to the Zarr store, relative to the filesystem root.
|
||||
# It is used to construct the Zarr root object.
|
||||
if filesystem is None and self.paths[0].endswith(".zip"):
|
||||
import fsspec
|
||||
|
||||
self._fs = fsspec.filesystem("zip", fo=self.paths[0])
|
||||
self._store_path = ""
|
||||
elif filesystem is None:
|
||||
from fsspec.implementations.arrow import ArrowFSWrapper
|
||||
|
||||
from ray.data.datasource.path_util import (
|
||||
_resolve_paths_and_filesystem,
|
||||
)
|
||||
|
||||
resolved_paths, pa_fs = _resolve_paths_and_filesystem([self.paths[0]])
|
||||
self._fs = ArrowFSWrapper(pa_fs)
|
||||
self._store_path = resolved_paths[0].rstrip("/")
|
||||
else:
|
||||
from pyarrow.fs import FileSystem
|
||||
|
||||
if isinstance(filesystem, AbstractFileSystem):
|
||||
self._fs = filesystem
|
||||
elif isinstance(filesystem, FileSystem):
|
||||
from fsspec.implementations.arrow import ArrowFSWrapper
|
||||
|
||||
self._fs = ArrowFSWrapper(filesystem)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"filesystem must be pyarrow.fs.FileSystem or "
|
||||
f"fsspec.spec.AbstractFileSystem, got "
|
||||
f"{type(filesystem).__name__}"
|
||||
)
|
||||
from fsspec.implementations.zip import ZipFileSystem
|
||||
|
||||
if isinstance(self._fs, ZipFileSystem) and self.paths[0].endswith(".zip"):
|
||||
# An explicit archive filesystem: the store is the archive root,
|
||||
# not a ``.zip``-named entry inside it.
|
||||
self._store_path = ""
|
||||
else:
|
||||
from fsspec.core import split_protocol
|
||||
|
||||
_, store_path = split_protocol(self.paths[0])
|
||||
self._store_path = store_path.rstrip("/")
|
||||
|
||||
if chunk_shapes is not None and not isinstance(
|
||||
chunk_shapes, (tuple, list, dict)
|
||||
):
|
||||
raise ValueError(
|
||||
f"chunk_shapes must be a non-empty sequence of positive "
|
||||
f"integers (list or tuple), or a dict, got {chunk_shapes!r}"
|
||||
)
|
||||
|
||||
self.chunk_shapes: tuple[int, ...] | dict[str, tuple[int, ...]] | None = None
|
||||
if chunk_shapes is not None:
|
||||
if isinstance(chunk_shapes, dict):
|
||||
self.chunk_shapes = _validate_chunk_shapes_dict(chunk_shapes)
|
||||
else:
|
||||
if not chunk_shapes or not all(
|
||||
_is_positive_int(x) for x in chunk_shapes
|
||||
):
|
||||
raise ValueError(
|
||||
"chunk_shapes must be a non-empty sequence of positive integers "
|
||||
f"(list or tuple), got {chunk_shapes!r}"
|
||||
)
|
||||
|
||||
self.chunk_shapes = tuple(int(x) for x in chunk_shapes)
|
||||
|
||||
# Open the store with zarr (consolidated metadata when available).
|
||||
# Detect consolidation by *trying* ``open_consolidated``.
|
||||
store = self._fs.get_mapper(self._store_path)
|
||||
try:
|
||||
self.root = zarr.open_consolidated(store, mode="r")
|
||||
self._consolidated = True
|
||||
except KeyError:
|
||||
self.root = zarr.open(store, mode="r")
|
||||
self._consolidated = False
|
||||
|
||||
self._metadata_by_path = self._load_metadata(array_paths)
|
||||
if not self._metadata_by_path:
|
||||
raise ValueError(
|
||||
f"No arrays discovered in Zarr store at {self.paths[0]!r}."
|
||||
)
|
||||
|
||||
# Reject per-array overrides that do not correspond to any selected
|
||||
# array in this read.
|
||||
if isinstance(self.chunk_shapes, dict):
|
||||
unknown_chunk_shape_keys = sorted(
|
||||
set(self.chunk_shapes) - set(self._metadata_by_path)
|
||||
)
|
||||
if unknown_chunk_shape_keys:
|
||||
raise ValueError(
|
||||
f"Unknown array path(s) in chunk_shapes: {unknown_chunk_shape_keys}"
|
||||
)
|
||||
|
||||
if not align_axis_0:
|
||||
self._aligned_array_names = None
|
||||
else:
|
||||
scalar_arrays = sorted(
|
||||
name for name, meta in self._metadata_by_path.items() if not meta.shape
|
||||
)
|
||||
if scalar_arrays:
|
||||
raise ValueError(
|
||||
f"align_axis_0=True requires every selected array to have "
|
||||
f"at least one axis, but these are 0-D (scalar): "
|
||||
f"{scalar_arrays}. Drop them with array_paths=[...]."
|
||||
)
|
||||
shape0_by_array = {
|
||||
name: meta.shape[0] for name, meta in self._metadata_by_path.items()
|
||||
}
|
||||
if len(set(shape0_by_array.values())) > 1:
|
||||
raise ValueError(
|
||||
f"All selected arrays must share shape[0] when "
|
||||
f"align_axis_0=True. Got: {shape0_by_array}. Pass a "
|
||||
f"shape-compatible subset via array_paths=[...]."
|
||||
)
|
||||
self._aligned_array_names = list(self._metadata_by_path.keys())
|
||||
|
||||
# Validate overlap. Only meaningful when arrays are co-iterated as
|
||||
# wide rows, since the trailing lookahead is exposed via the
|
||||
# per-array column being longer than ``t_stop - t_start``.
|
||||
if not isinstance(overlap, int) or overlap < 0:
|
||||
raise ValueError(f"overlap must be a non-negative integer, got {overlap!r}")
|
||||
if overlap and self._aligned_array_names is None:
|
||||
raise ValueError(
|
||||
"overlap requires align_axis_0=True. In the default long-form "
|
||||
"(chunk-per-row) mode, there's no wide row to extend forward — "
|
||||
"the ``chunk_slices`` column on each chunk row already exposes "
|
||||
"the global axis-0 range."
|
||||
)
|
||||
self.overlap = overlap
|
||||
|
||||
# Resolve per-array chunk geometry. ``effective_chunks`` raises a
|
||||
# ``ValueError`` if a shared ``chunk_shapes`` prefix or any per-array
|
||||
# ``chunk_shapes`` override is longer than the target array's rank —
|
||||
# so this loop is also where rank validation happens.
|
||||
self._array_chunks: dict[str, tuple[int, ...]] = {}
|
||||
self._array_grids: dict[str, tuple[int, ...]] = {}
|
||||
for name, meta in self._metadata_by_path.items():
|
||||
chunks = meta.effective_chunks(name, self.chunk_shapes)
|
||||
self._array_chunks[name] = chunks
|
||||
self._array_grids[name] = meta.grid_shape(chunks)
|
||||
|
||||
# If aligned, all listed arrays must share the same axis-0 chunk size
|
||||
# so each wide row corresponds to one axis-0 step across every array.
|
||||
if self._aligned_array_names is not None:
|
||||
axis_0_chunks = {
|
||||
name: self._array_chunks[name][0] for name in self._aligned_array_names
|
||||
}
|
||||
unique = set(axis_0_chunks.values())
|
||||
if len(unique) > 1:
|
||||
raise ValueError(
|
||||
f"Aligned arrays must share the same axis-0 chunk size. "
|
||||
f"Got: {axis_0_chunks}. Pass chunk_shapes=[N] (or a "
|
||||
f"per-array chunk_shapes dict that resolves all aligned "
|
||||
f"arrays to the same axis-0 prefix) to re-tile them."
|
||||
)
|
||||
|
||||
@property
|
||||
def supports_distributed_reads(self) -> bool:
|
||||
return self._supports_distributed_reads
|
||||
|
||||
def estimate_inmemory_data_size(self) -> Optional[int]:
|
||||
"""Total bytes = sum over selected arrays of ``prod(shape) * itemsize``."""
|
||||
return sum(
|
||||
math.prod(meta.shape) * meta.itemsize
|
||||
for meta in self._metadata_by_path.values()
|
||||
)
|
||||
|
||||
def get_read_tasks(
|
||||
self,
|
||||
parallelism: int,
|
||||
per_task_row_limit: Optional[int] = None,
|
||||
data_context: Optional["DataContext"] = None,
|
||||
) -> List[ReadTask]:
|
||||
"""Enumerate every chunk and wrap it (or batches of chunks) in ReadTasks."""
|
||||
from ray.data.context import DataContext
|
||||
|
||||
retry_match = (data_context or DataContext.get_current()).retried_io_errors
|
||||
if self._aligned_array_names is not None:
|
||||
return self._get_aligned_read_tasks(
|
||||
parallelism, per_task_row_limit, retry_match
|
||||
)
|
||||
return self._get_long_form_read_tasks(
|
||||
parallelism, per_task_row_limit, retry_match
|
||||
)
|
||||
|
||||
def _get_long_form_read_tasks(
|
||||
self,
|
||||
parallelism: int,
|
||||
per_task_row_limit: Optional[int],
|
||||
retry_match: Optional[List[str]],
|
||||
) -> List[ReadTask]:
|
||||
read_tasks: List[ReadTask] = []
|
||||
for name, meta in self._metadata_by_path.items():
|
||||
chunks = self._array_chunks[name]
|
||||
grid = self._array_grids[name]
|
||||
n_chunks = math.prod(grid)
|
||||
if n_chunks == 0:
|
||||
continue
|
||||
# Split the chunk grid into contiguous flat-index ranges. This is
|
||||
# O(n_tasks), not O(n_chunks): we never materialize a per-chunk list
|
||||
# on the driver -- the read fn unravels chunks lazily on the worker.
|
||||
n_tasks = max(1, min(parallelism, n_chunks))
|
||||
batch_size = math.ceil(n_chunks / n_tasks)
|
||||
for flat_start in range(0, n_chunks, batch_size):
|
||||
flat_stop = min(flat_start + batch_size, n_chunks)
|
||||
chunk_range = _ChunkRange(
|
||||
name, meta, chunks, grid, flat_start, flat_stop
|
||||
)
|
||||
read_tasks.append(
|
||||
ReadTask(
|
||||
_create_read_fn(
|
||||
chunk_range, self.root, per_task_row_limit, retry_match
|
||||
),
|
||||
BlockMetadata(
|
||||
num_rows=flat_stop - flat_start,
|
||||
size_bytes=self._estimate_range_mem_size(chunk_range),
|
||||
input_files=(self.paths[0],),
|
||||
exec_stats=None,
|
||||
),
|
||||
per_task_row_limit=per_task_row_limit,
|
||||
)
|
||||
)
|
||||
return read_tasks
|
||||
|
||||
def _estimate_range_mem_size(self, chunk_range: _ChunkRange) -> int:
|
||||
"""Upper-bound in-memory bytes for a chunk-grid range.
|
||||
|
||||
Assumes a full-size chunk per index; trailing-edge chunks are smaller,
|
||||
so this slightly over-estimates. O(1) -- it does not enumerate the range.
|
||||
"""
|
||||
n = chunk_range.flat_stop - chunk_range.flat_start
|
||||
return n * math.prod(chunk_range.chunks) * chunk_range.meta.itemsize
|
||||
|
||||
def _get_aligned_read_tasks(
|
||||
self,
|
||||
parallelism: int,
|
||||
per_task_row_limit: Optional[int],
|
||||
retry_match: Optional[List[str]],
|
||||
) -> List[ReadTask]:
|
||||
"""Aligned read tasks. See :meth:`get_read_tasks` for semantics."""
|
||||
assert self._aligned_array_names is not None
|
||||
# All aligned arrays share the same axis-0 chunk size (validated in
|
||||
# ``__init__``) and the same shape[0]. Read the geometry off the first.
|
||||
first_name = self._aligned_array_names[0]
|
||||
axis_0_chunk = self._array_chunks[first_name][0]
|
||||
shape0 = self._metadata_by_path[first_name].shape[0]
|
||||
|
||||
descriptors = [
|
||||
_AlignedChunkDescriptor(
|
||||
chunk_index=i,
|
||||
t_start=i * axis_0_chunk,
|
||||
t_stop=min((i + 1) * axis_0_chunk, shape0),
|
||||
t_stop_data=min((i + 1) * axis_0_chunk + self.overlap, shape0),
|
||||
)
|
||||
for i in range(math.ceil(shape0 / axis_0_chunk))
|
||||
]
|
||||
if not descriptors:
|
||||
return []
|
||||
|
||||
n_tasks = max(1, min(parallelism, len(descriptors)))
|
||||
batch_size = math.ceil(len(descriptors) / n_tasks)
|
||||
|
||||
read_tasks: List[ReadTask] = []
|
||||
for start in range(0, len(descriptors), batch_size):
|
||||
batch = descriptors[start : start + batch_size]
|
||||
read_tasks.append(
|
||||
ReadTask(
|
||||
_create_aligned_read_fn(
|
||||
batch,
|
||||
self._aligned_array_names,
|
||||
self.root,
|
||||
per_task_row_limit,
|
||||
retry_match,
|
||||
),
|
||||
BlockMetadata(
|
||||
num_rows=len(batch),
|
||||
size_bytes=self._estimate_aligned_batch_mem_size(batch),
|
||||
input_files=(self.paths[0],),
|
||||
exec_stats=None,
|
||||
),
|
||||
per_task_row_limit=per_task_row_limit,
|
||||
)
|
||||
)
|
||||
return read_tasks
|
||||
|
||||
def _estimate_aligned_batch_mem_size(
|
||||
self, batch: list[_AlignedChunkDescriptor]
|
||||
) -> int:
|
||||
"""Sum bytes across all (row, aligned-array) pairs in a wide-row batch.
|
||||
|
||||
Accounts for the trailing overlap data each row carries: the row's
|
||||
per-array slice covers ``[t_start, t_stop_data)``, not just
|
||||
``[t_start, t_stop)``.
|
||||
"""
|
||||
assert self._aligned_array_names is not None
|
||||
return sum(
|
||||
(desc.t_stop_data - desc.t_start)
|
||||
* (math.prod(meta.shape[1:]) if len(meta.shape) > 1 else 1)
|
||||
* meta.itemsize
|
||||
for desc in batch
|
||||
for meta in (
|
||||
self._metadata_by_path[name] for name in self._aligned_array_names
|
||||
)
|
||||
)
|
||||
|
||||
def _load_metadata(self, array_paths) -> dict[str, ZarrArrayMeta]:
|
||||
"""Read ``shape``/``chunks``/``dtype`` for the selected arrays off ``self.root``.
|
||||
|
||||
zarr validated the store's metadata when it was opened, so this only
|
||||
adapts the resulting ``zarr.Array`` objects. Discovery uses consolidated
|
||||
metadata when present, then explicit ``array_paths``, then an optional
|
||||
full scan (``allow_full_metadata_scan``). If ``array_paths`` is given,
|
||||
the discovered set is filtered down to it.
|
||||
"""
|
||||
import zarr
|
||||
from zarr.util import normalize_storage_path
|
||||
|
||||
root = self.root
|
||||
requested = (
|
||||
{normalize_storage_path(p) for p in array_paths} if array_paths else None
|
||||
)
|
||||
|
||||
if isinstance(root, zarr.Array):
|
||||
# A store that is itself an array exposes exactly one path: "" (root).
|
||||
# Reject any requested path that isn't the root so a bad ``array_paths``
|
||||
# fails loudly here instead of silently returning the root array.
|
||||
if requested is not None and requested != {""}:
|
||||
raise ValueError(
|
||||
f"This Zarr store is a single root-level array (path ''), "
|
||||
f"but array_paths={array_paths!r} requested other path(s). "
|
||||
f"Pass array_paths=[''] or omit it."
|
||||
)
|
||||
return {"": ZarrArrayMeta.from_zarr_array(root)}
|
||||
|
||||
if not self._consolidated and not self.allow_full_metadata_scan:
|
||||
if requested is None:
|
||||
raise ValueError(
|
||||
"No array_paths were provided and this Zarr store does not "
|
||||
"contain .zmetadata. Pass array_paths=[...] or set "
|
||||
"allow_full_metadata_scan=True."
|
||||
)
|
||||
out: dict[str, ZarrArrayMeta] = {}
|
||||
for raw in array_paths:
|
||||
name = normalize_storage_path(raw)
|
||||
try:
|
||||
arr = root[name]
|
||||
except KeyError as e:
|
||||
raise ValueError(
|
||||
f"Array path {raw!r} not found in Zarr store."
|
||||
) from e
|
||||
if not isinstance(arr, zarr.Array):
|
||||
raise ValueError(f"Array path {raw!r} is a group, not an array.")
|
||||
out[name] = ZarrArrayMeta.from_zarr_array(arr)
|
||||
return out
|
||||
|
||||
all_arrays: dict[str, ZarrArrayMeta] = {}
|
||||
|
||||
def _collect(name: str, obj) -> None:
|
||||
if isinstance(obj, zarr.Array):
|
||||
all_arrays[name] = ZarrArrayMeta.from_zarr_array(obj)
|
||||
|
||||
root.visititems(_collect)
|
||||
|
||||
if requested is not None:
|
||||
missing = sorted(requested - all_arrays.keys())
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Array(s) not found: {', '.join(repr(m) for m in missing)}. "
|
||||
f"Available: {', '.join(repr(a) for a in sorted(all_arrays))}"
|
||||
)
|
||||
all_arrays = {k: v for k, v in all_arrays.items() if k in requested}
|
||||
|
||||
return all_arrays
|
||||
Reference in New Issue
Block a user