Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changes/3983.bugfix.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Fixed field selection for Zarr v2 arrays with structured dtypes.
10 changes: 9 additions & 1 deletion src/zarr/core/array.py
Original file line number Diff line number Diff line change
Expand Up @@ -5397,6 +5397,7 @@ async def _get_selection(

# check fields are sensible
out_dtype = check_fields(fields, dtype)
has_fields = bool(fields)

# setup output buffer
if out is not None:
Expand All @@ -5414,6 +5415,11 @@ async def _get_selection(
dtype=out_dtype,
order=order,
)
read_buffer = (
prototype.nd_buffer.empty(shape=indexer.shape, dtype=dtype, order=order)
if has_fields
else out_buffer
)
if product(indexer.shape) > 0:
# need to use the order from the metadata for v2
_config = config
Expand Down Expand Up @@ -5446,7 +5452,7 @@ async def _get_selection(
)
for chunk_coords, chunk_selection, out_selection, is_complete_chunk in indexed_chunks
],
out_buffer,
read_buffer,
drop_axes=indexer.drop_axes,
)
if _config.read_missing_chunks is False:
Expand All @@ -5464,6 +5470,8 @@ async def _get_selection(
f"missing chunks with the fill value.\n"
f"Missing chunks:\n{chunks_str}"
)
if has_fields:
out_buffer[...] = read_buffer[fields]
if isinstance(indexer, BasicIndexer) and indexer.shape == ():
return out_buffer.as_scalar()
return out_buffer.as_ndarray_like()
Expand Down
18 changes: 17 additions & 1 deletion tests/test_v2.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import json
from pathlib import Path
from typing import Any, Literal
from typing import Any, Literal, cast

import numpy as np
import pytest
Expand Down Expand Up @@ -252,6 +252,22 @@ def test_structured_dtype_roundtrip(fill_value: float | bytes, tmp_path: Path) -
assert (a == za[:]).all()


def test_structured_dtype_field_selection(tmp_path: Path) -> None:
a = np.array(
[("Rex", 9, 81.0), ("Fido", 3, 27.0)],
dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
)
za = zarr.create_array(store=tmp_path / "data.zarr", data=a, zarr_format=2)

np.testing.assert_array_equal(cast(Any, za)["name"], a["name"])
np.testing.assert_array_equal(
za.get_basic_selection(slice(1, None), fields="age"), a["age"][1:]
)
np.testing.assert_array_equal(
za.get_basic_selection(Ellipsis, fields=["name", "age"]), a[["name", "age"]]
)


@pytest.mark.parametrize(
(
"fill_value",
Expand Down
Loading