From b14939b1b8fa15a210dc3ba67e4bf9bd3cb3908a Mon Sep 17 00:00:00 2001
From: Jeremy Maitin-Shepard <jbms@google.com>
Date: Thu, 15 Jul 2021 13:37:01 -0700
Subject: [PATCH] Add support for zarr dimension_separator metadata field

Historically zarr has named chunk files using dots by default,
e.g. "1.2.3", but has also supported slashes, e.g. "1/2/3".  However,
the choice was not indicated in the .zarray metadata file, and instead
needed to be specified out-of-band.  In TensorStore, that was
accomplished using the "key_encoding" spec member.

Zarr recently added a dimension_separator field to the metadata
format:

https://github.com/zarr-developers/zarr-python/pull/715

This commit adds support for that field.  Note that previously
TensorStore did not accept .zarray files that contain unknown
members, meaning that reading arrays with the dimension_separator
field would fail.

The existing "key_encoding" spec member is still supported, but
deprecated in favor of specifying "dimension_separator" as part of the
"metadata" field.  When creating a new array, TensorStore will always
write the "dimension_separator" field, in order to allow the array to
be re-opened without specifying the "dimension_separator".

Additionally, this change also makes TensorStore accept (and preserve)
extra .zarray members, as the zarr spec has now also changed to
require that.

PiperOrigin-RevId: 384998911
Change-Id: I485141b7786084d57c28005e3076f70462cd5b9e
---
 python/tensorstore/tensorstore_class.cc     |  10 ++
 tensorstore/driver/zarr/BUILD               |   1 +
 tensorstore/driver/zarr/driver.cc           |  69 ++++++---
 tensorstore/driver/zarr/driver_impl.h       |   2 +-
 tensorstore/driver/zarr/driver_impl_test.cc |   6 +-
 tensorstore/driver/zarr/driver_test.cc      | 159 ++++++++++++++++++--
 tensorstore/driver/zarr/metadata.cc         |  38 ++++-
 tensorstore/driver/zarr/metadata.h          |  23 +++
 tensorstore/driver/zarr/metadata_test.cc    |  20 +++
 tensorstore/driver/zarr/schema.yml          |  31 ++--
 tensorstore/driver/zarr/spec.cc             |  23 ++-
 tensorstore/driver/zarr/spec.h              |  16 --
 tensorstore/driver/zarr/spec_test.cc        |  40 ++---
 13 files changed, 327 insertions(+), 111 deletions(-)

diff --git a/python/tensorstore/tensorstore_class.cc b/python/tensorstore/tensorstore_class.cc
index 8e1835b33..e674a18f6 100644
--- a/python/tensorstore/tensorstore_class.cc
+++ b/python/tensorstore/tensorstore_class.cc
@@ -175,6 +175,7 @@ Asynchronous multi-dimensional array handle.
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<u4',
         'fill_value': None,
         'filters': None,
@@ -321,6 +322,7 @@ Spec that may be used to re-open or re-create the array.
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<u4',
         'fill_value': None,
         'filters': None,
@@ -1021,6 +1023,7 @@ that reflects the result of the indexing operation.  To read data, call
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<u4',
         'fill_value': None,
         'filters': None,
@@ -1148,6 +1151,7 @@ integer or boolean array indexing terms are applied orthogonally:
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<u4',
         'fill_value': None,
         'filters': None,
@@ -1257,6 +1261,7 @@ dimensions are unconditionally added as the first dimensions of the result
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<u4',
         'fill_value': None,
         'filters': None,
