Skip to content

Update to follow #287 and implement general feedback #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
BrianMichell merged 8 commits into from
May 8, 2026
Merged

Update to follow #287 and implement general feedback #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0b1b152
Update to follow #287 and implement general feedback
BrianMichell May 5, 2026
a0b56ef
Deslopping
BrianMichell May 6, 2026
fcacc9b
Deslop
BrianMichell May 6, 2026
7d35a8a
Deslop
BrianMichell May 6, 2026
46a8e04
Deslop
BrianMichell May 6, 2026
9bc62a0
Deslop
BrianMichell May 6, 2026
102b8cb
Deslop
BrianMichell May 6, 2026
923ce7b
deslop
BrianMichell May 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions tensorstore/driver/zarr3/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
177 changes: 15 additions & 162 deletions tensorstore/driver/zarr3/chunk_cache.cc
View file
Open in desktop

Large diffs are not rendered by default.

116 changes: 30 additions & 86 deletions tensorstore/driver/zarr3/chunk_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Index> field_shape,
internal::CachePool::WeakPtr data_cache_pool);

void Read(ZarrChunkCache::ReadRequest request,
AnyFlowReceiver<absl::Status, internal::ReadChunk,
Expand Down Expand Up @@ -210,16 +207,7 @@ class ZarrLeafChunkCache : public internal::KvsBackedChunkCache,

ZarrCodecChain::PreparedState::Ptr codec_state_;
ZarrDType dtype_;
bool open_as_void_;
bool original_is_structured_;
DataType original_dtype_; // Original dtype for void access encoding
bool grid_has_void_dimension_; // Whether grid().chunk_shape includes bytes dim

private:
Result<absl::InlinedVector<SharedArray<const void>, 1>> DecodeChunkAsVoid(
absl::Cord data);

Result<absl::Cord> EncodeChunkAsVoid(const SharedArray<const void>& byte_array);
std::vector<Index> field_shape_;
};

/// Chunk cache for a Zarr array where each chunk is a shard.
Expand All @@ -230,11 +218,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<Index> field_shape,
internal::CachePool::WeakPtr data_cache_pool);

const ZarrShardingCodec::PreparedState& sharding_codec_state() const {
return static_cast<const ZarrShardingCodec::PreparedState&>(
Expand Down Expand Up @@ -284,9 +269,7 @@ 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
std::vector<Index> field_shape_;

// Data cache pool, if it differs from `this->pool()` (which is equal to the
// metadata cache pool).
Expand All @@ -298,9 +281,6 @@ class ZarrShardedChunkCache : public internal::Cache, public ZarrChunkCache {
template <typename ChunkCacheImpl>
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).
class FieldKeyParserWrapper
: public internal::LexicographicalGridIndexKeyParser {
public:
Expand Down Expand Up @@ -339,63 +319,34 @@ 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<Index> 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<Index> 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<tensorstore::dtypes::byte_t>);
// chunked_to_cell_dimensions maps the grid dimensions to cell dimensions
// (the bytes dimension is unchunked, so not included here)
std::vector<DimensionIndex> 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<Index> 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<const DimensionIndex>(), // no inner_order available
nullptr)); // zero-init fill values

field_key_parser_.emplace(
sharding_state_->GetSubChunkStorageKeyParser(), full_rank);
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<DimensionIndex>(field_shape_ref.size());
const DimensionIndex spatial_rank = full_rank - extra_field_dims;

std::vector<Index> 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<const DimensionIndex>(), // no inner_order available
nullptr)); // zero-init fill values

if (extra_field_dims != 0) {
field_key_parser_.emplace(sharding_state_->GetSubChunkStorageKeyParser(),
full_rank);
}
}

Expand All @@ -408,21 +359,14 @@ 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_; }

internal::Cache& cache() final { return *this; }

ZarrShardingCodec::PreparedState::Ptr sharding_state_;
Executor executor_;
std::optional<internal::ChunkGridSpecification> void_grid_;
std::optional<internal::ChunkGridSpecification> field_grid_;
std::optional<FieldKeyParserWrapper> field_key_parser_;
};
Expand Down
12 changes: 12 additions & 0 deletions tensorstore/driver/zarr3/codec/codec_spec.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@

#include <array>
#include <optional>
#include <vector>

#include "absl/status/status.h"
#include "tensorstore/array.h"
Expand Down Expand Up @@ -136,10 +137,17 @@ struct ArrayDataTypeAndShapeInfo {
DataType dtype;

// 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 shape of the array on which the codec will operate.
// When present, has exactly `rank` valid entries.
std::optional<std::array<Index, kMaxRank>> shape;

// 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<Index> inner_shape;
};

// Specifies information about the chunk layout that must be propagated through
Expand All @@ -166,6 +174,7 @@ struct ArrayCodecResolveParameters {
DataType dtype;

// 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.
Expand All @@ -179,6 +188,9 @@ struct ArrayCodecResolveParameters {

// Specifies required inner order.
std::optional<std::array<DimensionIndex, kMaxRank>> inner_order;

// Inner trailing dimensions contributed by `dtype` (e.g. for `open_as_void`).
std::vector<Index> inner_shape;
};

// Spec for an "array -> array" codec.
Expand Down
Loading