From 0b1b1525e4c169a5d838415f470256b5231193b0 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Tue, 5 May 2026 19:57:37 +0000 Subject: [PATCH 1/8] Update to follow #287 and implement general feedback --- tensorstore/driver/zarr3/BUILD | 5 + tensorstore/driver/zarr3/chunk_cache.cc | 150 +---- tensorstore/driver/zarr3/chunk_cache.h | 139 ++--- tensorstore/driver/zarr3/codec/bytes.cc | 4 + tensorstore/driver/zarr3/codec/codec_spec.h | 39 +- .../driver/zarr3/codec/sharding_indexed.cc | 99 ++-- tensorstore/driver/zarr3/codec/transpose.cc | 23 +- tensorstore/driver/zarr3/driver.cc | 511 ++++++++---------- tensorstore/driver/zarr3/driver_test.cc | 63 ++- tensorstore/driver/zarr3/dtype.cc | 10 + tensorstore/driver/zarr3/dtype.h | 5 +- tensorstore/driver/zarr3/dtype_test.cc | 80 +++ tensorstore/driver/zarr3/metadata.cc | 416 +++++++++----- tensorstore/driver/zarr3/metadata.h | 97 +++- tensorstore/driver/zarr3/metadata_test.cc | 161 ++++++ 15 files changed, 1064 insertions(+), 738 deletions(-) diff --git a/tensorstore/driver/zarr3/BUILD b/tensorstore/driver/zarr3/BUILD index fdfe995f4..3a1ce4381 100644 --- a/tensorstore/driver/zarr3/BUILD +++ b/tensorstore/driver/zarr3/BUILD @@ -112,10 +112,12 @@ tensorstore_cc_library( "//tensorstore:rank", "//tensorstore:schema", "//tensorstore/driver/zarr3/codec", + "//tensorstore/driver/zarr3/codec:bytes", "//tensorstore/driver/zarr3/codec:codec_chain_spec", "//tensorstore/driver/zarr3/codec:sharding_indexed", "//tensorstore/index_space:dimension_units", "//tensorstore/index_space:index_transform", + "//tensorstore/internal:data_type_endian_conversion", "//tensorstore/internal:dimension_labels", "//tensorstore/internal:intrusive_ptr", "//tensorstore/internal:json_metadata_matching", @@ -129,6 +131,7 @@ tensorstore_cc_library( "//tensorstore/serialization", "//tensorstore/serialization:json", "//tensorstore/util:constant_vector", + "//tensorstore/util:endian", "//tensorstore/util:iterate", "//tensorstore/util:quote_string", "//tensorstore/util:result", @@ -307,10 +310,12 @@ tensorstore_cc_test( "//tensorstore:index", "//tensorstore:rank", "//tensorstore:schema", + "//tensorstore/driver/zarr3/codec:codec_chain_spec", "//tensorstore/internal/json_binding:bindable", "//tensorstore/internal/json_binding:gtest", "//tensorstore/internal/meta:integer_types", "//tensorstore/internal/testing:json_gtest", + "//tensorstore/util:endian", "//tensorstore/util:result", "//tensorstore/util:status", "//tensorstore/util:status_testutil", diff --git a/tensorstore/driver/zarr3/chunk_cache.cc b/tensorstore/driver/zarr3/chunk_cache.cc index 4a24f6297..8b3124c51 100644 --- a/tensorstore/driver/zarr3/chunk_cache.cc +++ b/tensorstore/driver/zarr3/chunk_cache.cc @@ -155,16 +155,12 @@ ZarrChunkCache::~ZarrChunkCache() = default; ZarrLeafChunkCache::ZarrLeafChunkCache( kvstore::DriverPtr store, ZarrCodecChain::PreparedState::Ptr codec_state, - ZarrDType dtype, internal::CachePool::WeakPtr /*data_cache_pool*/, - bool open_as_void, bool original_is_structured, DataType original_dtype, - bool grid_has_void_dimension) + ZarrDType dtype, std::vector field_shape, + internal::CachePool::WeakPtr /*data_cache_pool*/) : Base(std::move(store)), codec_state_(std::move(codec_state)), dtype_(std::move(dtype)), - open_as_void_(open_as_void), - original_is_structured_(original_is_structured), - original_dtype_(original_dtype), - grid_has_void_dimension_(grid_has_void_dimension) {} + field_shape_(std::move(field_shape)) {} void ZarrLeafChunkCache::Read(ZarrChunkCache::ReadRequest request, AnyFlowReceiver, 1>> -ZarrLeafChunkCache::DecodeChunkAsVoid(absl::Cord data) { - absl::InlinedVector, 1> field_arrays(1); - const auto& void_component_shape = grid().components[0].shape(); - - if (original_is_structured_) { - // Structured types: codec already expects bytes with extra dimension. - // Just decode directly to the void component shape. - TENSORSTORE_ASSIGN_OR_RETURN( - field_arrays[0], - codec_state_->DecodeArray(void_component_shape, std::move(data))); - return field_arrays; - } - - // Non-structured types: codec expects original dtype without extra - // dimension. Decode, then reinterpret as bytes. - // - // For top-level caches, grid().chunk_shape includes bytes dimension. - // For sub-chunk caches (inside sharding), grid() returns the sharding - // codec's sub_chunk_grid which doesn't have bytes dimension. - const Index bytes_per_element = dtype_.bytes_per_outer_element; - const auto& grid_chunk_shape = grid().chunk_shape; - - std::vector original_chunk_shape; - if (grid_has_void_dimension_) { - // Strip the bytes dimension to get original shape - original_chunk_shape.assign(grid_chunk_shape.begin(), - grid_chunk_shape.end() - 1); - } else { - // Sub-chunk cache: grid shape is already the original shape - original_chunk_shape.assign(grid_chunk_shape.begin(), - grid_chunk_shape.end()); - } - - // Decode using original codec shape - TENSORSTORE_ASSIGN_OR_RETURN( - auto decoded_array, - codec_state_->DecodeArray(original_chunk_shape, std::move(data))); - - // Verify decoded array is C-contiguous (codec chain should guarantee this) - assert(IsContiguousLayout(decoded_array.layout(), c_order, - decoded_array.dtype().size())); - - // Build the void output shape: original_shape + [bytes_per_element] - std::vector void_output_shape = original_chunk_shape; - void_output_shape.push_back(bytes_per_element); - - // Alias the decoded array's memory as bytes. - SharedElementPointer byte_element_pointer( - std::shared_ptr(decoded_array.element_pointer().pointer(), - decoded_array.data()), - dtype_v); - field_arrays[0] = - SharedArray(byte_element_pointer, void_output_shape); - return field_arrays; -} - Result, 1>> ZarrLeafChunkCache::DecodeChunk(span chunk_indices, absl::Cord data) { - if (open_as_void_) { - return DecodeChunkAsVoid(std::move(data)); - } - const size_t num_fields = dtype_.fields.size(); absl::InlinedVector, 1> field_arrays(num_fields); - // For single non-structured field, decode directly - if (num_fields == 1 && dtype_.fields[0].field_shape.empty()) { + // Fast path: scalar single-field arrays contribute no inner dimensions to + // the codec resolution shape, so the codec call is rank-aligned with the + // grid's component shape and no byte interleaving is required. + if (field_shape_.empty()) { + assert(num_fields == 1); TENSORSTORE_ASSIGN_OR_RETURN( field_arrays[0], codec_state_->DecodeArray(grid().components[0].shape(), std::move(data))); return field_arrays; } - // For structured types, decode byte array then extract fields - // Build decode shape: [chunk_dims..., bytes_per_outer_element] + // Otherwise the codec was resolved against `byte` data with the inner + // `field_shape_` dims appended -- multi-field structs, `rN` raw byte + // fields, and the `open_as_void` view. Decode the raw byte array, then + // split it into per-field components based purely on the field layout. const auto& chunk_shape = grid().chunk_shape; std::vector decode_shape(chunk_shape.begin(), chunk_shape.end()); - decode_shape.push_back(dtype_.bytes_per_outer_element); + decode_shape.insert(decode_shape.end(), field_shape_.begin(), + field_shape_.end()); TENSORSTORE_ASSIGN_OR_RETURN( auto byte_array, codec_state_->DecodeArray(decode_shape, std::move(data))); @@ -370,59 +311,30 @@ ZarrLeafChunkCache::DecodeChunk(span chunk_indices, return field_arrays; } -Result ZarrLeafChunkCache::EncodeChunkAsVoid( - const SharedArray& byte_array) { - if (original_is_structured_) { - // Structured types: codec already expects bytes with extra dimension. - return codec_state_->EncodeArray(byte_array); - } - - // Non-structured types: reinterpret bytes as original dtype/shape. - const Index bytes_per_element = dtype_.bytes_per_outer_element; - - // Build original chunk shape by stripping the bytes dimension - const auto& void_shape = byte_array.shape(); - std::vector original_shape(void_shape.begin(), void_shape.end() - 1); - - // Use the original dtype (stored during cache creation) for encoding. - // Create a view over the byte data with original dtype and layout. - // Use the aliasing constructor to share ownership with byte_array but - // interpret the data with the original dtype. - SharedArray encoded_array; - auto aliased_ptr = std::shared_ptr( - byte_array.pointer(), // Share ownership with byte_array - byte_array.data()); // But point to the raw data - encoded_array.element_pointer() = SharedElementPointer( - std::move(aliased_ptr), original_dtype_); - encoded_array.layout() = StridedLayout<>(c_order, bytes_per_element, - original_shape); - - return codec_state_->EncodeArray(encoded_array); -} - Result ZarrLeafChunkCache::EncodeChunk( span chunk_indices, span> component_arrays) { - if (open_as_void_) { - assert(component_arrays.size() == 1); - return EncodeChunkAsVoid(component_arrays[0]); - } - const size_t num_fields = dtype_.fields.size(); - // For single non-structured field, encode directly - if (num_fields == 1 && dtype_.fields[0].field_shape.empty()) { + // Fast path: see `DecodeChunk` -- when the metadata-level `field_shape_` + // is empty, the codec was resolved at the field's natural rank and dtype + // and we hand it the per-field array directly. + if (field_shape_.empty()) { + assert(num_fields == 1); assert(component_arrays.size() == 1); return codec_state_->EncodeArray(component_arrays[0]); } - // For structured types, combine multiple field arrays into a single byte array + // Otherwise pack each field's bytes into the interleaved byte buffer the + // codec chain expects (resolved against `byte` + the trailing `field_shape_` + // dims), then hand it off for encoding. assert(component_arrays.size() == num_fields); - // Build encode shape: [chunk_dims..., bytes_per_outer_element] + // Build encode shape: [chunk_dims..., field_shape...] const auto& chunk_shape = grid().chunk_shape; std::vector encode_shape(chunk_shape.begin(), chunk_shape.end()); - encode_shape.push_back(dtype_.bytes_per_outer_element); + encode_shape.insert(encode_shape.end(), field_shape_.begin(), + field_shape_.end()); // Allocate byte array for combined fields auto byte_array = AllocateArray(encode_shape, c_order, value_init); @@ -477,15 +389,12 @@ kvstore::Driver* ZarrLeafChunkCache::GetKvStoreDriver() { ZarrShardedChunkCache::ZarrShardedChunkCache( kvstore::DriverPtr store, ZarrCodecChain::PreparedState::Ptr codec_state, - ZarrDType dtype, internal::CachePool::WeakPtr data_cache_pool, - bool open_as_void, bool original_is_structured, DataType original_dtype, - bool /*grid_has_void_dimension*/) + ZarrDType dtype, std::vector field_shape, + internal::CachePool::WeakPtr data_cache_pool) : base_kvstore_(std::move(store)), codec_state_(std::move(codec_state)), dtype_(std::move(dtype)), - open_as_void_(open_as_void), - original_is_structured_(original_is_structured), - original_dtype_(original_dtype), + field_shape_(std::move(field_shape)), data_cache_pool_(std::move(data_cache_pool)) {} Result> TranslateCellToSourceTransformForShard( @@ -795,8 +704,7 @@ void ZarrShardedChunkCache::Entry::DoInitialize() { *sharding_state.sub_chunk_codec_chain, std::move(sharding_kvstore), cache.executor(), ZarrShardingCodec::PreparedState::Ptr(&sharding_state), - cache.dtype_, cache.data_cache_pool_, cache.open_as_void_, - cache.original_is_structured_, cache.original_dtype_); + cache.dtype_, cache.field_shape_, cache.data_cache_pool_); zarr_chunk_cache = new_cache.release(); return std::unique_ptr(&zarr_chunk_cache->cache()); }) diff --git a/tensorstore/driver/zarr3/chunk_cache.h b/tensorstore/driver/zarr3/chunk_cache.h index cc7238b76..17fdf0e7a 100644 --- a/tensorstore/driver/zarr3/chunk_cache.h +++ b/tensorstore/driver/zarr3/chunk_cache.h @@ -178,11 +178,8 @@ class ZarrLeafChunkCache : public internal::KvsBackedChunkCache, explicit ZarrLeafChunkCache(kvstore::DriverPtr store, ZarrCodecChain::PreparedState::Ptr codec_state, ZarrDType dtype, - internal::CachePool::WeakPtr data_cache_pool, - bool open_as_void, - bool original_is_structured, - DataType original_dtype, - bool grid_has_void_dimension = true); + std::vector field_shape, + internal::CachePool::WeakPtr data_cache_pool); void Read(ZarrChunkCache::ReadRequest request, AnyFlowReceiver, 1>> DecodeChunkAsVoid( - absl::Cord data); - - Result EncodeChunkAsVoid(const SharedArray& byte_array); + // Inner trailing dims contributed by the dtype to the codec resolution + // shape -- the same vector that lives on `ZarrMetadata::field_shape`. + // Empty for plain scalar single-field arrays; `{bytes_per_outer_element}` + // for multi-field structs and the byte-substituted void view; the per- + // field `field_shape` for `rN` raw byte fields. Used only to gate the + // "fast path" (no byte interleaving) vs the byte-packing path. + std::vector field_shape_; }; /// Chunk cache for a Zarr array where each chunk is a shard. @@ -230,11 +224,8 @@ class ZarrShardedChunkCache : public internal::Cache, public ZarrChunkCache { explicit ZarrShardedChunkCache(kvstore::DriverPtr store, ZarrCodecChain::PreparedState::Ptr codec_state, ZarrDType dtype, - internal::CachePool::WeakPtr data_cache_pool, - bool open_as_void, - bool original_is_structured, - DataType original_dtype, - bool grid_has_void_dimension = true); + std::vector field_shape, + internal::CachePool::WeakPtr data_cache_pool); const ZarrShardingCodec::PreparedState& sharding_codec_state() const { return static_cast( @@ -284,9 +275,11 @@ class ZarrShardedChunkCache : public internal::Cache, public ZarrChunkCache { kvstore::DriverPtr base_kvstore_; ZarrCodecChain::PreparedState::Ptr codec_state_; ZarrDType dtype_; - bool open_as_void_; - bool original_is_structured_; - DataType original_dtype_; // Original dtype for void access encoding + // Same semantics as `ZarrLeafChunkCache::field_shape_`: the metadata-level + // inner trailing dims used to gate the byte-packing decode/encode path + // and to determine how many trailing dims the sharding sub-chunk grid + // carries. + std::vector field_shape_; // Data cache pool, if it differs from `this->pool()` (which is equal to the // metadata cache pool). @@ -295,12 +288,21 @@ class ZarrShardedChunkCache : public internal::Cache, public ZarrChunkCache { /// Chunk cache mixin for a chunk cache where the entire chunk cache corresponds /// to a single shard. +/// +/// The sharding codec's `sub_chunk_grid` is a single-component grid whose +/// component shape includes any trailing `field_shape` dimensions contributed +/// by the dtype. This mixin replaces it with a field-level grid built +/// directly from the dtype (`CreateFieldGridSpecification`), so that the +/// surrounding code path operates on per-field typed arrays regardless of +/// whether the chunk is a multi-field struct, an `rN` raw byte field, or a +/// plain scalar. template class ZarrShardSubChunkCache : public ChunkCacheImpl { public: /// Wraps the sharding state's key parser to pad spatial-only cell indices /// (from a field-level grid) to the full dimensionality expected by the - /// shard index (which includes the bytes dimension with index 0). + /// shard index (which includes any trailing field_shape dimensions with + /// index 0). class FieldKeyParserWrapper : public internal::LexicographicalGridIndexKeyParser { public: @@ -339,63 +341,39 @@ class ZarrShardSubChunkCache : public ChunkCacheImpl { explicit ZarrShardSubChunkCache( kvstore::DriverPtr store, Executor executor, ZarrShardingCodec::PreparedState::Ptr sharding_state, - ZarrDType dtype, internal::CachePool::WeakPtr data_cache_pool, - bool open_as_void, bool original_is_structured, DataType original_dtype) + ZarrDType dtype, std::vector field_shape, + internal::CachePool::WeakPtr data_cache_pool) : ChunkCacheImpl(std::move(store), ZarrCodecChain::PreparedState::Ptr( sharding_state->sub_chunk_codec_state), - std::move(dtype), std::move(data_cache_pool), - open_as_void, original_is_structured, original_dtype, - /*grid_has_void_dimension=*/false), + std::move(dtype), std::move(field_shape), + std::move(data_cache_pool)), sharding_state_(std::move(sharding_state)), executor_(std::move(executor)) { - // For void access on non-structured types, create a modified grid - // with the bytes dimension added to the component shape. - // The grid's chunk_shape stays the same (determines cell layout). - if (ChunkCacheImpl::open_as_void_ && - !ChunkCacheImpl::original_is_structured_) { - const auto& original_grid = *sharding_state_->sub_chunk_grid; - const auto& orig_comp = original_grid.components[0]; - // Component chunk_shape gets bytes dimension, grid chunk_shape doesn't - std::vector void_comp_shape = orig_comp.chunk_shape; - void_comp_shape.push_back(ChunkCacheImpl::dtype_.bytes_per_outer_element); - // Create zero fill value with the void shape - auto fill_value = AllocateArray(void_comp_shape, c_order, value_init, - dtype_v); - // chunked_to_cell_dimensions maps the grid dimensions to cell dimensions - // (the bytes dimension is unchunked, so not included here) - std::vector chunked_to_cell(original_grid.chunk_shape.size()); - std::iota(chunked_to_cell.begin(), chunked_to_cell.end(), 0); - internal::ChunkGridSpecification::ComponentList components; - components.emplace_back( - internal::AsyncWriteArray::Spec{std::move(fill_value), - Box<>(void_comp_shape.size())}, - void_comp_shape, std::move(chunked_to_cell)); - void_grid_.emplace(std::move(components)); - } - - // For structured types (non-void), create a field-level grid where each - // component corresponds to a field with spatial-only shape. The sharding - // codec's sub_chunk_grid has a single byte component including the bytes - // dimension; this grid strips that dimension so that DecodeChunk/EncodeChunk - // can extract/combine per-field typed arrays. - if (ChunkCacheImpl::original_is_structured_ && - !ChunkCacheImpl::open_as_void_) { - const auto& original_grid = *sharding_state_->sub_chunk_grid; - const DimensionIndex full_rank = original_grid.chunk_shape.size(); - const DimensionIndex spatial_rank = full_rank - 1; - - std::vector spatial_chunk_shape( - original_grid.chunk_shape.begin(), - original_grid.chunk_shape.begin() + spatial_rank); - - field_grid_.emplace(CreateFieldGridSpecification( - spatial_chunk_shape, ChunkCacheImpl::dtype_, - span(), // no inner_order available - nullptr)); // zero-init fill values - - field_key_parser_.emplace( - sharding_state_->GetSubChunkStorageKeyParser(), full_rank); + // Strip any trailing inner dimensions contributed by `field_shape_` from + // the sharding codec's sub-chunk grid to get the spatial (chunked-only) + // rank. `field_shape_` is the metadata-level inner shape vector -- + // empty for plain scalars, single-element for current zarr v3 cases + // (`{bytes_per_outer_element}` for structs/void, `{N}` for `rN`). + const auto& field_shape_ref = ChunkCacheImpl::field_shape_; + const auto& original_grid = *sharding_state_->sub_chunk_grid; + const DimensionIndex full_rank = original_grid.chunk_shape.size(); + const DimensionIndex extra_field_dims = + static_cast(field_shape_ref.size()); + const DimensionIndex spatial_rank = full_rank - extra_field_dims; + + std::vector spatial_chunk_shape( + original_grid.chunk_shape.begin(), + original_grid.chunk_shape.begin() + spatial_rank); + + field_grid_.emplace(CreateFieldGridSpecification( + spatial_chunk_shape, ChunkCacheImpl::dtype_, + span(), // no inner_order available + nullptr)); // zero-init fill values + + if (extra_field_dims != 0) { + field_key_parser_.emplace(sharding_state_->GetSubChunkStorageKeyParser(), + full_rank); } } @@ -408,13 +386,7 @@ class ZarrShardSubChunkCache : public ChunkCacheImpl { } const internal::ChunkGridSpecification& grid() const override { - if (field_grid_) { - return *field_grid_; - } - if (void_grid_) { - return *void_grid_; - } - return *sharding_state_->sub_chunk_grid; + return *field_grid_; } const Executor& executor() const override { return executor_; } @@ -422,7 +394,6 @@ class ZarrShardSubChunkCache : public ChunkCacheImpl { ZarrShardingCodec::PreparedState::Ptr sharding_state_; Executor executor_; - std::optional void_grid_; std::optional field_grid_; std::optional field_key_parser_; }; diff --git a/tensorstore/driver/zarr3/codec/bytes.cc b/tensorstore/driver/zarr3/codec/bytes.cc index 17885008d..6d35eb866 100644 --- a/tensorstore/driver/zarr3/codec/bytes.cc +++ b/tensorstore/driver/zarr3/codec/bytes.cc @@ -79,6 +79,10 @@ absl::Status BytesCodecSpec::GetDecodedChunkLayout( !internal::IsTrivialDataType(array_info.dtype)) { return InvalidDataTypeError(array_info.dtype); } + // `array_info.rank` is the chunked rank only; any inner (`field_shape`) + // dims are reported via `array_info.inner_shape` and do not appear in the + // chunk layout info (they are pinned, identity-ordered, and carry no + // user-tunable layout). const DimensionIndex rank = array_info.rank; if (rank != dynamic_rank) { auto& inner_order = decoded.inner_order.emplace(); diff --git a/tensorstore/driver/zarr3/codec/codec_spec.h b/tensorstore/driver/zarr3/codec/codec_spec.h index fcf903610..bc4dfdd93 100644 --- a/tensorstore/driver/zarr3/codec/codec_spec.h +++ b/tensorstore/driver/zarr3/codec/codec_spec.h @@ -69,6 +69,7 @@ #include #include +#include #include "absl/status/status.h" #include "tensorstore/array.h" @@ -135,11 +136,30 @@ struct ArrayDataTypeAndShapeInfo { // Specifies the data type of the array on which the codec will operate. DataType dtype; - // Specifies the rank of the array on which the codec will operate. + // Specifies the *chunked* rank of the array on which the codec will operate. + // + // This is the user-visible rank: it does NOT include the inner trailing + // dimensions contributed by the dtype's `field_shape` (multi-field structs, + // `rN` raw byte fields, `open_as_void`'s byte substitution). Those inner + // dimensions, if any, are reported separately in `inner_shape` and are + // pinned at the trailing positions of the runtime array; "array -> array" + // codec specs cannot operate on them. DimensionIndex rank = dynamic_rank; - // Specifies the shape of the array on which the codec will operate. + // Specifies the chunked-rank shape of the array on which the codec will + // operate. When present, has exactly `rank` valid entries. std::optional> shape; + + // Inner trailing dimensions contributed by the dtype's `field_shape`. + // Empty (the default) when the dtype has no field_shape (i.e. plain scalar + // dtype with no codec substitution); non-empty for multi-field structs, + // `rN`, and the byte-substituted `open_as_void` view. + // + // "array -> array" codec specs must propagate this field unchanged from + // `decoded` to `encoded`; they cannot read it for resolution decisions and + // cannot reshape/permute it. The "array -> bytes" codec at the bottom of + // the chain consumes it (e.g. for sizing the encoded byte stream). + std::vector inner_shape; }; // Specifies information about the chunk layout that must be propagated through @@ -165,20 +185,27 @@ struct ArrayCodecResolveParameters { // Specifies the data type of the array on which the codec will operate. DataType dtype; - // Specifies the rank of the array on which the codec will operate. + // Specifies the *chunked* rank of the array on which the codec will operate + // (excludes any `inner_shape` dimensions; see `ArrayDataTypeAndShapeInfo`). DimensionIndex rank; // Specifies the fill value. SharedArray fill_value; - // Specifies requested read chunk shape. + // Specifies requested read chunk shape (chunked rank only). std::optional> read_chunk_shape; - // Specifies requested codec chunk shape. + // Specifies requested codec chunk shape (chunked rank only). std::optional> codec_chunk_shape; - // Specifies required inner order. + // Specifies required inner order (chunked rank only). std::optional> inner_order; + + // Inner trailing dimensions contributed by the dtype's `field_shape`; see + // the docstring on `ArrayDataTypeAndShapeInfo::inner_shape`. Propagated + // unchanged through "array -> array" codecs; consumed by the "array -> + // bytes" codec (e.g. for sizing the on-disk byte stream). + std::vector inner_shape; }; // Spec for an "array -> array" codec. diff --git a/tensorstore/driver/zarr3/codec/sharding_indexed.cc b/tensorstore/driver/zarr3/codec/sharding_indexed.cc index 6804f517e..4c3881891 100644 --- a/tensorstore/driver/zarr3/codec/sharding_indexed.cc +++ b/tensorstore/driver/zarr3/codec/sharding_indexed.cc @@ -175,35 +175,15 @@ class ShardingIndexedCodec : public ZarrShardingCodec { }; } // namespace -namespace { -// Custom merge function for sub_chunk_shape that handles structured types. -// For structured types (byte dtype with extra dimension), user-specified shapes -// may not include the bytes dimension. Shapes [4,4] and [4,4,3] are compatible -// if the first N elements match. -struct TryMergeSubChunkShape { - bool operator()(const std::optional>& a, - const std::optional>& b) const { - if (!a || !b) return true; - if (*a == *b) return true; - // Allow one shape to be an extension of the other by one dimension - // (the bytes dimension for structured types) - const auto& longer = a->size() > b->size() ? *a : *b; - const auto& shorter = a->size() > b->size() ? *b : *a; - if (longer.size() == shorter.size() + 1) { - return std::equal(shorter.begin(), shorter.end(), longer.begin()); - } - return false; - } -}; -} // namespace - absl::Status ShardingIndexedCodecSpec::MergeFrom(const ZarrCodecSpec& other, bool strict) { using Self = ShardingIndexedCodecSpec; const auto& other_options = static_cast(other).options; + // `sub_chunk_shape` lives at the chunked rank in both the on-disk + // zarr.json and the user-facing spec; the codec resolution path no longer + // extends it internally, so plain equality is sufficient. TENSORSTORE_RETURN_IF_ERROR(MergeConstraint<&Options::sub_chunk_shape>( - "chunk_shape", options, other_options, jb::DefaultBinder<>, - TryMergeSubChunkShape{})); + "chunk_shape", options, other_options)); TENSORSTORE_RETURN_IF_ERROR( internal_zarr3::MergeZarrCodecSpecs(options.index_codecs, other_options.index_codecs, strict)) @@ -237,17 +217,24 @@ const ZarrCodecChainSpec* ShardingIndexedCodecSpec::GetSubChunkCodecs() const { absl::Status ShardingIndexedCodecSpec::GetDecodedChunkLayout( const ArrayDataTypeAndShapeInfo& array_info, ArrayCodecChunkLayoutInfo& decoded) const { - ArrayDataTypeAndShapeInfo sub_chunk_info; + // `array_info.rank` is the chunked rank only; inner (`field_shape`) dims + // travel via `array_info.inner_shape` and are propagated unchanged through + // sub-chunk codecs without ever being permuted or reshaped. The codec + // spec's `sub_chunk_shape` is therefore at the same chunked rank as + // `array_info.rank`. if (options.sub_chunk_shape && - !RankConstraint::Implies(options.sub_chunk_shape->size(), - array_info.rank)) { + static_cast(options.sub_chunk_shape->size()) != + array_info.rank) { return SubChunkRankMismatch(*options.sub_chunk_shape, array_info.rank); } + ArrayDataTypeAndShapeInfo sub_chunk_info; sub_chunk_info.dtype = array_info.dtype; sub_chunk_info.rank = array_info.rank; + sub_chunk_info.inner_shape = array_info.inner_shape; if (options.sub_chunk_shape) { + auto& shape = sub_chunk_info.shape.emplace(); std::copy(options.sub_chunk_shape->begin(), options.sub_chunk_shape->end(), - sub_chunk_info.shape.emplace().begin()); + shape.begin()); } if (options.sub_chunk_codecs) { TENSORSTORE_RETURN_IF_ERROR(options.sub_chunk_codecs->GetDecodedChunkLayout( @@ -275,35 +262,39 @@ Result ShardingIndexedCodecSpec::Resolve( resolved_options = &resolved_spec_ptr->options; resolved_spec->reset(resolved_spec_ptr); } - std::vector extended_sub_chunk_shape; + // `decoded.rank` is the chunked rank only. Inner (`field_shape`) dims are + // carried separately on `decoded.inner_shape`, propagated through to the + // sub-chunk codec chain unchanged, and consumed at the leaf "array -> + // bytes" codec for byte-stream sizing. The codec spec's `sub_chunk_shape` + // is at the same chunked rank. span sub_chunk_shape; if (options.sub_chunk_shape) { sub_chunk_shape = *options.sub_chunk_shape; - // For structured types (byte dtype with extra dimension), the user may - // specify sub_chunk_shape without the bytes dimension. Extend it if needed. - if (sub_chunk_shape.size() + 1 == decoded.rank && - decoded.read_chunk_shape) { - extended_sub_chunk_shape.assign(sub_chunk_shape.begin(), - sub_chunk_shape.end()); - extended_sub_chunk_shape.push_back( - (*decoded.read_chunk_shape)[decoded.rank - 1]); - sub_chunk_shape = extended_sub_chunk_shape; - } } else if (decoded.read_chunk_shape) { sub_chunk_shape = span(decoded.read_chunk_shape->data(), decoded.rank); } else { return absl::InvalidArgumentError("\"chunk_shape\" must be specified"); } - if (sub_chunk_shape.size() != decoded.rank) { + if (static_cast(sub_chunk_shape.size()) != decoded.rank) { return SubChunkRankMismatch(sub_chunk_shape, decoded.rank); } + // The runtime ChunkGridSpecification partitions an extended-rank chunk + // (chunked + inner): chunked dims are partitioned by `sub_chunk_shape`, and + // inner dims are kept whole inside each sub-chunk. Build the runtime + // sub-chunk shape accordingly. + std::vector runtime_sub_chunk_shape(sub_chunk_shape.begin(), + sub_chunk_shape.end()); + runtime_sub_chunk_shape.insert(runtime_sub_chunk_shape.end(), + decoded.inner_shape.begin(), + decoded.inner_shape.end()); internal::IntrusivePtr codec; auto set_up_codecs = [&](const ZarrCodecChainSpec& sub_chunk_codecs) -> absl::Status { ArrayCodecResolveParameters sub_chunk_decoded; sub_chunk_decoded.dtype = decoded.dtype; sub_chunk_decoded.rank = decoded.rank; + sub_chunk_decoded.inner_shape = decoded.inner_shape; sub_chunk_decoded.fill_value = decoded.fill_value; if (decoded.read_chunk_shape) { std::copy_n(decoded.read_chunk_shape->begin(), decoded.rank, @@ -323,10 +314,11 @@ Result ShardingIndexedCodecSpec::Resolve( std::move(sub_chunk_decoded), encoded, resolved_options ? &resolved_options->sub_chunk_codecs.emplace() : nullptr)); - // Get sub-chunk codec chunk layout info. + // Get sub-chunk codec chunk layout info (chunked rank only). ArrayDataTypeAndShapeInfo array_info; array_info.dtype = decoded.dtype; array_info.rank = decoded.rank; + array_info.inner_shape = decoded.inner_shape; ArrayCodecChunkLayoutInfo layout_info; TENSORSTORE_RETURN_IF_ERROR( (resolved_options ? *resolved_options->sub_chunk_codecs @@ -338,16 +330,23 @@ Result ShardingIndexedCodecSpec::Resolve( layout_info.inner_order->begin() + decoded.rank, static_cast(0)); } + // Extend the inner_order to runtime rank with identity on inner dims. + const DimensionIndex inner_rank = + static_cast(decoded.inner_shape.size()); + for (DimensionIndex i = 0; i < inner_rank; ++i) { + (*layout_info.inner_order)[decoded.rank + i] = decoded.rank + i; + } + const DimensionIndex runtime_rank = decoded.rank + inner_rank; internal::ChunkGridSpecification::ComponentList components; TENSORSTORE_ASSIGN_OR_RETURN( auto broadcast_fill_value, - BroadcastArray(decoded.fill_value, BoxView<>(sub_chunk_shape.size()))); + BroadcastArray(decoded.fill_value, BoxView<>(runtime_rank))); auto& component = components.emplace_back( internal::AsyncWriteArray::Spec{ - std::move(broadcast_fill_value), Box<>(sub_chunk_shape.size()), + std::move(broadcast_fill_value), Box<>(runtime_rank), ContiguousLayoutPermutation<>(tensorstore::span( - layout_info.inner_order->data(), decoded.rank))}, - std::vector(sub_chunk_shape.begin(), sub_chunk_shape.end())); + layout_info.inner_order->data(), runtime_rank))}, + runtime_sub_chunk_shape); component.array_spec.fill_value_comparison_kind = EqualityComparisonKind::identical; codec = internal::MakeIntrusivePtr( @@ -356,7 +355,14 @@ Result ShardingIndexedCodecSpec::Resolve( options.index_location.value_or(ShardIndexLocation::kEnd); codec->sub_chunk_codec_chain_ = std::move(sub_chunk_codec_chain); if (resolved_options) { - resolved_options->sub_chunk_shape = codec->sub_chunk_grid_.chunk_shape; + // The user-form sub_chunk_shape is at the same chunked rank as the + // resolved spec; no internal extension to undo. + if (options.sub_chunk_shape) { + resolved_options->sub_chunk_shape = *options.sub_chunk_shape; + } else { + resolved_options->sub_chunk_shape.emplace(sub_chunk_shape.begin(), + sub_chunk_shape.end()); + } resolved_options->index_location = codec->index_location_; } return absl::OkStatus(); @@ -366,6 +372,7 @@ Result ShardingIndexedCodecSpec::Resolve( : ZarrCodecChainSpec{})) .Format("Error resolving sub-chunk codecs"); + // Index codecs index the chunked grid only, regardless of inner_shape. auto set_up_index_codecs = [&](const ZarrCodecChainSpec& index_codecs) -> absl::Status { TENSORSTORE_ASSIGN_OR_RETURN( diff --git a/tensorstore/driver/zarr3/codec/transpose.cc b/tensorstore/driver/zarr3/codec/transpose.cc index 0384d23f9..a5545ea32 100644 --- a/tensorstore/driver/zarr3/codec/transpose.cc +++ b/tensorstore/driver/zarr3/codec/transpose.cc @@ -232,11 +232,14 @@ Result> ResolveOrder( absl::Status TransposeCodecSpec::PropagateDataTypeAndShape( const ArrayDataTypeAndShapeInfo& decoded, ArrayDataTypeAndShapeInfo& encoded) const { + // The user permutation operates strictly on the chunked dimensions; any + // `inner_shape` contributed by the dtype is propagated as-is. DimensionIndex temp_perm[kMaxRank]; TENSORSTORE_ASSIGN_OR_RETURN( auto order, ResolveOrder(options.order, decoded.rank, temp_perm)); encoded.dtype = decoded.dtype; encoded.rank = order.size(); + encoded.inner_shape = decoded.inner_shape; if (decoded.shape) { auto& encoded_shape = encoded.shape.emplace(); const auto& decoded_shape = *decoded.shape; @@ -299,11 +302,14 @@ absl::Status TransposeCodecSpec::GetDecodedChunkLayout( Result TransposeCodecSpec::Resolve( ArrayCodecResolveParameters&& decoded, ArrayCodecResolveParameters& encoded, ZarrArrayToArrayCodecSpec::Ptr* resolved_spec) const { + // Spec-level resolution is at the chunked rank only; the user permutation + // never touches inner (`field_shape`) dimensions. DimensionIndex temp_perm[kMaxRank]; TENSORSTORE_ASSIGN_OR_RETURN( auto order, ResolveOrder(options.order, decoded.rank, temp_perm)); encoded.dtype = decoded.dtype; encoded.rank = decoded.rank; + encoded.inner_shape = decoded.inner_shape; assert(decoded.fill_value.rank() == 0); encoded.fill_value = std::move(decoded.fill_value); std::vector inverse_order(order.size()); @@ -318,7 +324,22 @@ Result TransposeCodecSpec::Resolve( resolved_spec->reset(new TransposeCodecSpec({TransposeCodecSpec::Order( std::vector(order.begin(), order.end()))})); } - return internal::MakeIntrusivePtr(std::move(inverse_order)); + // Build the runtime permutation at the *runtime* rank: chunked dims permuted + // as the user requested, inner (`field_shape`) dims pinned at the trailing + // positions with identity. The chunk cache hands extended-rank arrays to + // the codec at runtime; the runtime codec must therefore accept the extended + // rank without altering the inner dims. + const DimensionIndex chunked_rank = decoded.rank; + const DimensionIndex inner_rank = + static_cast(decoded.inner_shape.size()); + std::vector runtime_inverse_order(chunked_rank + inner_rank); + std::copy(inverse_order.begin(), inverse_order.end(), + runtime_inverse_order.begin()); + for (DimensionIndex i = 0; i < inner_rank; ++i) { + runtime_inverse_order[chunked_rank + i] = chunked_rank + i; + } + return internal::MakeIntrusivePtr( + std::move(runtime_inverse_order)); } TENSORSTORE_GLOBAL_INITIALIZER { diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index 7b5f4928a..7a6a12439 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -109,7 +109,7 @@ class ZarrDriverSpec ZarrMetadataConstraints metadata_constraints; std::string selected_field; - bool open_as_void; + bool open_as_void = false; constexpr static auto ApplyMembers = [](auto& x, auto f) { return f(internal::BaseCast(x), x.metadata_constraints, @@ -150,87 +150,50 @@ class ZarrDriverSpec return Base::ApplyOptions(std::move(options)); } - Result> GetDomainAsVoid() const { - const Index bytes_per_elem = - metadata_constraints.data_type->bytes_per_outer_element; - const DimensionIndex original_rank = metadata_constraints.shape->size(); - IndexDomainBuilder builder(original_rank + 1); - - // Set original dimensions from metadata (all origins are 0) - std::fill_n(builder.origin().begin(), original_rank + 1, Index{0}); - std::copy_n(metadata_constraints.shape->begin(), original_rank, - builder.shape().begin()); - builder.shape()[original_rank] = bytes_per_elem; - - // Set implicit bounds: array dims are implicit (resizable), bytes dim is explicit - builder.implicit_lower_bounds(DimensionSet(false)); - builder.implicit_upper_bounds(DimensionSet::UpTo(original_rank)); - - // Copy dimension names if available - if (metadata_constraints.dimension_names) { - for (DimensionIndex i = 0; i < original_rank; ++i) { - if (const auto& name = (*metadata_constraints.dimension_names)[i]; - name.has_value()) { - builder.labels()[i] = *name; + Result> GetDomain() const override { + // When the array contributes trailing field_shape dimensions + // (open_as_void or a selected rN/struct field), build the domain by + // appending those dimensions to the chunked dimensions stored in + // `metadata_constraints`. Both cases are uniform under the field_shape + // model: open_as_void is just `field_shape = {bytes_per_outer_element}` + // contributed by the synthetic void view of the data type. + if (metadata_constraints.data_type && metadata_constraints.shape) { + const Index byte_dim = + metadata_constraints.data_type->bytes_per_outer_element; + span field_shape; + if (open_as_void) { + field_shape = span(&byte_dim, 1); + } else if (!selected_field.empty()) { + const auto& dtype = *metadata_constraints.data_type; + for (const auto& field : dtype.fields) { + if (field.name == selected_field && !field.field_shape.empty()) { + field_shape = span(field.field_shape.data(), + field.field_shape.size()); + break; + } } } - } - - return builder.Finalize(); - } - - Result> GetDomain() const override { - // For void access with known dtype and shape, build domain directly - // to include the extra bytes dimension. - if (open_as_void && metadata_constraints.data_type && - metadata_constraints.shape) { - return GetDomainAsVoid(); - } - - // For fields with field_shape (like r16, r64), build domain to include - // the extra field dimensions. - if (!selected_field.empty() && metadata_constraints.data_type && - metadata_constraints.shape) { - const auto& dtype = *metadata_constraints.data_type; - // Find the selected field - for (size_t i = 0; i < dtype.fields.size(); ++i) { - if (dtype.fields[i].name == selected_field && - !dtype.fields[i].field_shape.empty()) { - const auto& field = dtype.fields[i]; - const DimensionIndex original_rank = metadata_constraints.shape->size(); - const DimensionIndex field_shape_rank = field.field_shape.size(); - const DimensionIndex total_rank = original_rank + field_shape_rank; - - IndexDomainBuilder builder(total_rank); - - // Set all origins to 0 - std::fill_n(builder.origin().begin(), total_rank, Index{0}); - - // Copy original dimensions - std::copy_n(metadata_constraints.shape->begin(), original_rank, - builder.shape().begin()); - - // Add field_shape dimensions - std::copy_n(field.field_shape.begin(), field_shape_rank, - builder.shape().begin() + original_rank); - - // Set implicit bounds: array dims are implicit (resizable), - // field_shape dims are explicit - builder.implicit_lower_bounds(DimensionSet(false)); - builder.implicit_upper_bounds(DimensionSet::UpTo(original_rank)); - - // Copy dimension names if available - if (metadata_constraints.dimension_names) { - for (DimensionIndex j = 0; j < original_rank; ++j) { - if (const auto& name = (*metadata_constraints.dimension_names)[j]; - name.has_value()) { - builder.labels()[j] = *name; - } + if (!field_shape.empty()) { + const DimensionIndex chunked_rank = + metadata_constraints.shape->size(); + const DimensionIndex total_rank = chunked_rank + field_shape.size(); + IndexDomainBuilder builder(total_rank); + std::fill_n(builder.origin().begin(), total_rank, Index{0}); + 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_upper_bounds(DimensionSet::UpTo(chunked_rank)); + if (metadata_constraints.dimension_names) { + for (DimensionIndex j = 0; j < chunked_rank; ++j) { + if (const auto& name = (*metadata_constraints.dimension_names)[j]; + name.has_value()) { + builder.labels()[j] = *name; } } - - return builder.Finalize(); } + return builder.Finalize(); } } @@ -258,57 +221,28 @@ class ZarrDriverSpec const ZarrDType& dtype = *constraints.data_type; - // Determine which field this spec refers to (or void access). - TENSORSTORE_ASSIGN_OR_RETURN( - size_t field_index, - GetFieldIndex(dtype, selected_field, open_as_void)); - - // ── Normal field access: just return that field's fill_value ─────────────── - if (field_index != kVoidFieldIndex) { - if (field_index < vec.size()) { - return vec[field_index]; - } - // Fallback to "no fill". - return SharedArray(); - } - // ── Void access: synthesize a byte-level fill value ──────────────────────── // - // We want a 1D byte array of length bytes_per_outer_element whose contents - // are exactly the Zarr-defined struct layout built from per-field fills. - - // Special case: "raw bytes" field (single byte_t field with flexible shape). - // In that case the existing fill array already has the correct bytes. - if (dtype.fields.size() == 1 && - dtype.fields[0].dtype.id() == DataTypeId::byte_t && - !dtype.fields[0].flexible_shape.empty()) { - // vec[0] should be a byte array of size bytes_per_outer_element. - return vec[0]; + // Defer to `MakeVoidFillValue` so the spec-level fill is byte-identical + // to the data-cache-level fill produced by `GetVoidMetadata` (codec + // endianness is honoured uniformly for non-struct, non-rN dtypes). + if (open_as_void) { + const ZarrCodecChainSpec empty_codec_specs; + const ZarrCodecChainSpec& codec_specs = + constraints.codec_specs ? *constraints.codec_specs + : empty_codec_specs; + return SharedArray( + MakeVoidFillValue(dtype, codec_specs, vec)); } - const Index nbytes = dtype.bytes_per_outer_element; - - auto byte_arr = AllocateArray( - span({nbytes}), c_order, default_init, - dtype_v); - auto* dst = static_cast(byte_arr.data()); - std::memset(dst, 0, static_cast(nbytes)); - - // Pack each field's scalar fill into its byte_offset region. - for (size_t i = 0; i < dtype.fields.size() && i < vec.size(); ++i) { - const auto& field = dtype.fields[i]; - const auto& field_fill = vec[i]; - if (!field_fill.valid()) continue; - - // We assume a single outer element per field here (which is exactly how - // FillValueJsonBinder constructs per-field fill values). - std::memcpy( - dst + field.byte_offset, - static_cast(field_fill.data()), - static_cast(field.num_bytes)); + // ── Normal field access: just return that field's fill_value ─────────────── + TENSORSTORE_ASSIGN_OR_RETURN( + size_t field_index, GetFieldIndex(dtype, selected_field)); + if (field_index < vec.size()) { + return vec[field_index]; } - - return byte_arr; + // Fallback to "no fill". + return SharedArray(); } Result GetDimensionUnits() const override { @@ -394,8 +328,18 @@ class DataCacheBase virtual ZarrChunkCache& zarr_chunk_cache() = 0; - /// Returns true if this cache was opened with open_as_void=true. - virtual bool is_void_access() const = 0; + /// Returns the fill value to expose to user-visible reads for the given + /// component (field) index. By default this is just the per-field fill + /// value stored in the natural metadata; the `open_as_void` data cache + /// overrides this to return the packed single-byte fill value computed at + /// open time. + virtual SharedArray GetEffectiveFillValue(size_t component_index) { + const auto& m = this->metadata(); + if (component_index >= m.fill_value.size()) { + return SharedArray(); + } + return m.fill_value[component_index]; + } absl::Status ValidateMetadataCompatibility( const void* existing_metadata_ptr, @@ -426,18 +370,26 @@ class DataCacheBase i < static_cast(metadata.shape.size()); ++i) { implicit_upper_bounds[i] = true; } - // Handle extra dimensions from field_shape (e.g., the bytes dimension added - // by open_as_void). This only applies when there's exactly one field since - // multiple fields could have different shapes, making bounds ambiguous. + // The user-visible domain is `metadata.shape` extended by the + // metadata-level `field_shape` (the same vector that drives the codec + // resolution shape). This covers `rN` raw byte fields, the + // sharding-substituted dtype, and the synthetic single-byte view + // produced by `GetVoidMetadata` for `open_as_void`. Multi-field structs + // have `metadata.field_shape == {bytes_per_outer_element}` for the + // codec view, but the *user* must select a field first; here we only + // extend bounds when the dtype has a single field, so the trailing + // dim(s) are unambiguously user-visible. if (bounds.rank() > static_cast(metadata.shape.size()) && - metadata.data_type.fields.size() == 1) { - const auto& field = metadata.data_type.fields[0]; - if (static_cast(metadata.shape.size() + - field.field_shape.size()) == - bounds.rank()) { - for (size_t i = 0; i < field.field_shape.size(); ++i) { - bounds.shape()[metadata.shape.size() + i] = field.field_shape[i]; - bounds.origin()[metadata.shape.size() + i] = 0; + metadata.data_type.fields.size() == 1 && + !metadata.field_shape.empty()) { + const DimensionIndex chunked_rank = + static_cast(metadata.shape.size()); + const DimensionIndex field_rank = + static_cast(metadata.field_shape.size()); + if (chunked_rank + field_rank == bounds.rank()) { + for (DimensionIndex i = 0; i < field_rank; ++i) { + bounds.shape()[chunked_rank + i] = metadata.field_shape[i]; + bounds.origin()[chunked_rank + i] = 0; } } } @@ -461,49 +413,13 @@ class DataCacheBase } static internal::ChunkGridSpecification GetChunkGridSpecification( - const ZarrMetadata& metadata, size_t field_index = 0) { + const ZarrMetadata& metadata) { assert(!metadata.fill_value.empty()); - internal::ChunkGridSpecification::ComponentList components; - - // Special case: void access - create single component for entire struct - if (field_index == kVoidFieldIndex) { - // For void access, create a zero-filled byte array as the fill value - const Index bytes_per_element = metadata.data_type.bytes_per_outer_element; - auto base_fill_value = AllocateArray( - span({bytes_per_element}), c_order, value_init, - dtype_v); - - // Broadcast to shape [unbounded, unbounded, ..., struct_size] - std::vector target_shape(metadata.rank, kInfIndex); - target_shape.push_back(bytes_per_element); - auto chunk_fill_value = - BroadcastArray(base_fill_value, BoxView<>(target_shape)).value(); - - // Add extra dimension for struct size in bytes - std::vector chunk_shape_with_bytes = metadata.chunk_shape; - chunk_shape_with_bytes.push_back(bytes_per_element); - - // Create permutation: copy existing inner_order and add the new dimension - std::vector void_permutation(metadata.rank + 1); - std::copy_n(metadata.inner_order.data(), metadata.rank, - void_permutation.begin()); - void_permutation[metadata.rank] = metadata.rank; // Add the bytes dimension - - auto& component = components.emplace_back( - internal::AsyncWriteArray::Spec{ - std::move(chunk_fill_value), - // Since all dimensions are resizable, just - // specify unbounded `valid_data_bounds`. - Box<>(metadata.rank + 1), - ContiguousLayoutPermutation<>( - span(void_permutation.data(), metadata.rank + 1))}, - chunk_shape_with_bytes); - component.array_spec.fill_value_comparison_kind = - EqualityComparisonKind::identical; - return internal::ChunkGridSpecification(std::move(components)); - } - - // Create one component per field using the shared helper + // The cache always sees a `metadata` whose `data_type` correctly + // describes the field structure -- including the synthetic single-byte + // field with `field_shape = {bytes_per_outer_element}` that + // `GetVoidMetadata` produces for `open_as_void`. `CreateFieldGridSpecification` + // therefore handles scalar, struct, rN, and open-as-void modes uniformly. return CreateFieldGridSpecification( metadata.chunk_shape, metadata.data_type, span(metadata.inner_order.data(), metadata.rank), &metadata.fill_value); @@ -669,15 +585,23 @@ using internal_kvs_backed_chunk_driver::DataCacheInitializer; template class ZarrDataCache : public ChunkCacheImpl, public DataCacheBase { public: + // `effective_metadata` is the metadata view the cache operates on for chunk + // decoding, grid layout, fills, transforms, and bounds. It is always + // populated -- when no substitution applies it aliases the persisted + // (natural) metadata; under `open_as_void` it is the byte-substituted view + // produced by `GetVoidMetadata`. `is_substituted` records whether a + // substitution was applied; it is the single bit required to round-trip + // `open_as_void` back into the spec. template explicit ZarrDataCache(DataCacheInitializer&& initializer, - std::string key_prefix, U&&... arg) + std::string key_prefix, + std::shared_ptr effective_metadata, + bool is_substituted, U&&... arg) : ChunkCacheImpl(std::move(initializer.store), std::forward(arg)...), DataCacheBase(std::move(initializer), std::move(key_prefix)), - grid_(DataCacheBase::GetChunkGridSpecification( - metadata(), - // Check if this is void access by examining the dtype - ChunkCacheImpl::open_as_void_ ? kVoidFieldIndex : false)) {} + effective_metadata_(std::move(effective_metadata)), + is_substituted_(is_substituted), + grid_(DataCacheBase::GetChunkGridSpecification(*effective_metadata_)) {} const internal::LexicographicalGridIndexKeyParser& GetChunkStorageKeyParser() final { @@ -688,8 +612,6 @@ class ZarrDataCache : public ChunkCacheImpl, public DataCacheBase { ZarrChunkCache& zarr_chunk_cache() final { return *this; } - bool is_void_access() const final { return ChunkCacheImpl::open_as_void_; } - const internal::ChunkGridSpecification& grid() const override { return grid_; } @@ -705,81 +627,55 @@ class ZarrDataCache : public ChunkCacheImpl, public DataCacheBase { return DataCacheBase::executor(); } - // Override to handle void access - check the dtype to see if this is void + // The framework hands `metadata_ptr` from the metadata cache (i.e. the + // current persisted form). Cache-internal layout/transform logic drives + // off `effective_metadata_` only when a substitution was applied at open + // time -- otherwise we forward the live `metadata_ptr` so that operations + // like `Resize` correctly observe the post-resize shape. Either way the + // base implementation reads `field_shape` off whichever metadata it gets + // and produces the right user-visible rank with no `open_as_void`-aware + // branching beyond the substitution site itself. + const void* RouteMetadata(const void* metadata_ptr) const { + return is_substituted_ ? effective_metadata_.get() : metadata_ptr; + } + Result> GetExternalToInternalTransform( const void* metadata_ptr, size_t component_index) override { - if (ChunkCacheImpl::open_as_void_) { - const auto& metadata = *static_cast(metadata_ptr); - // For void access, create transform with extra bytes dimension - const DimensionIndex rank = metadata.rank; - const Index bytes_per_element = metadata.data_type.bytes_per_outer_element; - const DimensionIndex total_rank = rank + 1; - - std::string_view normalized_dimension_names[kMaxRank]; - for (DimensionIndex i = 0; i < rank; ++i) { - if (const auto& name = metadata.dimension_names[i]; name.has_value()) { - normalized_dimension_names[i] = *name; - } - } - - auto builder = - tensorstore::IndexTransformBuilder<>(total_rank, total_rank); - std::vector full_shape = metadata.shape; - full_shape.push_back(bytes_per_element); - builder.input_shape(full_shape); - builder.input_labels(span(&normalized_dimension_names[0], total_rank)); - - // Set implicit bounds: array dims have implicit upper bounds (resizable), - // bytes dim has explicit bounds (fixed size). - builder.implicit_lower_bounds(DimensionSet(false)); - builder.implicit_upper_bounds(DimensionSet::UpTo(rank)); - - for (DimensionIndex i = 0; i < total_rank; ++i) { - builder.output_single_input_dimension(i, i); - } - return builder.Finalize(); - } - - // Not void access - delegate to base implementation - return DataCacheBase::GetExternalToInternalTransform(metadata_ptr, - component_index); + return DataCacheBase::GetExternalToInternalTransform( + RouteMetadata(metadata_ptr), component_index); } void GetChunkGridBounds(const void* metadata_ptr, MutableBoxView<> bounds, DimensionSet& implicit_lower_bounds, DimensionSet& implicit_upper_bounds) override { - if (ChunkCacheImpl::open_as_void_) { - const auto& metadata = *static_cast(metadata_ptr); - const DimensionIndex rank = metadata.rank; - const Index bytes_per_element = metadata.data_type.bytes_per_outer_element; - - assert(bounds.rank() == rank + 1); - std::fill(bounds.origin().begin(), bounds.origin().end(), Index(0)); - std::copy(metadata.shape.begin(), metadata.shape.end(), - bounds.shape().begin()); - bounds.shape()[rank] = bytes_per_element; - - implicit_lower_bounds = false; - implicit_upper_bounds = DimensionSet::UpTo(rank); - return; - } - - // Not void access - delegate to base implementation - DataCacheBase::GetChunkGridBounds(metadata_ptr, bounds, implicit_lower_bounds, + DataCacheBase::GetChunkGridBounds(RouteMetadata(metadata_ptr), bounds, + implicit_lower_bounds, implicit_upper_bounds); } absl::Status GetBoundSpecData(KvsDriverSpec& spec_base, const void* metadata_ptr, size_t component_index) override { + // The persisted (natural) metadata drives the spec body so the user sees + // their original dtype on round-trip. `open_as_void` is the only piece + // of substitution-time state the cache must surface back into the spec. TENSORSTORE_RETURN_IF_ERROR( DataCacheBase::GetBoundSpecData(spec_base, metadata_ptr, component_index)); auto& spec = static_cast(spec_base); - // Preserve the open_as_void flag so spec round-trips correctly - spec.open_as_void = ChunkCacheImpl::open_as_void_; + spec.open_as_void = is_substituted_; return absl::OkStatus(); } + SharedArray GetEffectiveFillValue( + size_t component_index) override { + if (component_index >= effective_metadata_->fill_value.size()) { + return SharedArray(); + } + return effective_metadata_->fill_value[component_index]; + } + + std::shared_ptr effective_metadata_; + bool is_substituted_; internal::ChunkGridSpecification grid_; }; @@ -799,42 +695,12 @@ class ZarrDriver : public ZarrDriverBase { Result> GetFillValue( IndexTransformView<> transform) override { - const auto& metadata = this->metadata(); - if (metadata.fill_value.empty()) { - return SharedArray(); - } - - // For void access, convert fill_value to byte array representation. - // This is similar to v2's CreateVoidMetadata fill_value handling. - // In zarr3, endianness is handled by the codec chain, so we just copy - // the raw bytes from each field's fill_value. - if (static_cast(cache())->is_void_access()) { - const Index nbytes = metadata.data_type.bytes_per_outer_element; - auto byte_fill = AllocateArray({nbytes}, c_order, value_init); - - // Copy bytes from each field's fill_value at their respective offsets - for (size_t field_i = 0; field_i < metadata.data_type.fields.size(); - ++field_i) { - const auto& field = metadata.data_type.fields[field_i]; - if (field_i >= metadata.fill_value.size() || - !metadata.fill_value[field_i].valid()) { - continue; - } - const auto& fill_value = metadata.fill_value[field_i]; - // Copy the raw bytes from the fill_value to the byte array - std::memcpy(byte_fill.data() + field.byte_offset, - fill_value.data(), - field.num_bytes); - } - - return byte_fill; - } - - size_t index = this->component_index(); - if (index >= metadata.fill_value.size()) { - return absl::OutOfRangeError("Component index out of bounds"); - } - return metadata.fill_value[index]; + // The data cache routes this through any substituted view metadata, so + // `open_as_void` returns the packed single-byte fill value at index 0 + // while normal access returns the per-component fill from the natural + // metadata. + return static_cast(cache())->GetEffectiveFillValue( + this->component_index()); } Future GetStorageStatistics( @@ -916,6 +782,35 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase { public: using ZarrDriver::OpenStateBase::OpenStateBase; + // Lazily computed (and memoised across `GetComponentIndex` / + // `GetDataCache`) byte-substituted view of the metadata. Only non-null + // when `spec().open_as_void` is true and a successful substitution has + // already been performed against the framework-supplied `metadata_ptr`. + // Both call sites pass the same pointer (the persisted metadata), so a + // single computation suffices for both validation and cache construction. + Result> GetOrComputeVoidMetadata( + const ZarrMetadata& metadata) { + if (!cached_void_metadata_for_) { + TENSORSTORE_ASSIGN_OR_RETURN(auto void_metadata, + internal_zarr3::GetVoidMetadata(metadata)); + cached_void_metadata_ = std::move(void_metadata); + cached_void_metadata_for_ = &metadata; + } + // Defensive: if the framework ever started passing a different metadata + // object across calls, recompute rather than serve stale state. In the + // current code path this branch is unreachable. + if (cached_void_metadata_for_ != &metadata) { + TENSORSTORE_ASSIGN_OR_RETURN(auto void_metadata, + internal_zarr3::GetVoidMetadata(metadata)); + cached_void_metadata_ = std::move(void_metadata); + cached_void_metadata_for_ = &metadata; + } + return cached_void_metadata_; + } + + std::shared_ptr cached_void_metadata_; + const ZarrMetadata* cached_void_metadata_for_ = nullptr; + std::string GetPrefixForDeleteExisting() override { return spec().store.path; } @@ -964,23 +859,36 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase { DataCacheInitializer&& initializer) override { const auto& metadata = *static_cast(initializer.metadata.get()); - ZarrDType dtype = spec().open_as_void ? metadata.GetVoidAccessDType() - : metadata.data_type; - // Determine if original dtype is structured (multiple fields). - // This affects how void access handles codec operations. - const bool original_is_structured = metadata.data_type.fields.size() > 1; - - // Get the original dtype for void access encoding (needed by EncodeChunk). - // For non-structured types, this is the single field's dtype. - DataType original_dtype = !metadata.data_type.fields.empty() - ? metadata.data_type.fields[0].dtype - : DataType{}; + + // Build the cache-side `effective_metadata`: under `open_as_void` this is + // the byte-substituted view produced by `GetVoidMetadata`; otherwise it + // aliases the persisted metadata. Everything downstream -- chunk cache, + // grid spec, codec state -- runs against that single view, with no + // `open_as_void`-specific branching beyond the substitution itself. + std::shared_ptr effective_metadata; + bool is_substituted = false; + if (spec().open_as_void) { + // Reuse the substitution that `GetComponentIndex` already performed; + // `GetOrComputeVoidMetadata` memoises by metadata pointer and avoids + // re-resolving the codec chain. `CHECK` here because validation has + // already succeeded by this point in the open flow. + TENSORSTORE_CHECK_OK_AND_ASSIGN(effective_metadata, + GetOrComputeVoidMetadata(metadata)); + is_substituted = true; + } else { + effective_metadata = std::static_pointer_cast( + initializer.metadata); + } + ZarrDType dtype = effective_metadata->data_type; + ZarrCodecChain::PreparedState::Ptr codec_state = + effective_metadata->codec_state; + std::vector field_shape = effective_metadata->field_shape; return internal_zarr3::MakeZarrChunkCache( *metadata.codecs, std::move(initializer), spec().store.path, - metadata.codec_state, dtype, - /*data_cache_pool=*/*cache_pool(), - spec().open_as_void, original_is_structured, original_dtype); + std::move(effective_metadata), is_substituted, std::move(codec_state), + std::move(dtype), std::move(field_shape), + /*data_cache_pool=*/*cache_pool()); } Result GetComponentIndex(const void* metadata_ptr, @@ -988,15 +896,26 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase { const auto& metadata = *static_cast(metadata_ptr); TENSORSTORE_RETURN_IF_ERROR( ValidateMetadata(metadata, spec().metadata_constraints)); + if (spec().open_as_void) { + if (!spec().selected_field.empty()) { + return absl::InvalidArgumentError( + "\"field\" and \"open_as_void\" are mutually exclusive"); + } + // Compute the void substitution once and stash it on the open state; + // `GetDataCache` reuses the same shared_ptr below to avoid re-resolving + // the codec chain. + TENSORSTORE_ASSIGN_OR_RETURN(auto void_metadata, + GetOrComputeVoidMetadata(metadata)); + TENSORSTORE_RETURN_IF_ERROR(ValidateMetadataSchema( + *void_metadata, /*field_index=*/0, spec().schema)); + // The void view exposes a single component (the synthetic byte field). + return static_cast(0); + } TENSORSTORE_ASSIGN_OR_RETURN( auto field_index, - GetFieldIndex(metadata.data_type, spec().selected_field, spec().open_as_void)); + GetFieldIndex(metadata.data_type, spec().selected_field)); TENSORSTORE_RETURN_IF_ERROR( ValidateMetadataSchema(metadata, field_index, spec().schema)); - if (field_index == kVoidFieldIndex) { - // Void access uses component 0 (raw bytes). - return static_cast(0); - } return field_index; } }; diff --git a/tensorstore/driver/zarr3/driver_test.cc b/tensorstore/driver/zarr3/driver_test.cc index 6c8d0659b..2987636b2 100644 --- a/tensorstore/driver/zarr3/driver_test.cc +++ b/tensorstore/driver/zarr3/driver_test.cc @@ -1883,6 +1883,10 @@ ::nlohmann::json GetOpenAsVoidSpec(std::string_view kvstore_path) { } TEST(Zarr3OpenAsVoidTest, SimpleType) { + // Per the zarr v3 open_as_void spec, opening a plain scalar dtype with + // `open_as_void` succeeds and yields a `byte` array whose innermost + // dimension equals the scalar's byte width. For int16 (2 bytes), an + // array of shape [4, 4] becomes byte-shape [4, 4, 2]. auto context = Context::Default(); ::nlohmann::json create_spec{ @@ -1909,12 +1913,17 @@ TEST(Zarr3OpenAsVoidTest, SimpleType) { {"open_as_void", true}, }; - EXPECT_THAT(tensorstore::Open(void_spec, context, tensorstore::OpenMode::open, - tensorstore::ReadWriteMode::read) - .result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - ".*open_as_void is only supported for structured dtypes.*")); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto void_store, + tensorstore::Open(void_spec, context, tensorstore::OpenMode::open, + tensorstore::ReadWriteMode::read) + .result()); + + EXPECT_EQ(tensorstore::dtype_v, + void_store.dtype()); + EXPECT_EQ(3, void_store.rank()); + EXPECT_THAT(void_store.domain().shape(), + ::testing::ElementsAre(4, 4, 2)); } TEST(Zarr3OpenAsVoidTest, StructuredType) { @@ -2716,11 +2725,12 @@ TEST(Zarr3OpenAsVoidTest, IncompatibleMetadata) { StatusIs(absl::StatusCode::kFailedPrecondition)); } -TEST(Zarr3OpenAsVoidTest, WithShardingRejectsSimpleType) { - // Test that open_as_void with sharding correctly rejects simple dtypes. +TEST(Zarr3OpenAsVoidTest, WithShardingSimpleType) { + // Per the zarr v3 open_as_void spec, simple dtypes are also supported even + // when wrapped in a sharding_indexed codec. For int32 (4 bytes) shape + // [8,8] with sub-chunks [4,4], the void view has shape [8,8,4] of `byte`. auto context = Context::Default(); - // Create a sharded array with simple dtype ::nlohmann::json create_spec{ {"driver", "zarr3"}, {"kvstore", {{"driver", "memory"}, {"path", "prefix/"}}}, @@ -2746,7 +2756,6 @@ TEST(Zarr3OpenAsVoidTest, WithShardingRejectsSimpleType) { tensorstore::ReadWriteMode::read_write) .result()); - // Write some data auto data = tensorstore::MakeArray( {{1, 2, 0, 0, 0, 0, 0, 0}, {3, 4, 0, 0, 0, 0, 0, 0}, @@ -2758,19 +2767,23 @@ TEST(Zarr3OpenAsVoidTest, WithShardingRejectsSimpleType) { {0, 0, 0, 0, 0, 0, 0, 0}}); TENSORSTORE_EXPECT_OK(tensorstore::Write(data, store).result()); - // Attempt to open with open_as_void=true - should fail for simple dtype ::nlohmann::json void_spec{ {"driver", "zarr3"}, {"kvstore", {{"driver", "memory"}, {"path", "prefix/"}}}, {"open_as_void", true}, }; - EXPECT_THAT(tensorstore::Open(void_spec, context, tensorstore::OpenMode::open, - tensorstore::ReadWriteMode::read) - .result(), - tensorstore::MatchesStatus( - absl::StatusCode::kInvalidArgument, - ".*open_as_void is only supported for structured dtypes.*")); + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto void_store, + tensorstore::Open(void_spec, context, tensorstore::OpenMode::open, + tensorstore::ReadWriteMode::read) + .result()); + + EXPECT_EQ(tensorstore::dtype_v, + void_store.dtype()); + EXPECT_EQ(3, void_store.rank()); + EXPECT_THAT(void_store.domain().shape(), + ::testing::ElementsAre(8, 8, 4)); } TEST(Zarr3OpenAsVoidTest, InvalidSchema) { @@ -2899,11 +2912,16 @@ TEST(Zarr3OpenAsVoidTest, StructBigEndian) { EXPECT_THAT(byte_array, MatchesArray(expected_array)); } -TEST(Zarr3OpenAsVoidTest, ShardedSubChunkShapeExtension) { +TEST(Zarr3OpenAsVoidTest, ShardedSubChunkShapeIsUserForm) { + // The sharding_indexed `chunk_shape` exposed in the round-tripped spec is + // the user-facing form (the chunked dimensions only), not the codec's + // internally extended form. This keeps zarr.json round-trips and + // user-provided spec constraints in a single canonical representation. + // For a struct{x: uint8, y: int16} with sub-chunk shape [4, 4], the spec + // round-trips as [4, 4] regardless of whether opened normally or via + // open_as_void. auto context = Context::Default(); - // Create sharded structured array. - // Struct: x (uint8), y (int16) -> 3 bytes. ::nlohmann::json create_spec{ {"driver", "zarr3"}, {"kvstore", {{"driver", "memory"}, {"path", "prefix_shard/"}}}, @@ -2933,7 +2951,6 @@ TEST(Zarr3OpenAsVoidTest, ShardedSubChunkShapeExtension) { tensorstore::ReadWriteMode::read_write) .result()); - // Open as void. ::nlohmann::json void_spec{ {"driver", "zarr3"}, {"kvstore", {{"driver", "memory"}, {"path", "prefix_shard/"}}}, @@ -2946,15 +2963,13 @@ TEST(Zarr3OpenAsVoidTest, ShardedSubChunkShapeExtension) { tensorstore::ReadWriteMode::read) .result()); - // Spec should have sub_chunk_shape extended with the bytes dimension [4, 4, - // 3] TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec, void_store.spec()); TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto spec_json, spec.ToJson()); EXPECT_THAT( spec_json["metadata"]["codecs"], ::testing::Contains(JsonSubValuesMatch( {{"/name", "sharding_indexed"}, - {"/configuration/chunk_shape", ::nlohmann::json({4, 4, 3})}}))); + {"/configuration/chunk_shape", ::nlohmann::json({4, 4})}}))); } // Helper: returns a JSON spec for creating a sharded structured array. diff --git a/tensorstore/driver/zarr3/dtype.cc b/tensorstore/driver/zarr3/dtype.cc index f0c0812fb..0935abc09 100644 --- a/tensorstore/driver/zarr3/dtype.cc +++ b/tensorstore/driver/zarr3/dtype.cc @@ -366,6 +366,16 @@ Result ParseDTypeNoDerived(const nlohmann::json& value) { } // namespace absl::Status ValidateDType(ZarrDType& dtype) { + // Invariant: every valid zarr v3 dtype has at least one field. All JSON + // parser entry points enforce this (`ValidateFieldsArrayNotEmpty` for + // struct/structured, single-field `resize(1)` for scalar / raw_bytes). + // This guard catches future programmatic constructions of `ZarrDType` + // that bypass the parser, so that consumers indexing `fields[0]` get a + // clean error rather than out-of-bounds access. + if (dtype.fields.empty()) { + return absl::FailedPreconditionError( + "zarr3 data type must have at least one field"); + } dtype.bytes_per_outer_element = 0; for (size_t field_i = 0; field_i < dtype.fields.size(); ++field_i) { auto& field = dtype.fields[field_i]; diff --git a/tensorstore/driver/zarr3/dtype.h b/tensorstore/driver/zarr3/dtype.h index 31db6a2c0..f22ac7a78 100644 --- a/tensorstore/driver/zarr3/dtype.h +++ b/tensorstore/driver/zarr3/dtype.h @@ -79,7 +79,10 @@ struct ZarrDType { std::string name; /// Inner array dimensions of this field, derived from `flexible_shape`. - /// Used for raw_bytes dtype and open_as_void byte dimensions. + /// Inner array shape contributed by this field to the full TensorStore + /// rank. For scalar fields this is empty; for raw byte fields (e.g. + /// `r24`, structured-as-void access) this is the inner shape that is + /// surfaced as trailing dimensions in the user-visible array. std::vector field_shape; /// Product of `field_shape` dimensions (derived value). diff --git a/tensorstore/driver/zarr3/dtype_test.cc b/tensorstore/driver/zarr3/dtype_test.cc index 924dc8e70..72a1f2694 100644 --- a/tensorstore/driver/zarr3/dtype_test.cc +++ b/tensorstore/driver/zarr3/dtype_test.cc @@ -39,6 +39,7 @@ using ::tensorstore::StatusIs; using ::tensorstore::internal_zarr3::ChooseBaseDType; using ::tensorstore::internal_zarr3::ParseBaseDType; using ::tensorstore::internal_zarr3::ParseDType; +using ::tensorstore::internal_zarr3::ValidateDType; using ::tensorstore::internal_zarr3::ZarrDType; using ::testing::HasSubstr; using ::tensorstore::IsOkAndHolds; @@ -228,6 +229,17 @@ TEST(ParseDType, DuplicateFieldName) { HasSubstr("Field name \"x\" occurs more than once"))); } +TEST(ValidateDType, RejectsEmptyFields) { + // Defensive guard: a default-constructed `ZarrDType` has no fields, which + // is an invariant violation for downstream consumers indexing `fields[0]`. + // The JSON parsers all enforce non-empty fields, so this only fires for + // programmatic constructions that bypass the parser. + ZarrDType empty; + EXPECT_THAT(ValidateDType(empty), + StatusIs(absl::StatusCode::kFailedPrecondition, + HasSubstr("must have at least one field"))); +} + TEST(ParseDType, NonStringFieldBaseDType) { EXPECT_THAT(ParseDType(::nlohmann::json::array_t{{"x", 3}}), StatusIs(absl::StatusCode::kInvalidArgument, @@ -583,6 +595,74 @@ TEST(ParseDType, StructuredWithFieldShape) { EXPECT_THAT(ParseDType(input), IsOkAndHolds(expected)); } +TEST(ParseBaseDType, R24FieldShape) { + // r24 = 24 bits = 3 bytes; field_shape encodes that as {3}. + EXPECT_THAT("r24", ParsesAsBaseDType(dtype_v, + std::vector{3})); +} + +TEST(ParseDType, R24SingleField) { + // r24 as a top-level dtype yields a single field with field_shape {3}. + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto dtype, ParseDType("r24")); + ASSERT_EQ(dtype.fields.size(), 1); + EXPECT_FALSE(dtype.has_fields); + EXPECT_EQ(dtype.bytes_per_outer_element, 3); + EXPECT_EQ(dtype.fields[0].dtype, dtype_v); + EXPECT_EQ(dtype.fields[0].num_bytes, 3); + EXPECT_EQ(dtype.fields[0].num_inner_elements, 3); + EXPECT_THAT(dtype.fields[0].field_shape, ::testing::ElementsAre(3)); +} + +TEST(ParseDType, R24InStruct) { + // r24 nested inside a struct: bytes pack normally, field_shape preserved. + ::nlohmann::json input = { + {"name", "struct"}, + {"configuration", + {{"fields", ::nlohmann::json::array( + {{{"name", "tag"}, {"data_type", "uint8"}}, + {{"name", "rgb"}, {"data_type", "r24"}}})}}}}; + + ZarrDType expected{}; + AddFieldToZarrDType(expected, + {{/*.encoded_dtype=*/"uint8", + /*.dtype=*/dtype_v, + /*.flexible_shape=*/{}}, + /*.name=*/"tag", + /*.field_shape=*/{}, + /*.num_inner_elements=*/1, + /*.byte_offset=*/0, + /*.num_bytes=*/1}); + AddFieldToZarrDType(expected, + {{/*.encoded_dtype=*/"r24", + /*.dtype=*/dtype_v, + /*.flexible_shape=*/{3}}, + /*.name=*/"rgb", + /*.field_shape=*/{3}, + /*.num_inner_elements=*/3, + /*.byte_offset=*/0, + /*.num_bytes=*/3}); + + EXPECT_THAT(ParseDType(input), IsOkAndHolds(expected)); + EXPECT_EQ(expected.bytes_per_outer_element, 4); + EXPECT_EQ(expected.fields[1].byte_offset, 1); +} + +TEST(ParseDType, R24SerializationRoundTrip) { + ::nlohmann::json input = { + {"name", "struct"}, + {"configuration", + {{"fields", ::nlohmann::json::array( + {{{"name", "rgb"}, {"data_type", "r24"}}})}}}}; + + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto dtype, ParseDType(input)); + ::nlohmann::json output = dtype; + // Field's data_type round-trips as the literal "r24" string. + EXPECT_EQ(output["configuration"]["fields"][0]["data_type"], "r24"); + // Re-parsing the serialized form yields the same dtype. + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto reparsed, ParseDType(output)); + EXPECT_EQ(dtype, reparsed); +} + TEST(ParseDType, ManyFieldsOffsets) { // Verify that many fields are supported and byte offsets are computed correctly. // There is no explicit limit on the number of fields; the limit is bounded only diff --git a/tensorstore/driver/zarr3/metadata.cc b/tensorstore/driver/zarr3/metadata.cc index 03c7f9fc0..6ba5f0a92 100644 --- a/tensorstore/driver/zarr3/metadata.cc +++ b/tensorstore/driver/zarr3/metadata.cc @@ -48,9 +48,12 @@ #include "tensorstore/codec_spec.h" #include "tensorstore/contiguous_layout.h" #include "tensorstore/data_type.h" +#include "tensorstore/driver/zarr3/codec/bytes.h" #include "tensorstore/driver/zarr3/codec/codec_chain_spec.h" #include "tensorstore/driver/zarr3/codec/codec_spec.h" #include "tensorstore/driver/zarr3/codec/sharding_indexed.h" +#include "tensorstore/internal/data_type_endian_conversion.h" +#include "tensorstore/util/endian.h" #include "tensorstore/driver/zarr3/default_nan.h" #include "tensorstore/driver/zarr3/dtype.h" #include "tensorstore/driver/zarr3/name_configuration_json_binder.h" @@ -648,47 +651,218 @@ std::string ZarrMetadata::GetCompatibilityKey() const { .dump(); } -ZarrDType ZarrMetadata::GetVoidAccessDType() const { +namespace { +// Walks `codec_specs` (transparently descending through any +// `sharding_indexed` layers) and returns the endianness selected by the +// innermost `bytes` codec. Falls back to `endian::native` when no `bytes` +// codec is reachable, when the codec was resolved against an endian-invariant +// data type, or when the user did not pin an explicit endian. +endian GetBytesCodecEndianImpl(const ZarrCodecChainSpec& codec_specs) { + const ZarrArrayToBytesCodecSpec* a2b = codec_specs.array_to_bytes.get(); + while (a2b != nullptr) { + if (auto* sharding = dynamic_cast(a2b)) { + const ZarrCodecChainSpec* sub = sharding->GetSubChunkCodecs(); + if (sub == nullptr) break; + a2b = sub->array_to_bytes.get(); + continue; + } + if (auto* bytes_codec = dynamic_cast(a2b)) { + return bytes_codec->options.endianness.value_or(endian::native); + } + break; + } + return endian::native; +} + +// Returns the byte order in which a per-field typed fill must be written +// when packed into the synthetic byte view consumed under `open_as_void`. +// +// The void byte buffer must mirror what would be read from a present chunk +// after the original codec chain encoded it: +// +// * Single scalar field (`field_shape` empty, dtype != byte): the chunk +// cache hands the typed array to the bytes codec, which performs endian +// conversion to the codec's configured endian. The fill must follow. +// * Multi-field struct, or `rN` raw byte field: the chunk cache packs +// bytes via `CopyArray` (native-endian memcpy) and the bytes codec at +// the bottom is resolved against `byte` (endian-invariant). Fill must +// therefore stay native-endian. +endian SelectFillTargetEndian(const ZarrDType& dtype, + const ZarrCodecChainSpec& codec_specs) { + if (dtype.fields.size() == 1 && dtype.fields[0].field_shape.empty()) { + return GetBytesCodecEndianImpl(codec_specs); + } + return endian::native; +} +} // namespace + +SharedArray MakeVoidFillValue( + const ZarrDType& dtype, const ZarrCodecChainSpec& codec_specs, + span> per_field_fill) { + const Index nbytes = dtype.bytes_per_outer_element; + auto byte_fill = AllocateArray(span({nbytes}), c_order, + value_init, + dtype_v); + auto* dst = static_cast(byte_fill.data()); + const endian target_endian = SelectFillTargetEndian(dtype, codec_specs); + const size_t num_fields = std::min( + dtype.fields.size(), static_cast(per_field_fill.size())); + for (size_t i = 0; i < num_fields; ++i) { + const auto& field = dtype.fields[i]; + const auto& fill = per_field_fill[i]; + if (!fill.valid()) continue; + ArrayView dst_view( + ElementPointer{static_cast(dst + field.byte_offset), + field.dtype}, + fill.layout()); + internal::EncodeArray(fill, dst_view, target_endian); + } + return byte_fill; +} + +ZarrDType MakeVoidDType(Index bytes_per_outer_element) { return ZarrDType{ /*.has_fields=*/false, /*.fields=*/{ZarrDType::Field{ ZarrDType::BaseDType{"", dtype_v, - {data_type.bytes_per_outer_element}}, + {bytes_per_outer_element}}, /*.name=*/"", - /*.field_shape=*/{data_type.bytes_per_outer_element}, - /*.num_inner_elements=*/data_type.bytes_per_outer_element, + /*.field_shape=*/{bytes_per_outer_element}, + /*.num_inner_elements=*/bytes_per_outer_element, /*.byte_offset=*/0, - /*.num_bytes=*/data_type.bytes_per_outer_element}}, - /*.bytes_per_outer_element=*/data_type.bytes_per_outer_element}; + /*.num_bytes=*/bytes_per_outer_element}}, + /*.bytes_per_outer_element=*/bytes_per_outer_element}; } -absl::Status ValidateMetadata(ZarrMetadata& metadata) { - // Determine if this is a structured type with multiple fields - const bool is_structured = - metadata.data_type.fields.size() > 1 || - (metadata.data_type.fields.size() == 1 && - !metadata.data_type.fields[0].field_shape.empty()); +absl::Status ValidateVoidCodecChain(const ZarrCodecChainSpec& codec_specs) { + // After the inner-shape architecture change, array-to-array codecs operate + // strictly on the chunked dimensions and propagate the inner dims via the + // `inner_shape` resolve parameter; they cannot rearrange or reshape the + // byte dim, and the chain Resolve enforces that they preserve the dtype. + // The only structural precondition this function still checks is that the + // innermost array-to-bytes codec (after unwinding any `sharding_indexed` + // layers) is the `bytes` codec, since other a-to-b codecs would alter the + // on-disk byte layout in ways `open_as_void` cannot reproduce. + const ZarrArrayToBytesCodecSpec* a2b = codec_specs.array_to_bytes.get(); + while (auto* sharding = + dynamic_cast(a2b)) { + const ZarrCodecChainSpec* sub = sharding->GetSubChunkCodecs(); + if (sub == nullptr) { + return absl::InvalidArgumentError( + "open_as_void: nested sharding_indexed codec is missing its " + "sub-chunk codec specification"); + } + a2b = sub->array_to_bytes.get(); + } + if (dynamic_cast(a2b) == nullptr) { + return absl::InvalidArgumentError( + "open_as_void requires the innermost array-to-bytes codec to be the " + "`bytes` codec (after unwrapping any sharding_indexed layers). " + "Codecs that alter the byte representation of the chunk (e.g. " + "blosc-direct, future bitround) are not supported under " + "open_as_void."); + } + return absl::OkStatus(); +} + +Result> GetVoidMetadata( + const ZarrMetadata& metadata) { + // Per the zarr v3 open_as_void spec, raw byte access is supported for any + // data type; the only structural precondition is that the codec chain is + // re-resolvable under the substituted `byte` data type. + TENSORSTORE_RETURN_IF_ERROR(ValidateVoidCodecChain(metadata.codec_specs)); + + auto void_metadata = std::make_shared(metadata); + + // Replace the data type with the synthetic single-field byte view. This is + // what every downstream consumer (chunk cache, codec resolution, schema + // validation) sees. The persisted on-disk metadata is untouched: this view + // exists only in memory. + void_metadata->data_type = + MakeVoidDType(metadata.data_type.bytes_per_outer_element); + + // Pack the per-field fill values into a single byte array following the + // struct's byte_offset layout, mirroring how a chunk is laid out on disk. + // Endian conversion (when applicable) is handled by `MakeVoidFillValue`. + void_metadata->fill_value = { + MakeVoidFillValue(metadata.data_type, metadata.codec_specs, + metadata.fill_value)}; + + // Cleared so that `ValidateMetadata` rederives `field_shape` from the + // newly-substituted dtype rather than carrying over the natural metadata's + // value. In every supported case this rederives to the same numeric value + // (the per-field field_shape on the synthetic byte field), but staging it + // through the standard derivation keeps the void path uniform with parse- + // time metadata. + void_metadata->field_shape.clear(); + + // Force re-resolution of the codec chain against the byte data type. The + // codec chain spec itself (endianness selections, sharding parameters, etc.) + // is unchanged; only the resolved codec instances and prepared state are + // recomputed. + void_metadata->codecs.reset(); + void_metadata->codec_state.reset(); + TENSORSTORE_RETURN_IF_ERROR(ValidateMetadata(*void_metadata)); + return void_metadata; +} - // Build the codec shape - for structured types, include bytes dimension - std::vector codec_shape(metadata.chunk_shape.begin(), - metadata.chunk_shape.end()); - if (is_structured) { - codec_shape.push_back(metadata.data_type.bytes_per_outer_element); +namespace { +// Populate `metadata.field_shape` (the codec-view inner trailing dimensions) +// from the data type, unless the caller pre-populated it (e.g. +// `GetVoidMetadata` pins it to `{bytes_per_outer_element}` after substituting +// the dtype). For single-field rN-style dtypes we hoist the per-field +// `field_shape`; for multi-field structs we use the single flat byte trailing +// dim. Single-field scalar dtypes contribute no inner codec dimensions. +// +// This is the single source of truth for `field_shape`. Call sites that need +// it (both at parse time via `ValidateMetadata` and at create time via +// `GetNewMetadata`) defer to this helper rather than recomputing locally. +void DeriveFieldShape(ZarrMetadata& metadata) { + if (!metadata.field_shape.empty()) return; + const auto& dtype = metadata.data_type; + if (dtype.fields.size() == 1 && !dtype.fields[0].field_shape.empty()) { + metadata.field_shape.assign(dtype.fields[0].field_shape.begin(), + dtype.fields[0].field_shape.end()); + } else if (dtype.fields.size() > 1) { + metadata.field_shape.push_back(dtype.bytes_per_outer_element); } +} +} // namespace + +absl::Status ValidateMetadata(ZarrMetadata& metadata) { + DeriveFieldShape(metadata); + + // The codec chain is resolved at the *chunked* rank only. Inner trailing + // dimensions contributed by the dtype's field_shape (multi-field structs, + // `rN` raw byte fields, `open_as_void`'s byte substitution) travel via + // `decoded.inner_shape`: array-to-array codecs propagate this field + // unchanged and cannot operate on it; the leaf array-to-bytes codec + // consumes it for byte-stream sizing. + // + // At runtime the chunk cache still hands extended-rank arrays to the + // codec chain, so each codec's runtime state is built at runtime rank + // (chunked + inner) inside its `Resolve` step. + const bool has_field_shape = !metadata.field_shape.empty(); + + std::vector runtime_shape(metadata.chunk_shape.begin(), + metadata.chunk_shape.end()); + runtime_shape.insert(runtime_shape.end(), metadata.field_shape.begin(), + metadata.field_shape.end()); if (!metadata.codecs) { ArrayCodecResolveParameters decoded; - if (!is_structured) { - decoded.dtype = metadata.data_type.fields[0].dtype; - decoded.rank = metadata.rank; - } else { - // For structured types, use byte dtype with extra dimension - decoded.dtype = dtype_v; - decoded.rank = metadata.rank + 1; + decoded.dtype = + has_field_shape ? dtype_v : metadata.data_type.fields[0].dtype; + decoded.rank = metadata.rank; + decoded.inner_shape = metadata.field_shape; + // `read_chunk_shape` is at the chunked rank only; sharding_indexed's + // internal `sub_chunk_shape` is at this same rank now too. + { + auto& read_chunk_shape = decoded.read_chunk_shape.emplace(); + std::copy_n(metadata.chunk_shape.begin(), metadata.rank, + read_chunk_shape.begin()); } - // Fill value for codec resolve might be complex. - // For structured types, create a byte fill value - if (metadata.fill_value.size() == 1 && !is_structured) { + if (metadata.fill_value.size() == 1 && !has_field_shape) { decoded.fill_value = metadata.fill_value[0]; } @@ -698,20 +872,14 @@ absl::Status ValidateMetadata(ZarrMetadata& metadata) { metadata.codec_specs.Resolve(std::move(decoded), encoded)); } - // Get codec chunk layout info. + // Get codec chunk layout info at the chunked rank. ArrayDataTypeAndShapeInfo array_info; - if (!is_structured) { - array_info.dtype = metadata.data_type.fields[0].dtype; - array_info.rank = metadata.rank; - std::copy_n(metadata.chunk_shape.begin(), metadata.rank, - array_info.shape.emplace().begin()); - } else { - array_info.dtype = dtype_v; - array_info.rank = metadata.rank + 1; - auto& shape = array_info.shape.emplace(); - std::copy_n(metadata.chunk_shape.begin(), metadata.rank, shape.begin()); - shape[metadata.rank] = metadata.data_type.bytes_per_outer_element; - } + array_info.dtype = + has_field_shape ? dtype_v : metadata.data_type.fields[0].dtype; + array_info.rank = metadata.rank; + array_info.inner_shape = metadata.field_shape; + std::copy_n(metadata.chunk_shape.begin(), metadata.rank, + array_info.shape.emplace().begin()); ArrayCodecChunkLayoutInfo layout_info; TENSORSTORE_RETURN_IF_ERROR( @@ -726,7 +894,7 @@ absl::Status ValidateMetadata(ZarrMetadata& metadata) { } TENSORSTORE_ASSIGN_OR_RETURN(metadata.codec_state, - metadata.codecs->Prepare(codec_shape)); + metadata.codecs->Prepare(runtime_shape)); return absl::OkStatus(); } @@ -823,31 +991,7 @@ std::string GetFieldNames(const ZarrDType& dtype) { } // namespace Result GetFieldIndex(const ZarrDType& dtype, - std::string_view selected_field, - bool open_as_void) { - // Special case: open_as_void requests raw byte access. - // Only allowed for structured dtypes (multiple fields or single field with - // field_shape like raw_bytes). Simple scalar types like int16 are not - // supported - use the normal typed access instead. - if (open_as_void) { - if (dtype.fields.empty()) { - return absl::FailedPreconditionError( - "Requested void access but dtype has no fields"); - } - // Check if dtype is structured: multiple fields OR single field with - // field_shape (e.g., raw_bytes) - const bool is_structured = - dtype.fields.size() > 1 || - (dtype.fields.size() == 1 && !dtype.fields[0].field_shape.empty()); - if (!is_structured) { - return absl::InvalidArgumentError( - "open_as_void is only supported for structured dtypes (multiple " - "fields or raw_bytes). For simple scalar types, use normal typed " - "access instead."); - } - return kVoidFieldIndex; - } - + std::string_view selected_field) { if (selected_field.empty()) { if (dtype.fields.size() != 1) { return absl::FailedPreconditionError(absl::StrFormat( @@ -873,14 +1017,8 @@ SpecRankAndFieldInfo GetSpecRankAndFieldInfo(const ZarrMetadata& metadata, size_t field_index) { SpecRankAndFieldInfo info; info.chunked_rank = metadata.rank; - if (field_index == kVoidFieldIndex) { - // Void access: no field, rank is chunked_rank + 1 (extra bytes dimension) - info.field = nullptr; - info.field_rank = 1; - } else { - info.field = &metadata.data_type.fields[field_index]; - info.field_rank = info.field->field_shape.size(); - } + info.field = &metadata.data_type.fields[field_index]; + info.field_rank = info.field->field_shape.size(); info.full_rank = info.chunked_rank + info.field_rank; return info; } @@ -1196,7 +1334,6 @@ CodecSpec GetCodecFromMetadata(const ZarrMetadata& metadata) { absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, size_t field_index, const Schema& schema) { auto info = GetSpecRankAndFieldInfo(metadata, field_index); - const bool is_void_access = (field_index == kVoidFieldIndex); if (!RankConstraint::EqualOrUnspecified(schema.rank(), info.full_rank)) { return absl::FailedPreconditionError(absl::StrFormat( @@ -1205,20 +1342,11 @@ absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, schema.rank(), info.full_rank)); } - // Dtype validation: for void access, dtype must be byte; otherwise use - // field's dtype. - const DataType expected_dtype = - is_void_access ? dtype_v : info.field->dtype; if (auto dtype = schema.dtype(); - !IsPossiblySameDataType(expected_dtype, dtype)) { + !IsPossiblySameDataType(info.field->dtype, dtype)) { return absl::FailedPreconditionError(absl::StrFormat( "data_type from metadata (%v) does not match dtype in schema (%v)", - expected_dtype, dtype)); - } - - // The following validations only apply to field access, not void access. - if (is_void_access) { - return absl::OkStatus(); + info.field->dtype, dtype)); } if (schema.domain().valid()) { @@ -1319,15 +1447,22 @@ Result> GetNewMetadata( return absl::InvalidArgumentError("dtype must be specified"); } - TENSORSTORE_ASSIGN_OR_RETURN( - size_t field_index, GetFieldIndex(metadata->data_type, selected_field, open_as_void)); - SpecRankAndFieldInfo info; - // For void access, field_index is kVoidFieldIndex (sentinel value) - if (field_index == kVoidFieldIndex) { - info.field = nullptr; + // Resolve the selected field: void mode does not pick a "real" field (the + // void substitution happens later, via `GetVoidMetadata`), but we still need + // a placeholder field index/info to drive shape/rank computation here. + size_t field_index; + if (open_as_void) { + if (!selected_field.empty()) { + return absl::InvalidArgumentError( + "\"field\" and \"open_as_void\" are mutually exclusive"); + } + field_index = 0; } else { - info.field = &metadata->data_type.fields[field_index]; + TENSORSTORE_ASSIGN_OR_RETURN( + field_index, GetFieldIndex(metadata->data_type, selected_field)); } + SpecRankAndFieldInfo info; + info.field = &metadata->data_type.fields[field_index]; info.chunked_rank = metadata_constraints.rank; if (info.chunked_rank == dynamic_rank && metadata_constraints.shape) { info.chunked_rank = metadata_constraints.shape->size(); @@ -1336,17 +1471,20 @@ Result> GetNewMetadata( schema.rank().rank != dynamic_rank) { info.chunked_rank = schema.rank().rank; } - // For void access on structured types, add the bytes dimension to chunked_rank - if (open_as_void && info.chunked_rank != dynamic_rank) { - info.chunked_rank += 1; - } - // For fields with field_shape (like r16, r64), add those dimensions to chunked_rank - if (info.field && !info.field->field_shape.empty() && - info.chunked_rank != dynamic_rank) { - info.chunked_rank += info.field->field_shape.size(); + // Number of synthetic trailing dimensions contributed by either an explicit + // field_shape (rN / struct field) or by `open_as_void` (the bytes + // dimension). Both are field_shape in the laramiel sense. + const DimensionIndex extra_field_dims = + open_as_void ? 1 + : (info.field ? info.field->field_shape.size() : 0); + if (extra_field_dims != 0 && info.chunked_rank != dynamic_rank) { + info.chunked_rank += extra_field_dims; } - // Set domain + // Set domain. When the field contributes trailing field_shape dimensions + // (rN, struct field, or the synthetic void byte dimension), extend the + // user-provided chunked shape with those dimensions before merging with the + // schema, so that the schema and metadata describe the same full rank. bool dimension_names_used = false; std::vector extended_shape; std::optional> constraint_shape_span; @@ -1383,15 +1521,11 @@ Result> GetNewMetadata( if (!domain.valid() || !IsFinite(domain.box())) { return absl::InvalidArgumentError("domain must be specified"); } - // For void access or fields with field_shape, domain includes extra dimensions, - // but metadata stores only the logical (base array) dimensions. - DimensionIndex extra_dims = 0; - if (open_as_void) { - extra_dims = 1; - } else if (info.field && !info.field->field_shape.empty()) { - extra_dims = info.field->field_shape.size(); - } - const DimensionIndex logical_rank = domain.rank() - extra_dims; + // The user-visible domain may include trailing dimensions contributed by + // either an explicit field_shape (rN / struct field) or by `open_as_void` + // (the bytes dimension), but the persisted metadata stores only the + // logical chunked dimensions. + const DimensionIndex logical_rank = domain.rank() - extra_field_dims; metadata->rank = logical_rank; info.chunked_rank = domain.rank(); // Keep extended rank for codec processing metadata->shape.assign(domain.shape().begin(), @@ -1458,24 +1592,26 @@ Result> GetNewMetadata( TENSORSTORE_ASSIGN_OR_RETURN(auto codec_spec, GetEffectiveCodec(metadata_constraints, schema)); - // Determine if this is a structured type (multiple fields or inner shape) - const bool is_structured = - metadata->data_type.fields.size() > 1 || - (metadata->data_type.fields.size() == 1 && - !metadata->data_type.fields[0].field_shape.empty()); + // Derive `metadata->field_shape` up front so the rest of this function (and + // the eventual `ValidateMetadata` re-check) drives off a single source of + // truth, identical to the parse-time code path. + DeriveFieldShape(*metadata); + const bool has_field_shape = !metadata->field_shape.empty(); + // Codec resolution is at the chunked rank; inner (`field_shape`) dims are + // carried via `decoded.inner_shape` and never appear in any rank/shape + // field that array-to-array codecs see. ArrayCodecResolveParameters decoded; - if (!is_structured) { - decoded.dtype = metadata->data_type.fields[0].dtype; - decoded.rank = metadata->rank; + decoded.dtype = + has_field_shape ? dtype_v : metadata->data_type.fields[0].dtype; + decoded.rank = metadata->rank; + decoded.inner_shape = metadata->field_shape; + if (!has_field_shape) { if (metadata->fill_value.size() == 1) { decoded.fill_value = metadata->fill_value[0]; } } else { - // For structured types, use byte dtype with extra dimension for the bytes - decoded.dtype = dtype_v; - decoded.rank = metadata->rank + 1; - // Create a zero-filled scalar byte as fill_value (gets broadcast to chunk shape) + // Zero-filled scalar byte fill_value (broadcast to chunk shape at use). decoded.fill_value = AllocateArray( span{}, c_order, value_init, dtype_v); } @@ -1487,34 +1623,22 @@ Result> GetNewMetadata( if (auto inner_order = chunk_layout.inner_order(); inner_order.valid()) { auto& dest = decoded.inner_order.emplace(); std::copy(inner_order.begin(), inner_order.end(), dest.begin()); - // For structured types, add the bytes dimension at the end - if (is_structured) { - dest[metadata->rank] = metadata->rank; - } } - // For structured types, read_chunk_shape needs decoded.rank dimensions - // (metadata->rank + 1 for the bytes dimension) + // `read_chunk_shape` is at the chunked rank only. span read_chunk_shape(decoded.read_chunk_shape.emplace().data(), - decoded.rank); + metadata->rank); - // ChooseReadWriteChunkShapes only operates on the logical dimensions. - // For void access, domain includes the bytes dimension; restrict to logical. TENSORSTORE_RETURN_IF_ERROR(internal::ChooseReadWriteChunkShapes( chunk_layout.read_chunk(), chunk_layout.write_chunk(), - SubBoxView(domain.box(), 0, metadata->rank), - read_chunk_shape.first(metadata->rank), metadata->chunk_shape)); + SubBoxView(domain.box(), 0, metadata->rank), read_chunk_shape, + metadata->chunk_shape)); - // For structured types, set the bytes dimension - if (is_structured) { - read_chunk_shape[metadata->rank] = - metadata->data_type.bytes_per_outer_element; - } - - // Compare only the logical dimensions to decide if sharding is needed if (!internal::RangesEqual(span(metadata->chunk_shape), - read_chunk_shape.first(metadata->rank))) { + read_chunk_shape)) { if (!codec_spec->codecs || codec_spec->codecs->sharding_height() == 0) { + // sub_chunk_shape is at the chunked rank now -- same as the + // user-facing form and the on-disk zarr.json representation. auto sharding_codec = internal::MakeIntrusivePtr( ShardingIndexedCodecSpec::Options{ @@ -1537,7 +1661,13 @@ Result> GetNewMetadata( TENSORSTORE_RETURN_IF_ERROR(set_up_codecs( codec_spec->codecs ? *codec_spec->codecs : ZarrCodecChainSpec{})); TENSORSTORE_RETURN_IF_ERROR(ValidateMetadata(*metadata)); - if (field_index != kVoidFieldIndex) { + if (open_as_void) { + // The user-supplied schema dtype is `byte` and rank is chunked + 1, neither + // of which match the (natural) `*metadata`. The driver re-validates against + // the void-substituted view via `GetVoidMetadata` after this returns, so we + // skip the non-applicable per-field schema check here. + TENSORSTORE_RETURN_IF_ERROR(ValidateVoidCodecChain(metadata->codec_specs)); + } else { TENSORSTORE_RETURN_IF_ERROR( ValidateMetadataSchema(*metadata, field_index, schema)); } diff --git a/tensorstore/driver/zarr3/metadata.h b/tensorstore/driver/zarr3/metadata.h index 27e2ac7b7..c63ff9033 100644 --- a/tensorstore/driver/zarr3/metadata.h +++ b/tensorstore/driver/zarr3/metadata.h @@ -53,10 +53,6 @@ namespace tensorstore { namespace internal_zarr3 { -// Sentinel value for field_index indicating access to the raw byte "void" view -// of a structured array. -constexpr size_t kVoidFieldIndex = size_t(-1); - // Defines how chunks map to keys in the underlying kvstore. // // https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#chunk-key-encoding @@ -130,15 +126,27 @@ struct ZarrMetadata { std::string GetCompatibilityKey() const; - /// Returns a synthetic ZarrDType for void access, representing the raw bytes - /// of the original data type. The returned dtype has a single unnamed field - /// of byte type with shape equal to bytes_per_outer_element. - ZarrDType GetVoidAccessDType() const; - ZarrCodecChain::Ptr codecs; ZarrCodecChain::PreparedState::Ptr codec_state; std::array inner_order; + // Inner trailing dimensions appended to `chunk_shape` to form the codec + // resolution shape. Empty for plain scalar single-field arrays. For + // single-field arrays whose dtype carries an `rN`-style `field_shape`, this + // mirrors that per-field `field_shape`. For multi-field structs and for + // the byte-substituted view produced by `GetVoidMetadata`, this is + // `{bytes_per_outer_element}` -- the single trailing byte dimension that + // pins the struct's interleaved byte layout into the codec chain. + // + // Conceptually `chunk_shape` describes the array's chunked dimensions and + // `field_shape` describes any inner (per-element) dimensions; the two are + // concatenated to form the rank the codec chain is resolved against. This + // replaces the previous implicit "shape prefix" handling where the trailing + // byte dimension was ad-hoc-appended at every consumer site. + // + // Not part of the persisted zarr.json -- derived during `ValidateMetadata`. + std::vector field_shape; + TENSORSTORE_DECLARE_JSON_DEFAULT_BINDER(ZarrMetadata, internal_json_binding::NoOptions, tensorstore::IncludeDefaults) @@ -231,9 +239,6 @@ Result> GetEffectiveCodec( CodecSpec GetCodecFromMetadata(const ZarrMetadata& metadata); /// Validates that `schema` is compatible with `metadata`. -/// -/// \param field_index The field index, or `kVoidFieldIndex` for void access. -/// For void access, only rank and dtype are validated (dtype must be byte). absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, size_t field_index, const Schema& schema); absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, @@ -241,6 +246,14 @@ absl::Status ValidateMetadataSchema(const ZarrMetadata& metadata, /// Converts `metadata_constraints` to a full metadata object. /// +/// When `open_as_void` is true, the persisted metadata is constructed for the +/// natural data type representation (i.e. unchanged from `metadata_constraints`), +/// but the user-visible domain rank is extended by one trailing byte dimension. +/// `GetNewMetadata` strips that trailing dimension from the user-supplied schema +/// shape/chunk_layout, and skips the dtype/domain compatibility check against +/// the schema (the caller is expected to re-validate against the void-substituted +/// metadata). +/// /// \error `absl::StatusCode::kInvalidArgument` if any required fields are /// unspecified. Result> GetNewMetadata( @@ -250,9 +263,59 @@ Result> GetNewMetadata( absl::Status ValidateDataType(DataType dtype); +/// Returns the index of `selected_field` within `dtype.fields`. +/// +/// If `selected_field` is empty, requires the dtype to have a single field and +/// returns 0. Result GetFieldIndex(const ZarrDType& dtype, - std::string_view selected_field, - bool open_as_void); + std::string_view selected_field); + +/// Returns a synthetic single-field "void view" `ZarrDType` whose only field +/// has dtype `byte` and `field_shape = {bytes_per_outer_element}`. This is +/// the representation handed to the chunk cache when `open_as_void` is in use, +/// and is treated by all downstream code as a normal `field_shape`-bearing +/// field (no separate `open_as_void` plumbing required). +ZarrDType MakeVoidDType(Index bytes_per_outer_element); + +/// Packs `per_field_fill` (one fill value per field of `dtype`) into a single +/// 1-D byte `SharedArray` of length `dtype.bytes_per_outer_element`, following +/// the struct's `byte_offset` layout. +/// +/// The byte order of each typed field in the packed buffer matches the byte +/// order that the bytes codec would write for a chunk encoded under the +/// supplied `codec_specs` -- for a single non-byte scalar field this means +/// the codec's configured endian; for multi-field structs and `rN` raw byte +/// fields it is native-endian (matching the chunk cache's CopyArray-based +/// per-field packing path). This guarantees that the synthetic void fill +/// returned for a missing chunk is byte-identical to what `read` returns for +/// a present chunk. +SharedArray MakeVoidFillValue( + const ZarrDType& dtype, const ZarrCodecChainSpec& codec_specs, + span> per_field_fill); + +/// Returns a void-access view of `metadata`: a new `ZarrMetadata` whose +/// `data_type` is `MakeVoidDType(...)`, whose `fill_value` is a single byte +/// array packed from each original field's per-field fill, and whose codec +/// chain has been re-resolved against the substituted byte data type with an +/// extra innermost dimension of `bytes_per_outer_element`. Persisted state +/// (`shape`, `chunk_shape`, `chunk_key_encoding`, etc.) is unchanged. +/// +/// Also validates that the resulting codec chain is acceptable for raw byte +/// access (see `ValidateVoidCodecChain`). +Result> GetVoidMetadata( + const ZarrMetadata& metadata); + +/// Validates that `codec_specs` is compatible with `open_as_void` access. +/// +/// The only structural precondition is that, after unwinding any +/// `sharding_indexed` layers, the innermost array-to-bytes codec is the +/// `bytes` codec. The previous "preserves dtype" and "preserves trailing +/// byte dim" rules on array-to-array codecs are now structurally guaranteed +/// by the codec resolution architecture: the chain exposes inner +/// (`field_shape`) dims to a-to-a codecs only via `inner_shape` (which they +/// must propagate verbatim), and `Resolve` forces every a-to-a codec to +/// preserve the substituted byte dtype. +absl::Status ValidateVoidCodecChain(const ZarrCodecChainSpec& codec_specs); struct SpecRankAndFieldInfo { /// Full rank of the TensorStore, if known. Equal to the chunked rank plus @@ -262,7 +325,7 @@ struct SpecRankAndFieldInfo { /// Number of chunked dimensions (the array's original rank). DimensionIndex chunked_rank = dynamic_rank; - /// Number of field dimensions (from field_shape, or 1 for void access). + /// Number of field dimensions contributed by `field->field_shape`. DimensionIndex field_rank = dynamic_rank; /// Data type field, or `nullptr` if unknown. @@ -274,7 +337,9 @@ absl::Status ValidateSpecRankAndFieldInfo(SpecRankAndFieldInfo& info); /// Gets spec rank and field info from metadata constraints. /// -/// When `open_as_void` is true, `info.field_rank` is 1 (the bytes dimension). +/// When `open_as_void` is true, `info.field_rank` is 1 (the synthetic bytes +/// dimension contributed by `MakeVoidDType`). Otherwise `info.field_rank` is +/// derived from the selected field's `field_shape`. /// /// \param metadata Metadata constraints. /// \param selected_field The field to access. Must be empty if `open_as_void` diff --git a/tensorstore/driver/zarr3/metadata_test.cc b/tensorstore/driver/zarr3/metadata_test.cc index 94f10a941..7cc8880ac 100644 --- a/tensorstore/driver/zarr3/metadata_test.cc +++ b/tensorstore/driver/zarr3/metadata_test.cc @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -34,6 +35,8 @@ #include "tensorstore/chunk_layout.h" #include "tensorstore/codec_spec.h" #include "tensorstore/data_type.h" +#include "tensorstore/driver/zarr3/codec/codec_chain_spec.h" +#include "tensorstore/driver/zarr3/dtype.h" #include "tensorstore/index.h" #include "tensorstore/internal/json_binding/bindable.h" #include "tensorstore/internal/json_binding/gtest.h" @@ -41,6 +44,7 @@ #include "tensorstore/internal/testing/json_gtest.h" #include "tensorstore/rank.h" #include "tensorstore/schema.h" +#include "tensorstore/util/endian.h" #include "tensorstore/util/result.h" #include "tensorstore/util/status.h" #include "tensorstore/util/status_testutil.h" @@ -70,6 +74,11 @@ using ::tensorstore::dtypes::float64_t; using ::tensorstore::internal::uint_t; using ::tensorstore::internal_zarr3::FillValueJsonBinder; using ::tensorstore::internal_zarr3::GetSpecRankAndFieldInfo; +using ::tensorstore::internal_zarr3::GetVoidMetadata; +using ::tensorstore::internal_zarr3::MakeVoidDType; +using ::tensorstore::internal_zarr3::ParseDType; +using ::tensorstore::internal_zarr3::ValidateVoidCodecChain; +using ::tensorstore::internal_zarr3::ZarrCodecChainSpec; using ::tensorstore::internal_zarr3::ZarrDType; using ::tensorstore::internal_zarr3::ZarrMetadata; using ::tensorstore::internal_zarr3::ZarrMetadataConstraints; @@ -954,4 +963,156 @@ TEST(TrySetMetadataConstraintsOnSchemaTest, SetsRankWithFieldShape) { EXPECT_EQ(schema.rank(), 4); } +TEST(MakeVoidDTypeTest, BuildsSingleByteFieldWithFieldShape) { + auto dtype = MakeVoidDType(/*bytes_per_outer_element=*/5); + EXPECT_FALSE(dtype.has_fields); + ASSERT_EQ(dtype.fields.size(), 1u); + EXPECT_EQ(dtype.bytes_per_outer_element, 5); + const auto& field = dtype.fields[0]; + EXPECT_EQ(field.dtype, dtype_v); + EXPECT_EQ(field.byte_offset, 0); + EXPECT_EQ(field.num_bytes, 5); + EXPECT_EQ(field.num_inner_elements, 5); + EXPECT_THAT(field.field_shape, ::testing::ElementsAre(5)); +} + +TEST(ValidateVoidCodecChainTest, AcceptsBareBytesCodec) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto chain, ZarrCodecChainSpec::FromJson( + ::nlohmann::json::array_t{{{"name", "bytes"}}})); + TENSORSTORE_EXPECT_OK(ValidateVoidCodecChain(chain)); +} + +TEST(ValidateVoidCodecChainTest, AcceptsBytesWithExplicitEndianness) { + // Endianness on the bytes codec is irrelevant for validation: we only care + // that the innermost array->bytes codec is `bytes`, regardless of its + // configured endianness (codec resolution will treat byte_t as + // endian-invariant). + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto chain, + ZarrCodecChainSpec::FromJson(::nlohmann::json::array_t{ + {{"name", "bytes"}, {"configuration", {{"endian", "big"}}}}})); + TENSORSTORE_EXPECT_OK(ValidateVoidCodecChain(chain)); +} + +TEST(ValidateVoidCodecChainTest, AcceptsShardingIndexedWrappingBytes) { + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto chain, + ZarrCodecChainSpec::FromJson(::nlohmann::json::array_t{ + {{"name", "sharding_indexed"}, + {"configuration", + {{"chunk_shape", {2, 2}}, + {"codecs", ::nlohmann::json::array_t{{{"name", "bytes"}}}}, + {"index_codecs", + ::nlohmann::json::array_t{{{"name", "bytes"}}, + {{"name", "crc32c"}}}}}}}})); + TENSORSTORE_EXPECT_OK(ValidateVoidCodecChain(chain)); +} + +TEST(ValidateVoidCodecChainTest, AcceptsTransposeAtChunkedRank) { + // After the inner-shape architecture change, transpose's permutation is + // strictly at the chunked rank; the inner (byte) dim is structurally + // unreachable, so this is the only well-formed shape for transpose + + // byte-substituted access. + TENSORSTORE_ASSERT_OK_AND_ASSIGN( + auto chain, + ZarrCodecChainSpec::FromJson(::nlohmann::json::array_t{ + {{"name", "transpose"}, {"configuration", {{"order", {1, 0}}}}}, + {{"name", "bytes"}}})); + TENSORSTORE_EXPECT_OK(ValidateVoidCodecChain(chain)); +} + +TEST(GetVoidMetadataTest, RoundTripsStructuredFillValue) { + ::nlohmann::json metadata_json = GetBasicMetadata(); + metadata_json["data_type"] = { + {"name", "struct"}, + {"configuration", + {{"fields", + ::nlohmann::json::array({{{"name", "x"}, {"data_type", "uint8"}}, + {{"name", "y"}, {"data_type", "int16"}}})}}}}; + metadata_json["fill_value"] = {{"x", 0xAA}, {"y", 0x1234}}; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, + ZarrMetadata::FromJson(metadata_json)); + + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto void_metadata, + GetVoidMetadata(metadata)); + + // The view metadata's data_type is the synthetic byte view: a single field + // with `field_shape = {bytes_per_outer_element}`. + EXPECT_FALSE(void_metadata->data_type.has_fields); + ASSERT_EQ(void_metadata->data_type.fields.size(), 1u); + EXPECT_EQ(void_metadata->data_type.bytes_per_outer_element, + metadata.data_type.bytes_per_outer_element); + EXPECT_THAT( + void_metadata->data_type.fields[0].field_shape, + ::testing::ElementsAre(metadata.data_type.bytes_per_outer_element)); + + // Persisted on-disk fields are unchanged. + EXPECT_EQ(void_metadata->shape, metadata.shape); + EXPECT_EQ(void_metadata->chunk_shape, metadata.chunk_shape); + + // Fill value is the per-field bytes packed at their byte_offsets, in the + // same byte order the bytes codec would produce on disk. Multi-field + // structs follow the chunk cache's native-endian CopyArray convention. + // Layout (native-endian): + // x (uint8) at offset 0 -> 0xAA + // y (int16) at offset 1 -> {0x34, 0x12} on little-endian hosts, + // {0x12, 0x34} on big-endian hosts. + ASSERT_EQ(void_metadata->fill_value.size(), 1u); + const auto& packed = void_metadata->fill_value[0]; + ASSERT_EQ(packed.rank(), 1); + ASSERT_EQ(packed.shape()[0], 3); + const auto* bytes = static_cast(packed.data()); + EXPECT_EQ(bytes[0], std::byte{0xAA}); + if (tensorstore::endian::native == tensorstore::endian::little) { + EXPECT_EQ(bytes[1], std::byte{0x34}); + EXPECT_EQ(bytes[2], std::byte{0x12}); + } else { + EXPECT_EQ(bytes[1], std::byte{0x12}); + EXPECT_EQ(bytes[2], std::byte{0x34}); + } +} + +TEST(GetVoidMetadataTest, ScalarFillEncodedInCodecEndian) { + // For a single non-byte scalar dtype, the bytes codec performs endian + // conversion when encoding chunks. The void fill must therefore reflect + // the codec's configured endian, NOT the host endian -- otherwise a + // missing-chunk void read returns different bytes than a present-chunk + // void read on a host whose native endian disagrees with the codec. + ::nlohmann::json metadata_json = GetBasicMetadata(); + metadata_json["data_type"] = "uint16"; + metadata_json["codecs"] = { + {{"name", "bytes"}, {"configuration", {{"endian", "big"}}}}}; + metadata_json["fill_value"] = 0x1234; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, + ZarrMetadata::FromJson(metadata_json)); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto void_metadata, + GetVoidMetadata(metadata)); + ASSERT_EQ(void_metadata->fill_value.size(), 1u); + const auto& packed = void_metadata->fill_value[0]; + ASSERT_EQ(packed.rank(), 1); + ASSERT_EQ(packed.shape()[0], 2); + const auto* bytes = static_cast(packed.data()); + // Big-endian uint16(0x1234) = {0x12, 0x34} regardless of host endianness. + EXPECT_EQ(bytes[0], std::byte{0x12}); + EXPECT_EQ(bytes[1], std::byte{0x34}); +} + +TEST(GetVoidMetadataTest, AcceptsScalar) { + // GetVoidMetadata for a plain scalar dtype yields a synthetic single-field + // byte view whose `field_shape` equals the scalar's byte width. + ::nlohmann::json metadata_json = GetBasicMetadata(); + metadata_json["data_type"] = "int32"; + metadata_json["fill_value"] = 0; + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto metadata, + ZarrMetadata::FromJson(metadata_json)); + TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto void_metadata, + GetVoidMetadata(metadata)); + EXPECT_FALSE(void_metadata->data_type.has_fields); + ASSERT_EQ(void_metadata->data_type.fields.size(), 1u); + EXPECT_EQ(void_metadata->data_type.bytes_per_outer_element, 4); + EXPECT_THAT(void_metadata->data_type.fields[0].field_shape, + ::testing::ElementsAre(4)); +} + } // namespace From a0b56efc0782939d809f2a41fbb60fc4ad1edf08 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 6 May 2026 16:23:04 +0000 Subject: [PATCH 2/8] Deslopping --- .../driver/zarr3/codec/sharding_indexed.cc | 22 +------------------ 1 file changed, 1 insertion(+), 21 deletions(-) diff --git a/tensorstore/driver/zarr3/codec/sharding_indexed.cc b/tensorstore/driver/zarr3/codec/sharding_indexed.cc index 4c3881891..4565b8d85 100644 --- a/tensorstore/driver/zarr3/codec/sharding_indexed.cc +++ b/tensorstore/driver/zarr3/codec/sharding_indexed.cc @@ -179,9 +179,6 @@ absl::Status ShardingIndexedCodecSpec::MergeFrom(const ZarrCodecSpec& other, bool strict) { using Self = ShardingIndexedCodecSpec; const auto& other_options = static_cast(other).options; - // `sub_chunk_shape` lives at the chunked rank in both the on-disk - // zarr.json and the user-facing spec; the codec resolution path no longer - // extends it internally, so plain equality is sufficient. TENSORSTORE_RETURN_IF_ERROR(MergeConstraint<&Options::sub_chunk_shape>( "chunk_shape", options, other_options)); TENSORSTORE_RETURN_IF_ERROR( @@ -217,11 +214,6 @@ const ZarrCodecChainSpec* ShardingIndexedCodecSpec::GetSubChunkCodecs() const { absl::Status ShardingIndexedCodecSpec::GetDecodedChunkLayout( const ArrayDataTypeAndShapeInfo& array_info, ArrayCodecChunkLayoutInfo& decoded) const { - // `array_info.rank` is the chunked rank only; inner (`field_shape`) dims - // travel via `array_info.inner_shape` and are propagated unchanged through - // sub-chunk codecs without ever being permuted or reshaped. The codec - // spec's `sub_chunk_shape` is therefore at the same chunked rank as - // `array_info.rank`. if (options.sub_chunk_shape && static_cast(options.sub_chunk_shape->size()) != array_info.rank) { @@ -262,11 +254,6 @@ Result ShardingIndexedCodecSpec::Resolve( resolved_options = &resolved_spec_ptr->options; resolved_spec->reset(resolved_spec_ptr); } - // `decoded.rank` is the chunked rank only. Inner (`field_shape`) dims are - // carried separately on `decoded.inner_shape`, propagated through to the - // sub-chunk codec chain unchanged, and consumed at the leaf "array -> - // bytes" codec for byte-stream sizing. The codec spec's `sub_chunk_shape` - // is at the same chunked rank. span sub_chunk_shape; if (options.sub_chunk_shape) { sub_chunk_shape = *options.sub_chunk_shape; @@ -279,10 +266,6 @@ Result ShardingIndexedCodecSpec::Resolve( if (static_cast(sub_chunk_shape.size()) != decoded.rank) { return SubChunkRankMismatch(sub_chunk_shape, decoded.rank); } - // The runtime ChunkGridSpecification partitions an extended-rank chunk - // (chunked + inner): chunked dims are partitioned by `sub_chunk_shape`, and - // inner dims are kept whole inside each sub-chunk. Build the runtime - // sub-chunk shape accordingly. std::vector runtime_sub_chunk_shape(sub_chunk_shape.begin(), sub_chunk_shape.end()); runtime_sub_chunk_shape.insert(runtime_sub_chunk_shape.end(), @@ -314,7 +297,7 @@ Result ShardingIndexedCodecSpec::Resolve( std::move(sub_chunk_decoded), encoded, resolved_options ? &resolved_options->sub_chunk_codecs.emplace() : nullptr)); - // Get sub-chunk codec chunk layout info (chunked rank only). + // Get sub-chunk codec chunk layout info. ArrayDataTypeAndShapeInfo array_info; array_info.dtype = decoded.dtype; array_info.rank = decoded.rank; @@ -355,8 +338,6 @@ Result ShardingIndexedCodecSpec::Resolve( options.index_location.value_or(ShardIndexLocation::kEnd); codec->sub_chunk_codec_chain_ = std::move(sub_chunk_codec_chain); if (resolved_options) { - // The user-form sub_chunk_shape is at the same chunked rank as the - // resolved spec; no internal extension to undo. if (options.sub_chunk_shape) { resolved_options->sub_chunk_shape = *options.sub_chunk_shape; } else { @@ -372,7 +353,6 @@ Result ShardingIndexedCodecSpec::Resolve( : ZarrCodecChainSpec{})) .Format("Error resolving sub-chunk codecs"); - // Index codecs index the chunked grid only, regardless of inner_shape. auto set_up_index_codecs = [&](const ZarrCodecChainSpec& index_codecs) -> absl::Status { TENSORSTORE_ASSIGN_OR_RETURN( From fcacc9b8f8ac859b84a5958e10ff04fceb3f8c2d Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 6 May 2026 16:24:21 +0000 Subject: [PATCH 3/8] Deslop --- tensorstore/driver/zarr3/codec/bytes.cc | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tensorstore/driver/zarr3/codec/bytes.cc b/tensorstore/driver/zarr3/codec/bytes.cc index 6d35eb866..17885008d 100644 --- a/tensorstore/driver/zarr3/codec/bytes.cc +++ b/tensorstore/driver/zarr3/codec/bytes.cc @@ -79,10 +79,6 @@ absl::Status BytesCodecSpec::GetDecodedChunkLayout( !internal::IsTrivialDataType(array_info.dtype)) { return InvalidDataTypeError(array_info.dtype); } - // `array_info.rank` is the chunked rank only; any inner (`field_shape`) - // dims are reported via `array_info.inner_shape` and do not appear in the - // chunk layout info (they are pinned, identity-ordered, and carry no - // user-tunable layout). const DimensionIndex rank = array_info.rank; if (rank != dynamic_rank) { auto& inner_order = decoded.inner_order.emplace(); From 7d35a8a373754b8b118d5a3b464769843a98293b Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 6 May 2026 16:28:29 +0000 Subject: [PATCH 4/8] Deslop --- tensorstore/driver/zarr3/codec/codec_spec.h | 41 +++++++-------------- 1 file changed, 13 insertions(+), 28 deletions(-) diff --git a/tensorstore/driver/zarr3/codec/codec_spec.h b/tensorstore/driver/zarr3/codec/codec_spec.h index bc4dfdd93..61567113c 100644 --- a/tensorstore/driver/zarr3/codec/codec_spec.h +++ b/tensorstore/driver/zarr3/codec/codec_spec.h @@ -136,29 +136,17 @@ struct ArrayDataTypeAndShapeInfo { // Specifies the data type of the array on which the codec will operate. DataType dtype; - // Specifies the *chunked* rank of the array on which the codec will operate. - // - // This is the user-visible rank: it does NOT include the inner trailing - // dimensions contributed by the dtype's `field_shape` (multi-field structs, - // `rN` raw byte fields, `open_as_void`'s byte substitution). Those inner - // dimensions, if any, are reported separately in `inner_shape` and are - // pinned at the trailing positions of the runtime array; "array -> array" - // codec specs cannot operate on them. + // Specifies the rank of the array on which the codec will operate. + // This excludes any inner dimensions contributed by `dtype`. DimensionIndex rank = dynamic_rank; - // Specifies the chunked-rank shape of the array on which the codec will - // operate. When present, has exactly `rank` valid entries. + // Specifies the shape of the array on which the codec will operate. + // When present, has exactly `rank` valid entries. std::optional> shape; - // Inner trailing dimensions contributed by the dtype's `field_shape`. - // Empty (the default) when the dtype has no field_shape (i.e. plain scalar - // dtype with no codec substitution); non-empty for multi-field structs, - // `rN`, and the byte-substituted `open_as_void` view. - // - // "array -> array" codec specs must propagate this field unchanged from - // `decoded` to `encoded`; they cannot read it for resolution decisions and - // cannot reshape/permute it. The "array -> bytes" codec at the bottom of - // the chain consumes it (e.g. for sizing the encoded byte stream). + // Inner trailing dimensions contributed by `dtype` (e.g. for `open_as_void`). + // These are propagated unchanged through "array -> array" codecs and consumed + // by the "array -> bytes" codec. std::vector inner_shape; }; @@ -185,26 +173,23 @@ struct ArrayCodecResolveParameters { // Specifies the data type of the array on which the codec will operate. DataType dtype; - // Specifies the *chunked* rank of the array on which the codec will operate - // (excludes any `inner_shape` dimensions; see `ArrayDataTypeAndShapeInfo`). + // Specifies the rank of the array on which the codec will operate. + // This excludes any inner dimensions contributed by `dtype`. DimensionIndex rank; // Specifies the fill value. SharedArray fill_value; - // Specifies requested read chunk shape (chunked rank only). + // Specifies requested read chunk shape. std::optional> read_chunk_shape; - // Specifies requested codec chunk shape (chunked rank only). + // Specifies requested codec chunk shape. std::optional> codec_chunk_shape; - // Specifies required inner order (chunked rank only). + // Specifies required inner order. std::optional> inner_order; - // Inner trailing dimensions contributed by the dtype's `field_shape`; see - // the docstring on `ArrayDataTypeAndShapeInfo::inner_shape`. Propagated - // unchanged through "array -> array" codecs; consumed by the "array -> - // bytes" codec (e.g. for sizing the on-disk byte stream). + // Inner trailing dimensions contributed by `dtype` (e.g. for `open_as_void`). std::vector inner_shape; }; From 46a8e04c61051391719b658b863c029b7c45c639 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 6 May 2026 18:08:19 +0000 Subject: [PATCH 5/8] Deslop --- tensorstore/driver/zarr3/chunk_cache.cc | 55 --------------------- tensorstore/driver/zarr3/chunk_cache.h | 27 ---------- tensorstore/driver/zarr3/codec/transpose.cc | 9 ---- tensorstore/driver/zarr3/driver_test.cc | 19 +------ 4 files changed, 1 insertion(+), 109 deletions(-) diff --git a/tensorstore/driver/zarr3/chunk_cache.cc b/tensorstore/driver/zarr3/chunk_cache.cc index 8b3124c51..58a317360 100644 --- a/tensorstore/driver/zarr3/chunk_cache.cc +++ b/tensorstore/driver/zarr3/chunk_cache.cc @@ -89,7 +89,6 @@ internal::ChunkGridSpecification CreateFieldGridSpecification( const size_t field_rank = field.field_shape.size(); const DimensionIndex total_rank = chunked_rank + field_rank; - // Get or create fill value for this field SharedArray fill_value; if (fill_values && field_i < fill_values->size()) { fill_value = (*fill_values)[field_i]; @@ -99,7 +98,6 @@ internal::ChunkGridSpecification CreateFieldGridSpecification( field.dtype); } - // Construct target shape for broadcasting: [unbounded..., field_shape...] std::vector target_shape(chunked_rank, kInfIndex); target_shape.insert(target_shape.end(), field.field_shape.begin(), field.field_shape.end()); @@ -107,15 +105,12 @@ internal::ChunkGridSpecification CreateFieldGridSpecification( auto chunk_fill_value = BroadcastArray(fill_value, BoxView<>(target_shape)).value(); - // Construct component chunk shape: [chunk_shape..., field_shape...] std::vector component_chunk_shape(chunk_shape.begin(), chunk_shape.end()); component_chunk_shape.insert(component_chunk_shape.end(), field.field_shape.begin(), field.field_shape.end()); - // Construct permutation: copy inner_order (if available), then identity - // for field dimensions std::vector component_permutation(total_rank); if (!inner_order.empty()) { assert(inner_order.size() == chunked_rank); @@ -128,14 +123,12 @@ internal::ChunkGridSpecification CreateFieldGridSpecification( std::iota(component_permutation.begin() + chunked_rank, component_permutation.end(), chunked_rank); - // Construct bounds: chunked dims unbounded, field dims fixed Box<> valid_data_bounds(total_rank); for (size_t i = 0; i < field_rank; ++i) { valid_data_bounds[chunked_rank + i] = IndexInterval::UncheckedSized(0, field.field_shape[i]); } - // chunked_to_cell_dimensions maps chunked grid dims to cell dims std::vector chunked_to_cell(chunked_rank); std::iota(chunked_to_cell.begin(), chunked_to_cell.end(), 0); @@ -238,9 +231,6 @@ ZarrLeafChunkCache::DecodeChunk(span chunk_indices, const size_t num_fields = dtype_.fields.size(); absl::InlinedVector, 1> field_arrays(num_fields); - // Fast path: scalar single-field arrays contribute no inner dimensions to - // the codec resolution shape, so the codec call is rank-aligned with the - // grid's component shape and no byte interleaving is required. if (field_shape_.empty()) { assert(num_fields == 1); TENSORSTORE_ASSIGN_OR_RETURN( @@ -249,10 +239,6 @@ ZarrLeafChunkCache::DecodeChunk(span chunk_indices, return field_arrays; } - // Otherwise the codec was resolved against `byte` data with the inner - // `field_shape_` dims appended -- multi-field structs, `rN` raw byte - // fields, and the `open_as_void` view. Decode the raw byte array, then - // split it into per-field components based purely on the field layout. const auto& chunk_shape = grid().chunk_shape; std::vector decode_shape(chunk_shape.begin(), chunk_shape.end()); decode_shape.insert(decode_shape.end(), field_shape_.begin(), @@ -261,33 +247,19 @@ ZarrLeafChunkCache::DecodeChunk(span chunk_indices, TENSORSTORE_ASSIGN_OR_RETURN( auto byte_array, codec_state_->DecodeArray(decode_shape, std::move(data))); - // Extract each field from the byte array. - // We create a strided view into the source that maps to each field's - // position within the interleaved struct layout, then use CopyArray which - // safely handles any layout differences via IterateOverArrays. for (size_t field_i = 0; field_i < num_fields; ++field_i) { const auto& field = dtype_.fields[field_i]; - // Use the component's shape (from the grid) for the result array const auto& component_shape = grid().components[field_i].shape(); auto result_array = AllocateArray(component_shape, c_order, default_init, field.dtype); - // Build the full view shape: [chunk_shape..., field_shape...] - // For fields with field_shape (like r16, r64), we need to include those - // dimensions in the view. std::vector view_shape(chunk_shape.begin(), chunk_shape.end()); view_shape.insert(view_shape.end(), field.field_shape.begin(), field.field_shape.end()); - // Build strides for the source view: - // - Outer dimensions (chunk_shape): each element separated by - // bytes_per_outer_element - // - Inner dimensions (field_shape): contiguous bytes within each element std::vector src_byte_strides(view_shape.size()); - // First compute strides for chunk dimensions ComputeStrides(c_order, dtype_.bytes_per_outer_element, chunk_shape, tensorstore::span(src_byte_strides.data(), chunk_shape.size())); - // Then compute strides for field_shape dimensions (contiguous within element) if (!field.field_shape.empty()) { ComputeStrides(c_order, static_cast(field.dtype.size()), field.field_shape, @@ -295,15 +267,12 @@ ZarrLeafChunkCache::DecodeChunk(span chunk_indices, field.field_shape.size())); } - // Create source ArrayView pointing to this field's offset within - // the interleaved byte array, with strides that skip over other fields. ArrayView src_field_view( {static_cast( static_cast(byte_array.data()) + field.byte_offset), field.dtype}, StridedLayoutView<>(view_shape, src_byte_strides)); - // Use CopyArray which safely handles any layout differences CopyArray(src_field_view, result_array); field_arrays[field_i] = std::move(result_array); } @@ -316,53 +285,32 @@ Result ZarrLeafChunkCache::EncodeChunk( span> component_arrays) { const size_t num_fields = dtype_.fields.size(); - // Fast path: see `DecodeChunk` -- when the metadata-level `field_shape_` - // is empty, the codec was resolved at the field's natural rank and dtype - // and we hand it the per-field array directly. if (field_shape_.empty()) { assert(num_fields == 1); assert(component_arrays.size() == 1); return codec_state_->EncodeArray(component_arrays[0]); } - // Otherwise pack each field's bytes into the interleaved byte buffer the - // codec chain expects (resolved against `byte` + the trailing `field_shape_` - // dims), then hand it off for encoding. assert(component_arrays.size() == num_fields); - // Build encode shape: [chunk_dims..., field_shape...] const auto& chunk_shape = grid().chunk_shape; std::vector encode_shape(chunk_shape.begin(), chunk_shape.end()); encode_shape.insert(encode_shape.end(), field_shape_.begin(), field_shape_.end()); - // Allocate byte array for combined fields auto byte_array = AllocateArray(encode_shape, c_order, value_init); - // Copy each field's data into the byte array at their respective offsets. - // We create a strided view into the destination that maps to each field's - // position within the interleaved struct layout, then use CopyArray which - // safely handles any source array strides via IterateOverArrays. for (size_t field_i = 0; field_i < num_fields; ++field_i) { const auto& field = dtype_.fields[field_i]; const auto& field_array = component_arrays[field_i]; - // Build the full view shape: [chunk_shape..., field_shape...] - // For fields with field_shape (like r16, r64), we need to include those - // dimensions in the view. std::vector view_shape(chunk_shape.begin(), chunk_shape.end()); view_shape.insert(view_shape.end(), field.field_shape.begin(), field.field_shape.end()); - // Build strides for the destination view: - // - Outer dimensions (chunk_shape): each element separated by - // bytes_per_outer_element - // - Inner dimensions (field_shape): contiguous bytes within each element std::vector dest_byte_strides(view_shape.size()); - // First compute strides for chunk dimensions ComputeStrides(c_order, dtype_.bytes_per_outer_element, chunk_shape, tensorstore::span(dest_byte_strides.data(), chunk_shape.size())); - // Then compute strides for field_shape dimensions (contiguous within element) if (!field.field_shape.empty()) { ComputeStrides(c_order, static_cast(field.dtype.size()), field.field_shape, @@ -370,13 +318,10 @@ Result ZarrLeafChunkCache::EncodeChunk( field.field_shape.size())); } - // Create destination ArrayView pointing to this field's offset within - // the interleaved byte array, with strides that skip over other fields. ArrayView dest_field_view( {static_cast(byte_array.data() + field.byte_offset), field.dtype}, StridedLayoutView<>(view_shape, dest_byte_strides)); - // Use CopyArray which safely handles any source strides via IterateOverArrays CopyArray(field_array, dest_field_view); } diff --git a/tensorstore/driver/zarr3/chunk_cache.h b/tensorstore/driver/zarr3/chunk_cache.h index 17fdf0e7a..71a314b74 100644 --- a/tensorstore/driver/zarr3/chunk_cache.h +++ b/tensorstore/driver/zarr3/chunk_cache.h @@ -207,12 +207,6 @@ class ZarrLeafChunkCache : public internal::KvsBackedChunkCache, ZarrCodecChain::PreparedState::Ptr codec_state_; ZarrDType dtype_; - // Inner trailing dims contributed by the dtype to the codec resolution - // shape -- the same vector that lives on `ZarrMetadata::field_shape`. - // Empty for plain scalar single-field arrays; `{bytes_per_outer_element}` - // for multi-field structs and the byte-substituted void view; the per- - // field `field_shape` for `rN` raw byte fields. Used only to gate the - // "fast path" (no byte interleaving) vs the byte-packing path. std::vector field_shape_; }; @@ -275,10 +269,6 @@ class ZarrShardedChunkCache : public internal::Cache, public ZarrChunkCache { kvstore::DriverPtr base_kvstore_; ZarrCodecChain::PreparedState::Ptr codec_state_; ZarrDType dtype_; - // Same semantics as `ZarrLeafChunkCache::field_shape_`: the metadata-level - // inner trailing dims used to gate the byte-packing decode/encode path - // and to determine how many trailing dims the sharding sub-chunk grid - // carries. std::vector field_shape_; // Data cache pool, if it differs from `this->pool()` (which is equal to the @@ -288,21 +278,9 @@ class ZarrShardedChunkCache : public internal::Cache, public ZarrChunkCache { /// Chunk cache mixin for a chunk cache where the entire chunk cache corresponds /// to a single shard. -/// -/// The sharding codec's `sub_chunk_grid` is a single-component grid whose -/// component shape includes any trailing `field_shape` dimensions contributed -/// by the dtype. This mixin replaces it with a field-level grid built -/// directly from the dtype (`CreateFieldGridSpecification`), so that the -/// surrounding code path operates on per-field typed arrays regardless of -/// whether the chunk is a multi-field struct, an `rN` raw byte field, or a -/// plain scalar. template class ZarrShardSubChunkCache : public ChunkCacheImpl { public: - /// Wraps the sharding state's key parser to pad spatial-only cell indices - /// (from a field-level grid) to the full dimensionality expected by the - /// shard index (which includes any trailing field_shape dimensions with - /// index 0). class FieldKeyParserWrapper : public internal::LexicographicalGridIndexKeyParser { public: @@ -350,11 +328,6 @@ class ZarrShardSubChunkCache : public ChunkCacheImpl { std::move(data_cache_pool)), sharding_state_(std::move(sharding_state)), executor_(std::move(executor)) { - // Strip any trailing inner dimensions contributed by `field_shape_` from - // the sharding codec's sub-chunk grid to get the spatial (chunked-only) - // rank. `field_shape_` is the metadata-level inner shape vector -- - // empty for plain scalars, single-element for current zarr v3 cases - // (`{bytes_per_outer_element}` for structs/void, `{N}` for `rN`). const auto& field_shape_ref = ChunkCacheImpl::field_shape_; const auto& original_grid = *sharding_state_->sub_chunk_grid; const DimensionIndex full_rank = original_grid.chunk_shape.size(); diff --git a/tensorstore/driver/zarr3/codec/transpose.cc b/tensorstore/driver/zarr3/codec/transpose.cc index a5545ea32..b67f51664 100644 --- a/tensorstore/driver/zarr3/codec/transpose.cc +++ b/tensorstore/driver/zarr3/codec/transpose.cc @@ -232,8 +232,6 @@ Result> ResolveOrder( absl::Status TransposeCodecSpec::PropagateDataTypeAndShape( const ArrayDataTypeAndShapeInfo& decoded, ArrayDataTypeAndShapeInfo& encoded) const { - // The user permutation operates strictly on the chunked dimensions; any - // `inner_shape` contributed by the dtype is propagated as-is. DimensionIndex temp_perm[kMaxRank]; TENSORSTORE_ASSIGN_OR_RETURN( auto order, ResolveOrder(options.order, decoded.rank, temp_perm)); @@ -302,8 +300,6 @@ absl::Status TransposeCodecSpec::GetDecodedChunkLayout( Result TransposeCodecSpec::Resolve( ArrayCodecResolveParameters&& decoded, ArrayCodecResolveParameters& encoded, ZarrArrayToArrayCodecSpec::Ptr* resolved_spec) const { - // Spec-level resolution is at the chunked rank only; the user permutation - // never touches inner (`field_shape`) dimensions. DimensionIndex temp_perm[kMaxRank]; TENSORSTORE_ASSIGN_OR_RETURN( auto order, ResolveOrder(options.order, decoded.rank, temp_perm)); @@ -324,11 +320,6 @@ Result TransposeCodecSpec::Resolve( resolved_spec->reset(new TransposeCodecSpec({TransposeCodecSpec::Order( std::vector(order.begin(), order.end()))})); } - // Build the runtime permutation at the *runtime* rank: chunked dims permuted - // as the user requested, inner (`field_shape`) dims pinned at the trailing - // positions with identity. The chunk cache hands extended-rank arrays to - // the codec at runtime; the runtime codec must therefore accept the extended - // rank without altering the inner dims. const DimensionIndex chunked_rank = decoded.rank; const DimensionIndex inner_rank = static_cast(decoded.inner_shape.size()); diff --git a/tensorstore/driver/zarr3/driver_test.cc b/tensorstore/driver/zarr3/driver_test.cc index 2987636b2..e31d7d4ec 100644 --- a/tensorstore/driver/zarr3/driver_test.cc +++ b/tensorstore/driver/zarr3/driver_test.cc @@ -1883,10 +1883,6 @@ ::nlohmann::json GetOpenAsVoidSpec(std::string_view kvstore_path) { } TEST(Zarr3OpenAsVoidTest, SimpleType) { - // Per the zarr v3 open_as_void spec, opening a plain scalar dtype with - // `open_as_void` succeeds and yields a `byte` array whose innermost - // dimension equals the scalar's byte width. For int16 (2 bytes), an - // array of shape [4, 4] becomes byte-shape [4, 4, 2]. auto context = Context::Default(); ::nlohmann::json create_spec{ @@ -2574,7 +2570,7 @@ TEST(Zarr3OpenAsVoidTest, GetSpecInfoRankConsistency) { EXPECT_EQ(4, void_spec.rank()); } -TEST(Zarr3OpenAsVoidTest, FillValue) { // TODO: We need to define behavior for whether fill_value is required always for struct dtype. +TEST(Zarr3OpenAsVoidTest, FillValue) { // Test that fill_value is correctly obtained from metadata when using // open_as_void. The void access should get the fill_value representing // the raw bytes of the original fill_value. @@ -2726,9 +2722,6 @@ TEST(Zarr3OpenAsVoidTest, IncompatibleMetadata) { } TEST(Zarr3OpenAsVoidTest, WithShardingSimpleType) { - // Per the zarr v3 open_as_void spec, simple dtypes are also supported even - // when wrapped in a sharding_indexed codec. For int32 (4 bytes) shape - // [8,8] with sub-chunks [4,4], the void view has shape [8,8,4] of `byte`. auto context = Context::Default(); ::nlohmann::json create_spec{ @@ -2913,13 +2906,6 @@ TEST(Zarr3OpenAsVoidTest, StructBigEndian) { } TEST(Zarr3OpenAsVoidTest, ShardedSubChunkShapeIsUserForm) { - // The sharding_indexed `chunk_shape` exposed in the round-tripped spec is - // the user-facing form (the chunked dimensions only), not the codec's - // internally extended form. This keeps zarr.json round-trips and - // user-provided spec constraints in a single canonical representation. - // For a struct{x: uint8, y: int16} with sub-chunk shape [4, 4], the spec - // round-trips as [4, 4] regardless of whether opened normally or via - // open_as_void. auto context = Context::Default(); ::nlohmann::json create_spec{ @@ -3680,9 +3666,6 @@ TEST(Zarr3StructuredTest, ShardedOpenAsVoidNoFieldCreate) { } TEST(Zarr3StructuredTest, ShardedWiderFieldRoundtrip) { - // Use {uint8, int32} so bytes_per_element=5, which is distinct from rank=3. - // This breaks the coincidental alignment in other tests where - // bytes_per_element=3 and the void rank is also 3. auto context = Context::Default(); ::nlohmann::json base_spec{ From 9bc62a0ec55e56d88591be54485a1ee58798c91f Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 6 May 2026 19:00:49 +0000 Subject: [PATCH 6/8] Deslop --- tensorstore/driver/zarr3/dtype.cc | 62 +++++++------------------- tensorstore/driver/zarr3/dtype.h | 7 ++- tensorstore/driver/zarr3/dtype_test.cc | 6 +-- 3 files changed, 21 insertions(+), 54 deletions(-) diff --git a/tensorstore/driver/zarr3/dtype.cc b/tensorstore/driver/zarr3/dtype.cc index 0935abc09..e006c499d 100644 --- a/tensorstore/driver/zarr3/dtype.cc +++ b/tensorstore/driver/zarr3/dtype.cc @@ -27,7 +27,6 @@ #include #include "absl/base/optimization.h" #include "absl/status/status.h" -#include "absl/strings/ascii.h" #include "absl/strings/numbers.h" #include "absl/strings/str_format.h" #include "absl/strings/str_join.h" @@ -44,7 +43,6 @@ #include "tensorstore/util/span.h" #include "tensorstore/util/status.h" - namespace tensorstore { namespace internal_zarr3 { @@ -76,12 +74,10 @@ Result ParseBaseDType(std::string_view dtype) { if (dtype == "complex128") return make_dtype(dtype_v<::tensorstore::dtypes::complex128_t>); - // Handle r raw bits type where N is number of bits (must be multiple of 8). - // Parse N as uint64_t so values above 2^31-1 (e.g. r8589934592) are accepted; - // (std::numeric_limits::max() / 8) < std::numeric_limits::max(), - // so num_bits / 8 always fits in Index. - if (!dtype.empty() && dtype[0] == 'r' && dtype.size() > 1 && - absl::ascii_isdigit(dtype[1])) { + // Handle r raw bits type, where N is a positive multiple of 8. Parse N + // as uint64_t to accept values above INT32_MAX (e.g. `r8589934592`); + // `uint64_t::max() / 8` fits in `Index`. + if (!dtype.empty() && dtype[0] == 'r') { std::string_view suffix = dtype.substr(1); uint64_t num_bits = 0; if (!absl::SimpleAtoi(suffix, &num_bits) || num_bits == 0 || @@ -93,16 +89,8 @@ Result ParseBaseDType(std::string_view dtype) { } Index num_bytes = static_cast(num_bits / 8); return ZarrDType::BaseDType{std::string(dtype), - dtype_v<::tensorstore::dtypes::byte_t>, - {num_bytes}}; - } - - // Handle bare "r" - must have a number after it - if (!dtype.empty() && dtype[0] == 'r') { - return absl::InvalidArgumentError(absl::StrFormat( - "%s data type is invalid; expected r where N is a positive " - "multiple of 8", - dtype)); + dtype_v<::tensorstore::dtypes::byte_t>, + {num_bytes}}; } constexpr std::string_view kSupported = @@ -115,9 +103,8 @@ Result ParseBaseDType(std::string_view dtype) { namespace { -/// Validates that a fields array contains at least one field. -/// -/// Per the Zarr v3 struct extension, the "fields" array MUST contain at least one field. +/// Validates that a fields array contains at least one field, as required by +/// the Zarr v3 struct extension. /// /// \param size The number of fields in the array. /// \param type_name The data type name for error messages ("struct" or @@ -132,17 +119,12 @@ absl::Status ValidateFieldsArrayNotEmpty(const ptrdiff_t size, return absl::OkStatus(); } -/// Parses a single struct field. -/// -/// Expected format: {"name": "field_name", "data_type": "float32"} -/// -/// Note: Nested struct types and extension data types with configuration -/// (e.g., numpy.datetime64) are valid per the Zarr v3 spec but are not -/// currently supported by TensorStore. +/// Parses a single struct field of the form +/// `{"name": "field_name", "data_type": "float32"}`. /// /// \param field_json The JSON object representing a single field. /// \param field[out] Filled with the parsed field on success. -/// \error `absl::StatusCode::kInvalidArgument` if `field_json` is not valid +/// \error `absl::StatusCode::kInvalidArgument` if `field_json` is not valid. absl::Status ParseObjectField(const nlohmann::json& field_json, ZarrDType::Field& field) { if (!field_json.is_object()) { @@ -366,12 +348,8 @@ Result ParseDTypeNoDerived(const nlohmann::json& value) { } // namespace absl::Status ValidateDType(ZarrDType& dtype) { - // Invariant: every valid zarr v3 dtype has at least one field. All JSON - // parser entry points enforce this (`ValidateFieldsArrayNotEmpty` for - // struct/structured, single-field `resize(1)` for scalar / raw_bytes). - // This guard catches future programmatic constructions of `ZarrDType` - // that bypass the parser, so that consumers indexing `fields[0]` get a - // clean error rather than out-of-bounds access. + // JSON parsers always produce at least one field; this guards programmatic + // constructions, since downstream consumers index `fields[0]` directly. if (dtype.fields.empty()) { return absl::FailedPreconditionError( "zarr3 data type must have at least one field"); @@ -472,20 +450,12 @@ Result ChooseBaseDType(DataType dtype) { return make_dtype("complex64"); if (dtype == dtype_v<::tensorstore::dtypes::complex128_t>) return make_dtype("complex128"); + // `byte_t` and `char_t` both encode as `r8`; `r8` parses back to `byte_t`. if (dtype == dtype_v<::tensorstore::dtypes::byte_t>) { - ZarrDType::BaseDType base_dtype; - base_dtype.dtype = dtype; - base_dtype.encoded_dtype = "r8"; - base_dtype.flexible_shape = {1}; - return base_dtype; + return D{"r8", dtype, {1}}; } if (dtype == dtype_v<::tensorstore::dtypes::char_t>) { - // char_t encodes as r8, which parses back to byte_t - ZarrDType::BaseDType base_dtype; - base_dtype.dtype = dtype_v<::tensorstore::dtypes::byte_t>; - base_dtype.encoded_dtype = "r8"; - base_dtype.flexible_shape = {1}; - return base_dtype; + return D{"r8", dtype_v<::tensorstore::dtypes::byte_t>, {1}}; } return absl::InvalidArgumentError( absl::StrFormat("Data type not supported: %v", dtype)); diff --git a/tensorstore/driver/zarr3/dtype.h b/tensorstore/driver/zarr3/dtype.h index f22ac7a78..cdff4d1b2 100644 --- a/tensorstore/driver/zarr3/dtype.h +++ b/tensorstore/driver/zarr3/dtype.h @@ -78,10 +78,9 @@ struct ZarrDType { /// specified as an array. Otherwise, is empty. std::string name; - /// Inner array dimensions of this field, derived from `flexible_shape`. - /// Inner array shape contributed by this field to the full TensorStore - /// rank. For scalar fields this is empty; for raw byte fields (e.g. - /// `r24`, structured-as-void access) this is the inner shape that is + /// Inner array dimensions of this field (derived from `flexible_shape`). + /// + /// Empty for scalar fields. For raw byte fields (e.g. `r24`) this is /// surfaced as trailing dimensions in the user-visible array. std::vector field_shape; diff --git a/tensorstore/driver/zarr3/dtype_test.cc b/tensorstore/driver/zarr3/dtype_test.cc index 72a1f2694..e1d2bb1ae 100644 --- a/tensorstore/driver/zarr3/dtype_test.cc +++ b/tensorstore/driver/zarr3/dtype_test.cc @@ -230,10 +230,8 @@ TEST(ParseDType, DuplicateFieldName) { } TEST(ValidateDType, RejectsEmptyFields) { - // Defensive guard: a default-constructed `ZarrDType` has no fields, which - // is an invariant violation for downstream consumers indexing `fields[0]`. - // The JSON parsers all enforce non-empty fields, so this only fires for - // programmatic constructions that bypass the parser. + // Bypassing the JSON parsers leaves `fields` empty; downstream code + // assumes `fields[0]` exists. ZarrDType empty; EXPECT_THAT(ValidateDType(empty), StatusIs(absl::StatusCode::kFailedPrecondition, From 102b8cb7791555158d4b06ee0ec767386034b805 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Wed, 6 May 2026 19:05:56 +0000 Subject: [PATCH 7/8] Deslop --- tensorstore/driver/zarr3/metadata.cc | 21 --------------------- tensorstore/driver/zarr3/metadata_test.cc | 14 +------------- 2 files changed, 1 insertion(+), 34 deletions(-) diff --git a/tensorstore/driver/zarr3/metadata.cc b/tensorstore/driver/zarr3/metadata.cc index 6ba5f0a92..3e7539018 100644 --- a/tensorstore/driver/zarr3/metadata.cc +++ b/tensorstore/driver/zarr3/metadata.cc @@ -287,10 +287,8 @@ absl::Status FillValueJsonBinder::operator()( std::vector>* obj, ::nlohmann::json* j) const { obj->resize(dtype.fields.size()); if (dtype.fields.size() == 1) { - // Special case: raw_bytes (single field with byte_t and flexible shape) if (dtype.fields[0].dtype.id() == DataTypeId::byte_t && !dtype.fields[0].flexible_shape.empty()) { - // Handle base64-encoded fill value for raw_bytes if (!j->is_string()) { return absl::InvalidArgumentError( "Expected base64-encoded string for raw_bytes fill_value"); @@ -301,7 +299,6 @@ absl::Status FillValueJsonBinder::operator()( "Expected valid base64-encoded fill value, but received: %s", j->dump())); } - // Verify size matches expected byte array size Index expected_size = dtype.fields[0].num_inner_elements; if (static_cast(b64_decoded.size()) != expected_size) { return absl::InvalidArgumentError(absl::StrFormat( @@ -309,7 +306,6 @@ absl::Status FillValueJsonBinder::operator()( "%d bytes", expected_size, b64_decoded.size())); } - // Create fill value array auto fill_arr = AllocateArray(dtype.fields[0].field_shape, c_order, default_init, dtype.fields[0].dtype); std::memcpy(fill_arr.data(), b64_decoded.data(), b64_decoded.size()); @@ -319,9 +315,7 @@ absl::Status FillValueJsonBinder::operator()( DecodeSingle(*j, dtype.fields[0].dtype, (*obj)[0])); } } else { - // For structured types, handle object, array, and base64-encoded string if (j->is_object()) { - // Zarr v3 struct extension format: {"field_name": value, ...} for (size_t i = 0; i < dtype.fields.size(); ++i) { const auto& field_name = dtype.fields[i].name; if (j->contains(field_name)) { @@ -333,14 +327,12 @@ absl::Status FillValueJsonBinder::operator()( } } } else if (j->is_string()) { - // Legacy: decode base64-encoded fill value for entire struct std::string b64_decoded; if (!absl::Base64Unescape(j->get(), &b64_decoded)) { return absl::InvalidArgumentError(absl::StrFormat( "Expected valid base64-encoded fill value, but received: %s", j->dump())); } - // Verify size matches expected struct size if (static_cast(b64_decoded.size()) != dtype.bytes_per_outer_element) { return absl::InvalidArgumentError(absl::StrFormat( @@ -348,7 +340,6 @@ absl::Status FillValueJsonBinder::operator()( "%d bytes", dtype.bytes_per_outer_element, b64_decoded.size())); } - // Extract per-field fill values from decoded bytes for (size_t i = 0; i < dtype.fields.size(); ++i) { const auto& field = dtype.fields[i]; auto arr = AllocateArray(span{}, c_order, default_init, @@ -358,7 +349,6 @@ absl::Status FillValueJsonBinder::operator()( (*obj)[i] = std::move(arr); } } else if (j->is_array()) { - // Legacy: array format [value1, value2, ...] if (j->size() != dtype.fields.size()) { return internal_json::ExpectedError( *j, absl::StrFormat("array of size %d", dtype.fields.size())); @@ -382,7 +372,6 @@ absl::Status FillValueJsonBinder::operator()( if (dtype.fields.size() == 1) { return EncodeSingle((*obj)[0], dtype.fields[0].dtype, *j); } - // Structured fill value - use object format per spec *j = ::nlohmann::json::object(); for (size_t i = 0; i < dtype.fields.size(); ++i) { ::nlohmann::json item; @@ -408,7 +397,6 @@ absl::Status FillValueJsonBinder::DecodeSingle(::nlohmann::json& j, AllocateArray(span{}, c_order, default_init, data_type); void* data = arr.data(); out = std::move(arr); - // Special handling for byte_t: use uint8_t functions since they're binary compatible auto type_id = data_type.id(); if (type_id == DataTypeId::byte_t) { type_id = DataTypeId::uint8_t; @@ -434,7 +422,6 @@ absl::Status FillValueJsonBinder::EncodeSingle( return absl::InvalidArgumentError( "data_type must be specified before fill_value"); } - // Special handling for byte_t: use uint8_t functions since they're binary compatible auto type_id = data_type.id(); if (type_id == DataTypeId::byte_t) { type_id = DataTypeId::uint8_t; @@ -902,7 +889,6 @@ absl::Status ValidateMetadata(const ZarrMetadata& metadata, const ZarrMetadataConstraints& constraints) { using internal::MetadataMismatchError; if (constraints.data_type) { - // Compare ZarrDType if (::nlohmann::json(*constraints.data_type) != ::nlohmann::json(metadata.data_type)) { return MetadataMismatchError( @@ -911,7 +897,6 @@ absl::Status ValidateMetadata(const ZarrMetadata& metadata, } } if (constraints.fill_value) { - // Compare vector of arrays if (constraints.fill_value->size() != metadata.fill_value.size()) { return MetadataMismatchError("fill_value size", constraints.fill_value->size(), @@ -1264,9 +1249,6 @@ absl::Status SetChunkLayoutFromMetadata( Result GetEffectiveChunkLayout( const ZarrMetadataConstraints& metadata_constraints, const Schema& schema) { SpecRankAndFieldInfo info; - // Use metadata_constraints.rank when available: it represents the logical - // rank (matching chunk_shape dimensions). schema.rank() may be larger when - // open_as_void adds a bytes dimension. info.chunked_rank = metadata_constraints.rank; if (info.chunked_rank == dynamic_rank) { info.chunked_rank = schema.rank().rank; @@ -1277,9 +1259,6 @@ Result GetEffectiveChunkLayout( if (info.chunked_rank == dynamic_rank && metadata_constraints.chunk_shape) { info.chunked_rank = metadata_constraints.chunk_shape->size(); } - // We can't easily know field info from constraints unless we parse data_type. - // If data_type is present and has 1 field, we can check it. - // For now, basic implementation. ChunkLayout chunk_layout = schema.chunk_layout(); std::optional> chunk_shape_span; diff --git a/tensorstore/driver/zarr3/metadata_test.cc b/tensorstore/driver/zarr3/metadata_test.cc index 7cc8880ac..c9da94ffe 100644 --- a/tensorstore/driver/zarr3/metadata_test.cc +++ b/tensorstore/driver/zarr3/metadata_test.cc @@ -659,7 +659,6 @@ TEST(FillValueTest, Float64) { } TEST(FillValueTest, StructuredObjectFormat) { - // Create a structured dtype with two fields ZarrDType dtype; dtype.has_fields = true; dtype.fields.resize(2); @@ -679,7 +678,6 @@ TEST(FillValueTest, StructuredObjectFormat) { FillValueJsonBinder binder(dtype); - // Test parsing object format ::nlohmann::json object_json = {{"x", 42}, {"y", -100}}; TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto fill_values, @@ -688,7 +686,6 @@ TEST(FillValueTest, StructuredObjectFormat) { EXPECT_EQ(*static_cast(fill_values[0].data()), 42); EXPECT_EQ(*static_cast(fill_values[1].data()), -100); - // Test serialization uses object format TENSORSTORE_ASSERT_OK_AND_ASSIGN(auto output_json, jb::ToJson(fill_values, binder)); EXPECT_TRUE(output_json.is_object()); @@ -697,7 +694,6 @@ TEST(FillValueTest, StructuredObjectFormat) { } TEST(FillValueTest, StructuredLegacyArrayFormat) { - // Create a structured dtype with two fields ZarrDType dtype; dtype.has_fields = true; dtype.fields.resize(2); @@ -717,7 +713,6 @@ TEST(FillValueTest, StructuredLegacyArrayFormat) { FillValueJsonBinder binder(dtype); - // Test parsing legacy array format ::nlohmann::json array_json = {123, 1.5}; TENSORSTORE_ASSERT_OK_AND_ASSIGN( auto fill_values, @@ -728,7 +723,6 @@ TEST(FillValueTest, StructuredLegacyArrayFormat) { } TEST(FillValueTest, StructuredObjectOmittedFieldsError) { - // Create a structured dtype with two fields ZarrDType dtype; dtype.has_fields = true; dtype.fields.resize(2); @@ -748,7 +742,6 @@ TEST(FillValueTest, StructuredObjectOmittedFieldsError) { FillValueJsonBinder binder(dtype); - // Test parsing object format with only one field specified ::nlohmann::json partial_json = {{"present", 42}}; EXPECT_THAT( jb::FromJson>>(partial_json, binder), @@ -827,7 +820,6 @@ TEST(GetSpecRankAndFieldInfoTest, SelectedFieldWithFieldShape) { TEST(GetSpecRankAndFieldInfoTest, DeriveRankFromSchema) { ZarrMetadataConstraints constraints; - // No rank in constraints constraints.data_type.emplace(); constraints.data_type->has_fields = false; constraints.data_type->fields.resize(1); @@ -842,7 +834,6 @@ TEST(GetSpecRankAndFieldInfoTest, DeriveRankFromSchema) { GetSpecRankAndFieldInfo(constraints, /*selected_field=*/"", schema, /*open_as_void=*/false)); EXPECT_EQ(info.full_rank, 3); - // chunked_rank derived from full_rank - field_rank (0) EXPECT_EQ(info.chunked_rank, 3); EXPECT_EQ(info.field_rank, 0); } @@ -872,7 +863,6 @@ TEST(ValidateSpecRankAndFieldInfoTest, DerivesFullRank) { tensorstore::internal_zarr3::SpecRankAndFieldInfo info; info.chunked_rank = 2; info.field_rank = 1; - // full_rank is dynamic TENSORSTORE_ASSERT_OK( tensorstore::internal_zarr3::ValidateSpecRankAndFieldInfo(info)); @@ -883,7 +873,6 @@ TEST(ValidateSpecRankAndFieldInfoTest, DerivesChunkedRank) { tensorstore::internal_zarr3::SpecRankAndFieldInfo info; info.full_rank = 5; info.field_rank = 2; - // chunked_rank is dynamic TENSORSTORE_ASSERT_OK( tensorstore::internal_zarr3::ValidateSpecRankAndFieldInfo(info)); @@ -894,7 +883,7 @@ TEST(ValidateSpecRankAndFieldInfoTest, InconsistentRanksError) { tensorstore::internal_zarr3::SpecRankAndFieldInfo info; info.full_rank = 5; info.chunked_rank = 2; - info.field_rank = 1; // 2 + 1 != 5 + info.field_rank = 1; EXPECT_THAT( tensorstore::internal_zarr3::ValidateSpecRankAndFieldInfo(info), @@ -959,7 +948,6 @@ TEST(TrySetMetadataConstraintsOnSchemaTest, SetsRankWithFieldShape) { constraints, /*selected_field=*/"matrix", /*open_as_void=*/false, schema)); - // Rank should be metadata_rank (2) + field_shape rank (2) = 4 EXPECT_EQ(schema.rank(), 4); } From 923ce7b872a63337bed0ae465ecc25cacac764f2 Mon Sep 17 00:00:00 2001 From: BrianMichell Date: Fri, 8 May 2026 13:36:44 +0000 Subject: [PATCH 8/8] deslop --- tensorstore/driver/zarr3/driver.cc | 168 +++++++---------------------- 1 file changed, 41 insertions(+), 127 deletions(-) diff --git a/tensorstore/driver/zarr3/driver.cc b/tensorstore/driver/zarr3/driver.cc index 7a6a12439..6a67c20b0 100644 --- a/tensorstore/driver/zarr3/driver.cc +++ b/tensorstore/driver/zarr3/driver.cc @@ -151,53 +151,50 @@ class ZarrDriverSpec } Result> GetDomain() const override { - // When the array contributes trailing field_shape dimensions - // (open_as_void or a selected rN/struct field), build the domain by - // appending those dimensions to the chunked dimensions stored in - // `metadata_constraints`. Both cases are uniform under the field_shape - // model: open_as_void is just `field_shape = {bytes_per_outer_element}` - // contributed by the synthetic void view of the data type. - if (metadata_constraints.data_type && metadata_constraints.shape) { - const Index byte_dim = - metadata_constraints.data_type->bytes_per_outer_element; - span field_shape; - if (open_as_void) { - field_shape = span(&byte_dim, 1); - } else if (!selected_field.empty()) { - const auto& dtype = *metadata_constraints.data_type; - for (const auto& field : dtype.fields) { - if (field.name == selected_field && !field.field_shape.empty()) { - field_shape = span(field.field_shape.data(), - field.field_shape.size()); - break; - } + if (!metadata_constraints.data_type || !metadata_constraints.shape) { + return GetEffectiveDomain(metadata_constraints, schema); + } + + const Index byte_dim = + metadata_constraints.data_type->bytes_per_outer_element; + span field_shape; + if (open_as_void) { + field_shape = span(&byte_dim, 1); + } else if (!selected_field.empty()) { + const auto& dtype = *metadata_constraints.data_type; + for (const auto& field : dtype.fields) { + if (field.name == selected_field && !field.field_shape.empty()) { + field_shape = span(field.field_shape.data(), + field.field_shape.size()); + break; } } - if (!field_shape.empty()) { - const DimensionIndex chunked_rank = - metadata_constraints.shape->size(); - const DimensionIndex total_rank = chunked_rank + field_shape.size(); - IndexDomainBuilder builder(total_rank); - std::fill_n(builder.origin().begin(), total_rank, Index{0}); - 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_upper_bounds(DimensionSet::UpTo(chunked_rank)); - if (metadata_constraints.dimension_names) { - for (DimensionIndex j = 0; j < chunked_rank; ++j) { - if (const auto& name = (*metadata_constraints.dimension_names)[j]; - name.has_value()) { - builder.labels()[j] = *name; - } - } + } + + if (field_shape.empty()) { + return GetEffectiveDomain(metadata_constraints, schema); + } + + const DimensionIndex chunked_rank = + metadata_constraints.shape->size(); + const DimensionIndex total_rank = chunked_rank + field_shape.size(); + IndexDomainBuilder builder(total_rank); + std::fill_n(builder.origin().begin(), total_rank, Index{0}); + 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_upper_bounds(DimensionSet::UpTo(chunked_rank)); + if (metadata_constraints.dimension_names) { + for (DimensionIndex j = 0; j < chunked_rank; ++j) { + if (const auto& name = (*metadata_constraints.dimension_names)[j]; + name.has_value()) { + builder.labels()[j] = *name; } - return builder.Finalize(); } } - - return GetEffectiveDomain(metadata_constraints, schema); + return builder.Finalize(); } Result> GetFillValue( @@ -206,14 +203,12 @@ class ZarrDriverSpec const auto& constraints = metadata_constraints; - // If constraints don't specify a fill value, just use the schema's. if (!constraints.fill_value || constraints.fill_value->empty()) { return fill_value; } const auto& vec = *constraints.fill_value; - // If we don't have dtype information, we can't do field-aware logic. if (!constraints.data_type) { if (!vec.empty()) return vec[0]; return fill_value; @@ -221,11 +216,6 @@ class ZarrDriverSpec const ZarrDType& dtype = *constraints.data_type; - // ── Void access: synthesize a byte-level fill value ──────────────────────── - // - // Defer to `MakeVoidFillValue` so the spec-level fill is byte-identical - // to the data-cache-level fill produced by `GetVoidMetadata` (codec - // endianness is honoured uniformly for non-struct, non-rN dtypes). if (open_as_void) { const ZarrCodecChainSpec empty_codec_specs; const ZarrCodecChainSpec& codec_specs = @@ -235,13 +225,11 @@ class ZarrDriverSpec MakeVoidFillValue(dtype, codec_specs, vec)); } - // ── Normal field access: just return that field's fill_value ─────────────── TENSORSTORE_ASSIGN_OR_RETURN( size_t field_index, GetFieldIndex(dtype, selected_field)); if (field_index < vec.size()) { return vec[field_index]; } - // Fallback to "no fill". return SharedArray(); } @@ -328,11 +316,6 @@ class DataCacheBase virtual ZarrChunkCache& zarr_chunk_cache() = 0; - /// Returns the fill value to expose to user-visible reads for the given - /// component (field) index. By default this is just the per-field fill - /// value stored in the natural metadata; the `open_as_void` data cache - /// overrides this to return the packed single-byte fill value computed at - /// open time. virtual SharedArray GetEffectiveFillValue(size_t component_index) { const auto& m = this->metadata(); if (component_index >= m.fill_value.size()) { @@ -370,15 +353,6 @@ class DataCacheBase i < static_cast(metadata.shape.size()); ++i) { implicit_upper_bounds[i] = true; } - // The user-visible domain is `metadata.shape` extended by the - // metadata-level `field_shape` (the same vector that drives the codec - // resolution shape). This covers `rN` raw byte fields, the - // sharding-substituted dtype, and the synthetic single-byte view - // produced by `GetVoidMetadata` for `open_as_void`. Multi-field structs - // have `metadata.field_shape == {bytes_per_outer_element}` for the - // codec view, but the *user* must select a field first; here we only - // extend bounds when the dtype has a single field, so the trailing - // dim(s) are unambiguously user-visible. if (bounds.rank() > static_cast(metadata.shape.size()) && metadata.data_type.fields.size() == 1 && !metadata.field_shape.empty()) { @@ -415,11 +389,6 @@ class DataCacheBase static internal::ChunkGridSpecification GetChunkGridSpecification( const ZarrMetadata& metadata) { assert(!metadata.fill_value.empty()); - // The cache always sees a `metadata` whose `data_type` correctly - // describes the field structure -- including the synthetic single-byte - // field with `field_shape = {bytes_per_outer_element}` that - // `GetVoidMetadata` produces for `open_as_void`. `CreateFieldGridSpecification` - // therefore handles scalar, struct, rN, and open-as-void modes uniformly. return CreateFieldGridSpecification( metadata.chunk_shape, metadata.data_type, span(metadata.inner_order.data(), metadata.rank), &metadata.fill_value); @@ -508,7 +477,6 @@ class DataCacheBase Result> GetExternalToInternalTransform( const void* metadata_ptr, size_t component_index) override { - // component_index corresponds to the selected field index const auto& metadata = *static_cast(metadata_ptr); const auto& field = metadata.data_type.fields[component_index]; const DimensionIndex rank = metadata.rank; @@ -547,7 +515,6 @@ class DataCacheBase auto& spec = static_cast(spec_base); const auto& metadata = *static_cast(metadata_ptr); spec.metadata_constraints = ZarrMetadataConstraints(metadata); - // Encode selected_field from component_index if (metadata.data_type.has_fields && component_index < metadata.data_type.fields.size()) { spec.selected_field = metadata.data_type.fields[component_index].name; @@ -585,13 +552,6 @@ using internal_kvs_backed_chunk_driver::DataCacheInitializer; template class ZarrDataCache : public ChunkCacheImpl, public DataCacheBase { public: - // `effective_metadata` is the metadata view the cache operates on for chunk - // decoding, grid layout, fills, transforms, and bounds. It is always - // populated -- when no substitution applies it aliases the persisted - // (natural) metadata; under `open_as_void` it is the byte-substituted view - // produced by `GetVoidMetadata`. `is_substituted` records whether a - // substitution was applied; it is the single bit required to round-trip - // `open_as_void` back into the spec. template explicit ZarrDataCache(DataCacheInitializer&& initializer, std::string key_prefix, @@ -627,14 +587,6 @@ class ZarrDataCache : public ChunkCacheImpl, public DataCacheBase { return DataCacheBase::executor(); } - // The framework hands `metadata_ptr` from the metadata cache (i.e. the - // current persisted form). Cache-internal layout/transform logic drives - // off `effective_metadata_` only when a substitution was applied at open - // time -- otherwise we forward the live `metadata_ptr` so that operations - // like `Resize` correctly observe the post-resize shape. Either way the - // base implementation reads `field_shape` off whichever metadata it gets - // and produces the right user-visible rank with no `open_as_void`-aware - // branching beyond the substitution site itself. const void* RouteMetadata(const void* metadata_ptr) const { return is_substituted_ ? effective_metadata_.get() : metadata_ptr; } @@ -656,9 +608,6 @@ class ZarrDataCache : public ChunkCacheImpl, public DataCacheBase { absl::Status GetBoundSpecData(KvsDriverSpec& spec_base, const void* metadata_ptr, size_t component_index) override { - // The persisted (natural) metadata drives the spec body so the user sees - // their original dtype on round-trip. `open_as_void` is the only piece - // of substitution-time state the cache must surface back into the spec. TENSORSTORE_RETURN_IF_ERROR( DataCacheBase::GetBoundSpecData(spec_base, metadata_ptr, component_index)); auto& spec = static_cast(spec_base); @@ -695,10 +644,6 @@ class ZarrDriver : public ZarrDriverBase { Result> GetFillValue( IndexTransformView<> transform) override { - // The data cache routes this through any substituted view metadata, so - // `open_as_void` returns the packed single-byte fill value at index 0 - // while normal access returns the per-component fill from the natural - // metadata. return static_cast(cache())->GetEffectiveFillValue( this->component_index()); } @@ -782,34 +727,16 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase { public: using ZarrDriver::OpenStateBase::OpenStateBase; - // Lazily computed (and memoised across `GetComponentIndex` / - // `GetDataCache`) byte-substituted view of the metadata. Only non-null - // when `spec().open_as_void` is true and a successful substitution has - // already been performed against the framework-supplied `metadata_ptr`. - // Both call sites pass the same pointer (the persisted metadata), so a - // single computation suffices for both validation and cache construction. Result> GetOrComputeVoidMetadata( const ZarrMetadata& metadata) { - if (!cached_void_metadata_for_) { - TENSORSTORE_ASSIGN_OR_RETURN(auto void_metadata, - internal_zarr3::GetVoidMetadata(metadata)); - cached_void_metadata_ = std::move(void_metadata); - cached_void_metadata_for_ = &metadata; - } - // Defensive: if the framework ever started passing a different metadata - // object across calls, recompute rather than serve stale state. In the - // current code path this branch is unreachable. - if (cached_void_metadata_for_ != &metadata) { - TENSORSTORE_ASSIGN_OR_RETURN(auto void_metadata, + if (!cached_void_metadata_) { + TENSORSTORE_ASSIGN_OR_RETURN(cached_void_metadata_, internal_zarr3::GetVoidMetadata(metadata)); - cached_void_metadata_ = std::move(void_metadata); - cached_void_metadata_for_ = &metadata; } return cached_void_metadata_; } std::shared_ptr cached_void_metadata_; - const ZarrMetadata* cached_void_metadata_for_ = nullptr; std::string GetPrefixForDeleteExisting() override { return spec().store.path; @@ -860,18 +787,9 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase { const auto& metadata = *static_cast(initializer.metadata.get()); - // Build the cache-side `effective_metadata`: under `open_as_void` this is - // the byte-substituted view produced by `GetVoidMetadata`; otherwise it - // aliases the persisted metadata. Everything downstream -- chunk cache, - // grid spec, codec state -- runs against that single view, with no - // `open_as_void`-specific branching beyond the substitution itself. std::shared_ptr effective_metadata; bool is_substituted = false; if (spec().open_as_void) { - // Reuse the substitution that `GetComponentIndex` already performed; - // `GetOrComputeVoidMetadata` memoises by metadata pointer and avoids - // re-resolving the codec chain. `CHECK` here because validation has - // already succeeded by this point in the open flow. TENSORSTORE_CHECK_OK_AND_ASSIGN(effective_metadata, GetOrComputeVoidMetadata(metadata)); is_substituted = true; @@ -901,14 +819,10 @@ class ZarrDriver::OpenState : public ZarrDriver::OpenStateBase { return absl::InvalidArgumentError( "\"field\" and \"open_as_void\" are mutually exclusive"); } - // Compute the void substitution once and stash it on the open state; - // `GetDataCache` reuses the same shared_ptr below to avoid re-resolving - // the codec chain. TENSORSTORE_ASSIGN_OR_RETURN(auto void_metadata, GetOrComputeVoidMetadata(metadata)); TENSORSTORE_RETURN_IF_ERROR(ValidateMetadataSchema( *void_metadata, /*field_index=*/0, spec().schema)); - // The void view exposes a single component (the synthetic byte field). return static_cast(0); } TENSORSTORE_ASSIGN_OR_RETURN(