@@ -1378,6 +1383,7 @@ Computes a virtual view using an explicit :ref:`index transform<index-transform>
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<u4',
         'fill_value': None,
         'filters': None,
@@ -1840,6 +1846,7 @@ Opens or creates a :py:class:`TensorStore` from a :py:class:`Spec`.
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<i4',
         'fill_value': None,
         'filters': None,
@@ -1950,6 +1957,7 @@ properties that are left unconstrained:
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<f4',
         'fill_value': 42.0,
         'filters': None,
@@ -1996,6 +2004,7 @@ determine a matching chunk layout automatically:
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '<f4',
         'fill_value': None,
         'filters': None,
@@ -2080,6 +2089,7 @@ schema constraints:
           'id': 'blosc',
           'shuffle': -1,
         },
+        'dimension_separator': '.',
         'dtype': '>f4',
         'fill_value': None,
         'filters': None,
diff --git a/tensorstore/driver/zarr/BUILD b/tensorstore/driver/zarr/BUILD
index 3e52a5740..8ab72f90e 100644
--- a/tensorstore/driver/zarr/BUILD
+++ b/tensorstore/driver/zarr/BUILD
@@ -262,6 +262,7 @@ tensorstore_cc_library(
         "//tensorstore:schema",
         "//tensorstore/index_space:index_transform_builder",
         "//tensorstore/internal:json",
+        "//tensorstore/internal:json_metadata_matching",
         "//tensorstore/util:quote_string",
         "//tensorstore/util:result",
         "@com_github_nlohmann_json//:nlohmann_json",
diff --git a/tensorstore/driver/zarr/driver.cc b/tensorstore/driver/zarr/driver.cc
index 9994e80a2..3875c71d4 100644
--- a/tensorstore/driver/zarr/driver.cc
+++ b/tensorstore/driver/zarr/driver.cc
@@ -41,8 +41,18 @@ namespace internal_zarr {
 namespace {
 constexpr const char kZarrMetadataKey[] = ".zarray";
 
-inline char GetChunkKeyEncodingSeparator(ChunkKeyEncoding key_encoding) {
-  return key_encoding == ChunkKeyEncoding::kDotSeparated ? '.' : '/';
+inline char GetDimensionSeparatorChar(DimensionSeparator dimension_separator) {
+  return dimension_separator == DimensionSeparator::kDotSeparated ? '.' : '/';
+}
+
+DimensionSeparator GetDimensionSeparator(
+    const ZarrPartialMetadata& partial_metadata, const ZarrMetadata& metadata) {
+  if (metadata.dimension_separator) {
+    return *metadata.dimension_separator;
+  } else if (partial_metadata.dimension_separator) {
+    return *partial_metadata.dimension_separator;
+  }
+  return DimensionSeparator::kDotSeparated;
 }
 
 Result<ZarrMetadataPtr> ParseEncodedMetadata(std::string_view encoded_value) {
@@ -95,15 +105,13 @@ class ZarrDriver
   template <template <typename> class MaybeBound = internal::ContextUnbound>
   struct SpecT : public internal_kvs_backed_chunk_driver::SpecT<MaybeBound> {
     std::string key_prefix;
-    ChunkKeyEncoding key_encoding;
     ZarrPartialMetadata partial_metadata;
     SelectedField selected_field;
 
     constexpr static auto ApplyMembers = [](auto& x, auto f) {
       return f(internal::BaseCast<
                    internal_kvs_backed_chunk_driver::SpecT<MaybeBound>>(x),
-               x.key_prefix, x.key_encoding, x.partial_metadata,
-               x.selected_field);
+               x.key_prefix, x.partial_metadata, x.selected_field);
     };
   };
 
@@ -123,12 +131,25 @@ class ZarrDriver
       internal_kvs_backed_chunk_driver::SpecJsonBinder,
       jb::Member("path", jb::Projection(&SpecT<>::key_prefix,
                                         jb::DefaultInitializedValue())),
-      jb::Member("key_encoding",
-                 jb::Projection(
-                     &SpecT<>::key_encoding,
-                     jb::DefaultInitializedValue(ChunkKeyEncodingJsonBinder))),
       jb::Member("metadata", jb::Projection(&SpecT<>::partial_metadata,
                                             jb::DefaultInitializedValue())),
+      // Deprecated `key_encoding` property.
+      jb::LoadSave(jb::OptionalMember(
+          "key_encoding",
+          jb::Compose<DimensionSeparator>(
+              [](auto is_loading, const auto& options, auto* obj,
+                 DimensionSeparator* value) {
+                auto& sep = obj->partial_metadata.dimension_separator;
+                if (sep && *sep != *value) {
+                  return absl::InvalidArgumentError(tensorstore::StrCat(
+                      "value (", ::nlohmann::json(*value).dump(),
+                      ") does not match value in metadata (",
+                      ::nlohmann::json(*sep).dump(), ")"));
+                }
+                sep = *value;
+                return absl::OkStatus();
+              },
+              DimensionSeparatorJsonBinder))),
       jb::Member("field",
                  jb::Projection(&SpecT<>::selected_field,
                                 jb::DefaultValue<jb::kNeverIncludeDefaults>(
@@ -224,12 +245,12 @@ class DataCache : public internal_kvs_backed_chunk_driver::DataCache {
 
  public:
   explicit DataCache(Initializer initializer, std::string key_prefix,
-                     ChunkKeyEncoding key_encoding)
+                     DimensionSeparator dimension_separator)
       : Base(initializer,
              GetChunkGridSpecification(*static_cast<const ZarrMetadata*>(
                  initializer.metadata.get()))),
         key_prefix_(std::move(key_prefix)),
-        key_encoding_(key_encoding) {}
+        dimension_separator_(dimension_separator) {}
 
   Status ValidateMetadataCompatibility(const void* existing_metadata_ptr,
                                        const void* new_metadata_ptr) override {
@@ -345,8 +366,8 @@ class DataCache : public internal_kvs_backed_chunk_driver::DataCache {
 
   std::string GetChunkStorageKey(const void* metadata,
                                  span<const Index> cell_indices) override {
-    return internal::JoinPath(key_prefix_,
-                              EncodeChunkIndices(cell_indices, key_encoding_));
+    return internal::JoinPath(
+        key_prefix_, EncodeChunkIndices(cell_indices, dimension_separator_));
   }
 
   Status GetBoundSpecData(
@@ -358,7 +379,6 @@ class DataCache : public internal_kvs_backed_chunk_driver::DataCache {
     const auto& metadata = *static_cast<const ZarrMetadata*>(metadata_ptr);
     spec.key_prefix = key_prefix_;
     spec.selected_field = EncodeSelectedField(component_index, metadata.dtype);
-    spec.key_encoding = key_encoding_;
     auto& pm = spec.partial_metadata;
     pm.rank = metadata.rank;
     pm.zarr_format = metadata.zarr_format;
@@ -369,6 +389,7 @@ class DataCache : public internal_kvs_backed_chunk_driver::DataCache {
     pm.order = metadata.order;
     pm.dtype = metadata.dtype;
     pm.fill_value = metadata.fill_value;
+    pm.dimension_separator = dimension_separator_;
     return absl::OkStatus();
   }
 
@@ -391,7 +412,7 @@ class DataCache : public internal_kvs_backed_chunk_driver::DataCache {
 
  private:
   std::string key_prefix_;
-  ChunkKeyEncoding key_encoding_;
+  DimensionSeparator dimension_separator_;
 };
 
 class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase {
@@ -427,15 +448,21 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase {
   std::string GetDataCacheKey(const void* metadata) override {
     std::string result;
     const auto& spec = this->spec();
-    internal::EncodeCacheKey(&result, spec.key_prefix, spec.key_encoding,
-                             *static_cast<const ZarrMetadata*>(metadata));
+    const auto& zarr_metadata = *static_cast<const ZarrMetadata*>(metadata);
+    internal::EncodeCacheKey(
+        &result, spec.key_prefix,
+        GetDimensionSeparator(spec.partial_metadata, zarr_metadata),
+        zarr_metadata);
     return result;
   }
 
   std::unique_ptr<internal_kvs_backed_chunk_driver::DataCache> GetDataCache(
       DataCache::Initializer initializer) override {
-    return std::make_unique<DataCache>(std::move(initializer),
-                                       spec().key_prefix, spec().key_encoding);
+    const auto& metadata =
+        *static_cast<const ZarrMetadata*>(initializer.metadata.get());
+    return std::make_unique<DataCache>(
+        std::move(initializer), spec().key_prefix,
+        GetDimensionSeparator(spec().partial_metadata, metadata));
   }
 
   Result<std::size_t> GetComponentIndex(const void* metadata_ptr,
@@ -456,8 +483,8 @@ const internal::DriverRegistration<ZarrDriver> registration;
 }  // namespace
 
 std::string EncodeChunkIndices(span<const Index> indices,
-                               ChunkKeyEncoding key_encoding) {
-  const char separator = GetChunkKeyEncodingSeparator(key_encoding);
+                               DimensionSeparator dimension_separator) {
+  const char separator = GetDimensionSeparatorChar(dimension_separator);
   std::string key;
   for (DimensionIndex i = 0; i < indices.size(); ++i) {
     if (i != 0) {
diff --git a/tensorstore/driver/zarr/driver_impl.h b/tensorstore/driver/zarr/driver_impl.h
index 6785299d9..cc6d1d6ef 100644
--- a/tensorstore/driver/zarr/driver_impl.h
+++ b/tensorstore/driver/zarr/driver_impl.h
@@ -36,7 +36,7 @@ namespace internal_zarr {
 
 /// Encodes a chunk grid index vector as a storage key suffix.
 std::string EncodeChunkIndices(span<const Index> indices,
-                               ChunkKeyEncoding key_encoding);
+                               DimensionSeparator dimension_separator);
 
 }  // namespace internal_zarr
 }  // namespace tensorstore
diff --git a/tensorstore/driver/zarr/driver_impl_test.cc b/tensorstore/driver/zarr/driver_impl_test.cc
index 00538f389..514c18595 100644
--- a/tensorstore/driver/zarr/driver_impl_test.cc
+++ b/tensorstore/driver/zarr/driver_impl_test.cc
@@ -40,7 +40,7 @@ using tensorstore::Result;
 using tensorstore::span;
 using tensorstore::TransactionMode;
 using tensorstore::internal_kvs_backed_chunk_driver::ResizeParameters;
-using tensorstore::internal_zarr::ChunkKeyEncoding;
+using tensorstore::internal_zarr::DimensionSeparator;
 using tensorstore::internal_zarr::ZarrMetadata;
 
 Result<tensorstore::IndexTransform<>> ResolveBoundsFromMetadata(
@@ -88,12 +88,12 @@ Result<ResizeParameters> GetResizeParameters(
 
 TEST(EncodeChunkIndicesTest, DotSeparated) {
   EXPECT_EQ("1.2.3", EncodeChunkIndices(span<const Index>({1, 2, 3}),
-                                        ChunkKeyEncoding::kDotSeparated));
+                                        DimensionSeparator::kDotSeparated));
 }
 
 TEST(EncodeChunkIndicesTest, SlashSeparated) {
   EXPECT_EQ("1/2/3", EncodeChunkIndices(span<const Index>({1, 2, 3}),
-                                        ChunkKeyEncoding::kSlashSeparated));
+                                        DimensionSeparator::kSlashSeparated));
 }
 
 TEST(ResolveBoundsFromMetadataTest, Basic) {
diff --git a/tensorstore/driver/zarr/driver_test.cc b/tensorstore/driver/zarr/driver_test.cc
index 29fe68e6b..acac118df 100644
--- a/tensorstore/driver/zarr/driver_test.cc
+++ b/tensorstore/driver/zarr/driver_test.cc
@@ -234,6 +234,7 @@ TEST(ZarrDriverTest, Create) {
                    {"dtype", "<i2"},
                    {"shape", {100, 100}},
                    {"chunks", {3, 2}},
+                   {"dimension_separator", "."},
                }))),
           Pair("prefix/3.4",    //
                DecodedMatches(  //
@@ -639,6 +640,7 @@ TEST(ZarrDriverTest, CreateBigEndian) {
                    {"dtype", ">i2"},
                    {"shape", {100, 100}},
                    {"chunks", {3, 2}},
+                   {"dimension_separator", "."},
                }))),
           Pair("prefix/3.4",    //
                DecodedMatches(  //
@@ -701,6 +703,7 @@ TEST(ZarrDriverTest, CreateBfloat16) {
                    {"dtype", "bfloat16"},
                    {"shape", {100, 100}},
                    {"chunks", {3, 2}},
+                   {"dimension_separator", "."},
                }))),
           Pair("prefix/1.1",  //
                ::testing::MatcherCast<absl::Cord>(
@@ -753,6 +756,7 @@ TEST(ZarrDriverTest, CreateBigEndianUnaligned) {
                     ::nlohmann::json::array_t{{"x", "|b1"}, {"y", ">i2"}}},
                    {"shape", {100, 100}},
                    {"chunks", {3, 2}},
+                   {"dimension_separator", "."},
                }))),
           Pair(
               "prefix/3.4",    //
@@ -824,6 +828,7 @@ TEST(ZarrDriverTest, CreateLittleEndianUnaligned) {
                     ::nlohmann::json::array_t{{"x", "|b1"}, {"y", "<i2"}}},
                    {"shape", {100, 100}},
                    {"chunks", {3, 2}},
+                   {"dimension_separator", "."},
                }))),
           Pair(
               "prefix/3.4",    //
@@ -866,9 +871,11 @@ TEST(ZarrDriverTest, CreateComplexWithFillValue) {
 
 ::nlohmann::json GetBasicResizeMetadata() {
   return {
-      {"zarr_format", 2},      {"order", "C"},          {"filters", nullptr},
-      {"fill_value", nullptr}, {"compressor", nullptr}, {"dtype", "|i1"},
-      {"shape", {100, 100}},   {"chunks", {3, 2}},
+      {"zarr_format", 2},           {"order", "C"},
+      {"filters", nullptr},         {"fill_value", nullptr},
+      {"compressor", nullptr},      {"dtype", "|i1"},
+      {"shape", {100, 100}},        {"chunks", {3, 2}},
+      {"dimension_separator", "."},
   };
 }
 
@@ -877,6 +884,7 @@ TEST(ZarrDriverTest, KeyEncodingWithSlash) {
   // Create the store.
   ::nlohmann::json storage_spec{{"driver", "memory"}};
   ::nlohmann::json zarr_metadata_json = GetBasicResizeMetadata();
+  zarr_metadata_json.erase("dimension_separator");
   ::nlohmann::json json_spec{
       {"driver", "zarr"},
       {"kvstore", storage_spec},
@@ -884,6 +892,7 @@ TEST(ZarrDriverTest, KeyEncodingWithSlash) {
       {"key_encoding", "/"},
       {"metadata", zarr_metadata_json},
   };
+  zarr_metadata_json["dimension_separator"] = "/";
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(
       auto store,
       tensorstore::Open(json_spec, context, tensorstore::OpenMode::create,
@@ -1996,19 +2005,22 @@ TENSORSTORE_GLOBAL_INITIALIZER {
       {"driver", "zarr"},
       {"path", "prefix"},
       {"metadata",
-       {{"chunks", {3, 2}},
-        {"compressor",
-         {{"blocksize", 0},
-          {"clevel", 5},
-          {"cname", "lz4"},
-          {"id", "blosc"},
-          {"shuffle", -1}}},
-        {"dtype", "<i2"},
-        {"fill_value", nullptr},
-        {"filters", nullptr},
-        {"order", "C"},
-        {"shape", {100, 100}},
-        {"zarr_format", 2}}},
+       {
+           {"chunks", {3, 2}},
+           {"compressor",
+            {{"blocksize", 0},
+             {"clevel", 5},
+             {"cname", "lz4"},
+             {"id", "blosc"},
+             {"shuffle", -1}}},
+           {"dtype", "<i2"},
+           {"fill_value", nullptr},
+           {"filters", nullptr},
+           {"order", "C"},
+           {"shape", {100, 100}},
+           {"zarr_format", 2},
+           {"dimension_separator", "."},
+       }},
       {"kvstore", {{"driver", "memory"}}},
       {"transform",
        {{"input_exclusive_max", {{100}, {100}}},
@@ -2074,6 +2086,7 @@ TENSORSTORE_GLOBAL_INITIALIZER {
              {"shape", bounds.shape()},
              {"chunks", {4, 5}},
              {"filters", nullptr},
+             {"dimension_separator", "."},
          }},
         {"transform",
          {
@@ -2094,12 +2107,14 @@ TEST(DriverTest, NoPrefix) {
   // Create the store.
   ::nlohmann::json storage_spec{{"driver", "memory"}};
   ::nlohmann::json zarr_metadata_json = GetBasicResizeMetadata();
+  zarr_metadata_json.erase("dimension_separator");
   ::nlohmann::json json_spec{
       {"driver", "zarr"},
       {"kvstore", storage_spec},
       {"key_encoding", "/"},
       {"metadata", zarr_metadata_json},
   };
+  zarr_metadata_json["dimension_separator"] = "/";
   TENSORSTORE_ASSERT_OK_AND_ASSIGN(
       auto store,
       tensorstore::Open(json_spec, context, tensorstore::OpenMode::create,
@@ -2122,6 +2137,37 @@ TEST(DriverTest, NoPrefix) {
           Pair("1/1", Bytes({5, 6, 0, 0, 0, 0}))));
 }
 
+TEST(DriverTest, DimensionSeparatorMatch) {
+  ::nlohmann::json storage_spec{{"driver", "memory"}};
+  ::nlohmann::json zarr_metadata_json = GetBasicResizeMetadata();
+  zarr_metadata_json["dimension_separator"] = "/";
+  ::nlohmann::json json_spec{
+      {"driver", "zarr"},
+      {"kvstore", storage_spec},
+      {"key_encoding", "/"},
+      {"metadata", zarr_metadata_json},
+  };
+  TENSORSTORE_EXPECT_OK(
+      tensorstore::Open(json_spec, tensorstore::OpenMode::create));
+}
+
+TEST(DriverTest, DimensionSeparatorMismatch) {
+  ::nlohmann::json storage_spec{{"driver", "memory"}};
+  ::nlohmann::json zarr_metadata_json = GetBasicResizeMetadata();
+  ::nlohmann::json json_spec{
+      {"driver", "zarr"},
+      {"kvstore", storage_spec},
+      {"key_encoding", "/"},
+      {"metadata", zarr_metadata_json},
+  };
+  EXPECT_THAT(
+      tensorstore::Open(json_spec, tensorstore::OpenMode::create).result(),
+      MatchesStatus(
+          absl::StatusCode::kInvalidArgument,
+          "Error parsing object member \"key_encoding\": "
+          "value \\(\"/\"\\) does not match value in metadata \\(\"\\.\"\\)"));
+}
+
 TEST(DriverTest, ChunkLayout) {
   ::nlohmann::json json_spec{
       {"driver", "zarr"},
@@ -2498,4 +2544,85 @@ TEST(DriverCreateCheckSchemaTest, Basic) {
       });
 }
 
+void TestReadWriteWithDimensionSeparator(std::string dimension_separator) {
+  std::string other_dimension_separator =
+      dimension_separator == "." ? "/" : ".";
+  ::nlohmann::json storage_spec{{"driver", "memory"}};
+  auto context = tensorstore::Context::Default();
+
+  // Write array
+  {
+    TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+        auto store,
+        tensorstore::Open({{"driver", "zarr"},
+                           {"kvstore", storage_spec},
+                           {
+                               "metadata",
+                               {{"dimension_separator", dimension_separator}},
+                           }},
+                          context, tensorstore::OpenMode::create,
+                          dtype_v<uint32_t>, Schema::Shape({1, 1}))
+            .result());
+    TENSORSTORE_EXPECT_OK(
+        tensorstore::Write(tensorstore::MakeScalarArray<uint32_t>(42), store));
+  }
+
+  const auto perform_read = [&](::nlohmann::json metadata)
+      -> tensorstore::Result<tensorstore::SharedArray<const void>> {
+    TENSORSTORE_ASSIGN_OR_RETURN(auto store,
+                                 tensorstore::Open({{"driver", "zarr"},
+                                                    {"kvstore", storage_spec},
+                                                    {"metadata", metadata}},
+                                                   context)
+                                     .result());
+    return tensorstore::Read<tensorstore::zero_origin>(store).result();
+  };
+
+  const auto matches_write =
+      ::testing::Optional(tensorstore::MakeArray<uint32_t>({{42}}));
+
+  // Read array with explicit dimension_separator.
+  EXPECT_THAT(perform_read({{"dimension_separator", dimension_separator}}),
+              matches_write);
+
+  // Read array without explicit dimension_separator.  This relies on the
+  // dimension_separator in the `.zarray` file.
+  EXPECT_THAT(perform_read(::nlohmann::json::object_t()), matches_write);
+
+  // Read array with wrong dimension_separator.
+  EXPECT_THAT(
+      perform_read({{"dimension_separator", other_dimension_separator}}),
+      MatchesStatus(absl::StatusCode::kFailedPrecondition,
+                    ".*: Expected \"dimension_separator\" of \"[./]\" but "
+                    "received: \"[./]\""));
+
+  // Remove dimension_separator field from metadata.
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(
+      auto kv_store, KeyValueStore::Open(context, storage_spec, {}).result());
+  TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata_json_encoded,
+                                   kv_store->Read(".zarray").result());
+  auto metadata_json =
+      tensorstore::internal::ParseJson(metadata_json_encoded.value.Flatten());
+  metadata_json.erase("dimension_separator");
+  TENSORSTORE_ASSERT_OK(
+      kv_store->Write(".zarray", absl::Cord(metadata_json.dump())));
+
+  // Read array with explicit dimension_separator.
+  EXPECT_THAT(perform_read({{"dimension_separator", dimension_separator}}),
+              matches_write);
+
+  if (dimension_separator != ".") {
+    // Read array with wrong dimension_separator.  The fill value is returned
+    // since the wrong chunk key is used.
+    EXPECT_THAT(perform_read(::nlohmann::json::object_t()),
+                ::testing::Optional(tensorstore::MakeArray<uint32_t>({{0}})));
+  }
+}
+
+TEST(DimensionSeparatorTest, Dot) { TestReadWriteWithDimensionSeparator("."); }
+
+TEST(DimensionSeparatorTest, Slash) {
+  TestReadWriteWithDimensionSeparator("/");
+}
+
 }  // namespace
diff --git a/tensorstore/driver/zarr/metadata.cc b/tensorstore/driver/zarr/metadata.cc
index 46eae4ce7..d6bd37183 100644
--- a/tensorstore/driver/zarr/metadata.cc
+++ b/tensorstore/driver/zarr/metadata.cc
@@ -36,6 +36,19 @@ TENSORSTORE_DEFINE_JSON_BINDER(
     jb::Enum<ContiguousLayoutOrder, std::string_view>({{c_order, "C"},
                                                        {fortran_order, "F"}}))
 
+TENSORSTORE_DEFINE_JSON_BINDER(DimensionSeparatorJsonBinder,
+                               jb::Enum<DimensionSeparator, std::string_view>(
+                                   {{DimensionSeparator::kDotSeparated, "."},
+                                    {DimensionSeparator::kSlashSeparated,
+                                     "/"}}))
+
+void to_json(::nlohmann::json& out, DimensionSeparator value) {
+  DimensionSeparatorJsonBinder(/*is_loading=*/std::false_type{},
+                               /*options=*/internal_json_binding::NoOptions{},
+                               &value, &out)
+      .IgnoreError();
+}
+
 namespace {
 
 Result<double> DecodeFloat(const nlohmann::json& j) {
@@ -301,8 +314,18 @@ constexpr auto MetadataJsonBinder = [](auto maybe_optional) {
                        }))),
         jb::Member("order",
                    jb::Projection(&T::order, maybe_optional(OrderJsonBinder))),
-        jb::Member("filters", jb::Projection(&T::filters)))(is_loading, options,
-                                                            obj, j);
+        jb::Member("filters", jb::Projection(&T::filters)),
+        jb::Member("dimension_separator",
+                   jb::Projection(&T::dimension_separator,
+                                  jb::Optional(DimensionSeparatorJsonBinder))),
+        [](auto is_loading, const auto& options, auto* obj, auto* j) {
+          if constexpr (std::is_same_v<T, ZarrMetadata>) {
+            return jb::DefaultBinder<>(is_loading, options, &obj->extra_members,
+                                       j);
+          } else {
+            return absl::OkStatus();
+          }
+        })(is_loading, options, obj, j);
   };
 };
 
@@ -399,11 +422,22 @@ Result<absl::Cord> EncodeChunk(const ZarrMetadata& metadata,
 bool IsMetadataCompatible(const ZarrMetadata& a, const ZarrMetadata& b) {
   // Rank must be the same.
   if (a.shape.size() != b.shape.size()) return false;
+
   auto a_json = ::nlohmann::json(a);
   auto b_json = ::nlohmann::json(b);
+
   // Shape is allowed to differ.
   a_json.erase("shape");
   b_json.erase("shape");
+
+  // Extra members are allowed to differ.
+  for (const auto& [key, value] : a.extra_members) {
+    a_json.erase(key);
+  }
+  for (const auto& [key, value] : b.extra_members) {
+    b_json.erase(key);
+  }
+
   return a_json == b_json;
 }
 
diff --git a/tensorstore/driver/zarr/metadata.h b/tensorstore/driver/zarr/metadata.h
index e9b43fc63..0d04c7be2 100644
--- a/tensorstore/driver/zarr/metadata.h
+++ b/tensorstore/driver/zarr/metadata.h
@@ -78,6 +78,22 @@ TENSORSTORE_DECLARE_JSON_BINDER(OrderJsonBinder, ContiguousLayoutOrder,
                                 internal_json_binding::NoOptions,
                                 internal_json_binding::NoOptions)
 
+/// Specifies how chunk index vectors are encoded as keys.
+///
+/// Index vectors are encoded in order as their base-10 ASCII representation,
+/// separated by either "." or "/".
+enum class DimensionSeparator {
+  kDotSeparated = 0,
+  kSlashSeparated = 1,
+};
+
+TENSORSTORE_DECLARE_JSON_BINDER(DimensionSeparatorJsonBinder,
+                                DimensionSeparator,
+                                internal_json_binding::NoOptions,
+                                internal_json_binding::NoOptions)
+
+void to_json(::nlohmann::json& out, DimensionSeparator value);
+
 /// Parsed representation of a zarr `.zarray` metadata JSON file.
 struct ZarrMetadata {
   // The following members are common to both `ZarrMetadata` and
@@ -105,6 +121,11 @@ struct ZarrMetadata {
   /// `dtype.fields`.
   std::vector<SharedArray<const void>> fill_value;
 
+  /// If not specified, the open-time option is used instead.
+  std::optional<DimensionSeparator> dimension_separator;
+
+  ::nlohmann::json::object_t extra_members;
+
   // Derived information computed from `dtype`, `order`, and `chunks`.
 
   ZarrChunkLayout chunk_layout;
@@ -151,6 +172,8 @@ struct ZarrPartialMetadata {
   /// `dtype.fields`.
   std::optional<std::vector<SharedArray<const void>>> fill_value;
 
+  std::optional<DimensionSeparator> dimension_separator;
+
   TENSORSTORE_DECLARE_JSON_DEFAULT_BINDER(ZarrPartialMetadata,
                                           internal_json_binding::NoOptions,
                                           tensorstore::IncludeDefaults)
diff --git a/tensorstore/driver/zarr/metadata_test.cc b/tensorstore/driver/zarr/metadata_test.cc
index ff05547a6..537459c67 100644
--- a/tensorstore/driver/zarr/metadata_test.cc
+++ b/tensorstore/driver/zarr/metadata_test.cc
@@ -35,6 +35,8 @@ using tensorstore::float16_t;
 using tensorstore::MakeArray;
 using tensorstore::MakeScalarArray;
 using tensorstore::MatchesStatus;
+using tensorstore::internal_zarr::DimensionSeparator;
+using tensorstore::internal_zarr::DimensionSeparatorJsonBinder;
 using tensorstore::internal_zarr::EncodeFillValue;
 using tensorstore::internal_zarr::OrderJsonBinder;
 using tensorstore::internal_zarr::ParseDType;
@@ -692,4 +694,22 @@ TEST(ParseMetadataTest, InvalidRank) {
   });
 }
 
+TEST(DimensionSeparatorTest, JsonBinderTest) {
+  tensorstore::TestJsonBinderRoundTrip<DimensionSeparator>(
+      {
+          {DimensionSeparator::kDotSeparated, "."},
+          {DimensionSeparator::kSlashSeparated, "/"},
+      },
+      DimensionSeparatorJsonBinder);
+}
+
+TEST(DimensionSeparatorTest, JsonBinderTestInvalid) {
+  tensorstore::TestJsonBinderFromJson<DimensionSeparator>(
+      {
+          {"x", MatchesStatus(absl::StatusCode::kInvalidArgument)},
+          {3, MatchesStatus(absl::StatusCode::kInvalidArgument)},
+      },
+      DimensionSeparatorJsonBinder);
+}
+
 }  // namespace
diff --git a/tensorstore/driver/zarr/schema.yml b/tensorstore/driver/zarr/schema.yml
index 7cd7f1870..1a93ca5f8 100644
--- a/tensorstore/driver/zarr/schema.yml
+++ b/tensorstore/driver/zarr/schema.yml
@@ -15,17 +15,6 @@ allOf:
         Must be specified if the `.metadata.dtype` specified in the array
         metadata has more than one field.
       default: null
-    key_encoding:
-      enum:
-      - .
-      - /
-      type: string
-      default: .
-      title: Specifies the encoding of chunk indices into key-value store keys.
-      description: |
-        The default value of :json:`"."` corresponds to the default encoding
-        used by Zarr, while :json:`"/"` corresponds to the encoding used by
-        :py:obj:`zarr.storage.NestedDirectoryStore`.
     metadata:
       title: Zarr array metadata.
       description: |
@@ -93,6 +82,26 @@ allOf:
               :json:`"C"` for C order, :json:`"F"` for Fortran order.  When
               creating a new array, defaults to :json:`"C"`.
       - $ref: "#codec-properties"
+      - type: object
+        properties:
+          dimension_separator:
+            enum:
+              - "."
+              - "/"
+            title: Specifies the encoding of chunk indices into key-value store keys.
+            description: |
+              The default value of :json:`"."` corresponds to the default encoding
+              used by Zarr, while :json:`"/"` corresponds to the encoding used by
+              :py:obj:`zarr.storage.NestedDirectoryStore`.
+    key_encoding:
+      enum:
+      - .
+      - /
+      type: string
+      default: .
+      title: Specifies the encoding of chunk indices into key-value store keys.
+      description: |
+        Deprecated.  Equivalent to specifying `.metadata.dimension_separator`.
 examples:
 - driver: zarr
   "kvstore":
diff --git a/tensorstore/driver/zarr/spec.cc b/tensorstore/driver/zarr/spec.cc
index 7df982d38..d835db236 100644
--- a/tensorstore/driver/zarr/spec.cc
+++ b/tensorstore/driver/zarr/spec.cc
@@ -17,11 +17,13 @@
 #include "tensorstore/codec_spec_registry.h"
 #include "tensorstore/index_space/index_domain_builder.h"
 #include "tensorstore/internal/json.h"
+#include "tensorstore/internal/json_metadata_matching.h"
 #include "tensorstore/util/quote_string.h"
 
 namespace tensorstore {
 namespace internal_zarr {
 
+using internal::MetadataMismatchError;
 namespace jb = tensorstore::internal_json_binding;
 
 CodecSpec::Ptr ZarrCodecSpec::Clone() const {
@@ -58,14 +60,6 @@ TENSORSTORE_DEFINE_JSON_DEFAULT_BINDER(
 namespace {
 const internal::CodecSpecRegistration<ZarrCodecSpec> encoding_registration;
 
-template <typename T>
-Status MetadataMismatchError(const char* name, const T& expected,
-                             const T& actual) {
-  return absl::FailedPreconditionError(StrCat(
-      "Expected ", QuoteString(name), " of ", ::nlohmann::json(expected).dump(),
-      " but received: ", ::nlohmann::json(actual).dump()));
-}
-
 }  // namespace
 
 void GetChunkInnerOrder(DimensionIndex chunked_rank,
@@ -109,6 +103,12 @@ Status ValidateMetadata(const ZarrMetadata& metadata,
       return MetadataMismatchError("fill_value", a, b);
     }
   }
+  if (constraints.dimension_separator && metadata.dimension_separator &&
+      *constraints.dimension_separator != *metadata.dimension_separator) {
+    return MetadataMismatchError("dimension_separator",
+                                 *constraints.dimension_separator,
+                                 *metadata.dimension_separator);
+  }
   return absl::OkStatus();
 }
 
@@ -117,6 +117,8 @@ Result<ZarrMetadataPtr> GetNewMetadata(
     const SelectedField& selected_field, const Schema& schema) {
   ZarrMetadataPtr metadata = std::make_shared<ZarrMetadata>();
   metadata->zarr_format = partial_metadata.zarr_format.value_or(2);
+  metadata->dimension_separator = partial_metadata.dimension_separator.value_or(
+      DimensionSeparator::kDotSeparated);
 
   // Determine the new zarr metadata based on `partial_metadata` and `schema`.
   // Note that zarr dtypes can be NumPy "structured data types"
@@ -512,10 +514,5 @@ SelectedField EncodeSelectedField(std::size_t field_index,
   return field.name;
 }
 
-TENSORSTORE_DEFINE_JSON_BINDER(ChunkKeyEncodingJsonBinder,
-                               jb::Enum<ChunkKeyEncoding, std::string_view>(
-                                   {{ChunkKeyEncoding::kDotSeparated, "."},
-                                    {ChunkKeyEncoding::kSlashSeparated, "/"}}))
-
 }  // namespace internal_zarr
 }  // namespace tensorstore
diff --git a/tensorstore/driver/zarr/spec.h b/tensorstore/driver/zarr/spec.h
index 3e9e1bf4d..a4ec10053 100644
--- a/tensorstore/driver/zarr/spec.h
+++ b/tensorstore/driver/zarr/spec.h
@@ -55,22 +55,6 @@ class ZarrCodecSpec : public CodecSpec {
 Status ValidateMetadata(const ZarrMetadata& metadata,
                         const ZarrPartialMetadata& constraints);
 
-/// Specifies how chunk index vectors are encoded as keys.
-///
-/// Index vectors are encoded in order as their base-10 ASCII representation,
-/// separated by either "." or "/".
-///
-/// The zarr metadata format does not specify the key encoding; therefore, it
-/// must be specified separately when opening the driver.
-enum class ChunkKeyEncoding {
-  kDotSeparated = 0,
-  kSlashSeparated = 1,
-};
-
-TENSORSTORE_DECLARE_JSON_BINDER(ChunkKeyEncodingJsonBinder, ChunkKeyEncoding,
-                                internal_json_binding::NoOptions,
-                                internal_json_binding::NoOptions)
-
 /// An empty string indicates the singleton field if the dtype does not have
 /// fields.
 using SelectedField = std::string;
diff --git a/tensorstore/driver/zarr/spec_test.cc b/tensorstore/driver/zarr/spec_test.cc
index 8a32c2ac3..45fb63f68 100644
--- a/tensorstore/driver/zarr/spec_test.cc
+++ b/tensorstore/driver/zarr/spec_test.cc
@@ -32,8 +32,6 @@ using tensorstore::dtype_v;
 using tensorstore::MatchesJson;
 using tensorstore::MatchesStatus;
 using tensorstore::Schema;
-using tensorstore::internal_zarr::ChunkKeyEncoding;
-using tensorstore::internal_zarr::ChunkKeyEncodingJsonBinder;
 using tensorstore::internal_zarr::GetFieldIndex;
 using tensorstore::internal_zarr::ParseDType;
 using tensorstore::internal_zarr::ParseSelectedField;
@@ -41,14 +39,6 @@ using tensorstore::internal_zarr::SelectedField;
 using tensorstore::internal_zarr::ZarrMetadata;
 using tensorstore::internal_zarr::ZarrPartialMetadata;
 
-TEST(ParsePartialMetadataTest, ExtraMember) {
-  tensorstore::TestJsonBinderFromJson<ZarrPartialMetadata>({
-      {{{"foo", "x"}},
-       MatchesStatus(absl::StatusCode::kInvalidArgument,
-                     "Object includes extra members: \"foo\"")},
-  });
-}
-
 TEST(ParsePartialMetadataTest, InvalidZarrFormat) {
   tensorstore::TestJsonBinderFromJson<ZarrPartialMetadata>({
       {{{"zarr_format", "2"}},
@@ -158,24 +148,6 @@ TEST(ParsePartialMetadataTest, Complete) {
   EXPECT_THAT(*result.chunks, ::testing::ElementsAre(3, 2));
 }
 
-TEST(ChunkKeyEncodingTest, JsonBinderTest) {
-  tensorstore::TestJsonBinderRoundTrip<ChunkKeyEncoding>(
-      {
-          {ChunkKeyEncoding::kDotSeparated, "."},
-          {ChunkKeyEncoding::kSlashSeparated, "/"},
-      },
-      ChunkKeyEncodingJsonBinder);
-}
-
-TEST(ChunkKeyEncodingTest, JsonBinderTestInvalid) {
-  tensorstore::TestJsonBinderFromJson<ChunkKeyEncoding>(
-      {
-          {"x", MatchesStatus(absl::StatusCode::kInvalidArgument)},
-          {3, MatchesStatus(absl::StatusCode::kInvalidArgument)},
-      },
-      ChunkKeyEncodingJsonBinder);
-}
-
 TEST(ParseSelectedFieldTest, Null) {
   EXPECT_EQ(SelectedField(), ParseSelectedField(nullptr));
 }
@@ -281,6 +253,7 @@ TEST(GetNewMetadataTest, FullMetadata) {
                   {"order", "C"},
                   {"shape", {5, 6}},
                   {"zarr_format", 2},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -307,6 +280,7 @@ TEST(GetNewMetadataTest, AutomaticChunks) {
           {"order", "C"},
           {"shape", {2, 3}},
           {"zarr_format", 2},
+          {"dimension_separator", "."},
       })));
 }
 
@@ -339,6 +313,7 @@ TEST(GetNewMetadataTest, NoCompressor) {
                        {"blocksize", 0},
                        {"shuffle", -1},
                    }},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -380,6 +355,7 @@ TEST(GetNewMetadataTest, SchemaDomainDtype) {
                        {"blocksize", 0},
                        {"shuffle", -1},
                    }},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -409,6 +385,7 @@ TEST(GetNewMetadataTest, SchemaDomainDtypeFillValue) {
                        {"blocksize", 0},
                        {"shuffle", -1},
                    }},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -439,6 +416,7 @@ TEST(GetNewMetadataTest, SchemaObjectWithDomainDtypeFillValue) {
                        {"blocksize", 0},
                        {"shuffle", -1},
                    }},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -459,6 +437,7 @@ TEST(GetNewMetadataTest, SchemaDtypeShapeCodec) {
                   {"chunks", {100, 200}},
                   {"dtype", "<i4"},
                   {"compressor", nullptr},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -479,6 +458,7 @@ TEST(GetNewMetadataTest, SchemaDtypeInnerOrderC) {
                   {"chunks", {100, 200}},
                   {"dtype", "<i4"},
                   {"compressor", nullptr},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -499,6 +479,7 @@ TEST(GetNewMetadataTest, SchemaDtypeInnerOrderFortran) {
                   {"chunks", {100, 200}},
                   {"dtype", "<i4"},
                   {"compressor", nullptr},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -519,6 +500,7 @@ TEST(GetNewMetadataTest, SchemaDtypeInnerOrderFortranFieldShape) {
                   {"chunks", {100, 200}},
                   {"dtype", {{"x", "<u4", {2, 3}}}},
                   {"compressor", nullptr},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -547,6 +529,7 @@ TEST(GetNewMetadataTest, SchemaDtypeInnerOrderInvalidSoft) {
                   {"chunks", {100, 102, 102}},
                   {"dtype", "<i4"},
                   {"compressor", nullptr},
+                  {"dimension_separator", "."},
               })));
 }
 
@@ -601,6 +584,7 @@ TEST(GetNewMetadataTest, SchemaFillValueRedundant) {
           {"chunks", {100, 200}},
           {"dtype", "<u4"},
           {"compressor", nullptr},
+          {"dimension_separator", "."},
       })));
 }