diff --git a/changes/4503.bugfix.md b/changes/4503.bugfix.md new file mode 100644 index 0000000000..245ef25820 --- /dev/null +++ b/changes/4503.bugfix.md @@ -0,0 +1 @@ +Writing one field of a structured array with `fields=` no longer overwrites the other fields of each selected record. A tuple naming several fields is now rejected for writes, like a list. diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 06f7338755..0ced32074c 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -5846,6 +5846,8 @@ async def _set_selection( # check fields are sensible check_fields(fields, dtype) fields = check_no_multi_fields(fields) + # the dtype of the values being written: a single field, or the whole record + value_dtype = check_fields(fields, dtype) # check value shape if np.isscalar(value): @@ -5854,21 +5856,39 @@ async def _set_selection( # TODO: need to handle array types that don't support __array_function__ # like PyTorch and JAX array_like_ = cast("np._typing._SupportsArrayFunc", array_like) - value = np.asanyarray(value, dtype=dtype, like=array_like_) + value = np.asanyarray(value, dtype=value_dtype, like=array_like_) else: if not hasattr(value, "shape"): - value = np.asarray(value, dtype) + value = np.asarray(value, value_dtype) # assert ( # value.shape == indexer.shape # ), f"shape of value doesn't match indexer shape. Expected {indexer.shape}, got {value.shape}" - if not hasattr(value, "dtype") or value.dtype.name != dtype.name: + if not hasattr(value, "dtype") or value.dtype.name != value_dtype.name: if hasattr(value, "astype"): # Handle things that are already NDArrayLike more efficiently - value = value.astype(dtype=dtype, order="A") + value = value.astype(dtype=value_dtype, order="A") else: - value = np.array(value, dtype=dtype, order="A") + value = np.array(value, dtype=value_dtype, order="A") value = cast("NDArrayLike", value) + if fields: + # A field write must preserve the other fields of each record, so read the + # selected records, replace the field, and write the whole records back. + # Missing chunks are read as the fill value, as for any partial chunk write. + records = prototype.nd_buffer.empty(shape=indexer.shape, dtype=dtype) + await _get_selection( + store_path, + metadata, + codec_pipeline, + replace(config, read_missing_chunks=True), + chunk_grid, + indexer, + prototype=prototype, + out=records, + ) + records[fields] = value + value = records.as_ndarray_like() + # We accept any ndarray like object from the user and convert it # to an NDBuffer (or subclass). From this point onwards, we only pass # Buffer and NDBuffer between components. diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 8ea59eb960..4a7553d06f 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -1538,12 +1538,12 @@ def check_fields(fields: Fields | None, dtype: np.dtype[Any]) -> np.dtype[Any]: return dtype -def check_no_multi_fields(fields: Fields | None) -> Fields | None: - if isinstance(fields, list): - if len(fields) == 1: - return fields[0] - elif len(fields) > 1: +def check_no_multi_fields(fields: Fields | None) -> str | None: + if isinstance(fields, list | tuple): + if len(fields) > 1: raise IndexError("multiple fields are not supported for this operation") + # an empty sequence selects no fields, i.e. the whole record + return fields[0] if fields else None return fields diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 785f089513..0ca9fde64e 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -2044,9 +2044,9 @@ def test_get_selections_with_fields(store: StorePath) -> None: z.get_basic_selection(Ellipsis, fields=slice(None)) # type: ignore[arg-type] -@pytest.mark.xfail(reason="fields are not supported in v3") +@pytest.mark.filterwarnings("ignore::zarr.errors.UnstableSpecificationWarning") def test_set_selections_with_fields(store: StorePath) -> None: - """Would verify that basic, orthogonal, coordinate, and mask set-selections with structured-array `fields` correctly write individual fields and reject multi-field assignment (xfail: fields unsupported in v3).""" + """Basic, orthogonal, coordinate, and mask set-selections with structured-array `fields` write only the named field and reject multi-field assignment.""" v = np.array( [("aaa", 1, 4.2), ("bbb", 2, 8.4), ("ccc", 3, 12.6)], dtype=[("foo", "S3"), ("bar", "i4"), ("baz", "f8")], @@ -2073,11 +2073,11 @@ def test_set_selections_with_fields(store: StorePath) -> None: with pytest.raises(IndexError): z.set_basic_selection(Ellipsis, v, fields=fields) with pytest.raises(IndexError): - z.set_orthogonal_selection([0, 2], v, fields=fields) # type: ignore[arg-type] + z.set_orthogonal_selection([0, 2], v[[0, 2]], fields=fields) # type: ignore[arg-type] with pytest.raises(IndexError): - z.set_coordinate_selection([0, 2], v, fields=fields) + z.set_coordinate_selection([0, 2], v[[0, 2]], fields=fields) with pytest.raises(IndexError): - z.set_mask_selection([True, False, True], v, fields=fields) # type: ignore[arg-type] + z.set_mask_selection([True, False, True], v[[0, 2]], fields=fields) # type: ignore[arg-type] else: if isinstance(fields, list) and len(fields) == 1: @@ -2131,6 +2131,55 @@ def test_set_selections_with_fields(store: StorePath) -> None: assert_array_equal(a, z[:]) +@pytest.mark.parametrize( + ("zarr_format", "shards"), [(2, None), (3, None), (3, (6,))], ids=["v2", "v3", "v3-sharded"] +) +@pytest.mark.parametrize("read_missing_chunks", [True, False]) +def test_set_field_preserves_other_fields( + zarr_format: ZarrFormat, shards: tuple[int, ...] | None, read_missing_chunks: bool +) -> None: + """Writing one field of a structured array leaves the other fields of each record unchanged, including in chunks that were never written.""" + dtype = np.dtype([("a", "i4"), ("b", "f8"), ("c", "i2")]) + data = np.array([(i, i * 1.5, -i) for i in range(10)], dtype=dtype) + z = zarr.create_array( + {}, + shape=(12,), + chunks=(3,), + shards=shards, + dtype=dtype, + zarr_format=zarr_format, + fill_value=None if zarr_format == 2 else 0, + config={"read_missing_chunks": read_missing_chunks}, + ) + # the last chunk (elements 9-11) is only partly written, the last shard is missing + z[:10] = data + expected = np.zeros(12, dtype=dtype) + expected[:10] = data + + z.set_basic_selection(slice(2, 8), np.full(6, 7.0), fields="b") + expected["b"][2:8] = 7.0 + z.set_basic_selection(4, 9, fields=["a"]) + expected["a"][4] = 9 + z.set_orthogonal_selection([0, 11], np.array([-5, -6], dtype="i2"), fields="c") + expected["c"][[0, 11]] = [-5, -6] + z.set_coordinate_selection([10, 1], 3.25, fields="b") + expected["b"][[10, 1]] = 3.25 + mask = np.zeros(12, dtype=bool) + mask[[5, 9]] = True + z.set_mask_selection(mask, 8, fields="a") + expected["a"][mask] = 8 + + assert_array_equal(z[:], expected) + + +def test_set_multiple_fields_as_tuple_raises() -> None: + """A tuple naming several fields is rejected for writes, like a list.""" + dtype = np.dtype([("a", "i4"), ("b", "f8")]) + z = zarr.create_array({}, shape=(4,), chunks=(2,), dtype=dtype, fill_value=0) + with pytest.raises(IndexError, match="multiple fields are not supported"): + z.set_basic_selection(slice(None), 1, fields=("a", "b")) + + def test_slice_selection_uints() -> None: """make_slice_selection accepts unsigned integer indices without error and produces correct shape.""" arr = np.arange(24).reshape((4, 6))