zarr3: support the regular_bounded chunk grid

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",

tensorstore/driver/zarr3/chunk_cache.cc

@@ -104,12 +104,18 @@ internal::ChunkGridSpecification CreateFieldGridSpecification(
                                  field.dtype);
     }
 
-    std::vector<Index> 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<Index> component_chunk_shape(chunk_shape.begin(),
                                              chunk_shape.end());

tensorstore/driver/zarr3/driver.cc

@@ -20,6 +20,7 @@
 #include <cassert>
 #include <cstring>
 #include <memory>
+#include <nlohmann/json.hpp>
 #include <optional>
 #include <string>
 #include <string_view>
@@ -34,7 +35,6 @@
 #include "absl/strings/str_format.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
-#include <nlohmann/json.hpp>
 #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<DimensionIndex>(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<DimensionIndex>(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<const ZarrMetadata*>(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<Index> full_shape = metadata.shape;
     full_shape.insert(full_shape.end(), field.field_shape.begin(),
                       field.field_shape.end());
+    std::vector<Index> 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();
   }

tensorstore/driver/zarr3/driver_test.cc

@@ -14,28 +14,29 @@
 
 /// End-to-end tests of the zarr3 driver.
 
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
 #include <stddef.h>
 #include <stdint.h>
 
 #include <algorithm>
 #include <cstring>
 #include <functional>
+#include <nlohmann/json.hpp>
 #include <optional>
 #include <ostream>
 #include <string>
 #include <string_view>
 #include <utility>
 #include <vector>
 
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
 #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 <nlohmann/json.hpp>
 #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<std::string> ChunkKeys(const tensorstore::KvStore& kvs) {
+  std::vector<std::string> 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<uint8_t>({-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<uint8_t>, Schema::Shape({6}),
+                        ChunkLayout::ChunkShape({3}),
+                        tensorstore::OpenMode::create)
+          .result());
+  TENSORSTORE_ASSERT_OK(tensorstore::Write(
+      tensorstore::MakeArray<uint8_t>({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<uint8_t>(
+                  {-2}, {0, 0, 10, 11, 12, 13, 14, 15})));
+  TENSORSTORE_ASSERT_OK(
+      tensorstore::Write(tensorstore::MakeOffsetArray<uint8_t>({-2}, {1, 2}),
+                         grown | tensorstore::Dims(0).HalfOpenInterval(-2, 0)));
+  EXPECT_THAT(tensorstore::Read(grown).result(),
+              ::testing::Optional(tensorstore::MakeOffsetArray<uint8_t>(
+                  {-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<uint8_t>({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<uint8_t>({-4}, {1, 2, 3, 4}),
+      grown | tensorstore::Dims(0).HalfOpenInterval(-4, 0)));
+  EXPECT_THAT(tensorstore::Read(grown).result(),
+              ::testing::Optional(tensorstore::MakeOffsetArray<uint8_t>(
+                  {-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<uint8_t>(
+                  {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<uint8_t>(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";