Sub-chunk slicing for uncompressed ManifestArrays (part of #86)

conftest.py

@@ -80,6 +80,20 @@ def local_registry():
     return ObjectStoreRegistry({"file://": LocalStore()})
 
 
+@pytest.fixture
+def netcdf3_file(tmp_path: Path):
+    """Factory for writing a temporary netCDF3 file with a caller-supplied Dataset."""
+
+    def _make(ds: xr.Dataset | None = None, name: str = "file.nc") -> Path:
+        if ds is None:
+            ds = xr.Dataset({"foo": ("x", np.array([1, 2, 3]))})
+        filepath = tmp_path / name
+        ds.to_netcdf(filepath, format="NETCDF3_CLASSIC")
+        return filepath
+
+    return _make
+
+
 @pytest.fixture(params=["int8", "uint8", "float32"])
 def zarr_array_fill_value(request):
     store = zarr.storage.MemoryStore()

docs/data_structures.md

@@ -246,6 +246,7 @@ The whole point is to manipulate references to the data without actually loading
 !!! note
     You can index into a `ManifestArray` as long as the selection aligns with chunk boundaries — slicing through the interior of a chunk would require loading the chunk's bytes, which a virtual array deliberately cannot do.
     Chunk-aligned integer and slice indexing is supported, including mixed integer + slice indexers; integer indexers drop the indexed axis as in numpy. Misaligned selections raise `SubChunkIndexingError`.
+    As a special case, an **uncompressed** array supports slicing _within_ a single source chunk along its largest-stride storage axis — the result is a new chunk reference with a bumped byte offset and a smaller length, no data loaded. This works for `[BytesCodec]` arrays (C-order; the axis is axis 0) and `[TransposeCodec(order=...), BytesCodec]` arrays (the axis is `order[0]` — e.g. the last axis for F-order). Useful for picking out a row from a multi-row chunk produced by a parser like the netCDF3 one.
     Arbitrary fancy indexing (e.g. with a boolean mask or integer array) is not supported, since it would generally require loading data.
 
 ## Zarr Groups

docs/releases.md

@@ -6,6 +6,8 @@
 
 - `ManifestArray` now supports chunk-aligned integer and slice indexing along each axis, including multi-chunk slices, mixed integer + slice indexers, and selections that include a partial final chunk. Integer indexers drop the indexed axis (numpy / array-API semantics) and are legal only when `chunk_size == 1` along that axis; slice indexers preserve the axis. This makes `xarray.Dataset.isel` work end-to-end on virtual datasets for any chunk-aligned selection. Indexers that would split individual chunks raise a new `SubChunkIndexingError` (a `ValueError` subclass) — a permanent constraint of a virtual array, not a missing feature. Previously slice misalignment silently no-op'd while integer indexing unconditionally raised `NotImplementedError`. Closes [#51](https://github.com/zarr-developers/VirtualiZarr/issues/51), supersedes [#499](https://github.com/zarr-developers/VirtualiZarr/pull/499).
   By [Tom Nicholas](https://github.com/TomNicholas).
+- Slicing along the largest-stride storage axis of an uncompressed `ManifestArray` can now sub-divide a chunk — the result is a new reference into the same source file with a bumped byte offset and a smaller length. Eligible codec stacks are `[BytesCodec]` (C-order; the axis is axis 0) and `[TransposeCodec(order=...), BytesCodec]` (e.g. F-order with `order=(N-1, ..., 0)`, where the axis is `order[0]` — typically the last axis). Useful for picking a single timestep from a multi-row chunk produced by, e.g., the netCDF3 parser, without rechunking. Limited to slices that fit within one source chunk; multi-chunk-spanning sub-chunk slices, `step != 1`, and sub-chunk indexing on other axes still raise `SubChunkIndexingError`. Addresses part of [#86](https://github.com/zarr-developers/VirtualiZarr/issues/86).
+  By [Tom Nicholas](https://github.com/TomNicholas).
 
 ### Bug fixes
 

virtualizarr/manifests/array.py

@@ -242,6 +242,12 @@ def __getitem__(
         - ``int`` — drops the indexed axis, following numpy / array-API semantics. Only legal
           when ``chunk_size == 1`` along that axis; otherwise picking a single element would
           require splitting a chunk.
+        - Slice along the largest-stride storage axis of an **uncompressed** array that fits
+          entirely within one source chunk — handled by rewriting the chunk reference's byte
+          offset/length rather than splitting bytes. Useful for picking a single timestep from
+          a multi-row chunk on a parser like the netCDF3 one. The eligible-axis is axis 0 for
+          a plain ``[BytesCodec]`` array (C-order) or axis ``order[0]`` of a prepended
+          ``[TransposeCodec(order=...), BytesCodec]`` (e.g. the last axis for F-order).
 
         Anything else — fancy indexing with arrays, misaligned slices, ``step != 1`` —
         raises ``SubChunkIndexingError`` or ``NotImplementedError``.

virtualizarr/manifests/indexing.py

@@ -1,7 +1,9 @@
 from types import EllipsisType
-from typing import TYPE_CHECKING, Any, TypeAlias, cast
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeGuard, cast
 
 import numpy as np
+from zarr.codecs import BytesCodec, TransposeCodec
+from zarr.core.metadata.v3 import ArrayV3Metadata
 
 from virtualizarr.manifests.array_api import expand_dims
 from virtualizarr.manifests.manifest import ChunkManifest
@@ -172,33 +174,62 @@ def apply_selection(
             raise TypeError(f"Invalid indexer type: {indexer_1d}")
         narrowed_indexers.append(indexer_1d)
 
+    sub_chunk_axis = _uncompressed_sub_chunk_axis(marr.metadata)
+
     new_shape: list[int] = []
     new_chunks: list[int] = []
     chunk_grid_selectors: list[int | slice] = []
     kept_axes: list[int] = []
+    # At most one sub-chunk axis (whichever axis has the largest byte stride in storage).
+    # The byte adjustment is uniform across every surviving chunk, since chunks share layout.
+    sub_chunk_byte_adjust: tuple[int, int] | None = None
     for axis, (axis_length, chunk_size, indexer_1d) in enumerate(
         zip(marr.shape, marr.chunks, narrowed_indexers, strict=True)
     ):
-        chunk_grid_selector, new_axis_length = _compute_chunk_aligned_selection_1d(
-            indexer_1d, axis_length=axis_length, chunk_size=chunk_size
-        )
+        chunk_grid_selector: int | slice
+        if axis == sub_chunk_axis and _is_sub_chunk_slice(
+            indexer_1d, axis_length, chunk_size
+        ):
+            chunk_grid_selector, new_axis_length, sub_chunk_byte_adjust = (
+                _compute_sub_chunk_axis_selection(
+                    indexer_1d,
+                    axis_length=axis_length,
+                    chunk_size=chunk_size,
+                    other_axis_chunks=tuple(
+                        c for i, c in enumerate(marr.chunks) if i != axis
+                    ),
+                    itemsize=marr.dtype.itemsize,
+                )
+            )
+            new_chunks_for_axis = new_axis_length
+        else:
+            chunk_grid_selector, new_axis_length = _compute_chunk_aligned_selection_1d(
+                indexer_1d, axis_length=axis_length, chunk_size=chunk_size
+            )
+            new_chunks_for_axis = chunk_size
+
         chunk_grid_selectors.append(chunk_grid_selector)
-        # int selectors drop the axis from the output array
+        # int indexers drop the axis from the output array; slices preserve it (including the
+        # sub-chunk path, which uses a length-1 chunk-grid slice selector).
         if not isinstance(indexer_1d, int):
             new_shape.append(new_axis_length)
-            new_chunks.append(chunk_size)
+            new_chunks.append(new_chunks_for_axis)
             kept_axes.append(axis)
 
     chunk_grid_selectors_tuple = tuple(chunk_grid_selectors)
 
-    # short-circuit if every axis selects the whole chunk grid via a slice (a no-op)
-    if all(
+    # short-circuit if every axis selects the whole chunk grid via a slice (a no-op).
+    # A pending sub-chunk byte adjustment is real work even if its single source chunk
+    # happens to span the whole chunk grid along that axis, so don't short-circuit then.
+    if sub_chunk_byte_adjust is None and all(
         isinstance(cgs, slice) and cgs == slice(0, dim, 1)
         for cgs, dim in zip(chunk_grid_selectors_tuple, marr.manifest.shape_chunk_grid)
     ):
         return marr
 
     new_manifest = _subset_manifest(marr.manifest, chunk_grid_selectors_tuple)
+    if sub_chunk_byte_adjust is not None:
+        new_manifest = _shift_manifest_byte_ranges(new_manifest, *sub_chunk_byte_adjust)
     old_dimension_names = marr.metadata.dimension_names
     # zarr's dimension_names is tuple[str | None, ...] but copy_and_replace_metadata's
     # type hint says Iterable[str]; the runtime handles None entries fine, so cast through.
@@ -265,6 +296,38 @@ def _compute_chunk_aligned_selection_1d(
     return slice(chunk_start, chunk_stop, 1), stop - start
 
 
+def _compute_sub_chunk_axis_selection(
+    indexer_1d: slice,
+    axis_length: int,
+    chunk_size: int,
+    other_axis_chunks: tuple[int, ...],
+    itemsize: int,
+) -> tuple[slice, int, tuple[int, int]]:
+    """
+    Translate a sub-chunk slice along the eligible (largest-stride) storage axis into a
+    chunk-grid selector, an output axis length, and a uniform byte adjustment
+    ``(offset_delta, new_chunk_byte_length)`` applied to every surviving chunk reference.
+
+    Callers must have already confirmed that this slice is sub-chunk-eligible via
+    ``_is_sub_chunk_slice`` and that the array is uncompressed via
+    ``_uncompressed_sub_chunk_axis``.
+    """
+    start, stop, _ = indexer_1d.indices(axis_length)
+    chunk_index = start // chunk_size
+    new_axis_length = stop - start
+    # Bytes per index step along this axis within one chunk is the product of every
+    # *other* axis's chunk size, times itemsize. Order doesn't matter since the product
+    # is commutative.
+    stride_bytes = int(np.prod(other_axis_chunks)) * itemsize
+    inner_offset_bytes = (start - chunk_index * chunk_size) * stride_bytes
+    sub_chunk_byte_adjust = (inner_offset_bytes, new_axis_length * stride_bytes)
+    return (
+        slice(chunk_index, chunk_index + 1, 1),
+        new_axis_length,
+        sub_chunk_byte_adjust,
+    )
+
+
 def _subset_manifest(
     manifest: ChunkManifest, chunk_grid_selectors: tuple[int | slice, ...]
 ) -> ChunkManifest:
@@ -323,3 +386,86 @@ def _subset_manifest(
         inlined=new_inlined,
         validate_paths=False,
     )
+
+
+def _uncompressed_sub_chunk_axis(metadata: ArrayV3Metadata) -> int | None:
+    """
+    Return the axis along which sub-chunk slicing is implementable for this array, or
+    ``None`` if the codec stack disqualifies it.
+
+    Sub-chunk slicing rewrites an existing chunk reference's byte offset and length,
+    so it only works when chunk bytes are raw element values in a fixed memory order —
+    i.e., no compression, no value transforms, no checksums. The eligible codec stacks
+    are:
+
+    - ``[BytesCodec]`` — C-order layout; the axis with the largest byte stride is axis 0.
+    - ``[TransposeCodec(order=perm), BytesCodec]`` — stored layout is the logical array
+      permuted by ``perm``; the axis with the largest byte stride in storage is logical
+      axis ``perm[0]``. For the F-order case ``perm = (n-1, n-2, ..., 0)`` this picks out
+      the last axis.
+    """
+    codecs = metadata.codecs
+    if len(codecs) == 1 and isinstance(codecs[0], BytesCodec):
+        return 0
+    if (
+        len(codecs) == 2
+        and isinstance(codecs[0], TransposeCodec)
+        and isinstance(codecs[1], BytesCodec)
+    ):
+        return int(codecs[0].order[0])
+    return None
+
+
+def _is_sub_chunk_slice(
+    indexer_1d: int | slice, axis_length: int, chunk_size: int
+) -> TypeGuard[slice]:
+    """
+    True iff this is a slice that should take the sub-chunk path: step == 1, non-empty,
+    fits entirely within one source chunk, and is NOT already chunk-aligned (chunk-aligned
+    slices go through the simpler aligned path).
+
+    Typed as ``TypeGuard[slice]`` so callers can pass the narrowed indexer straight into
+    helpers that take a ``slice``.
+    """
+    if not isinstance(indexer_1d, slice):
+        return False
+    start, stop, step = indexer_1d.indices(axis_length)
+    if step != 1 or start >= stop:
+        return False
+    # chunk-aligned slices are handled by _compute_chunk_aligned_selection_1d
+    aligned = start % chunk_size == 0 and (
+        stop == axis_length or stop % chunk_size == 0
+    )
+    if aligned:
+        return False
+    # contained in a single source chunk?
+    return start // chunk_size == (stop - 1) // chunk_size
+
+
+def _shift_manifest_byte_ranges(
+    manifest: ChunkManifest, offset_delta: int, new_length: int
+) -> ChunkManifest:
+    """
+    Return a new ``ChunkManifest`` whose virtual chunk references point to a uniform
+    sub-range of each original chunk: ``offset += offset_delta`` and ``length = new_length``.
+
+    Used by the uncompressed-axis-0 sub-chunk path, where every surviving chunk shares the
+    same byte layout and therefore the same byte adjustment.
+    """
+    new_offsets = cast(
+        "np.ndarray[Any, np.dtype[np.uint64]]",
+        manifest._offsets + np.uint64(offset_delta),
+    )
+    new_lengths = cast(
+        "np.ndarray[Any, np.dtype[np.uint64]]",
+        np.full_like(manifest._lengths, np.uint64(new_length)),
+    )
+    # paths and any inlined-chunk dict carry through unchanged: inlined chunks aren't
+    # involved here (this path is only taken for uncompressed virtual references).
+    return ChunkManifest.from_arrays(
+        paths=manifest._paths,
+        offsets=new_offsets,
+        lengths=new_lengths,
+        inlined=dict(manifest._inlined),
+        validate_paths=False,
+    )

virtualizarr/tests/__init__.py

@@ -38,7 +38,6 @@ def _importorskip(
 has_tifffile, requires_tifffile = _importorskip("tifffile")
 has_imagecodecs, requires_imagecodecs = _importorskip("imagecodecs")
 has_hdf5plugin, requires_hdf5plugin = _importorskip("hdf5plugin")
-has_zarr_python, requires_zarr_python = _importorskip("zarr")
 has_dask, requires_dask = _importorskip("dask")
 has_obstore, requires_obstore = _importorskip("obstore")
 has_tiff, requires_tiff = _importorskip("virtual_tiff")

virtualizarr/tests/test_integration.py

@@ -18,15 +18,16 @@
     ManifestStore,
 )
 from virtualizarr.manifests.utils import create_v3_array_metadata
-from virtualizarr.parsers import HDFParser, ZarrParser
+from virtualizarr.parsers import HDFParser, NetCDF3Parser, ZarrParser
 from virtualizarr.parsers.kerchunk.translator import manifestgroup_from_kerchunk_refs
 from virtualizarr.tests import (
     has_fastparquet,
     has_icechunk,
     has_kerchunk,
+    requires_icechunk,
     requires_kerchunk,
     requires_network,
-    requires_zarr_python,
+    requires_scipy,
     slow_test,
 )
 from virtualizarr.tests.utils import PYTEST_TMP_DIRECTORY_URL_PREFIX
@@ -179,7 +180,6 @@ def roundtrip_as_in_memory_icechunk(
     )
 
 
-@requires_zarr_python
 @pytest.mark.parametrize(
     "roundtrip_func",
     [
@@ -551,3 +551,49 @@ def test_roundtrip_dataset_with_multiple_compressors():
         ) as observed,
     ):
         xr.testing.assert_allclose(expected, observed)
+
+
+@requires_scipy
+@requires_icechunk
+def test_subchunk_slice_netcdf3_through_icechunk_roundtrip(
+    netcdf3_file, local_registry
+):
+    # End-to-end check of the uncompressed sub-chunk slicing path: parse a netCDF3
+    # file (whose variables become single multi-row chunks), .isel within a chunk,
+    # write to icechunk, read back, and confirm bytes match a direct netCDF3 read +
+    # slice on the same file.
+    data = np.arange(32, dtype=np.float64).reshape(8, 4)
+    nc_path = netcdf3_file(xr.Dataset({"foo": (["time", "x"], data)}))
+    nc_url = f"file://{nc_path}"
+
+    with open_virtual_dataset(
+        url=nc_url, parser=NetCDF3Parser(), registry=local_registry
+    ) as vds:
+        # netCDF3 puts all rows in one source chunk; this slice is sub-chunk on axis 0.
+        sliced_vds = vds.isel(time=slice(1, 3))
+
+        storage = icechunk.Storage.new_in_memory()
+        config = icechunk.RepositoryConfig.default()
+        container = icechunk.VirtualChunkContainer(
+            url_prefix=PYTEST_TMP_DIRECTORY_URL_PREFIX,
+            store=icechunk.local_filesystem_store(PYTEST_TMP_DIRECTORY_URL_PREFIX),
+        )
+        config.set_virtual_chunk_container(container)
+        repo = icechunk.Repository.create(
+            storage=storage,
+            config=config,
+            authorize_virtual_chunk_access={PYTEST_TMP_DIRECTORY_URL_PREFIX: None},
+        )
+        session = repo.writable_session("main")
+        sliced_vds.vz.to_icechunk(session.store)
+        session.commit("sub-chunk slice")
+
+    read_session = repo.readonly_session("main")
+    with (
+        xr.open_zarr(
+            read_session.store, zarr_format=3, consolidated=False
+        ) as roundtripped,
+        xr.open_dataset(nc_path) as direct,
+    ):
+        expected = direct.isel(time=slice(1, 3))
+        xrt.assert_identical(roundtripped.load(), expected.load())