From 03c6f77e551da21c54f440ca03c84e4b9650e74d Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 21:39:46 +0000 Subject: [PATCH 1/6] zarr3: support the offset chunk key encoding The `offset` encoding stores a chunk under the `default` key of its grid index plus a fixed per-dimension integer offset, so that an array grows in the negative direction of a dimension by whole chunks as a metadata-only change (zarr-extensions chunk-key-encodings/offset). Adds the kind to `ChunkKeyEncoding` with its `offset` vector (required for this kind, rejected for the others, length checked against the rank in ValidateMetadata), applies it in FormatKey and GetChunkStorageKey, and accepts a leading `-` in ParseKey. MinGridIndexForLexicographicalOrder splits off every grid index whose key is negative, since negative base-10 components do not sort numerically, and shifts the digit-count rule for the rest. Schema docs and tests, including a growth that rewrites only zarr.json and a shrink that deletes through key ranges. --- tensorstore/driver/zarr3/driver.cc | 46 ++++-- tensorstore/driver/zarr3/driver_test.cc | 174 +++++++++++++++++++++- tensorstore/driver/zarr3/index.rst | 2 + tensorstore/driver/zarr3/metadata.cc | 80 +++++++--- tensorstore/driver/zarr3/metadata.h | 16 +- tensorstore/driver/zarr3/metadata_test.cc | 98 ++++++++++-- tensorstore/driver/zarr3/schema.yml | 44 ++++++ 7 files changed, 410 insertions(+), 50 deletions(-) diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index c4f1763e6..70cb370a0 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -34,7 +35,6 @@ #include "absl/strings/str_format.h" #include "absl/time/clock.h" #include "absl/time/time.h" -#include #include "tensorstore/array.h" #include "tensorstore/array_storage_statistics.h" #include "tensorstore/box.h" @@ -402,7 +402,7 @@ class DataCacheBase *static_cast(initial_metadata().get()); const DimensionIndex rank = metadata.rank; char separator = metadata.chunk_key_encoding.separator; - if (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kDefault) { + if (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2) { key = absl::StrCat( key_prefix_, "c", rank != 0 ? std::string_view(&separator, 1) : std::string_view()); @@ -414,8 +414,10 @@ class DataCacheBase } internal::FormatGridIndexKeyWithDimensionSeparator( key, separator, - [](std::string& out, DimensionIndex dim, Index grid_index) { - absl::StrAppend(&out, grid_index); + [&](std::string& out, DimensionIndex dim, Index grid_index) { + absl::StrAppend( + &out, + grid_index + metadata.chunk_key_encoding.GridIndexOffset(dim)); }, rank, grid_indices.subspan(0, @@ -430,16 +432,24 @@ class DataCacheBase assert(!grid_indices.empty()); key.remove_prefix( key_prefix_.size() + - (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kDefault ? 2 - : 0)); + (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2 ? 2 : 0)); if (!internal::ParseGridIndexKeyWithDimensionSeparator( metadata.chunk_key_encoding.separator, - [](std::string_view part, DimensionIndex dim, Index& grid_index) { - if (part.empty() || !absl::ascii_isdigit(part.front()) || - !absl::ascii_isdigit(part.back()) || - !absl::SimpleAtoi(part, &grid_index)) { + [&](std::string_view part, DimensionIndex dim, Index& grid_index) { + // Under the `offset` encoding a key component may be negative; + // `-0` is never produced and is rejected. + const bool negative = metadata.chunk_key_encoding.kind == + ChunkKeyEncoding::kOffset && + part.size() > 1 && part.front() == '-'; + const std::string_view digits = negative ? part.substr(1) : part; + if (digits.empty() || !absl::ascii_isdigit(digits.front()) || + !absl::ascii_isdigit(digits.back()) || + !absl::SimpleAtoi(digits, &grid_index) || + (negative && grid_index == 0)) { return false; } + if (negative) grid_index = -grid_index; + grid_index -= metadata.chunk_key_encoding.GridIndexOffset(dim); return true; }, key, grid_indices.subspan(0, metadata.rank))) { @@ -451,18 +461,28 @@ class DataCacheBase Index MinGridIndexForLexicographicalOrder( DimensionIndex dim, IndexInterval grid_interval) const final { - return internal::MinValueWithMaxBase10Digits(grid_interval.exclusive_max()); + const auto& metadata = + *static_cast(initial_metadata().get()); + // Keys hold `grid_index + offset`. A negative value sorts before every + // non-negative one but negative values do not sort numerically among + // themselves, so every grid index with a negative key is split off on + // its own; the non-negative keys follow the base-10 digit-count rule, + // shifted back into grid coordinates. + const Index offset = metadata.chunk_key_encoding.GridIndexOffset(dim); + return std::max(-offset, internal::MinValueWithMaxBase10Digits( + grid_interval.exclusive_max() + offset) - + offset); } std::string GetChunkStorageKey(span cell_indices) override { const auto& metadata = *static_cast(initial_metadata().get()); - if (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kDefault) { + if (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2) { std::string key = absl::StrCat(key_prefix_, "c"); for (DimensionIndex i = 0; i < cell_indices.size(); ++i) { absl::StrAppend( &key, std::string_view(&metadata.chunk_key_encoding.separator, 1), - cell_indices[i]); + cell_indices[i] + metadata.chunk_key_encoding.GridIndexOffset(i)); } return key; } diff --git a/tensorstore/driver/zarr3/driver_test.cc b/tensorstore/driver/zarr3/driver_test.cc index c0c03b4ee..87dff62d6 100644 --- a/tensorstore/driver/zarr3/driver_test.cc +++ b/tensorstore/driver/zarr3/driver_test.cc @@ -14,12 +14,15 @@ /// End-to-end tests of the zarr3 driver. +#include +#include #include #include #include #include #include +#include #include #include #include @@ -27,14 +30,11 @@ #include #include -#include -#include #include "absl/status/status.h" #include "absl/strings/cord.h" #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/time/clock.h" -#include #include "tensorstore/array.h" #include "tensorstore/array_testutil.h" #include "tensorstore/batch.h" @@ -1219,6 +1219,174 @@ TEST(ChunkKeyEncodingTest, V2Rank2Slash) { ::testing::Optional(::testing::ElementsAre("0/0", "0/1", "1/0", "1/1"))); } +TEST(ChunkKeyEncodingTest, OffsetRank1) { + EXPECT_THAT(GetKeysForChunkKeyEncoding( + {{"name", "offset"}, {"configuration", {{"offset", {-1}}}}}, + Schema::Shape{{5}}, ChunkLayout::ChunkShape({3})), + ::testing::Optional(::testing::ElementsAre("c/-1", "c/0"))); +} + +TEST(ChunkKeyEncodingTest, OffsetRank2Dot) { + EXPECT_THAT( + GetKeysForChunkKeyEncoding( + {{"name", "offset"}, + {"configuration", {{"separator", "."}, {"offset", {1, -1}}}}}, + Schema::Shape{{5, 6}}, ChunkLayout::ChunkShape({3, 3})), + ::testing::Optional( + ::testing::ElementsAre("c.1.-1", "c.1.0", "c.2.-1", "c.2.0"))); +} + +TEST(ChunkKeyEncodingTest, OffsetZeroMatchesDefault) { + EXPECT_THAT(GetKeysForChunkKeyEncoding( + {{"name", "offset"}, {"configuration", {{"offset", {0, 0}}}}}, + Schema::Shape{{5, 6}}, ChunkLayout::ChunkShape({3, 3})), + ::testing::Optional( + ::testing::ElementsAre("c/0/0", "c/0/1", "c/1/0", "c/1/1"))); +} + +TEST(ChunkKeyEncodingTest, OffsetRankMismatch) { + EXPECT_THAT(GetKeysForChunkKeyEncoding( + {{"name", "offset"}, {"configuration", {{"offset", {0}}}}}, + Schema::Shape{{5, 6}}, ChunkLayout::ChunkShape({3, 3})), + tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*\"offset\".*")); +} + +TEST(ChunkKeyEncodingTest, OffsetReadsBackWhatWasWritten) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open( + { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"metadata", + {{"chunk_key_encoding", + {{"name", "offset"}, {"configuration", {{"offset", {-2}}}}}}}}, + }, + dtype_v, Schema::Shape({7}), ChunkLayout::ChunkShape({3}), + tensorstore::OpenMode::create) + .result()); + TENSORSTORE_ASSERT_OK(tensorstore::Write( + tensorstore::MakeArray({1, 2, 3, 4, 5, 6, 7}), store)); + EXPECT_THAT(tensorstore::Read(store).result(), + ::testing::Optional( + tensorstore::MakeArray({1, 2, 3, 4, 5, 6, 7}))); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto entries, tensorstore::kvstore::ListFuture(store.kvstore()).result()); + std::vector keys; + for (const auto& e : entries) keys.push_back(e.key); + EXPECT_THAT(keys, ::testing::UnorderedElementsAre("zarr.json", "c/-2", "c/-1", + "c/0")); +} + +// The point of the encoding: an array grows in the negative direction by +// editing `zarr.json` alone. Chunks written under the `default` encoding +// are read back at their new indices without having moved. +TEST(ChunkKeyEncodingTest, + OffsetGrowsInTheNegativeDirectionWithoutMovingChunks) { + auto context = Context::Default(); + ::nlohmann::json spec{ + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open(spec, context, dtype_v, Schema::Shape({6}), + ChunkLayout::ChunkShape({3}), + tensorstore::OpenMode::create) + .result()); + TENSORSTORE_ASSERT_OK(tensorstore::Write( + tensorstore::MakeArray({10, 11, 12, 13, 14, 15}), store)); + auto kvs = store.kvstore(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto read_result, tensorstore::kvstore::Read(kvs, "zarr.json").result()); + auto metadata = ::nlohmann::json::parse(std::string(read_result.value)); + const auto keys_before = [&] { + std::vector keys; + auto entries = tensorstore::kvstore::ListFuture(kvs).value(); + for (const auto& e : entries) { + if (e.key != "zarr.json") keys.push_back(e.key); + } + std::sort(keys.begin(), keys.end()); + return keys; + }(); + EXPECT_THAT(keys_before, ::testing::ElementsAre("c/0", "c/1")); + + // Two chunks in front: the whole operation. + metadata["shape"] = {12}; + metadata["chunk_key_encoding"] = {{"name", "offset"}, + {"configuration", {{"offset", {-2}}}}}; + TENSORSTORE_ASSERT_OK( + tensorstore::kvstore::Write(kvs, "zarr.json", absl::Cord(metadata.dump())) + .result()); + + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto grown, + tensorstore::Open(spec, context, tensorstore::OpenMode::open).result()); + EXPECT_THAT(grown.domain().shape(), ::testing::ElementsAre(12)); + EXPECT_THAT(tensorstore::Read(grown).result(), + ::testing::Optional(tensorstore::MakeArray( + {0, 0, 0, 0, 0, 0, 10, 11, 12, 13, 14, 15}))); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeArray({1, 2, 3}), + grown | tensorstore::Dims(0).SizedInterval(0, 3))); + EXPECT_THAT(tensorstore::Read(grown).result(), + ::testing::Optional(tensorstore::MakeArray( + {1, 2, 3, 0, 0, 0, 10, 11, 12, 13, 14, 15}))); + std::vector keys_after; + auto entries_after = tensorstore::kvstore::ListFuture(kvs).value(); + for (const auto& e : entries_after) { + if (e.key != "zarr.json") keys_after.push_back(e.key); + } + EXPECT_THAT(keys_after, + ::testing::UnorderedElementsAre("c/-2", "c/0", "c/1")); +} + +// Shrinking deletes chunks through key ranges, which must cope with the +// negative key components not sorting numerically. +TEST(ChunkKeyEncodingTest, OffsetResizeDeletesOnlyTheChunksOutsideTheNewShape) { + auto context = Context::Default(); + ::nlohmann::json spec{ + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"metadata", + {{"chunk_key_encoding", + {{"name", "offset"}, {"configuration", {{"offset", {-2, 0}}}}}}}}, + }; + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open(spec, context, dtype_v, Schema::Shape({12, 2}), + ChunkLayout::ChunkShape({3, 1}), + tensorstore::OpenMode::create) + .result()); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeScalarArray(7), store)); + const auto keys = [&] { + std::vector keys; + auto entries = tensorstore::kvstore::ListFuture(store.kvstore()).value(); + for (const auto& e : entries) { + if (e.key != "zarr.json") keys.push_back(e.key); + } + std::sort(keys.begin(), keys.end()); + return keys; + }; + EXPECT_THAT(keys(), + ::testing::ElementsAre("c/-1/0", "c/-1/1", "c/-2/0", "c/-2/1", + "c/0/0", "c/0/1", "c/1/0", "c/1/1")); + TENSORSTORE_ASSERT_OK( + tensorstore::Resize(store, + {{tensorstore::kImplicit, tensorstore::kImplicit}}, + {{6, tensorstore::kImplicit}}) + .result()); + EXPECT_THAT(keys(), + ::testing::ElementsAre("c/-1/0", "c/-1/1", "c/-2/0", "c/-2/1")); + TENSORSTORE_ASSERT_OK( + tensorstore::Resize( + store, {{tensorstore::kImplicit, tensorstore::kImplicit}}, {{3, 1}}) + .result()); + EXPECT_THAT(keys(), ::testing::ElementsAre("c/-2/0")); +} + TENSORSTORE_GLOBAL_INITIALIZER { tensorstore::internal::TestTensorStoreDriverSpecRoundtripOptions options; options.test_name = "zarr3"; diff --git a/tensorstore/driver/zarr3/index.rst b/tensorstore/driver/zarr3/index.rst index 0acfc92b1..d67e33d1a 100644 --- a/tensorstore/driver/zarr3/index.rst +++ b/tensorstore/driver/zarr3/index.rst @@ -74,6 +74,8 @@ metadata. .. json:schema:: driver/zarr3/ChunkKeyEncoding.v2 +.. json:schema:: driver/zarr3/ChunkKeyEncoding.offset + Mapping to TensorStore Schema ----------------------------- diff --git a/tensorstore/driver/zarr3/metadata.cc b/tensorstore/driver/zarr3/metadata.cc index a0977bfa3..402944cd4 100644 --- a/tensorstore/driver/zarr3/metadata.cc +++ b/tensorstore/driver/zarr3/metadata.cc @@ -456,26 +456,59 @@ TENSORSTORE_DEFINE_JSON_DEFAULT_BINDER(ChunkKeyEncoding, [](auto is_loading, jb::Projection<&Self::kind>(jb::Enum({ {Self::kDefault, "default"}, {Self::kV2, "v2"}, + {Self::kOffset, "offset"}, })), - jb::Object(jb::Member("separator", - jb::Projection<&Self::separator>(jb::DefaultValue( - [outer = obj](auto* obj) -> absl::Status { - switch (outer->kind) { - case Self::kDefault: - *obj = '/'; - break; - case Self::kV2: - *obj = '.'; - break; - default: - ABSL_UNREACHABLE(); - } - return absl::OkStatus(); - }, - jb::Enum({ - {'.', "."}, - {'/', "/"}, - })))))) // + jb::Object( + jb::Member("separator", + jb::Projection<&Self::separator>(jb::DefaultValue( + [outer = obj](auto* obj) -> absl::Status { + switch (outer->kind) { + case Self::kDefault: + case Self::kOffset: + *obj = '/'; + break; + case Self::kV2: + *obj = '.'; + break; + default: + ABSL_UNREACHABLE(); + } + return absl::OkStatus(); + }, + jb::Enum({ + {'.', "."}, + {'/', "/"}, + })))), + // Required for `offset`, and only valid there. The length is + // checked against the rank in `ValidateMetadata`, where the rank + // is known. + jb::Member("offset", + jb::Projection<&Self::offset>( + [outer = obj](auto is_loading, const auto& options, + auto* obj, auto* j) -> absl::Status { + if (outer->kind != Self::kOffset) { + if constexpr (is_loading) { + if (!j->is_discarded()) { + return absl::InvalidArgumentError( + "\"offset\" is only valid for the " + "\"offset\" chunk key encoding"); + } + obj->clear(); + } else { + *j = ::nlohmann::json::value_t::discarded; + } + return absl::OkStatus(); + } + if constexpr (is_loading) { + if (j->is_discarded()) { + return absl::InvalidArgumentError( + "\"offset\" must be specified for the " + "\"offset\" chunk key encoding"); + } + } + return jb::Array(jb::Integer())( + is_loading, options, obj, j); + })))) // )(is_loading, options, obj, j); }) @@ -838,6 +871,15 @@ absl::Status ValidateStructEndianness(const ZarrDType& dtype, absl::Status ValidateMetadata(ZarrMetadata& metadata) { DeriveFieldShape(metadata); + if (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kOffset && + static_cast(metadata.chunk_key_encoding.offset.size()) != + metadata.rank) { + return absl::InvalidArgumentError(absl::StrCat( + "Expected \"offset\" of the chunk key encoding to have length ", + metadata.rank, ", but received: ", + ::nlohmann::json(metadata.chunk_key_encoding.offset).dump())); + } + // Skip when codecs are already resolved: the create path runs the // pre-Resolve check itself in `GetNewMetadata`, and post-Resolve the // leaf bytes spec has its `endian` stripped. diff --git a/tensorstore/driver/zarr3/metadata.h b/tensorstore/driver/zarr3/metadata.h index 0bfb8071e..a4c1870c4 100644 --- a/tensorstore/driver/zarr3/metadata.h +++ b/tensorstore/driver/zarr3/metadata.h @@ -23,6 +23,7 @@ #include #include +#include #include #include #include @@ -30,7 +31,6 @@ #include #include "absl/status/status.h" -#include #include "tensorstore/array.h" #include "tensorstore/chunk_layout.h" #include "tensorstore/codec_spec.h" @@ -61,16 +61,28 @@ struct ChunkKeyEncoding { enum Kind { kDefault, kV2, + // `default` keys of `grid_index + offset`: + // https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset + kOffset, }; Kind kind; char separator; + // For `kOffset`, one value per dimension, added to the grid index before + // it is formatted. Empty for the other kinds. + std::vector offset; + + // Offset applied to the grid index along `dim` when formatting a key. + Index GridIndexOffset(DimensionIndex dim) const { + return kind == kOffset ? offset[dim] : 0; + } TENSORSTORE_DECLARE_JSON_DEFAULT_BINDER(ChunkKeyEncoding, internal_json_binding::NoOptions, tensorstore::IncludeDefaults) friend bool operator==(const ChunkKeyEncoding& a, const ChunkKeyEncoding& b) { - return a.kind == b.kind && a.separator == b.separator; + return a.kind == b.kind && a.separator == b.separator && + a.offset == b.offset; } friend bool operator!=(const ChunkKeyEncoding& a, const ChunkKeyEncoding& b) { return !(a == b); diff --git a/tensorstore/driver/zarr3/metadata_test.cc b/tensorstore/driver/zarr3/metadata_test.cc index 22b2850e2..d527c66cc 100644 --- a/tensorstore/driver/zarr3/metadata_test.cc +++ b/tensorstore/driver/zarr3/metadata_test.cc @@ -14,22 +14,22 @@ #include "tensorstore/driver/zarr3/metadata.h" +#include +#include #include #include #include #include +#include #include #include #include #include #include -#include -#include #include "absl/base/casts.h" #include "absl/status/status.h" -#include #include "tensorstore/array.h" #include "tensorstore/array_testutil.h" #include "tensorstore/chunk_layout.h" @@ -72,6 +72,7 @@ using ::tensorstore::dtypes::float16_t; using ::tensorstore::dtypes::float32_t; using ::tensorstore::dtypes::float64_t; using ::tensorstore::internal::uint_t; +using ::tensorstore::internal_zarr3::ChunkKeyEncoding; using ::tensorstore::internal_zarr3::FillValueJsonBinder; using ::tensorstore::internal_zarr3::GetSpecRankAndFieldInfo; using ::tensorstore::internal_zarr3::GetVoidMetadata; @@ -151,6 +152,75 @@ TEST(MetadataTest, ParseValidNoDimensionNames) { EXPECT_THAT(metadata.user_attributes, MatchesJson({{"a", "b"}, {"c", "d"}})); } +TEST(MetadataTest, OffsetChunkKeyEncoding) { + auto json = GetBasicMetadata(); + json["chunk_key_encoding"] = { + {"name", "offset"}, + {"configuration", {{"separator", "."}, {"offset", {-2, 0, 5}}}}}; + tensorstore::TestJsonBinderRoundTripJsonOnly({json}); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, ZarrMetadata::FromJson(json)); + EXPECT_EQ(metadata.chunk_key_encoding.kind, ChunkKeyEncoding::kOffset); + EXPECT_EQ(metadata.chunk_key_encoding.separator, '.'); + EXPECT_THAT(metadata.chunk_key_encoding.offset, + ::testing::ElementsAre(-2, 0, 5)); + EXPECT_EQ(metadata.chunk_key_encoding.GridIndexOffset(0), -2); + EXPECT_EQ(metadata.chunk_key_encoding.GridIndexOffset(2), 5); +} + +TEST(MetadataTest, OffsetChunkKeyEncodingSeparatorDefaultsToSlash) { + auto json = GetBasicMetadata(); + json["chunk_key_encoding"] = {{"name", "offset"}, + {"configuration", {{"offset", {0, 0, 0}}}}}; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, ZarrMetadata::FromJson(json)); + EXPECT_EQ(metadata.chunk_key_encoding.separator, '/'); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto out, metadata.ToJson()); + // The default separator is omitted on output, as for `default`. + EXPECT_THAT(out["chunk_key_encoding"], + MatchesJson({{"name", "offset"}, + {"configuration", {{"offset", {0, 0, 0}}}}})); +} + +TEST(MetadataTest, OffsetChunkKeyEncodingRankMismatch) { + auto json = GetBasicMetadata(); + json["chunk_key_encoding"] = {{"name", "offset"}, + {"configuration", {{"offset", {0, 0}}}}}; + EXPECT_THAT(ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*\"offset\".*length 3.*"))); +} + +TEST(MetadataTest, OffsetChunkKeyEncodingRequiresOffset) { + auto json = GetBasicMetadata(); + json["chunk_key_encoding"] = {{"name", "offset"}}; + EXPECT_THAT(ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*\"offset\" must be specified.*"))); + json["chunk_key_encoding"] = {{"name", "offset"}, + {"configuration", {{"separator", "."}}}}; + EXPECT_THAT(ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*\"offset\" must be specified.*"))); +} + +TEST(MetadataTest, OffsetIsOnlyValidForTheOffsetChunkKeyEncoding) { + auto json = GetBasicMetadata(); + json["chunk_key_encoding"] = {{"name", "default"}, + {"configuration", {{"offset", {0, 0, 0}}}}}; + EXPECT_THAT(ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*\"offset\" is only valid.*"))); +} + +TEST(MetadataTest, OffsetChunkKeyEncodingEquality) { + ChunkKeyEncoding a{ChunkKeyEncoding::kOffset, '/', {-1, 0}}; + ChunkKeyEncoding b{ChunkKeyEncoding::kOffset, '/', {-1, 0}}; + ChunkKeyEncoding c{ChunkKeyEncoding::kOffset, '/', {-2, 0}}; + ChunkKeyEncoding d{ChunkKeyEncoding::kDefault, '/', {}}; + EXPECT_EQ(a, b); + EXPECT_NE(a, c); + EXPECT_NE(a, d); +} + TEST(MetadataTest, UnknownExtensionAttribute) { auto json = GetBasicMetadata(); json["foo"] = 42; @@ -1169,9 +1239,10 @@ TEST(StructEndianTest, ModernStructRequiresExplicitEndianForMultiByteField) { {{"name", "y"}, {"data_type", "int16"}}})}}}}; metadata_json["fill_value"] = {{"x", 0}, {"y", 0}}; metadata_json["codecs"] = ::nlohmann::json::array_t{{{"name", "bytes"}}}; - EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - ".*specify an explicit `endian`.*")); + EXPECT_THAT( + ZarrMetadata::FromJson(metadata_json), + tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*specify an explicit `endian`.*"))); } // All single-byte fields: no byte order to disambiguate. @@ -1267,9 +1338,10 @@ TEST(StructEndianTest, ShardedModernStructWithoutLeafEndianRejected) { {{"name", "y"}, {"data_type", "int16"}}})}}}}; metadata_json["fill_value"] = {{"x", 0}, {"y", 0}}; metadata_json["codecs"] = ShardingWrappedBytes(std::nullopt); - EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - ".*specify an explicit `endian`.*")); + EXPECT_THAT( + ZarrMetadata::FromJson(metadata_json), + tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*specify an explicit `endian`.*"))); } // Legacy default must be injected at the leaf bytes codec, not the @@ -1305,8 +1377,8 @@ TEST(StructFillValueFormTest, ModernStructRejectsBase64Fill) { metadata_json["codecs"] = ::nlohmann::json::array_t{ {{"name", "bytes"}, {"configuration", {{"endian", "little"}}}}}; EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - ".*Expected object.*")); + tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*Expected object.*"))); } TEST(StructFillValueFormTest, ModernStructRejectsArrayFill) { @@ -1321,8 +1393,8 @@ TEST(StructFillValueFormTest, ModernStructRejectsArrayFill) { metadata_json["codecs"] = ::nlohmann::json::array_t{ {{"name", "bytes"}, {"configuration", {{"endian", "little"}}}}}; EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - ".*Expected object.*")); + tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, + MatchesRegex(".*Expected object.*"))); } // Verifies the deferred-parse path resolves endian before decoding diff --git a/tensorstore/driver/zarr3/schema.yml b/tensorstore/driver/zarr3/schema.yml index 4b1bc20ce..6f44b9f2e 100644 --- a/tensorstore/driver/zarr3/schema.yml +++ b/tensorstore/driver/zarr3/schema.yml @@ -295,6 +295,50 @@ definitions: - const: "." default: "." title: "Separator character between dimensions" + chunk_key_encoding_offset: + $id: driver/zarr3/ChunkKeyEncoding.offset + title: Offset chunk key encoding. + description: | + The `~driver/zarr3/ChunkKeyEncoding.default` encoding applied to the + chunk grid position plus a fixed per-dimension integer offset, so that + an array can grow in the negative direction of a dimension by whole + chunks without any existing chunk changing its key. Refer to `the + extension registry + `__ + for details. + allOf: + - $ref: driver/zarr3/ChunkKeyEncoding + - type: object + properties: + name: + const: "offset" + configuration: + type: object + properties: + separator: + type: string + oneOf: + - const: "/" + - const: "." + default: "/" + title: "Separator character between dimensions" + offset: + type: array + items: + type: integer + title: | + Offset added to the chunk grid position along each dimension + before it is formatted as a key. + description: | + Must have one element per dimension. An element may be + negative, in which case the corresponding key component carries + a leading ``-``. With every element ``0`` the keys are + identical to those of the + `~driver/zarr3/ChunkKeyEncoding.default` encoding. + examples: + - [-2, 0, 0] + required: + - offset codec_chain: $id: driver/zarr3/CodecChain title: Specifies a chain of codecs. From 3c5a7304cfe4e12f7652eec8d39d57ea4fa5a70c Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 21:45:42 +0000 Subject: [PATCH 2/6] zarr3: bound the key offset lookup to the array rank --- tensorstore/driver/zarr3/driver.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index 70cb370a0..e1b3a58b5 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -480,9 +480,12 @@ class DataCacheBase if (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2) { std::string key = absl::StrCat(key_prefix_, "c"); for (DimensionIndex i = 0; i < cell_indices.size(); ++i) { + const Index offset = + i < metadata.rank ? metadata.chunk_key_encoding.GridIndexOffset(i) + : 0; absl::StrAppend( &key, std::string_view(&metadata.chunk_key_encoding.separator, 1), - cell_indices[i] + metadata.chunk_key_encoding.GridIndexOffset(i)); + cell_indices[i] + offset); } return key; } From ead79b258432eccd2412060fc503c47fcaa3b652 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 22:32:12 +0000 Subject: [PATCH 3/6] zarr3: shorten the offset encoding comments --- tensorstore/driver/zarr3/driver.cc | 8 +------- tensorstore/driver/zarr3/driver_test.cc | 6 ------ tensorstore/driver/zarr3/metadata.cc | 6 ++---- tensorstore/driver/zarr3/metadata.h | 5 +---- tensorstore/driver/zarr3/metadata_test.cc | 1 - 5 files changed, 4 insertions(+), 22 deletions(-) diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index e1b3a58b5..b9a29707c 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -436,8 +436,6 @@ class DataCacheBase if (!internal::ParseGridIndexKeyWithDimensionSeparator( metadata.chunk_key_encoding.separator, [&](std::string_view part, DimensionIndex dim, Index& grid_index) { - // Under the `offset` encoding a key component may be negative; - // `-0` is never produced and is rejected. const bool negative = metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kOffset && part.size() > 1 && part.front() == '-'; @@ -463,11 +461,7 @@ class DataCacheBase DimensionIndex dim, IndexInterval grid_interval) const final { const auto& metadata = *static_cast(initial_metadata().get()); - // Keys hold `grid_index + offset`. A negative value sorts before every - // non-negative one but negative values do not sort numerically among - // themselves, so every grid index with a negative key is split off on - // its own; the non-negative keys follow the base-10 digit-count rule, - // shifted back into grid coordinates. + // Negative key components do not sort numerically: split them off singly. const Index offset = metadata.chunk_key_encoding.GridIndexOffset(dim); return std::max(-offset, internal::MinValueWithMaxBase10Digits( grid_interval.exclusive_max() + offset) - diff --git a/tensorstore/driver/zarr3/driver_test.cc b/tensorstore/driver/zarr3/driver_test.cc index 87dff62d6..20814f932 100644 --- a/tensorstore/driver/zarr3/driver_test.cc +++ b/tensorstore/driver/zarr3/driver_test.cc @@ -1279,9 +1279,6 @@ TEST(ChunkKeyEncodingTest, OffsetReadsBackWhatWasWritten) { "c/0")); } -// The point of the encoding: an array grows in the negative direction by -// editing `zarr.json` alone. Chunks written under the `default` encoding -// are read back at their new indices without having moved. TEST(ChunkKeyEncodingTest, OffsetGrowsInTheNegativeDirectionWithoutMovingChunks) { auto context = Context::Default(); @@ -1312,7 +1309,6 @@ TEST(ChunkKeyEncodingTest, }(); EXPECT_THAT(keys_before, ::testing::ElementsAre("c/0", "c/1")); - // Two chunks in front: the whole operation. metadata["shape"] = {12}; metadata["chunk_key_encoding"] = {{"name", "offset"}, {"configuration", {{"offset", {-2}}}}}; @@ -1342,8 +1338,6 @@ TEST(ChunkKeyEncodingTest, ::testing::UnorderedElementsAre("c/-2", "c/0", "c/1")); } -// Shrinking deletes chunks through key ranges, which must cope with the -// negative key components not sorting numerically. TEST(ChunkKeyEncodingTest, OffsetResizeDeletesOnlyTheChunksOutsideTheNewShape) { auto context = Context::Default(); ::nlohmann::json spec{ diff --git a/tensorstore/driver/zarr3/metadata.cc b/tensorstore/driver/zarr3/metadata.cc index 402944cd4..452b08d74 100644 --- a/tensorstore/driver/zarr3/metadata.cc +++ b/tensorstore/driver/zarr3/metadata.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -42,7 +43,6 @@ #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" -#include #include "tensorstore/array.h" #include "tensorstore/box.h" #include "tensorstore/chunk_layout.h" @@ -479,9 +479,7 @@ TENSORSTORE_DEFINE_JSON_DEFAULT_BINDER(ChunkKeyEncoding, [](auto is_loading, {'.', "."}, {'/', "/"}, })))), - // Required for `offset`, and only valid there. The length is - // checked against the rank in `ValidateMetadata`, where the rank - // is known. + // Length is checked against the rank in `ValidateMetadata`. jb::Member("offset", jb::Projection<&Self::offset>( [outer = obj](auto is_loading, const auto& options, diff --git a/tensorstore/driver/zarr3/metadata.h b/tensorstore/driver/zarr3/metadata.h index a4c1870c4..9651c5428 100644 --- a/tensorstore/driver/zarr3/metadata.h +++ b/tensorstore/driver/zarr3/metadata.h @@ -61,17 +61,14 @@ struct ChunkKeyEncoding { enum Kind { kDefault, kV2, - // `default` keys of `grid_index + offset`: // https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset kOffset, }; Kind kind; char separator; - // For `kOffset`, one value per dimension, added to the grid index before - // it is formatted. Empty for the other kinds. + // One value per dimension for `kOffset`, empty otherwise. std::vector offset; - // Offset applied to the grid index along `dim` when formatting a key. Index GridIndexOffset(DimensionIndex dim) const { return kind == kOffset ? offset[dim] : 0; } diff --git a/tensorstore/driver/zarr3/metadata_test.cc b/tensorstore/driver/zarr3/metadata_test.cc index d527c66cc..ce43b4c94 100644 --- a/tensorstore/driver/zarr3/metadata_test.cc +++ b/tensorstore/driver/zarr3/metadata_test.cc @@ -174,7 +174,6 @@ TEST(MetadataTest, OffsetChunkKeyEncodingSeparatorDefaultsToSlash) { TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, ZarrMetadata::FromJson(json)); EXPECT_EQ(metadata.chunk_key_encoding.separator, '/'); TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto out, metadata.ToJson()); - // The default separator is omitted on output, as for `default`. EXPECT_THAT(out["chunk_key_encoding"], MatchesJson({{"name", "offset"}, {"configuration", {{"offset", {0, 0, 0}}}}})); From 92e1fec29249be5181ffaf149e066a3dda5e96ee Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 23:12:25 +0000 Subject: [PATCH 4/6] zarr3: return early when every key in the interval is negative --- tensorstore/driver/zarr3/driver.cc | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index b9a29707c..bf45e5bf6 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -463,9 +463,13 @@ class DataCacheBase *static_cast(initial_metadata().get()); // Negative key components do not sort numerically: split them off singly. const Index offset = metadata.chunk_key_encoding.GridIndexOffset(dim); - return std::max(-offset, internal::MinValueWithMaxBase10Digits( - grid_interval.exclusive_max() + offset) - - offset); + const Index exclusive_max = grid_interval.exclusive_max(); + if (exclusive_max + offset <= 0) { + return -offset; + } + return std::max( + -offset, + internal::MinValueWithMaxBase10Digits(exclusive_max + offset) - offset); } std::string GetChunkStorageKey(span cell_indices) override { From 0e48a2fcfaf16483827c6d9d4b6b29fb9af8c578 Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Thu, 15 Oct 2026 23:27:52 +0000 Subject: [PATCH 5/6] zarr3: strip the right prefix length when parsing a rank-0 key --- tensorstore/driver/zarr3/driver.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index bf45e5bf6..49554467b 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -430,9 +430,10 @@ class DataCacheBase *static_cast(initial_metadata().get()); assert(key.size() > key_prefix_.size()); assert(!grid_indices.empty()); - key.remove_prefix( - key_prefix_.size() + - (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2 ? 2 : 0)); + key.remove_prefix(key_prefix_.size() + + (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2 + ? (metadata.rank != 0 ? 2 : 1) + : 0)); if (!internal::ParseGridIndexKeyWithDimensionSeparator( metadata.chunk_key_encoding.separator, [&](std::string_view part, DimensionIndex dim, Index& grid_index) { From 5ba73539c4cbf715d37f59e31ae627d1c5022abc Mon Sep 17 00:00:00 2001 From: Jeff Rhoades Date: Fri, 16 Oct 2026 23:01:26 +0000 Subject: [PATCH 6/6] zarr3: support the regular_bounded chunk grid instead of an offset key encoding Rework following zarr-extensions#79. The `offset` chunk key encoding is removed. `regular_bounded` is a regular chunk grid with an `origin` (the array coordinate where chunk 0 starts) and a `lower_bound` (the coordinate of the first element). The array's domain starts at `lower_bound`, and the internal grid space is the array coordinate minus `origin`, so chunk grid indices are the absolute ones and may be negative. The lower bound of a `regular_bounded` array is implicit, so `Resize` can move it; chunks that fall off the front are deleted through the existing box-difference path. `regular` arrays are unchanged. Key parsing accepts a leading `-` and rejects `-0`. The per-component fill value is broadcast over an unbounded domain, since chunks may now sit at negative internal coordinates. The chunk layout reports `origin` as the grid origin, and the compatibility key ignores `lower_bound`, as it ignores `shape`. --- tensorstore/driver/zarr3/BUILD | 1 + tensorstore/driver/zarr3/chunk_cache.cc | 14 +- tensorstore/driver/zarr3/driver.cc | 80 +++---- tensorstore/driver/zarr3/driver_test.cc | 253 ++++++++++++---------- tensorstore/driver/zarr3/index.rst | 2 - tensorstore/driver/zarr3/metadata.cc | 226 ++++++++++++------- tensorstore/driver/zarr3/metadata.h | 30 ++- tensorstore/driver/zarr3/metadata_test.cc | 113 ++++------ tensorstore/driver/zarr3/schema.yml | 69 ++---- 9 files changed, 432 insertions(+), 356 deletions(-) diff --git a/tensorstore/driver/zarr3/BUILD b/tensorstore/driver/zarr3/BUILD index 5a177ba7c..8aee5fbac 100644 --- a/tensorstore/driver/zarr3/BUILD +++ b/tensorstore/driver/zarr3/BUILD @@ -165,6 +165,7 @@ tensorstore_cc_test( ":zarr3", "//tensorstore", "//tensorstore:array", + "//tensorstore:array_storage_statistics", "//tensorstore:array_testutil", "//tensorstore:batch", "//tensorstore:box", diff --git a/tensorstore/driver/zarr3/chunk_cache.cc b/tensorstore/driver/zarr3/chunk_cache.cc index 708c4cdaa..e11ced95a 100644 --- a/tensorstore/driver/zarr3/chunk_cache.cc +++ b/tensorstore/driver/zarr3/chunk_cache.cc @@ -104,12 +104,18 @@ internal::ChunkGridSpecification CreateFieldGridSpecification( field.dtype); } - std::vector target_shape(chunked_rank, kInfIndex); - target_shape.insert(target_shape.end(), field.field_shape.begin(), - field.field_shape.end()); + // Unbounded in both directions: chunk grid indices may be negative. + Box<> fill_value_domain(total_rank); + for (DimensionIndex i = 0; i < chunked_rank; ++i) { + fill_value_domain[i] = IndexInterval::Infinite(); + } + for (size_t i = 0; i < field_rank; ++i) { + fill_value_domain[chunked_rank + i] = + IndexInterval::UncheckedSized(0, field.field_shape[i]); + } auto chunk_fill_value = - BroadcastArray(fill_value, BoxView<>(target_shape)).value(); + BroadcastArray(fill_value, BoxView<>(fill_value_domain)).value(); std::vector component_chunk_shape(chunk_shape.begin(), chunk_shape.end()); diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index 49554467b..875b9de5f 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -182,11 +182,17 @@ class ZarrDriverSpec const DimensionIndex total_rank = chunked_rank + field_shape.size(); IndexDomainBuilder builder(total_rank); std::fill_n(builder.origin().begin(), total_rank, Index{0}); + if (metadata_constraints.lower_bound) { + std::copy_n(metadata_constraints.lower_bound->begin(), chunked_rank, + builder.origin().begin()); + } std::copy_n(metadata_constraints.shape->begin(), chunked_rank, builder.shape().begin()); std::copy_n(field_shape.begin(), field_shape.size(), builder.shape().begin() + chunked_rank); - builder.implicit_lower_bounds(DimensionSet(false)); + builder.implicit_lower_bounds(metadata_constraints.lower_bound + ? DimensionSet::UpTo(chunked_rank) + : DimensionSet(false)); builder.implicit_upper_bounds(DimensionSet::UpTo(chunked_rank)); if (metadata_constraints.dimension_names) { for (DimensionIndex j = 0; j < chunked_rank; ++j) { @@ -353,6 +359,8 @@ class DataCacheBase implicit_upper_bounds = false; for (DimensionIndex i = 0; i < static_cast(metadata.shape.size()); ++i) { + bounds.origin()[i] = metadata.LowerBound(i) - metadata.GridOrigin(i); + implicit_lower_bounds[i] = metadata.chunk_grid_bounded; implicit_upper_bounds[i] = true; } if (bounds.rank() > static_cast(metadata.shape.size()) && @@ -380,10 +388,17 @@ class DataCacheBase assert(rank == new_inclusive_min.size()); assert(rank == new_exclusive_max.size()); for (DimensionIndex i = 0; i < rank; ++i) { - assert(ExplicitIndexOr(new_inclusive_min[i], 0) == 0); - const Index new_size = new_exclusive_max[i]; - if (new_size == kImplicit) continue; - new_metadata->shape[i] = new_size; + const Index origin = new_metadata->GridOrigin(i); + const Index old_min = new_metadata->LowerBound(i) - origin; + const Index new_min = ExplicitIndexOr(new_inclusive_min[i], old_min); + const Index new_max = ExplicitIndexOr(new_exclusive_max[i], + old_min + new_metadata->shape[i]); + if (new_metadata->chunk_grid_bounded) { + new_metadata->lower_bound[i] = new_min + origin; + } else { + assert(new_min == 0); + } + new_metadata->shape[i] = new_max - new_min; } return new_metadata; } @@ -402,7 +417,7 @@ class DataCacheBase *static_cast(initial_metadata().get()); const DimensionIndex rank = metadata.rank; char separator = metadata.chunk_key_encoding.separator; - if (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2) { + if (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kDefault) { key = absl::StrCat( key_prefix_, "c", rank != 0 ? std::string_view(&separator, 1) : std::string_view()); @@ -414,10 +429,8 @@ class DataCacheBase } internal::FormatGridIndexKeyWithDimensionSeparator( key, separator, - [&](std::string& out, DimensionIndex dim, Index grid_index) { - absl::StrAppend( - &out, - grid_index + metadata.chunk_key_encoding.GridIndexOffset(dim)); + [](std::string& out, DimensionIndex dim, Index grid_index) { + absl::StrAppend(&out, grid_index); }, rank, grid_indices.subspan(0, @@ -430,25 +443,22 @@ class DataCacheBase *static_cast(initial_metadata().get()); assert(key.size() > key_prefix_.size()); assert(!grid_indices.empty()); - key.remove_prefix(key_prefix_.size() + - (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2 - ? (metadata.rank != 0 ? 2 : 1) - : 0)); + key.remove_prefix(key_prefix_.size() + (metadata.chunk_key_encoding.kind == + ChunkKeyEncoding::kDefault + ? (metadata.rank != 0 ? 2 : 1) + : 0)); if (!internal::ParseGridIndexKeyWithDimensionSeparator( metadata.chunk_key_encoding.separator, - [&](std::string_view part, DimensionIndex dim, Index& grid_index) { - const bool negative = metadata.chunk_key_encoding.kind == - ChunkKeyEncoding::kOffset && - part.size() > 1 && part.front() == '-'; + [](std::string_view part, DimensionIndex dim, Index& grid_index) { + const bool negative = part.size() > 1 && part.front() == '-'; const std::string_view digits = negative ? part.substr(1) : part; if (digits.empty() || !absl::ascii_isdigit(digits.front()) || !absl::ascii_isdigit(digits.back()) || !absl::SimpleAtoi(digits, &grid_index) || - (negative && grid_index == 0)) { + (negative && digits.front() == '0')) { return false; } if (negative) grid_index = -grid_index; - grid_index -= metadata.chunk_key_encoding.GridIndexOffset(dim); return true; }, key, grid_indices.subspan(0, metadata.rank))) { @@ -460,31 +470,18 @@ class DataCacheBase Index MinGridIndexForLexicographicalOrder( DimensionIndex dim, IndexInterval grid_interval) const final { - const auto& metadata = - *static_cast(initial_metadata().get()); - // Negative key components do not sort numerically: split them off singly. - const Index offset = metadata.chunk_key_encoding.GridIndexOffset(dim); - const Index exclusive_max = grid_interval.exclusive_max(); - if (exclusive_max + offset <= 0) { - return -offset; - } - return std::max( - -offset, - internal::MinValueWithMaxBase10Digits(exclusive_max + offset) - offset); + return internal::MinValueWithMaxBase10Digits(grid_interval.exclusive_max()); } std::string GetChunkStorageKey(span cell_indices) override { const auto& metadata = *static_cast(initial_metadata().get()); - if (metadata.chunk_key_encoding.kind != ChunkKeyEncoding::kV2) { + if (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kDefault) { std::string key = absl::StrCat(key_prefix_, "c"); for (DimensionIndex i = 0; i < cell_indices.size(); ++i) { - const Index offset = - i < metadata.rank ? metadata.chunk_key_encoding.GridIndexOffset(i) - : 0; absl::StrAppend( &key, std::string_view(&metadata.chunk_key_encoding.separator, 1), - cell_indices[i] + offset); + cell_indices[i]); } return key; } @@ -517,17 +514,26 @@ class DataCacheBase std::vector full_shape = metadata.shape; full_shape.insert(full_shape.end(), field.field_shape.begin(), field.field_shape.end()); + std::vector full_origin(total_rank, 0); + for (DimensionIndex i = 0; i < rank; ++i) { + full_origin[i] = metadata.LowerBound(i); + } + builder.input_origin(full_origin); builder.input_shape(full_shape); builder.input_labels(span(&normalized_dimension_names[0], total_rank)); + DimensionSet implicit_lower_bounds(false); DimensionSet implicit_upper_bounds(false); for (DimensionIndex i = 0; i < rank; ++i) { + implicit_lower_bounds[i] = metadata.chunk_grid_bounded; implicit_upper_bounds[i] = true; } + builder.implicit_lower_bounds(implicit_lower_bounds); builder.implicit_upper_bounds(implicit_upper_bounds); for (DimensionIndex i = 0; i < total_rank; ++i) { - builder.output_single_input_dimension(i, i); + builder.output_single_input_dimension( + i, i < rank ? -metadata.GridOrigin(i) : 0, 1, i); } return builder.Finalize(); } diff --git a/tensorstore/driver/zarr3/driver_test.cc b/tensorstore/driver/zarr3/driver_test.cc index d1df45bfe..ce6c75e99 100644 --- a/tensorstore/driver/zarr3/driver_test.cc +++ b/tensorstore/driver/zarr3/driver_test.cc @@ -36,6 +36,7 @@ #include "absl/strings/str_format.h" #include "absl/time/clock.h" #include "tensorstore/array.h" +#include "tensorstore/array_storage_statistics.h" #include "tensorstore/array_testutil.h" #include "tensorstore/batch.h" #include "tensorstore/box.h" @@ -1278,73 +1279,62 @@ TEST(ChunkKeyEncodingTest, V2Rank2Slash) { ::testing::Optional(::testing::ElementsAre("0/0", "0/1", "1/0", "1/1"))); } -TEST(ChunkKeyEncodingTest, OffsetRank1) { - EXPECT_THAT(GetKeysForChunkKeyEncoding( - {{"name", "offset"}, {"configuration", {{"offset", {-1}}}}}, - Schema::Shape{{5}}, ChunkLayout::ChunkShape({3})), - ::testing::Optional(::testing::ElementsAre("c/-1", "c/0"))); -} - -TEST(ChunkKeyEncodingTest, OffsetRank2Dot) { - EXPECT_THAT( - GetKeysForChunkKeyEncoding( - {{"name", "offset"}, - {"configuration", {{"separator", "."}, {"offset", {1, -1}}}}}, - Schema::Shape{{5, 6}}, ChunkLayout::ChunkShape({3, 3})), - ::testing::Optional( - ::testing::ElementsAre("c.1.-1", "c.1.0", "c.2.-1", "c.2.0"))); -} - -TEST(ChunkKeyEncodingTest, OffsetZeroMatchesDefault) { - EXPECT_THAT(GetKeysForChunkKeyEncoding( - {{"name", "offset"}, {"configuration", {{"offset", {0, 0}}}}}, - Schema::Shape{{5, 6}}, ChunkLayout::ChunkShape({3, 3})), - ::testing::Optional( - ::testing::ElementsAre("c/0/0", "c/0/1", "c/1/0", "c/1/1"))); +::nlohmann::json RegularBoundedSpec(Index shape, Index origin, + Index lower_bound) { + return { + {"driver", "zarr3"}, + {"kvstore", "memory://"}, + {"metadata", + {{"shape", {shape}}, + {"data_type", "uint8"}, + {"chunk_grid", + {{"name", "regular_bounded"}, + {"configuration", + {{"chunk_shape", {3}}, + {"origin", {origin}}, + {"lower_bound", {lower_bound}}}}}}}}, + }; } -TEST(ChunkKeyEncodingTest, OffsetRankMismatch) { - EXPECT_THAT(GetKeysForChunkKeyEncoding( - {{"name", "offset"}, {"configuration", {{"offset", {0}}}}}, - Schema::Shape{{5, 6}}, ChunkLayout::ChunkShape({3, 3})), - tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, - ".*\"offset\".*")); +std::vector ChunkKeys(const tensorstore::KvStore& kvs) { + std::vector keys; + auto entries = tensorstore::kvstore::ListFuture(kvs).value(); + for (const auto& e : entries) { + if (e.key != "zarr.json") keys.push_back(e.key); + } + std::sort(keys.begin(), keys.end()); + return keys; } -TEST(ChunkKeyEncodingTest, OffsetReadsBackWhatWasWritten) { +TEST(RegularBoundedChunkGridTest, DomainStartsAtTheLowerBound) { + auto context = Context::Default(); TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto store, - tensorstore::Open( - { - {"driver", "zarr3"}, - {"kvstore", "memory://"}, - {"metadata", - {{"chunk_key_encoding", - {{"name", "offset"}, {"configuration", {{"offset", {-2}}}}}}}}, - }, - dtype_v, Schema::Shape({7}), ChunkLayout::ChunkShape({3}), - tensorstore::OpenMode::create) + tensorstore::Open(RegularBoundedSpec(/*shape=*/9, /*origin=*/1, + /*lower_bound=*/-4), + context, tensorstore::OpenMode::create) .result()); - TENSORSTORE_ASSERT_OK(tensorstore::Write( - tensorstore::MakeArray({1, 2, 3, 4, 5, 6, 7}), store)); - EXPECT_THAT(tensorstore::Read(store).result(), - ::testing::Optional( - tensorstore::MakeArray({1, 2, 3, 4, 5, 6, 7}))); + EXPECT_EQ(store.domain().box(), tensorstore::BoxView<1>({-4}, {9})); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto layout, store.chunk_layout()); + EXPECT_THAT(layout.grid_origin(), ::testing::ElementsAre(1)); + auto data = + tensorstore::MakeOffsetArray({-4}, {1, 2, 3, 4, 5, 6, 7, 8, 9}); + TENSORSTORE_ASSERT_OK(tensorstore::Write(data, store)); + EXPECT_THAT(tensorstore::Read(store).result(), ::testing::Optional(data)); + // Chunk -2 is [-5, -2): the first element sits one place into it. + EXPECT_THAT(ChunkKeys(store.kvstore()), + ::testing::ElementsAre("c/-1", "c/-2", "c/0", "c/1")); TENSORSTORE_ASSERT_OK_AND_ASSIGN( - auto entries, tensorstore::kvstore::ListFuture(store.kvstore()).result()); - std::vector keys; - for (const auto& e : entries) keys.push_back(e.key); - EXPECT_THAT(keys, ::testing::UnorderedElementsAre("zarr.json", "c/-2", "c/-1", - "c/0")); + auto reopened, + tensorstore::Open({{"driver", "zarr3"}, {"kvstore", "memory://"}}, + context, tensorstore::OpenMode::open) + .result()); + EXPECT_THAT(tensorstore::Read(reopened).result(), ::testing::Optional(data)); } -TEST(ChunkKeyEncodingTest, - OffsetGrowsInTheNegativeDirectionWithoutMovingChunks) { +TEST(RegularBoundedChunkGridTest, GrowsInFrontByEditingZarrJsonAlone) { auto context = Context::Default(); - ::nlohmann::json spec{ - {"driver", "zarr3"}, - {"kvstore", "memory://"}, - }; + ::nlohmann::json spec{{"driver", "zarr3"}, {"kvstore", "memory://"}}; TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto store, tensorstore::Open(spec, context, dtype_v, Schema::Shape({6}), @@ -1354,23 +1344,17 @@ TEST(ChunkKeyEncodingTest, TENSORSTORE_ASSERT_OK(tensorstore::Write( tensorstore::MakeArray({10, 11, 12, 13, 14, 15}), store)); auto kvs = store.kvstore(); + EXPECT_THAT(ChunkKeys(kvs), ::testing::ElementsAre("c/0", "c/1")); TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto read_result, tensorstore::kvstore::Read(kvs, "zarr.json").result()); auto metadata = ::nlohmann::json::parse(std::string(read_result.value)); - const auto keys_before = [&] { - std::vector keys; - auto entries = tensorstore::kvstore::ListFuture(kvs).value(); - for (const auto& e : entries) { - if (e.key != "zarr.json") keys.push_back(e.key); - } - std::sort(keys.begin(), keys.end()); - return keys; - }(); - EXPECT_THAT(keys_before, ::testing::ElementsAre("c/0", "c/1")); - - metadata["shape"] = {12}; - metadata["chunk_key_encoding"] = {{"name", "offset"}, - {"configuration", {{"offset", {-2}}}}}; + + // Two elements in front: less than a chunk. + metadata["shape"] = {8}; + metadata["chunk_grid"] = { + {"name", "regular_bounded"}, + {"configuration", + {{"chunk_shape", {3}}, {"origin", {0}}, {"lower_bound", {-2}}}}}; TENSORSTORE_ASSERT_OK( tensorstore::kvstore::Write(kvs, "zarr.json", absl::Cord(metadata.dump())) .result()); @@ -1378,66 +1362,99 @@ TEST(ChunkKeyEncodingTest, TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto grown, tensorstore::Open(spec, context, tensorstore::OpenMode::open).result()); - EXPECT_THAT(grown.domain().shape(), ::testing::ElementsAre(12)); + EXPECT_EQ(grown.domain().box(), tensorstore::BoxView<1>({-2}, {8})); EXPECT_THAT(tensorstore::Read(grown).result(), - ::testing::Optional(tensorstore::MakeArray( - {0, 0, 0, 0, 0, 0, 10, 11, 12, 13, 14, 15}))); + ::testing::Optional(tensorstore::MakeOffsetArray( + {-2}, {0, 0, 10, 11, 12, 13, 14, 15}))); TENSORSTORE_ASSERT_OK( - tensorstore::Write(tensorstore::MakeArray({1, 2, 3}), - grown | tensorstore::Dims(0).SizedInterval(0, 3))); + tensorstore::Write(tensorstore::MakeOffsetArray({-2}, {1, 2}), + grown | tensorstore::Dims(0).HalfOpenInterval(-2, 0))); EXPECT_THAT(tensorstore::Read(grown).result(), - ::testing::Optional(tensorstore::MakeArray( - {1, 2, 3, 0, 0, 0, 10, 11, 12, 13, 14, 15}))); - std::vector keys_after; - auto entries_after = tensorstore::kvstore::ListFuture(kvs).value(); - for (const auto& e : entries_after) { - if (e.key != "zarr.json") keys_after.push_back(e.key); - } - EXPECT_THAT(keys_after, - ::testing::UnorderedElementsAre("c/-2", "c/0", "c/1")); + ::testing::Optional(tensorstore::MakeOffsetArray( + {-2}, {1, 2, 10, 11, 12, 13, 14, 15}))); + EXPECT_THAT(ChunkKeys(kvs), ::testing::ElementsAre("c/-1", "c/0", "c/1")); } -TEST(ChunkKeyEncodingTest, OffsetResizeDeletesOnlyTheChunksOutsideTheNewShape) { +TEST(RegularBoundedChunkGridTest, ResizeMovesTheLowerBound) { auto context = Context::Default(); - ::nlohmann::json spec{ - {"driver", "zarr3"}, - {"kvstore", "memory://"}, - {"metadata", - {{"chunk_key_encoding", - {{"name", "offset"}, {"configuration", {{"offset", {-2, 0}}}}}}}}, - }; TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto store, - tensorstore::Open(spec, context, dtype_v, Schema::Shape({12, 2}), - ChunkLayout::ChunkShape({3, 1}), - tensorstore::OpenMode::create) + tensorstore::Open(RegularBoundedSpec(/*shape=*/6, /*origin=*/0, + /*lower_bound=*/0), + context, tensorstore::OpenMode::create) .result()); - TENSORSTORE_ASSERT_OK( - tensorstore::Write(tensorstore::MakeScalarArray(7), store)); - const auto keys = [&] { - std::vector keys; - auto entries = tensorstore::kvstore::ListFuture(store.kvstore()).value(); - for (const auto& e : entries) { - if (e.key != "zarr.json") keys.push_back(e.key); - } - std::sort(keys.begin(), keys.end()); - return keys; - }; - EXPECT_THAT(keys(), - ::testing::ElementsAre("c/-1/0", "c/-1/1", "c/-2/0", "c/-2/1", - "c/0/0", "c/0/1", "c/1/0", "c/1/1")); - TENSORSTORE_ASSERT_OK( - tensorstore::Resize(store, - {{tensorstore::kImplicit, tensorstore::kImplicit}}, - {{6, tensorstore::kImplicit}}) + TENSORSTORE_ASSERT_OK(tensorstore::Write( + tensorstore::MakeArray({10, 11, 12, 13, 14, 15}), store)); + + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto grown, + tensorstore::Resize(store, {{-4}}, {{tensorstore::kImplicit}}).result()); + EXPECT_EQ(grown.domain().box(), tensorstore::BoxView<1>({-4}, {10})); + TENSORSTORE_ASSERT_OK(tensorstore::Write( + tensorstore::MakeOffsetArray({-4}, {1, 2, 3, 4}), + grown | tensorstore::Dims(0).HalfOpenInterval(-4, 0))); + EXPECT_THAT(tensorstore::Read(grown).result(), + ::testing::Optional(tensorstore::MakeOffsetArray( + {-4}, {1, 2, 3, 4, 10, 11, 12, 13, 14, 15}))); + EXPECT_THAT(ChunkKeys(store.kvstore()), + ::testing::ElementsAre("c/-1", "c/-2", "c/0", "c/1")); + + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto shrunk, + tensorstore::Resize(grown, {{1}}, {{tensorstore::kImplicit}}).result()); + EXPECT_EQ(shrunk.domain().box(), tensorstore::BoxView<1>({1}, {5})); + EXPECT_THAT(tensorstore::Read(shrunk).result(), + ::testing::Optional(tensorstore::MakeOffsetArray( + {1}, {11, 12, 13, 14, 15}))); + EXPECT_THAT(ChunkKeys(store.kvstore()), ::testing::ElementsAre("c/0", "c/1")); +} + +TEST(RegularBoundedChunkGridTest, StorageStatisticsSeeNegativeKeys) { + auto context = Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open(RegularBoundedSpec(/*shape=*/30, /*origin=*/0, + /*lower_bound=*/-15), + context, tensorstore::OpenMode::create) + .result()); + using ::tensorstore::ArrayStorageStatistics; + TENSORSTORE_ASSERT_OK(tensorstore::Write( + tensorstore::MakeScalarArray(7), + store | tensorstore::Dims(0).HalfOpenInterval(-15, -9))); + EXPECT_THAT(tensorstore::GetStorageStatistics( + store | tensorstore::Dims(0).HalfOpenInterval(-15, -9), + ArrayStorageStatistics::query_fully_stored) + .result(), + ::testing::Optional(ArrayStorageStatistics{ + /*.mask=*/ArrayStorageStatistics::query_fully_stored, + /*.not_stored=*/false, /*.fully_stored=*/true})); + EXPECT_THAT(tensorstore::GetStorageStatistics( + store | tensorstore::Dims(0).HalfOpenInterval(-9, 15), + ArrayStorageStatistics::query_not_stored) + .result(), + ::testing::Optional(ArrayStorageStatistics{ + /*.mask=*/ArrayStorageStatistics::query_not_stored, + /*.not_stored=*/true})); +} + +TEST(ChunkKeyEncodingTest, NegativeZeroIsNotAKey) { + auto context = Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open(RegularBoundedSpec(/*shape=*/6, /*origin=*/0, + /*lower_bound=*/-3), + context, tensorstore::OpenMode::create) .result()); - EXPECT_THAT(keys(), - ::testing::ElementsAre("c/-1/0", "c/-1/1", "c/-2/0", "c/-2/1")); TENSORSTORE_ASSERT_OK( - tensorstore::Resize( - store, {{tensorstore::kImplicit, tensorstore::kImplicit}}, {{3, 1}}) + tensorstore::kvstore::Write(store.kvstore(), "c/-0", absl::Cord("x")) .result()); - EXPECT_THAT(keys(), ::testing::ElementsAre("c/-2/0")); + using ::tensorstore::ArrayStorageStatistics; + EXPECT_THAT(tensorstore::GetStorageStatistics( + store, ArrayStorageStatistics::query_not_stored) + .result(), + ::testing::Optional(ArrayStorageStatistics{ + /*.mask=*/ArrayStorageStatistics::query_not_stored, + /*.not_stored=*/true})); } TENSORSTORE_GLOBAL_INITIALIZER { diff --git a/tensorstore/driver/zarr3/index.rst b/tensorstore/driver/zarr3/index.rst index d67e33d1a..0acfc92b1 100644 --- a/tensorstore/driver/zarr3/index.rst +++ b/tensorstore/driver/zarr3/index.rst @@ -74,8 +74,6 @@ metadata. .. json:schema:: driver/zarr3/ChunkKeyEncoding.v2 -.. json:schema:: driver/zarr3/ChunkKeyEncoding.offset - Mapping to TensorStore Schema ----------------------------- diff --git a/tensorstore/driver/zarr3/metadata.cc b/tensorstore/driver/zarr3/metadata.cc index 82955fd19..24169cd66 100644 --- a/tensorstore/driver/zarr3/metadata.cc +++ b/tensorstore/driver/zarr3/metadata.cc @@ -23,7 +23,6 @@ #include #include #include -#include #include #include #include @@ -43,6 +42,7 @@ #include "absl/strings/str_cat.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" +#include #include "tensorstore/array.h" #include "tensorstore/box.h" #include "tensorstore/chunk_layout.h" @@ -456,57 +456,26 @@ TENSORSTORE_DEFINE_JSON_DEFAULT_BINDER(ChunkKeyEncoding, [](auto is_loading, jb::Projection<&Self::kind>(jb::Enum({ {Self::kDefault, "default"}, {Self::kV2, "v2"}, - {Self::kOffset, "offset"}, })), - jb::Object( - jb::Member("separator", - jb::Projection<&Self::separator>(jb::DefaultValue( - [outer = obj](auto* obj) -> absl::Status { - switch (outer->kind) { - case Self::kDefault: - case Self::kOffset: - *obj = '/'; - break; - case Self::kV2: - *obj = '.'; - break; - default: - ABSL_UNREACHABLE(); - } - return absl::OkStatus(); - }, - jb::Enum({ - {'.', "."}, - {'/', "/"}, - })))), - // Length is checked against the rank in `ValidateMetadata`. - jb::Member("offset", - jb::Projection<&Self::offset>( - [outer = obj](auto is_loading, const auto& options, - auto* obj, auto* j) -> absl::Status { - if (outer->kind != Self::kOffset) { - if constexpr (is_loading) { - if (!j->is_discarded()) { - return absl::InvalidArgumentError( - "\"offset\" is only valid for the " - "\"offset\" chunk key encoding"); - } - obj->clear(); - } else { - *j = ::nlohmann::json::value_t::discarded; - } - return absl::OkStatus(); - } - if constexpr (is_loading) { - if (j->is_discarded()) { - return absl::InvalidArgumentError( - "\"offset\" must be specified for the " - "\"offset\" chunk key encoding"); - } - } - return jb::Array(jb::Integer())( - is_loading, options, obj, j); - })))) // + jb::Object(jb::Member("separator", + jb::Projection<&Self::separator>(jb::DefaultValue( + [outer = obj](auto* obj) -> absl::Status { + switch (outer->kind) { + case Self::kDefault: + *obj = '/'; + break; + case Self::kV2: + *obj = '.'; + break; + default: + ABSL_UNREACHABLE(); + } + return absl::OkStatus(); + }, + jb::Enum({ + {'.', "."}, + {'/', "/"}, + })))))) // )(is_loading, options, obj, j); }) @@ -557,6 +526,24 @@ constexpr auto MetadataJsonBinder = [] { } }; + // In constraints, `chunk_grid` as a whole may be absent. + constexpr auto maybe_optional_chunk_grid = [](auto binder) { + return [=](auto is_loading, const auto& options, auto* obj, + auto* j) -> absl::Status { + if constexpr (Constraints) { + if constexpr (is_loading) { + if (j->is_discarded()) return absl::OkStatus(); + } else { + if (!obj->chunk_shape && !obj->chunk_grid_bounded) { + *j = ::nlohmann::json::value_t::discarded; + return absl::OkStatus(); + } + } + } + return binder(is_loading, options, obj, j); + }; + }; + return [=](auto is_loading, const auto& options, auto* obj, auto* j) { using Self = absl::remove_cvref_t; DimensionIndex* rank = nullptr; @@ -564,6 +551,35 @@ constexpr auto MetadataJsonBinder = [] { rank = &obj->rank; } + // A `regular_bounded` configuration member, rejected under `regular`. + auto bounded_grid_member = [&](auto name, auto member) { + return jb::Member( + name, + [=](auto is_loading, const auto& options, auto* obj, + auto* j) -> absl::Status { + bool bounded; + if constexpr (Constraints) { + bounded = obj->chunk_grid_bounded.value_or(false); + } else { + bounded = obj->chunk_grid_bounded; + } + if (!bounded) { + if constexpr (is_loading) { + if (!j->is_discarded()) { + return absl::InvalidArgumentError( + "Only valid for the \"regular_bounded\" chunk grid"); + } + } else { + *j = ::nlohmann::json::value_t::discarded; + } + return absl::OkStatus(); + } + return maybe_optional( + jb::DimensionIndexedVector(rank, jb::Integer()))( + is_loading, options, &(obj->*member), j); + }); + }; + auto ensure_zarr_dtype = [&]() -> Result { if constexpr (std::is_same_v) { return obj->zarr_dtype; @@ -651,13 +667,23 @@ constexpr auto MetadataJsonBinder = [] { maybe_optional(jb::DefaultBinder<>))), jb::Member( "chunk_grid", - jb::Projection<&Self::chunk_shape>(maybe_optional(jb::Object( - jb::Member("name", jb::Constant([] { return "regular"; })), + maybe_optional_chunk_grid(jb::Object( + jb::Member("name", + jb::Projection<&Self::chunk_grid_bounded>( + maybe_optional(jb::Enum({ + {false, "regular"}, + {true, "regular_bounded"}, + })))), jb::Member( "configuration", - jb::Object(jb::Member("chunk_shape", - jb::ChunkShapeVector(rank)))) // - )))), // + jb::Object( + jb::Member( + "chunk_shape", + jb::Projection<&Self::chunk_shape>( + maybe_optional(jb::ChunkShapeVector(rank)))), + bounded_grid_member("origin", &Self::chunk_grid_origin), + non_compatibility_field(bounded_grid_member( + "lower_bound", &Self::lower_bound))))))), jb::Member("codecs", jb::Projection<&Self::codec_specs>(maybe_optional( ZarrCodecChainJsonBinder))), // Allow empty storage_transformers list. @@ -869,15 +895,6 @@ absl::Status ValidateStructEndianness(const ZarrDType& dtype, absl::Status ValidateMetadata(ZarrMetadata& metadata) { DeriveFieldShape(metadata); - if (metadata.chunk_key_encoding.kind == ChunkKeyEncoding::kOffset && - static_cast(metadata.chunk_key_encoding.offset.size()) != - metadata.rank) { - return absl::InvalidArgumentError(absl::StrCat( - "Expected \"offset\" of the chunk key encoding to have length ", - metadata.rank, ", but received: ", - ::nlohmann::json(metadata.chunk_key_encoding.offset).dump())); - } - // Skip when codecs are already resolved: the create path runs the // pre-Resolve check itself in `GetNewMetadata`, and post-Resolve the // leaf bytes spec has its `endian` stripped. @@ -1014,6 +1031,23 @@ absl::Status ValidateMetadata(const ZarrMetadata& metadata, return MetadataMismatchError("chunk_shape", *constraints.chunk_shape, metadata.chunk_shape); } + if (constraints.chunk_grid_bounded && + *constraints.chunk_grid_bounded != metadata.chunk_grid_bounded) { + return MetadataMismatchError( + "chunk_grid", + *constraints.chunk_grid_bounded ? "regular_bounded" : "regular", + metadata.chunk_grid_bounded ? "regular_bounded" : "regular"); + } + if (constraints.chunk_grid_origin && + *constraints.chunk_grid_origin != metadata.chunk_grid_origin) { + return MetadataMismatchError("origin", *constraints.chunk_grid_origin, + metadata.chunk_grid_origin); + } + if (constraints.lower_bound && + *constraints.lower_bound != metadata.lower_bound) { + return MetadataMismatchError("lower_bound", *constraints.lower_bound, + metadata.lower_bound); + } if (constraints.chunk_key_encoding && *constraints.chunk_key_encoding != metadata.chunk_key_encoding) { return MetadataMismatchError("chunk_key_encoding", @@ -1208,7 +1242,8 @@ Result> GetEffectiveDomain( const SpecRankAndFieldInfo& info, std::optional> metadata_shape, std::optional>> dimension_names, - const Schema& schema, bool* dimension_names_used) { + const Schema& schema, bool* dimension_names_used, + std::optional> lower_bound = std::nullopt) { const DimensionIndex rank = info.chunked_rank; if (dimension_names_used) *dimension_names_used = false; auto domain = schema.domain(); @@ -1238,6 +1273,11 @@ Result> GetEffectiveDomain( builder.shape(*metadata_shape); builder.implicit_upper_bounds(true); } + if (lower_bound) { + std::fill(builder.origin().begin(), builder.origin().end(), Index{0}); + std::copy(lower_bound->begin(), lower_bound->end(), + builder.origin().begin()); + } } else { builder.origin(GetConstantVector(builder.rank())); } @@ -1262,7 +1302,11 @@ Result> GetEffectiveDomain( domain, MergeIndexDomains(domain, domain_from_metadata), _.Format("Mismatch between metadata and schema") .With(internal::ConvertInvalidArgumentToFailedPrecondition)); - return WithImplicitDimensions(domain, false, true); + return WithImplicitDimensions(domain, + lower_bound + ? DimensionSet::UpTo(lower_bound->size()) + : DimensionSet(false), + true); } Result> GetEffectiveDomain( @@ -1285,21 +1329,27 @@ Result> GetEffectiveDomain( metadata_constraints.dimension_names->size()); } + std::optional> lower_bound_span; + if (metadata_constraints.lower_bound) { + lower_bound_span.emplace(metadata_constraints.lower_bound->data(), + metadata_constraints.lower_bound->size()); + } return GetEffectiveDomain(info, shape_span, names_span, schema, - dimension_names_used); + dimension_names_used, lower_bound_span); } absl::Status SetChunkLayoutFromMetadata( const SpecRankAndFieldInfo& info, std::optional> chunk_shape, - const ZarrCodecChainSpec* codecs, ChunkLayout& chunk_layout) { + const ZarrCodecChainSpec* codecs, ChunkLayout& chunk_layout, + std::optional> grid_origin) { const DimensionIndex rank = info.chunked_rank; if (rank == dynamic_rank) { return absl::OkStatus(); } TENSORSTORE_RETURN_IF_ERROR(chunk_layout.Set(RankConstraint(rank))); - TENSORSTORE_RETURN_IF_ERROR(chunk_layout.Set( - ChunkLayout::GridOrigin(GetConstantVector(rank)))); + TENSORSTORE_RETURN_IF_ERROR(chunk_layout.Set(ChunkLayout::GridOrigin( + grid_origin ? *grid_origin : GetConstantVector(rank)))); if (chunk_shape) { assert(chunk_shape->size() == rank); @@ -1365,11 +1415,16 @@ Result GetEffectiveChunkLayout( chunk_shape_span.emplace(metadata_constraints.chunk_shape->data(), metadata_constraints.chunk_shape->size()); } + std::optional> grid_origin_span; + if (metadata_constraints.chunk_grid_origin) { + grid_origin_span.emplace(metadata_constraints.chunk_grid_origin->data(), + metadata_constraints.chunk_grid_origin->size()); + } TENSORSTORE_RETURN_IF_ERROR(SetChunkLayoutFromMetadata( info, chunk_shape_span, metadata_constraints.codec_specs ? &*metadata_constraints.codec_specs : nullptr, - chunk_layout)); + chunk_layout, grid_origin_span)); return chunk_layout; } @@ -1433,9 +1488,14 @@ absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, std::optional>> dimension_names_span; dimension_names_span.emplace(metadata.dimension_names.data(), metadata.dimension_names.size()); + std::optional> lower_bound_span; + if (metadata.chunk_grid_bounded) { + lower_bound_span.emplace(metadata.lower_bound.data(), + metadata.lower_bound.size()); + } TENSORSTORE_RETURN_IF_ERROR(GetEffectiveDomain( info, metadata_shape_span, dimension_names_span, schema, - /*dimension_names_used=*/nullptr)); + /*dimension_names_used=*/nullptr, lower_bound_span)); } if (schema.chunk_layout().rank() != dynamic_rank) { @@ -1443,8 +1503,14 @@ absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, std::optional> chunk_shape_span; chunk_shape_span.emplace(metadata.chunk_shape.data(), metadata.chunk_shape.size()); + std::optional> grid_origin_span; + if (metadata.chunk_grid_bounded) { + grid_origin_span.emplace(metadata.chunk_grid_origin.data(), + metadata.chunk_grid_origin.size()); + } TENSORSTORE_RETURN_IF_ERROR(SetChunkLayoutFromMetadata( - info, chunk_shape_span, &metadata.codec_specs, chunk_layout)); + info, chunk_shape_span, &metadata.codec_specs, chunk_layout, + grid_origin_span)); if (chunk_layout.codec_chunk_shape().hard_constraint) { return absl::InvalidArgumentError("codec_chunk_shape not supported"); } @@ -1613,6 +1679,15 @@ Result> GetNewMetadata( domain.shape().begin() + logical_rank); metadata->dimension_names.assign(domain.labels().begin(), domain.labels().begin() + logical_rank); + metadata->chunk_grid_bounded = + metadata_constraints.chunk_grid_bounded.value_or(false); + if (metadata->chunk_grid_bounded) { + metadata->chunk_grid_origin = + metadata_constraints.chunk_grid_origin.value_or( + std::vector(logical_rank, 0)); + metadata->lower_bound = metadata_constraints.lower_bound.value_or( + std::vector(logical_rank, 0)); + } for (DimensionIndex i = 0; i < logical_rank; ++i) { auto& name = metadata->dimension_names[i]; @@ -1808,6 +1883,9 @@ ZarrMetadataConstraints::ZarrMetadataConstraints(const ZarrMetadata& metadata) dimension_names(metadata.dimension_names), chunk_key_encoding(metadata.chunk_key_encoding), chunk_shape(metadata.chunk_shape), + chunk_grid_bounded(metadata.chunk_grid_bounded), + chunk_grid_origin(metadata.chunk_grid_origin), + lower_bound(metadata.lower_bound), codec_specs(metadata.codec_specs), fill_value(metadata.fill_value), unknown_extension_attributes(metadata.unknown_extension_attributes) {} diff --git a/tensorstore/driver/zarr3/metadata.h b/tensorstore/driver/zarr3/metadata.h index 9651c5428..b4e9f6678 100644 --- a/tensorstore/driver/zarr3/metadata.h +++ b/tensorstore/driver/zarr3/metadata.h @@ -61,25 +61,16 @@ struct ChunkKeyEncoding { enum Kind { kDefault, kV2, - // https://github.com/zarr-developers/zarr-extensions/tree/main/chunk-key-encodings/offset - kOffset, }; Kind kind; char separator; - // One value per dimension for `kOffset`, empty otherwise. - std::vector offset; - - Index GridIndexOffset(DimensionIndex dim) const { - return kind == kOffset ? offset[dim] : 0; - } TENSORSTORE_DECLARE_JSON_DEFAULT_BINDER(ChunkKeyEncoding, internal_json_binding::NoOptions, tensorstore::IncludeDefaults) friend bool operator==(const ChunkKeyEncoding& a, const ChunkKeyEncoding& b) { - return a.kind == b.kind && a.separator == b.separator && - a.offset == b.offset; + return a.kind == b.kind && a.separator == b.separator; } friend bool operator!=(const ChunkKeyEncoding& a, const ChunkKeyEncoding& b) { return !(a == b); @@ -133,10 +124,23 @@ struct ZarrMetadata { std::vector> dimension_names; ChunkKeyEncoding chunk_key_encoding; std::vector chunk_shape; + // `regular_bounded` chunk grid: `chunk_grid_origin` is the array coordinate + // where chunk 0 starts and `lower_bound` the coordinate of the first + // element. Both are empty for the `regular` chunk grid. + bool chunk_grid_bounded = false; + std::vector chunk_grid_origin; + std::vector lower_bound; ZarrCodecChainSpec codec_specs; std::vector> fill_value; ::nlohmann::json::object_t unknown_extension_attributes; + Index GridOrigin(DimensionIndex dim) const { + return chunk_grid_bounded ? chunk_grid_origin[dim] : 0; + } + Index LowerBound(DimensionIndex dim) const { + return chunk_grid_bounded ? lower_bound[dim] : 0; + } + std::string GetCompatibilityKey() const; ZarrCodecChain::Ptr codecs; @@ -186,6 +190,9 @@ struct ZarrMetadataConstraints { std::optional>> dimension_names; std::optional chunk_key_encoding; std::optional> chunk_shape; + std::optional chunk_grid_bounded; + std::optional> chunk_grid_origin; + std::optional> lower_bound; std::optional codec_specs; std::optional>> fill_value; ::nlohmann::json::object_t unknown_extension_attributes; @@ -219,7 +226,8 @@ Result> GetEffectiveDomain( absl::Status SetChunkLayoutFromMetadata( const SpecRankAndFieldInfo& info, std::optional> chunk_shape, - const ZarrCodecChainSpec* codecs, ChunkLayout& chunk_layout); + const ZarrCodecChainSpec* codecs, ChunkLayout& chunk_layout, + std::optional> grid_origin = std::nullopt); absl::Status SetChunkLayoutFromMetadata( DataType dtype, DimensionIndex rank, std::optional> chunk_shape, diff --git a/tensorstore/driver/zarr3/metadata_test.cc b/tensorstore/driver/zarr3/metadata_test.cc index ce43b4c94..5f045ac6f 100644 --- a/tensorstore/driver/zarr3/metadata_test.cc +++ b/tensorstore/driver/zarr3/metadata_test.cc @@ -72,7 +72,6 @@ using ::tensorstore::dtypes::float16_t; using ::tensorstore::dtypes::float32_t; using ::tensorstore::dtypes::float64_t; using ::tensorstore::internal::uint_t; -using ::tensorstore::internal_zarr3::ChunkKeyEncoding; using ::tensorstore::internal_zarr3::FillValueJsonBinder; using ::tensorstore::internal_zarr3::GetSpecRankAndFieldInfo; using ::tensorstore::internal_zarr3::GetVoidMetadata; @@ -152,72 +151,58 @@ TEST(MetadataTest, ParseValidNoDimensionNames) { EXPECT_THAT(metadata.user_attributes, MatchesJson({{"a", "b"}, {"c", "d"}})); } -TEST(MetadataTest, OffsetChunkKeyEncoding) { +TEST(MetadataTest, RegularBoundedChunkGrid) { auto json = GetBasicMetadata(); - json["chunk_key_encoding"] = { - {"name", "offset"}, - {"configuration", {{"separator", "."}, {"offset", {-2, 0, 5}}}}}; + json["chunk_grid"] = {{"name", "regular_bounded"}, + {"configuration", + {{"chunk_shape", {1, 2, 3}}, + {"origin", {0, -1, 5}}, + {"lower_bound", {-3, 0, 2}}}}}; tensorstore::TestJsonBinderRoundTripJsonOnly({json}); TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, ZarrMetadata::FromJson(json)); - EXPECT_EQ(metadata.chunk_key_encoding.kind, ChunkKeyEncoding::kOffset); - EXPECT_EQ(metadata.chunk_key_encoding.separator, '.'); - EXPECT_THAT(metadata.chunk_key_encoding.offset, - ::testing::ElementsAre(-2, 0, 5)); - EXPECT_EQ(metadata.chunk_key_encoding.GridIndexOffset(0), -2); - EXPECT_EQ(metadata.chunk_key_encoding.GridIndexOffset(2), 5); -} - -TEST(MetadataTest, OffsetChunkKeyEncodingSeparatorDefaultsToSlash) { - auto json = GetBasicMetadata(); - json["chunk_key_encoding"] = {{"name", "offset"}, - {"configuration", {{"offset", {0, 0, 0}}}}}; - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, ZarrMetadata::FromJson(json)); - EXPECT_EQ(metadata.chunk_key_encoding.separator, '/'); - TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto out, metadata.ToJson()); - EXPECT_THAT(out["chunk_key_encoding"], - MatchesJson({{"name", "offset"}, - {"configuration", {{"offset", {0, 0, 0}}}}})); + EXPECT_TRUE(metadata.chunk_grid_bounded); + EXPECT_THAT(metadata.chunk_shape, ::testing::ElementsAre(1, 2, 3)); + EXPECT_THAT(metadata.chunk_grid_origin, ::testing::ElementsAre(0, -1, 5)); + EXPECT_THAT(metadata.lower_bound, ::testing::ElementsAre(-3, 0, 2)); } -TEST(MetadataTest, OffsetChunkKeyEncodingRankMismatch) { +TEST(MetadataTest, RegularChunkGridRejectsBoundedMembers) { auto json = GetBasicMetadata(); - json["chunk_key_encoding"] = {{"name", "offset"}, - {"configuration", {{"offset", {0, 0}}}}}; + json["chunk_grid"]["configuration"]["origin"] = {0, 0, 0}; EXPECT_THAT(ZarrMetadata::FromJson(json), StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*\"offset\".*length 3.*"))); + HasSubstr("regular_bounded"))); } -TEST(MetadataTest, OffsetChunkKeyEncodingRequiresOffset) { +TEST(MetadataTest, RegularBoundedChunkGridRequiresBothMembers) { auto json = GetBasicMetadata(); - json["chunk_key_encoding"] = {{"name", "offset"}}; - EXPECT_THAT(ZarrMetadata::FromJson(json), - StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*\"offset\" must be specified.*"))); - json["chunk_key_encoding"] = {{"name", "offset"}, - {"configuration", {{"separator", "."}}}}; - EXPECT_THAT(ZarrMetadata::FromJson(json), - StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*\"offset\" must be specified.*"))); + json["chunk_grid"] = { + {"name", "regular_bounded"}, + {"configuration", {{"chunk_shape", {1, 2, 3}}, {"origin", {0, 0, 0}}}}}; + EXPECT_THAT( + ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("lower_bound"))); + json["chunk_grid"]["configuration"]["lower_bound"] = {0, 0}; + EXPECT_THAT( + ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("lower_bound"))); } -TEST(MetadataTest, OffsetIsOnlyValidForTheOffsetChunkKeyEncoding) { +TEST(MetadataTest, CompatibilityKeyIgnoresTheLowerBound) { auto json = GetBasicMetadata(); - json["chunk_key_encoding"] = {{"name", "default"}, - {"configuration", {{"offset", {0, 0, 0}}}}}; - EXPECT_THAT(ZarrMetadata::FromJson(json), - StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*\"offset\" is only valid.*"))); -} - -TEST(MetadataTest, OffsetChunkKeyEncodingEquality) { - ChunkKeyEncoding a{ChunkKeyEncoding::kOffset, '/', {-1, 0}}; - ChunkKeyEncoding b{ChunkKeyEncoding::kOffset, '/', {-1, 0}}; - ChunkKeyEncoding c{ChunkKeyEncoding::kOffset, '/', {-2, 0}}; - ChunkKeyEncoding d{ChunkKeyEncoding::kDefault, '/', {}}; - EXPECT_EQ(a, b); - EXPECT_NE(a, c); - EXPECT_NE(a, d); + json["chunk_grid"] = {{"name", "regular_bounded"}, + {"configuration", + {{"chunk_shape", {1, 2, 3}}, + {"origin", {0, 0, 0}}, + {"lower_bound", {0, 0, 0}}}}}; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto a, ZarrMetadata::FromJson(json)); + json["shape"] = {14, 11, 12}; + json["chunk_grid"]["configuration"]["lower_bound"] = {-4, 0, 0}; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto b, ZarrMetadata::FromJson(json)); + EXPECT_EQ(a.GetCompatibilityKey(), b.GetCompatibilityKey()); + json["chunk_grid"]["configuration"]["origin"] = {1, 0, 0}; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto c, ZarrMetadata::FromJson(json)); + EXPECT_NE(a.GetCompatibilityKey(), c.GetCompatibilityKey()); } TEST(MetadataTest, UnknownExtensionAttribute) { @@ -1238,10 +1223,9 @@ TEST(StructEndianTest, ModernStructRequiresExplicitEndianForMultiByteField) { {{"name", "y"}, {"data_type", "int16"}}})}}}}; metadata_json["fill_value"] = {{"x", 0}, {"y", 0}}; metadata_json["codecs"] = ::nlohmann::json::array_t{{{"name", "bytes"}}}; - EXPECT_THAT( - ZarrMetadata::FromJson(metadata_json), - tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*specify an explicit `endian`.*"))); + EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), + tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*specify an explicit `endian`.*")); } // All single-byte fields: no byte order to disambiguate. @@ -1337,10 +1321,9 @@ TEST(StructEndianTest, ShardedModernStructWithoutLeafEndianRejected) { {{"name", "y"}, {"data_type", "int16"}}})}}}}; metadata_json["fill_value"] = {{"x", 0}, {"y", 0}}; metadata_json["codecs"] = ShardingWrappedBytes(std::nullopt); - EXPECT_THAT( - ZarrMetadata::FromJson(metadata_json), - tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*specify an explicit `endian`.*"))); + EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), + tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*specify an explicit `endian`.*")); } // Legacy default must be injected at the leaf bytes codec, not the @@ -1376,8 +1359,8 @@ TEST(StructFillValueFormTest, ModernStructRejectsBase64Fill) { metadata_json["codecs"] = ::nlohmann::json::array_t{ {{"name", "bytes"}, {"configuration", {{"endian", "little"}}}}}; EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), - tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*Expected object.*"))); + tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*Expected object.*")); } TEST(StructFillValueFormTest, ModernStructRejectsArrayFill) { @@ -1392,8 +1375,8 @@ TEST(StructFillValueFormTest, ModernStructRejectsArrayFill) { metadata_json["codecs"] = ::nlohmann::json::array_t{ {{"name", "bytes"}, {"configuration", {{"endian", "little"}}}}}; EXPECT_THAT(ZarrMetadata::FromJson(metadata_json), - tensorstore::StatusIs(absl::StatusCode::kInvalidArgument, - MatchesRegex(".*Expected object.*"))); + tensorstore::MatchesStatus(absl::StatusCode::kInvalidArgument, + ".*Expected object.*")); } // Verifies the deferred-parse path resolves endian before decoding diff --git a/tensorstore/driver/zarr3/schema.yml b/tensorstore/driver/zarr3/schema.yml index 6f44b9f2e..2f02f0c69 100644 --- a/tensorstore/driver/zarr3/schema.yml +++ b/tensorstore/driver/zarr3/schema.yml @@ -87,7 +87,16 @@ definitions: type: object properties: name: - const: "regular" + oneOf: + - const: "regular" + - const: "regular_bounded" + title: Chunk grid. + description: | + ``"regular_bounded"`` is a regular grid whose chunk 0 starts at + array coordinate ``origin`` and whose domain starts at + ``lower_bound``, so chunk grid indices, and the keys formed from + them, may be negative and the first chunk may be partly filled. + The lower bound is resizable. configuration: type: object properties: @@ -104,6 +113,20 @@ definitions: to the `Schema.chunk_layout`. examples: - [64, 64, 64] + origin: + type: array + items: + type: integer + title: | + Array coordinate where chunk 0 starts, for each dimension. + ``"regular_bounded"`` only. + lower_bound: + type: array + items: + type: integer + title: | + Array coordinate of the first element, for each dimension. + ``"regular_bounded"`` only. chunk_key_encoding: $ref: driver/zarr3/ChunkKeyEncoding fill_value: @@ -295,50 +318,6 @@ definitions: - const: "." default: "." title: "Separator character between dimensions" - chunk_key_encoding_offset: - $id: driver/zarr3/ChunkKeyEncoding.offset - title: Offset chunk key encoding. - description: | - The `~driver/zarr3/ChunkKeyEncoding.default` encoding applied to the - chunk grid position plus a fixed per-dimension integer offset, so that - an array can grow in the negative direction of a dimension by whole - chunks without any existing chunk changing its key. Refer to `the - extension registry - `__ - for details. - allOf: - - $ref: driver/zarr3/ChunkKeyEncoding - - type: object - properties: - name: - const: "offset" - configuration: - type: object - properties: - separator: - type: string - oneOf: - - const: "/" - - const: "." - default: "/" - title: "Separator character between dimensions" - offset: - type: array - items: - type: integer - title: | - Offset added to the chunk grid position along each dimension - before it is formatted as a key. - description: | - Must have one element per dimension. An element may be - negative, in which case the corresponding key component carries - a leading ``-``. With every element ``0`` the keys are - identical to those of the - `~driver/zarr3/ChunkKeyEncoding.default` encoding. - examples: - - [-2, 0, 0] - required: - - offset codec_chain: $id: driver/zarr3/CodecChain title: Specifies a chain of codecs.