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 c4f1763e6..875b9de5f 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" @@ -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; } @@ -428,18 +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::kDefault ? 2 - : 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) { - if (part.empty() || !absl::ascii_isdigit(part.front()) || - !absl::ascii_isdigit(part.back()) || - !absl::SimpleAtoi(part, &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 && digits.front() == '0')) { return false; } + if (negative) grid_index = -grid_index; return true; }, key, grid_indices.subspan(0, metadata.rank))) { @@ -495,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 7d1bfeaac..ce6c75e99 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,15 +30,13 @@ #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_storage_statistics.h" #include "tensorstore/array_testutil.h" #include "tensorstore/batch.h" #include "tensorstore/box.h" @@ -1278,6 +1279,184 @@ TEST(ChunkKeyEncodingTest, V2Rank2Slash) { ::testing::Optional(::testing::ElementsAre("0/0", "0/1", "1/0", "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}}}}}}}}, + }; +} + +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(RegularBoundedChunkGridTest, DomainStartsAtTheLowerBound) { + auto context = Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open(RegularBoundedSpec(/*shape=*/9, /*origin=*/1, + /*lower_bound=*/-4), + context, tensorstore::OpenMode::create) + .result()); + 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 reopened, + tensorstore::Open({{"driver", "zarr3"}, {"kvstore", "memory://"}}, + context, tensorstore::OpenMode::open) + .result()); + EXPECT_THAT(tensorstore::Read(reopened).result(), ::testing::Optional(data)); +} + +TEST(RegularBoundedChunkGridTest, GrowsInFrontByEditingZarrJsonAlone) { + 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(); + 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)); + + // 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()); + + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto grown, + tensorstore::Open(spec, context, tensorstore::OpenMode::open).result()); + EXPECT_EQ(grown.domain().box(), tensorstore::BoxView<1>({-2}, {8})); + EXPECT_THAT(tensorstore::Read(grown).result(), + ::testing::Optional(tensorstore::MakeOffsetArray( + {-2}, {0, 0, 10, 11, 12, 13, 14, 15}))); + TENSORSTORE_ASSERT_OK( + tensorstore::Write(tensorstore::MakeOffsetArray({-2}, {1, 2}), + grown | tensorstore::Dims(0).HalfOpenInterval(-2, 0))); + EXPECT_THAT(tensorstore::Read(grown).result(), + ::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(RegularBoundedChunkGridTest, ResizeMovesTheLowerBound) { + auto context = Context::Default(); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto store, + tensorstore::Open(RegularBoundedSpec(/*shape=*/6, /*origin=*/0, + /*lower_bound=*/0), + context, tensorstore::OpenMode::create) + .result()); + 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()); + TENSORSTORE_ASSERT_OK( + tensorstore::kvstore::Write(store.kvstore(), "c/-0", absl::Cord("x")) + .result()); + 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 { tensorstore::internal::TestTensorStoreDriverSpecRoundtripOptions options; options.test_name = "zarr3"; diff --git a/tensorstore/driver/zarr3/metadata.cc b/tensorstore/driver/zarr3/metadata.cc index e1d69b950..24169cd66 100644 --- a/tensorstore/driver/zarr3/metadata.cc +++ b/tensorstore/driver/zarr3/metadata.cc @@ -526,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; @@ -533,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; @@ -620,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. @@ -974,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", @@ -1168,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(); @@ -1198,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())); } @@ -1222,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( @@ -1245,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); @@ -1325,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; } @@ -1393,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) { @@ -1403,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"); } @@ -1573,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]; @@ -1768,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 0bfb8071e..b4e9f6678 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" @@ -124,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; @@ -177,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; @@ -210,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 22b2850e2..5f045ac6f 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" @@ -151,6 +151,60 @@ TEST(MetadataTest, ParseValidNoDimensionNames) { EXPECT_THAT(metadata.user_attributes, MatchesJson({{"a", "b"}, {"c", "d"}})); } +TEST(MetadataTest, RegularBoundedChunkGrid) { + auto json = GetBasicMetadata(); + 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_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, RegularChunkGridRejectsBoundedMembers) { + auto json = GetBasicMetadata(); + json["chunk_grid"]["configuration"]["origin"] = {0, 0, 0}; + EXPECT_THAT(ZarrMetadata::FromJson(json), + StatusIs(absl::StatusCode::kInvalidArgument, + HasSubstr("regular_bounded"))); +} + +TEST(MetadataTest, RegularBoundedChunkGridRequiresBothMembers) { + auto json = GetBasicMetadata(); + 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, CompatibilityKeyIgnoresTheLowerBound) { + auto json = GetBasicMetadata(); + 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) { auto json = GetBasicMetadata(); json["foo"] = 42; diff --git a/tensorstore/driver/zarr3/schema.yml b/tensorstore/driver/zarr3/schema.yml index 4b1bc20ce..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: