chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,937 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import fsspec
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow.fs
|
||||
import pytest
|
||||
import zarr
|
||||
from pytest_lazy_fixtures import lf as lazy_fixture
|
||||
|
||||
import ray
|
||||
from ray.data._internal.datasource import zarrv2_datasource
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.tests.conftest import * # noqa: F401, F403
|
||||
|
||||
|
||||
def _execute_read_tasks(tasks) -> pd.DataFrame:
|
||||
frames = [
|
||||
BlockAccessor.for_block(block).to_pandas() for task in tasks for block in task()
|
||||
]
|
||||
return pd.concat(frames, ignore_index=True)
|
||||
|
||||
|
||||
def _reconstruct_array(df: pd.DataFrame, array_name: str) -> np.ndarray:
|
||||
"""Concatenate all chunks of one array from a long-form result frame."""
|
||||
sub = df[df["array"] == array_name].sort_values(
|
||||
"chunk_index", key=lambda col: col.map(tuple)
|
||||
)
|
||||
return np.concatenate(list(sub["chunk"]), axis=0)
|
||||
|
||||
|
||||
def _write_real_zarr_store(
|
||||
store_path: Path,
|
||||
arrays: dict, # {name: (data, chunks)}
|
||||
) -> Path:
|
||||
"""Write a real Zarr v2 store from numpy arrays and consolidate metadata."""
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
for name, (data, chunks) in arrays.items():
|
||||
root.create_dataset(name, data=data, chunks=chunks, dtype=data.dtype)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zarrv2_group_store(tmp_path) -> Path:
|
||||
"""Two arrays at the store root, both 2-D and 1-D, axis-0-aligned (shape[0]==5)."""
|
||||
return _write_real_zarr_store(
|
||||
tmp_path / "group.zarr",
|
||||
{
|
||||
"images": (np.arange(20, dtype="<i4").reshape(5, 4), (2, 4)),
|
||||
"nested": (np.arange(5, dtype="|u1"), (2,)),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zarrv2_root_store(tmp_path) -> Path:
|
||||
"""Single-array store with the array sitting directly at the store root."""
|
||||
store_path = tmp_path / "root.zarr"
|
||||
arr = zarr.open(
|
||||
str(store_path),
|
||||
mode="w",
|
||||
shape=(5, 4),
|
||||
chunks=(2, 4),
|
||||
dtype="<i4",
|
||||
)
|
||||
arr[:] = np.arange(20, dtype="<i4").reshape(5, 4)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_fsspec_fs():
|
||||
"""fsspec local filesystem (for parametrized cross-fs read tests)."""
|
||||
return fsspec.filesystem("file")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def heterogeneous_zarrv2_store(tmp_path) -> Path:
|
||||
"""A store mixing different ranks, shape[0]s, dtypes, and native chunk sizes.
|
||||
|
||||
Mirrors the UMI-style real-world layout where ``data/*`` arrays share an
|
||||
axis-0 timestep count but differ in everything else, and ``meta/*``
|
||||
arrays live in a separate axis-0 universe entirely. The chunk-per-row
|
||||
datasource handles all of these in one read; nothing has to align.
|
||||
"""
|
||||
store_path = tmp_path / "heterogeneous.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
# 4-D image tensor with tiny axis-0 chunks (1 image per chunk).
|
||||
root.create_dataset(
|
||||
"data/camera0_rgb",
|
||||
data=np.arange(20 * 2 * 2 * 3, dtype="|u1").reshape(20, 2, 2, 3),
|
||||
chunks=(1, 2, 2, 3),
|
||||
)
|
||||
# 2-D pose array, same shape[0]=20, much larger axis-0 chunks (10).
|
||||
root.create_dataset(
|
||||
"data/robot0_eef_pos",
|
||||
data=np.arange(20 * 3, dtype="<f4").reshape(20, 3),
|
||||
chunks=(10, 3),
|
||||
)
|
||||
# Episode-boundary metadata: separate axis-0 universe.
|
||||
root.create_dataset(
|
||||
"meta/episode_ends",
|
||||
data=np.array([5, 12, 20], dtype="<i8"),
|
||||
chunks=(3,),
|
||||
)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def unconsolidated_zarrv2_store(tmp_path) -> Path:
|
||||
"""Two arrays at the store root, no ``.zmetadata``.
|
||||
|
||||
Exercises the no-``.zmetadata`` code paths (per-array ``.zarray``
|
||||
discovery and full-store walk) — the common shape of real-world stores
|
||||
behind plain HTTPS or other listing-less filesystems.
|
||||
"""
|
||||
store_path = tmp_path / "unconsolidated.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
root.create_dataset(
|
||||
"images", data=np.arange(20, dtype="<i4").reshape(5, 4), chunks=(2, 4)
|
||||
)
|
||||
root.create_dataset("nested", data=np.arange(5, dtype="|u1"), chunks=(2,))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def aligned_zarrv2_store(tmp_path) -> Path:
|
||||
"""Three arrays sharing ``shape[0]=8``, different ranks and native chunks.
|
||||
|
||||
Models the UMI-style case where data arrays co-stride on the timestep
|
||||
axis but differ in everything else.
|
||||
"""
|
||||
store_path = tmp_path / "aligned.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
root.create_dataset(
|
||||
"img",
|
||||
data=np.arange(8 * 4 * 4 * 3, dtype="|u1").reshape(8, 4, 4, 3),
|
||||
chunks=(2, 4, 4, 3),
|
||||
)
|
||||
root.create_dataset(
|
||||
"state",
|
||||
data=np.arange(8 * 3, dtype="<f4").reshape(8, 3),
|
||||
chunks=(4, 3), # different native axis-0 chunks than img
|
||||
)
|
||||
root.create_dataset(
|
||||
"label",
|
||||
data=np.arange(8, dtype="<i8"),
|
||||
chunks=(8,),
|
||||
)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zarr_zip_store(tmp_path) -> Path:
|
||||
"""A small Zarr store packed into a ``.zip`` for URL-detection tests."""
|
||||
src = tmp_path / "src.zarr"
|
||||
_write_real_zarr_store(
|
||||
src,
|
||||
{
|
||||
"data": (np.arange(12, dtype="<i4").reshape(6, 2), (3, 2)),
|
||||
},
|
||||
)
|
||||
zip_path = tmp_path / "store.zarr.zip"
|
||||
import shutil
|
||||
|
||||
shutil.make_archive(
|
||||
base_name=str(tmp_path / "store.zarr"),
|
||||
format="zip",
|
||||
root_dir=str(src),
|
||||
)
|
||||
assert zip_path.exists()
|
||||
return zip_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Metadata discovery
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_normalizes_requested_root_array_path(zarrv2_root_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_root_store),
|
||||
array_paths=[""],
|
||||
)
|
||||
assert list(datasource._metadata_by_path) == [""]
|
||||
|
||||
|
||||
def test_normalizes_requested_array_paths(zarrv2_group_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store),
|
||||
array_paths=["images/", "nested"],
|
||||
)
|
||||
assert list(datasource._metadata_by_path) == ["images", "nested"]
|
||||
|
||||
|
||||
def test_rejects_missing_array_paths(zarrv2_group_store):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Array\(s\) not found: 'missing'\. Available: 'images', 'nested'",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store),
|
||||
array_paths=["missing"],
|
||||
)
|
||||
|
||||
|
||||
def test_loads_per_array_zarray_without_zmetadata(unconsolidated_zarrv2_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store),
|
||||
array_paths=["images", "nested"],
|
||||
)
|
||||
assert set(datasource._metadata_by_path) == {"images", "nested"}
|
||||
|
||||
|
||||
def test_full_scan_discovers_arrays_without_zmetadata(unconsolidated_zarrv2_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store),
|
||||
allow_full_metadata_scan=True,
|
||||
)
|
||||
assert set(datasource._metadata_by_path) == {"images", "nested"}
|
||||
|
||||
|
||||
def test_requires_array_paths_or_full_scan_when_unconsolidated(
|
||||
unconsolidated_zarrv2_store,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
r"No array_paths were provided and this Zarr store does not "
|
||||
r"contain \.zmetadata"
|
||||
),
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(unconsolidated_zarrv2_store))
|
||||
|
||||
|
||||
def test_array_paths_missing_zarray_file_raises_value_error(
|
||||
unconsolidated_zarrv2_store,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Array path 'missing' not found",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store),
|
||||
array_paths=["missing"],
|
||||
)
|
||||
|
||||
|
||||
def test_local_scheme_pins_reads_to_driver_node(zarrv2_group_store):
|
||||
"""``local://`` stores can't be read distributed; plain/cloud paths can."""
|
||||
local = zarrv2_datasource.ZarrV2Datasource("local://" + str(zarrv2_group_store))
|
||||
assert local.supports_distributed_reads is False
|
||||
|
||||
plain = zarrv2_datasource.ZarrV2Datasource(str(zarrv2_group_store))
|
||||
assert plain.supports_distributed_reads is True
|
||||
|
||||
|
||||
def test_consolidation_detected_via_open_consolidated(
|
||||
zarrv2_group_store, unconsolidated_zarrv2_store
|
||||
):
|
||||
"""``_consolidated`` reflects whether ``.zmetadata`` actually opened."""
|
||||
consolidated = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store), array_paths=["images"]
|
||||
)
|
||||
assert consolidated._consolidated is True
|
||||
|
||||
unconsolidated = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store), array_paths=["images"]
|
||||
)
|
||||
assert unconsolidated._consolidated is False
|
||||
|
||||
|
||||
def test_array_paths_rejects_group_path(tmp_path):
|
||||
"""Requesting a group path (not an array) on an unconsolidated store errors."""
|
||||
store_path = tmp_path / "withgroup.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
grp = root.create_group("grp")
|
||||
grp.create_dataset("inner", data=np.arange(4, dtype="<i4"), chunks=(2,))
|
||||
# Not consolidated -> the per-array ``.zarray`` lookup path.
|
||||
with pytest.raises(ValueError, match="is a group, not an array"):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(store_path), array_paths=["grp"])
|
||||
|
||||
|
||||
def test_root_array_rejects_non_root_array_paths(zarrv2_root_store):
|
||||
"""A single root-level array rejects array_paths that aren't the root ''."""
|
||||
with pytest.raises(ValueError, match="single root-level array"):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_root_store), array_paths=["missing"]
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# chunk_shapes validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk_shapes, match",
|
||||
[
|
||||
("invalid", "positive integers"),
|
||||
({"images": 1}, "positive integers"),
|
||||
({"does_not_exist": [2]}, "Unknown array path"),
|
||||
],
|
||||
)
|
||||
def test_rejects_invalid_chunk_shapes(zarrv2_group_store, chunk_shapes, match):
|
||||
with pytest.raises(ValueError, match=match):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store), chunk_shapes=chunk_shapes
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk_shapes,array_paths,expected",
|
||||
[
|
||||
# No chunk_shapes: every array reads at its native chunk size.
|
||||
# 4-D image with tiny chunks coexists with 2-D pose with big chunks —
|
||||
# nothing is forced into a shared min/max.
|
||||
(
|
||||
None,
|
||||
None,
|
||||
{
|
||||
"data/camera0_rgb": (1, 2, 2, 3),
|
||||
"data/robot0_eef_pos": (10, 3),
|
||||
"meta/episode_ends": (3,),
|
||||
},
|
||||
),
|
||||
# ``[5]`` prefix overrides axis 0 across arrays of all ranks at once.
|
||||
(
|
||||
[5],
|
||||
None,
|
||||
{
|
||||
"data/camera0_rgb": (5, 2, 2, 3),
|
||||
"data/robot0_eef_pos": (5, 3),
|
||||
"meta/episode_ends": (5,),
|
||||
},
|
||||
),
|
||||
# Length-2 prefix overrides axes 0+1; needs every selected array to
|
||||
# have rank >= 2, so we filter out ``meta/episode_ends`` (rank 1).
|
||||
(
|
||||
[5, 1],
|
||||
["data/camera0_rgb", "data/robot0_eef_pos"],
|
||||
{
|
||||
"data/camera0_rgb": (5, 1, 2, 3),
|
||||
"data/robot0_eef_pos": (5, 1),
|
||||
},
|
||||
),
|
||||
# Per-array overrides may retile only some arrays while others keep
|
||||
# their native chunks.
|
||||
(
|
||||
{
|
||||
"data/camera0_rgb": [5],
|
||||
"data/robot0_eef_pos": [7],
|
||||
},
|
||||
None,
|
||||
{
|
||||
"data/camera0_rgb": (5, 2, 2, 3),
|
||||
"data/robot0_eef_pos": (7, 3),
|
||||
"meta/episode_ends": (3,),
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_chunk_shapes_resolution_across_mixed_rank(
|
||||
heterogeneous_zarrv2_store, chunk_shapes, array_paths, expected
|
||||
):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(heterogeneous_zarrv2_store),
|
||||
chunk_shapes=chunk_shapes,
|
||||
array_paths=array_paths,
|
||||
)
|
||||
assert datasource._array_chunks == expected
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# align_axis_0 (wide-form mode)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_align_axis_0_emits_wide_rows(ray_start_regular_shared, aligned_zarrv2_store):
|
||||
"""Wide-row schema: ``t_start``, ``t_stop``, one column per selected array."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
)
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
|
||||
assert set(df.columns) == {"t_start", "t_stop", "img", "state", "label"}
|
||||
# shape[0]=8, chunk_shapes=[4] -> 2 rows.
|
||||
assert len(df) == 2
|
||||
# Reconstruct each array by concatenating slices in order.
|
||||
img_recon = np.concatenate(list(df["img"]), axis=0)
|
||||
assert img_recon.shape == (8, 4, 4, 3)
|
||||
state_recon = np.concatenate(list(df["state"]), axis=0)
|
||||
assert state_recon.shape == (8, 3)
|
||||
label_recon = np.concatenate(list(df["label"]), axis=0)
|
||||
assert label_recon.shape == (8,)
|
||||
# t_start/t_stop are correct.
|
||||
starts = sorted(df["t_start"].tolist())
|
||||
stops = sorted(df["t_stop"].tolist())
|
||||
assert starts == [0, 4]
|
||||
assert stops == [4, 8]
|
||||
|
||||
|
||||
def test_align_axis_0_column_set(ray_start_regular_shared, aligned_zarrv2_store):
|
||||
"""array_paths selects which arrays are read; aligned mode emits one column
|
||||
per selected array (plus t_start/t_stop)."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
array_paths=["img", "state"],
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
)
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
|
||||
assert set(df.columns) == {"t_start", "t_stop", "img", "state"}
|
||||
|
||||
|
||||
def test_align_axis_0_rejects_misaligned_shape0(heterogeneous_zarrv2_store):
|
||||
"""Misalignment raises with the per-array shape[0] breakdown."""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"All selected arrays must share shape\[0\]",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(heterogeneous_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[5],
|
||||
)
|
||||
|
||||
|
||||
def test_align_axis_0_rejects_divergent_axis_0_chunks(aligned_zarrv2_store):
|
||||
"""If aligned arrays end up with different axis-0 chunks, error clearly.
|
||||
|
||||
The native chunks differ (img=2, state=4, label=8) — without a
|
||||
``chunk_shapes`` re-tile they all stay at native, and the validator
|
||||
catches the mismatch.
|
||||
"""
|
||||
with pytest.raises(
|
||||
ValueError, match="Aligned arrays must share the same axis-0 chunk size"
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# overlap (aligned-mode lookahead)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_overlap_extends_chunk_data(ray_start_regular_shared, aligned_zarrv2_store):
|
||||
"""``overlap=N`` makes each row's per-array slice cover ``N`` extra timesteps.
|
||||
|
||||
Aligned store has shape[0]=8, ``chunk_shapes=[4]`` -> rows own [0,4) and [4,8).
|
||||
With ``overlap=2``, row 0's data covers [0,6) and row 1's data covers [4,8)
|
||||
(clipped at the store end since 4+4+2 > 8).
|
||||
"""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
overlap=2,
|
||||
)
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
|
||||
# Ownership unchanged: 2 rows of width 4 each.
|
||||
assert sorted(zip(df["t_start"], df["t_stop"])) == [(0, 4), (4, 8)]
|
||||
# Data extents: row 0 has 6 timesteps, row 1 has 4 (clipped at shape[0]=8).
|
||||
rows = sorted(df.to_dict("records"), key=lambda r: r["t_start"])
|
||||
assert rows[0]["img"].shape[0] == 6 # 4 owned + 2 overlap
|
||||
assert rows[0]["state"].shape[0] == 6
|
||||
assert rows[1]["img"].shape[0] == 4 # 4 owned + 0 overlap (clipped)
|
||||
assert rows[1]["state"].shape[0] == 4
|
||||
|
||||
|
||||
def test_overlap_requires_align_axis_0(aligned_zarrv2_store):
|
||||
"""``overlap`` in long-form (no ``align_axis_0``) is a clear error."""
|
||||
with pytest.raises(ValueError, match="overlap requires align_axis_0=True"):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
overlap=2,
|
||||
)
|
||||
|
||||
|
||||
def test_overlap_rejects_negative_and_non_int(aligned_zarrv2_store):
|
||||
bad_values: list[Any] = [-1, 1.5, "two"]
|
||||
|
||||
for bad in bad_values:
|
||||
with pytest.raises(ValueError, match="overlap must be a non-negative integer"):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
overlap=bad,
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_shapes_rejected_when_longer_than_smallest_array(
|
||||
heterogeneous_zarrv2_store,
|
||||
):
|
||||
"""A shared ``chunk_shapes`` override longer than a target rank is an error."""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"chunk_shapes override for array .* has 2 axes but array of shape .* has rank 1",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(heterogeneous_zarrv2_store),
|
||||
chunk_shapes=[2, 2], # OK for 2-D and 4-D, fails for 1-D episode_ends
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Filesystem handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_accepts_pyarrow_fs_filesystem(zarrv2_group_store):
|
||||
"""A pyarrow.fs.FileSystem passed in is wrapped into fsspec internally."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store),
|
||||
filesystem=pyarrow.fs.LocalFileSystem(),
|
||||
)
|
||||
from fsspec.spec import AbstractFileSystem
|
||||
|
||||
assert isinstance(datasource._fs, AbstractFileSystem)
|
||||
assert set(datasource._metadata_by_path) == {"images", "nested"}
|
||||
|
||||
|
||||
def test_rejects_unsupported_filesystem_type():
|
||||
"""Filesystem that's neither pyarrow.fs nor fsspec raises ``TypeError``."""
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=r"filesystem must be pyarrow\.fs\.FileSystem or",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
"/tmp/some.zarr",
|
||||
filesystem="not-a-filesystem",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# .zarr.zip URL support
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_reads_zarr_zip_local_path(ray_start_regular_shared, zarr_zip_store):
|
||||
"""A local ``.zarr.zip`` path auto-wires fsspec's ZipFileSystem."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store))
|
||||
# The store has one array "data" of shape (6, 2) chunks (3, 2) -> 2 chunks.
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
|
||||
assert len(df) == 2
|
||||
assert set(df["array"]) == {"data"}
|
||||
recon = _reconstruct_array(df, "data")
|
||||
np.testing.assert_array_equal(recon, np.arange(12, dtype="<i4").reshape(6, 2))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Read task generation and execution (end-to-end)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_read_tasks_batches_chunks_by_parallelism(tmp_path):
|
||||
"""5 chunks split across parallelism=3 produces batches [2, 2, 1]."""
|
||||
store_path = tmp_path / "store.zarr"
|
||||
_write_real_zarr_store(
|
||||
store_path,
|
||||
{"images": (np.arange(5 * 4, dtype="<i4").reshape(5, 4), (1, 4))},
|
||||
)
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
|
||||
read_tasks = datasource.get_read_tasks(parallelism=3)
|
||||
|
||||
assert len(read_tasks) == 3
|
||||
assert [task.metadata.num_rows for task in read_tasks] == [2, 2, 1]
|
||||
assert all(task.metadata.input_files == (str(store_path),) for task in read_tasks)
|
||||
|
||||
|
||||
def test_long_form_chunk_index_order_matches_grid(ray_start_regular_shared, tmp_path):
|
||||
"""Lazy grid-range tasks emit chunk_index in the same row-major order as a
|
||||
full grid enumeration (regression guard for the lazy-unravel refactor)."""
|
||||
from itertools import product
|
||||
|
||||
store_path = tmp_path / "order.zarr"
|
||||
# shape (6, 4), chunks (2, 2) -> grid (3, 2) = 6 chunks.
|
||||
_write_real_zarr_store(
|
||||
store_path, {"a": (np.arange(6 * 4, dtype="<i4").reshape(6, 4), (2, 2))}
|
||||
)
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
# parallelism=2 -> two flat-index ranges; concatenated they must be in order.
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
|
||||
got = [tuple(int(x) for x in ci) for ci in df["chunk_index"]]
|
||||
assert got == list(product(range(3), range(2)))
|
||||
|
||||
|
||||
def test_per_task_row_limit_caps_chunks_read(
|
||||
ray_start_regular_shared, tmp_path, monkeypatch
|
||||
):
|
||||
"""per_task_row_limit bounds how many chunks a task actually reads, so a
|
||||
downstream ``limit`` doesn't pull the whole batch's I/O."""
|
||||
store_path = tmp_path / "limit.zarr"
|
||||
_write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (1,))})
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
|
||||
reads = []
|
||||
real_read_chunk = zarrv2_datasource._read_chunk
|
||||
|
||||
def _spy(*args, **kwargs):
|
||||
reads.append(1)
|
||||
return real_read_chunk(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(zarrv2_datasource, "_read_chunk", _spy)
|
||||
|
||||
# parallelism=1 -> one task batching all 10 chunks; cap it at 3.
|
||||
tasks = datasource.get_read_tasks(parallelism=1, per_task_row_limit=3)
|
||||
blocks = [block for task in tasks for block in task()]
|
||||
|
||||
total_rows = sum(BlockAccessor.for_block(b).num_rows() for b in blocks)
|
||||
assert total_rows == 3
|
||||
# The fix: only 3 chunks were actually read (not all 10, then truncated).
|
||||
assert len(reads) == 3
|
||||
|
||||
|
||||
def test_read_chunk_retries_transient_io(monkeypatch):
|
||||
"""_read_chunk retries reads whose error matches retry_match (Ray Data's
|
||||
DataContext.retried_io_errors), then succeeds."""
|
||||
monkeypatch.setattr("time.sleep", lambda *_: None) # no backoff in the test
|
||||
|
||||
class _FlakyArray:
|
||||
attempts = 0
|
||||
|
||||
def __getitem__(self, _idx):
|
||||
type(self).attempts += 1
|
||||
if self.attempts < 3:
|
||||
raise OSError("Connection reset by peer")
|
||||
return np.arange(4, dtype="<i4")
|
||||
|
||||
class _Root:
|
||||
def __getitem__(self, _name):
|
||||
return _FlakyArray()
|
||||
|
||||
out = zarrv2_datasource._read_chunk(
|
||||
_Root(), # pyrefly: ignore[bad-argument-type]
|
||||
"x",
|
||||
((0, 4),),
|
||||
retry_match=["Connection reset"],
|
||||
)
|
||||
np.testing.assert_array_equal(out, np.arange(4, dtype="<i4"))
|
||||
assert _FlakyArray.attempts == 3 # failed twice, then succeeded
|
||||
|
||||
|
||||
def test_long_form_schema_and_materialization(ray_start_regular_shared, tmp_path):
|
||||
"""End-to-end: long-form rows are emitted with the expected columns and data."""
|
||||
store_path = tmp_path / "aligned.zarr"
|
||||
images_src = np.arange(20, dtype="<i4").reshape(5, 4)
|
||||
labels_src = np.arange(5, dtype="|u1")
|
||||
_write_real_zarr_store(
|
||||
store_path,
|
||||
{
|
||||
"images": (images_src, (2, 4)),
|
||||
"labels": (labels_src, (2,)),
|
||||
},
|
||||
)
|
||||
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
|
||||
|
||||
# Schema is the long-form quad.
|
||||
assert list(df.columns) == ["array", "chunk_index", "chunk_slices", "chunk"]
|
||||
# 3 chunks for images (5/2), 3 chunks for labels (5/2) = 6 rows total.
|
||||
assert len(df) == 6
|
||||
assert set(df["array"]) == {"images", "labels"}
|
||||
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)
|
||||
|
||||
# ``chunk_slices`` matches the actual chunk shape and indexes back to
|
||||
# the source array: arr[start:stop, ...] equals the chunk.
|
||||
for _, row in df.iterrows():
|
||||
slices = row["chunk_slices"]
|
||||
chunk = row["chunk"]
|
||||
assert len(slices) == chunk.ndim
|
||||
for axis, (start, stop) in enumerate(slices):
|
||||
assert stop - start == chunk.shape[axis]
|
||||
if row["array"] == "images":
|
||||
np.testing.assert_array_equal(
|
||||
chunk,
|
||||
images_src[slices[0][0] : slices[0][1], slices[1][0] : slices[1][1]],
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_shapes_override_changes_grid(ray_start_regular_shared, tmp_path):
|
||||
"""User-supplied chunk_shapes controls the chunk grid and row count."""
|
||||
store_path = tmp_path / "tile.zarr"
|
||||
src = np.arange(10, dtype="<i4")
|
||||
_write_real_zarr_store(store_path, {"data": (src, (2,))}) # native: 5 chunks
|
||||
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path), chunk_shapes=[5])
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
|
||||
assert sorted(chunk.shape[0] for chunk in df["chunk"]) == [5, 5]
|
||||
|
||||
|
||||
def test_heterogeneous_store_emits_one_row_per_chunk(
|
||||
ray_start_regular_shared, heterogeneous_zarrv2_store
|
||||
):
|
||||
"""Mixed-rank/shape/dtype arrays each contribute their chunk count to the output."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(heterogeneous_zarrv2_store))
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
|
||||
|
||||
# Expected chunk counts:
|
||||
# data/camera0_rgb shape=(20,2,2,3) chunks=(1,2,2,3) → 20 chunks
|
||||
# data/robot0_eef_pos shape=(20,3) chunks=(10,3) → 2 chunks
|
||||
# meta/episode_ends shape=(3,) chunks=(3,) → 1 chunk
|
||||
counts = df.groupby("array").size().to_dict()
|
||||
assert counts == {
|
||||
"data/camera0_rgb": 20,
|
||||
"data/robot0_eef_pos": 2,
|
||||
"meta/episode_ends": 1,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Estimator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_estimate_inmemory_data_size(tmp_path):
|
||||
"""Estimate = sum over arrays of numel * dtype.itemsize."""
|
||||
store_path = tmp_path / "est.zarr"
|
||||
_write_real_zarr_store(
|
||||
store_path,
|
||||
{
|
||||
"a": (np.zeros((5, 4), dtype="<i4"), (2, 4)),
|
||||
"b": (np.zeros(5, dtype="|u1"), (2,)),
|
||||
},
|
||||
)
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
# 5*4*4 (a) + 5*1 (b) = 80 + 5 = 85
|
||||
assert datasource.estimate_inmemory_data_size() == 85
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-filesystem end-to-end (Ray Data convention)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fs",
|
||||
[
|
||||
None,
|
||||
lazy_fixture("local_fs"), # pyarrow.fs (gets wrapped to fsspec internally)
|
||||
lazy_fixture("local_fsspec_fs"), # native fsspec
|
||||
],
|
||||
)
|
||||
def test_read_zarr_basic_across_filesystems(ray_start_regular_shared, fs, local_path):
|
||||
"""Round-trip a real Zarr store through read_zarr for each filesystem flavor.
|
||||
|
||||
Mirrors the parametrized read-path coverage other Ray Data datasources use
|
||||
(lance, parquet, json, hudi, …) — exercises None / pyarrow.fs / fsspec
|
||||
input shapes against the same store written to a local path.
|
||||
"""
|
||||
store_path = os.path.join(local_path, "data.zarr")
|
||||
images_src = np.arange(20, dtype="<i4").reshape(5, 4)
|
||||
labels_src = np.arange(5, dtype="|u1")
|
||||
_write_real_zarr_store(
|
||||
Path(store_path),
|
||||
{
|
||||
"images": (images_src, (2, 4)),
|
||||
"labels": (labels_src, (2,)),
|
||||
},
|
||||
)
|
||||
|
||||
ds = ray.data.read_zarr(store_path, filesystem=fs)
|
||||
|
||||
# 3 chunks each for images and labels (5/2 → ceil = 3) → 6 rows total.
|
||||
assert ds.count() == 6
|
||||
df = pd.DataFrame(ds.take_all())
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)
|
||||
|
||||
|
||||
def test_rejects_zarr_v3(tmp_path, monkeypatch):
|
||||
"""read_zarr targets zarr-python 2.x; an incompatible v3 install must raise a
|
||||
clear, actionable error at construction, not a cryptic ImportError mid-read."""
|
||||
monkeypatch.setattr(zarr, "__version__", "3.0.1")
|
||||
with pytest.raises(ImportError, match=r"zarr-python 2\.x"):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(tmp_path))
|
||||
|
||||
|
||||
def test_explicit_filesystem_strips_uri_scheme(ray_start_regular_shared, tmp_path):
|
||||
"""An explicit ``filesystem=`` plus a scheme-prefixed path must strip the
|
||||
scheme so the store path is backend-relative. Regression: pyarrow
|
||||
filesystems can't resolve a ``file://`` / ``gs://`` prefix in the path."""
|
||||
store_path = tmp_path / "scheme.zarr"
|
||||
_write_real_zarr_store(store_path, {"data": (np.arange(6, dtype="<i4"), (2,))})
|
||||
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(
|
||||
f"file://{store_path}", filesystem=pyarrow.fs.LocalFileSystem()
|
||||
)
|
||||
assert ds._store_path == str(store_path)
|
||||
df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
|
||||
assert len(df) == 3
|
||||
|
||||
|
||||
def test_get_read_tasks_parallelism_zero(tmp_path):
|
||||
"""parallelism=0 must not divide by zero; fall back to a single task."""
|
||||
store_path = tmp_path / "p0.zarr"
|
||||
_write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (2,))})
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
tasks = ds.get_read_tasks(parallelism=0)
|
||||
assert len(tasks) >= 1
|
||||
|
||||
|
||||
def test_align_axis_0_rejects_scalar_array(tmp_path):
|
||||
"""align_axis_0=True with a 0-D (scalar) array must raise a clear error
|
||||
rather than an IndexError when reading the (empty) axis-0 chunk size."""
|
||||
store_path = tmp_path / "scalar.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
root.create_dataset("vec", data=np.arange(8, dtype="<i4"), chunks=(4,))
|
||||
root.create_dataset("scalar", data=np.array(42, dtype="<i4")) # 0-D
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
|
||||
with pytest.raises(ValueError, match=r"0-D \(scalar\)"):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(store_path), align_axis_0=True)
|
||||
|
||||
|
||||
def test_reads_zarr_zip_with_explicit_zip_filesystem(
|
||||
ray_start_regular_shared, zarr_zip_store
|
||||
):
|
||||
"""A .zip path read through an explicitly-passed fsspec ZipFileSystem must
|
||||
resolve the store at the archive root (store path ``""``), not treat the
|
||||
``.zip`` name as an entry inside the archive."""
|
||||
zip_fs = fsspec.filesystem("zip", fo=str(zarr_zip_store))
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store), filesystem=zip_fs)
|
||||
assert ds._store_path == ""
|
||||
df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
|
||||
assert len(df) == 2
|
||||
|
||||
|
||||
def test_align_axis_0_columns_unify_across_blocks(
|
||||
ray_start_regular_shared, aligned_zarrv2_store
|
||||
):
|
||||
"""Wide-form gives each array its own column, so blocks combine cleanly
|
||||
across the dataset even with trailing edge chunks of differing shape -- the
|
||||
batch-safe schema for row-aligned arrays."""
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas
|
||||
from ray.data.block import BlockAccessor
|
||||
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store), align_axis_0=True, chunk_shapes=[3]
|
||||
)
|
||||
blocks = [block for task in ds.get_read_tasks(parallelism=64) for block in task()]
|
||||
assert len(blocks) > 1 # actually exercise cross-block unification
|
||||
schemas = [BlockAccessor.for_block(b).to_arrow().schema for b in blocks]
|
||||
unified = unify_schemas(schemas) # must not raise
|
||||
assert {"t_start", "t_stop", "img", "state", "label"}.issubset(set(unified.names))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Custom codec registration in Ray workers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_ray():
|
||||
"""A clean Ray for a test that needs its own ``ray.init`` (e.g. a custom
|
||||
``runtime_env``). Unlike ``shutdown_only`` (teardown only), it also shuts
|
||||
down any pre-existing cluster, so isolation doesn't depend on test order.
|
||||
"""
|
||||
if ray.is_initialized():
|
||||
ray.shutdown()
|
||||
yield
|
||||
if ray.is_initialized():
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def test_custom_codec_succeeds_with_worker_setup_hook(fresh_ray, tmp_path):
|
||||
"""Test that we successfully register a custom codec.
|
||||
|
||||
numcodecs' registry is process-local.
|
||||
"""
|
||||
import numcodecs
|
||||
|
||||
def _register_codec():
|
||||
import numcodecs
|
||||
import numpy as np
|
||||
|
||||
class _RayZarrTestCodec(numcodecs.abc.Codec):
|
||||
codec_id = "ray_zarr_test_codec"
|
||||
|
||||
def encode(self, buf):
|
||||
return bytes(buf)
|
||||
|
||||
def decode(self, buf, out=None):
|
||||
arr = np.frombuffer(buf, dtype=np.uint8)
|
||||
if out is not None:
|
||||
out[:] = arr.view(out.dtype)
|
||||
return out
|
||||
return arr.copy()
|
||||
|
||||
numcodecs.register_codec(_RayZarrTestCodec)
|
||||
|
||||
# Register driver-side so we can write the store.
|
||||
_register_codec()
|
||||
|
||||
store_path = tmp_path / "codec_test.zarr"
|
||||
arr = zarr.open(
|
||||
str(store_path),
|
||||
mode="w",
|
||||
shape=(8,),
|
||||
chunks=(4,),
|
||||
dtype="u1",
|
||||
compressor=numcodecs.get_codec({"id": "ray_zarr_test_codec"}),
|
||||
)
|
||||
arr[:] = np.arange(8, dtype="u1")
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
|
||||
ray.init(
|
||||
num_cpus=1,
|
||||
logging_level=logging.ERROR,
|
||||
log_to_driver=False,
|
||||
runtime_env={"worker_process_setup_hook": _register_codec},
|
||||
)
|
||||
ds = ray.data.read_zarr(str(store_path))
|
||||
rows = sorted(ds.take_all(), key=lambda r: tuple(r["chunk_index"]))
|
||||
recon = np.concatenate([r["chunk"] for r in rows])
|
||||
np.testing.assert_array_equal(recon, np.arange(8, dtype="u1"))
|
||||
Reference in New Issue
Block a user