array.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
 
#ifndef LIBZARR_ARRAY_HPP
#define LIBZARR_ARRAY_HPP
 
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
 
#include "libzarr/codecs.hpp"
#include "libzarr/detail/common.hpp"
#include "libzarr/metadata.hpp"
#include "libzarr/sharding.hpp"
#include "libzarr/store.hpp"
#include "libzarr/types.hpp"
#include "libzarr/v2.hpp"
#include "libzarr/v3.hpp"
 
 
namespace zarr {
 
namespace detail {
 
inline void validate_path(const std::string& path) {
  if (path.empty()) {
    return;
  }
  if (path.front() == '/' || path.back() == '/') {
    throw error("node path must not start or end with '/': '" + path + "'");
  }
  std::size_t start = 0;
  while (start <= path.size()) {
    const std::size_t slash = path.find('/', start);
    const std::size_t end = slash == std::string::npos ? path.size() : slash;
    if (end == start) {
      throw error("node path has an empty segment: '" + path + "'");
    }
    if (path[start] == '.') {
      throw error("node path segments must not start with '.': '" + path + "'");
    }
    if (slash == std::string::npos) {
      break;
    }
    start = slash + 1;
  }
}
 
inline bool next_index(std::vector<std::uint64_t>& index,
                       const std::vector<std::uint64_t>& extents) {
  std::size_t d = extents.size();
  while (d-- > 0) {
    if (++index[d] < extents[d]) {
      return true;
    }
    index[d] = 0;
  }
  return false;
}
 
inline bool next_index_box(std::vector<std::uint64_t>& index, const std::vector<std::uint64_t>& lo,
                           const std::vector<std::uint64_t>& hi) {
  std::size_t d = lo.size();
  while (d-- > 0) {
    if (index[d] < hi[d]) {
      ++index[d];
      return true;
    }
    index[d] = lo[d];
  }
  return false;
}
 
}  // namespace detail
 
struct ArraySpec {
  ZarrFormat format = ZarrFormat::v2;
  std::vector<std::uint64_t> shape;
  std::vector<std::uint64_t> chunks;
  DataType dtype;
  std::vector<CodecSpec> codecs;
  std::optional<Bytes> fill;
  json attributes = json::object();
  char dimension_separator = '.';
  json dimension_names;
  std::vector<std::uint64_t> shards;
};
 
class Array {
 public:
  static Array create(std::shared_ptr<Store> store, const std::string& path,
                      const ArraySpec& spec) {
    if (!store) {
      throw error("Array::create: null store");
    }
    detail::validate_path(path);
    const bool v3 = spec.format == ZarrFormat::v3;
    const std::string ctx = v3 ? v3::meta_key(path) : v2::meta_key(path, v2::kArraySuffix);
    if (spec.chunks.size() != spec.shape.size()) {
      throw error(ctx + ": chunks rank " + std::to_string(spec.chunks.size()) + " != shape rank " +
                  std::to_string(spec.shape.size()));
    }
    for (const std::uint64_t c : spec.chunks) {
      if (c == 0) {
        throw error(ctx + ": chunk extents must be positive");
      }
    }
    if (spec.dimension_separator != '.' && spec.dimension_separator != '/') {
      throw error(ctx + ": dimension_separator must be '.' or '/'");
    }
 
    ArrayMeta meta;
    meta.format = spec.format;
    meta.shape = spec.shape;
    meta.chunk_shape = spec.chunks;
    meta.dtype = spec.dtype;
    meta.attributes = spec.attributes;
    apply_format_members(spec, meta, ctx);
    if (spec.fill) {
      if (spec.fill->size() != spec.dtype.itemsize) {
        throw error(ctx + ": fill is " + std::to_string(spec.fill->size()) +
                    " bytes, dtype needs " + std::to_string(spec.dtype.itemsize));
      }
      meta.fill = spec.fill;
    } else {
      meta.fill = Bytes(spec.dtype.itemsize, 0);  // canonical default: zeros
    }
    meta.codecs.push_back({"bytes", {{"endian", "little"}}});
    for (const CodecSpec& codec : spec.codecs) {
      meta.codecs.push_back(codec);
    }
 
    Array array(std::move(store), path, std::move(meta));  // resolves + validates codecs
    if (v3) {
      array.store_->write(v3::meta_key(path),
                          canonical_json_bytes(v3::emit_array_meta(array.meta_)));
    } else {
      v2::write_meta_key(*array.store_, array.meta_store_key(), v2::emit_array_meta(array.meta_));
      array.write_attributes();
    }
    return array;
  }
 
  [[nodiscard]] static Array open(std::shared_ptr<Store> store, const std::string& path,
                                  OpenOptions options = {}) {
    return open_impl(std::move(store), path, options, nullptr);
  }
 
  [[nodiscard]] const ArrayMeta& meta() const { return meta_; }
  [[nodiscard]] const std::string& path() const { return path_; }
  [[nodiscard]] std::vector<std::uint64_t> grid_shape() const { return meta_.grid_shape(); }
  [[nodiscard]] std::uint64_t nbytes() const {
    return meta_.element_count() * meta_.dtype.itemsize;
  }
  [[nodiscard]] std::uint64_t chunk_nbytes() const { return pipeline_.decoded_chunk_bytes(); }
 
  [[nodiscard]] Bytes read_chunk(const std::vector<std::uint64_t>& index) const {
    auto stored = chunk_store_->read(chunk_store_key(index));
    if (!stored) {
      return filled_chunk();
    }
    return pipeline_.decode(std::move(*stored));
  }
 
  void write_chunk(const std::vector<std::uint64_t>& index, const void* data, std::size_t size) {
    Bytes chunk(static_cast<const std::uint8_t*>(data),
                static_cast<const std::uint8_t*>(data) + size);
    chunk_store_->write(chunk_store_key(index), pipeline_.encode(std::move(chunk)));
    chunk_store_->flush();
  }
 
  [[nodiscard]] Bytes read_chunk_range(const std::vector<std::uint64_t>& index,
                                       std::uint64_t element_offset,
                                       std::uint64_t element_count) const {
    if (!pipeline_.supports_partial_read()) {
      throw error("byte-range chunk reads need an uncompressed, untransposed layout");
    }
    const std::uint32_t itemsize = meta_.dtype.itemsize;
    const std::uint64_t chunk_elements = meta_.chunk_element_count();
    if (element_count > chunk_elements || element_offset > chunk_elements - element_count) {
      throw error("chunk range [" + std::to_string(element_offset) + ", +" +
                  std::to_string(element_count) + ") exceeds " + std::to_string(chunk_elements) +
                  " elements");
    }
    auto stored = chunk_store_->read_range(
        chunk_store_key(index),
        ByteRange::slice(element_offset * itemsize, element_count * itemsize));
    if (!stored) {
      Bytes out(detail::checked_size(element_count * itemsize, "chunk range"));
      detail::fill_elements(out.data(), element_count, meta_.fill ? meta_.fill->data() : nullptr,
                            itemsize);
      return out;
    }
    return pipeline_.decode_range(std::move(*stored));
  }
 
  void read(void* dst, std::size_t size) const {
    read_region(std::vector<std::uint64_t>(meta_.shape.size(), 0), meta_.shape, dst, size);
  }
 
  void write(const void* src, std::size_t size) {
    write_region(std::vector<std::uint64_t>(meta_.shape.size(), 0), meta_.shape, src, size);
  }
 
  void read_region(const std::vector<std::uint64_t>& origin,
                   const std::vector<std::uint64_t>& shape, void* dst, std::size_t size) const {
    validate_region(origin, shape, size, "read_region");
    if (size == 0) {
      return;
    }
    auto* out = static_cast<std::uint8_t*>(dst);
    for_each_region_chunk(origin, shape, [&](const RegionChunk& rc) {
      const Bytes chunk = read_chunk(rc.index);
      detail::copy_box(chunk.data(), meta_.chunk_shape, rc.origin_in_chunk, out, shape,
                       rc.origin_in_region, rc.box, meta_.dtype.itemsize);
    });
  }
 
  void write_region(const std::vector<std::uint64_t>& origin,
                    const std::vector<std::uint64_t>& shape, const void* src, std::size_t size) {
    validate_region(origin, shape, size, "write_region");
    if (size == 0) {
      return;
    }
    const auto* in = static_cast<const std::uint8_t*>(src);
    for_each_region_chunk(origin, shape, [&](const RegionChunk& rc) {
      Bytes chunk = rc.covered ? filled_chunk() : read_chunk(rc.index);
      detail::copy_box(in, shape, rc.origin_in_region, chunk.data(), meta_.chunk_shape,
                       rc.origin_in_chunk, rc.box, meta_.dtype.itemsize);
      chunk_store_->write(chunk_store_key(rc.index), pipeline_.encode(std::move(chunk)));
    });
    chunk_store_->flush();
  }
 
  [[nodiscard]] const json& attributes() const { return meta_.attributes; }
 
  void set_attributes(json attributes) {
    meta_.attributes = std::move(attributes);
    if (meta_.format == ZarrFormat::v3) {
      const std::string key = v3::meta_key(path_);
      const auto bytes = store_->read(key);
      if (!bytes) {
        throw error(key + ": metadata disappeared");
      }
      json doc = detail::parse_json(*bytes, key);
      if (meta_.attributes.empty()) {
        doc.erase("attributes");
      } else {
        doc["attributes"] = meta_.attributes;
      }
      store_->write(key, canonical_json_bytes(doc));
      return;
    }
    write_attributes();
  }
 
  [[nodiscard]] std::string chunk_store_key(const std::vector<std::uint64_t>& index) const {
    const auto grid = meta_.grid_shape();
    if (index.size() != grid.size()) {
      throw error("chunk index rank " + std::to_string(index.size()) + " != array rank " +
                  std::to_string(grid.size()));
    }
    for (std::size_t d = 0; d < grid.size(); ++d) {
      if (index[d] >= grid[d]) {
        throw error("chunk index " + std::to_string(index[d]) + " out of range for dimension " +
                    std::to_string(d) + " (grid extent " + std::to_string(grid[d]) + ")");
      }
    }
    const std::string relative = meta_.key_encoding == ChunkKeyKind::v3_default
                                     ? v3::chunk_key(index, meta_.dimension_separator)
                                     : v2::chunk_key(index, meta_.dimension_separator);
    return path_.empty() ? relative : path_ + "/" + relative;
  }
 
 private:
  friend class Group;
 
  static Array open_impl(std::shared_ptr<Store> store, const std::string& path, OpenOptions options,
                         const std::shared_ptr<const json>& consolidated) {
    if (!store) {
      throw error("Array::open: null store");
    }
    detail::validate_path(path);
    const auto read_doc = [&](const std::string& key) -> std::optional<json> {
      if (consolidated) {
        const auto it = consolidated->find(key);
        if (it == consolidated->end()) {
          return std::nullopt;
        }
        return *it;
      }
      const auto bytes = store->read(key);
      if (!bytes) {
        return std::nullopt;
      }
      return detail::parse_json(*bytes, key);
    };
 
    // Probe order: zarr.json first, so v3 opens cost one round-trip.
    const std::string v3_key = v3::meta_key(path);
    if (const auto doc = read_doc(v3_key)) {
      if (doc->is_object() && doc->value("node_type", "") == std::string("group")) {
        throw error("'" + path + "' is a group, not an array");
      }
      return {std::move(store), path, v3::parse_array_meta(*doc, v3_key, options.lenient)};
    }
 
    const std::string meta_key = v2::meta_key(path, v2::kArraySuffix);
    auto doc = read_doc(meta_key);
    if (!doc) {
      if (store->exists(v2::meta_key(path, v2::kGroupSuffix))) {
        throw error("'" + path + "' is a group, not an array");
      }
      throw error("no array at '" + path + "' (neither " + v3_key + " nor " + meta_key + " found)");
    }
    ArrayMeta meta = v2::parse_array_meta(*doc, meta_key);
    if (const auto attrs = read_doc(v2::meta_key(path, v2::kAttrsSuffix))) {
      meta.attributes = *attrs;
    }
    return {std::move(store), path, std::move(meta)};
  }
 
  static void apply_format_members(const ArraySpec& spec, ArrayMeta& meta, const std::string& ctx) {
    if (spec.format != ZarrFormat::v3) {
      meta.dimension_separator = spec.dimension_separator;
      if (!spec.dimension_names.is_null()) {
        throw error(ctx + ": dimension_names is a v3 feature");
      }
      if (!spec.shards.empty()) {
        throw error(ctx + ": sharding is a v3 feature");
      }
      return;
    }
    // Canonical v3 creation: the "default" chunk-key encoding with '/'.
    meta.key_encoding = ChunkKeyKind::v3_default;
    meta.dimension_separator = '/';
    if (spec.dimension_names.is_array()) {
      if (spec.dimension_names.size() != spec.shape.size()) {
        throw error(ctx + ": dimension_names must have rank length");
      }
      meta.dimension_names = spec.dimension_names;
    }
    if (spec.shards.empty()) {
      return;
    }
    if (spec.shards.size() != spec.chunks.size()) {
      throw error(ctx + ": shards rank must match chunks rank");
    }
    for (std::size_t d = 0; d < spec.shards.size(); ++d) {
      // v3 sharding spec: chunks must evenly divide the shard.
      if (spec.shards[d] == 0 || spec.shards[d] % spec.chunks[d] != 0) {
        throw error(ctx + ": each shard extent must be a positive multiple of the chunk extent");
      }
    }
    ShardLevel level;
    level.shard_shape = spec.shards;
    level.index_codecs = {{"bytes", {{"endian", "little"}}}, {"crc32c", {}}};
    meta.shard_levels.push_back(std::move(level));
  }
 
  Array(std::shared_ptr<Store> store, std::string path, ArrayMeta meta)
      : store_(std::move(store)),
        path_(std::move(path)),
        meta_(std::move(meta)),
        pipeline_(CodecPipeline::resolve(meta_)),
        chunk_store_(wrap_shards(store_, meta_, path_)) {
    // Materializing a chunk must be possible on this platform (wasm32!).
    detail::checked_size(pipeline_.decoded_chunk_bytes(), "chunk");
  }
 
  static std::shared_ptr<Store> wrap_shards(std::shared_ptr<Store> store, const ArrayMeta& meta,
                                            const std::string& path) {
    std::shared_ptr<Store> chunks = std::move(store);
    const std::string prefix = path.empty() ? "" : path + "/";
    for (std::size_t i = 0; i < meta.shard_levels.size(); ++i) {
      chunks = std::make_shared<detail_shard::ShardStore>(
          std::move(chunks), detail_shard::params_for_level(meta, i, prefix));
    }
    return chunks;
  }
 
  [[nodiscard]] std::string meta_store_key() const { return v2::meta_key(path_, v2::kArraySuffix); }
 
  void write_attributes() {
    const std::string key = v2::meta_key(path_, v2::kAttrsSuffix);
    if (meta_.attributes.empty()) {
      v2::erase_meta_key(*store_, key);  // canonical: no empty .zattrs documents
    } else {
      v2::write_meta_key(*store_, key, meta_.attributes);
    }
  }
 
  [[nodiscard]] Bytes filled_chunk() const {
    Bytes chunk(detail::checked_size(pipeline_.decoded_chunk_bytes(), "chunk"));
    detail::fill_elements(chunk.data(), meta_.chunk_element_count(),
                          meta_.fill ? meta_.fill->data() : nullptr, meta_.dtype.itemsize);
    return chunk;
  }
 
  struct RegionChunk {
    std::vector<std::uint64_t> index;             
    std::vector<std::uint64_t> origin_in_chunk;   
    std::vector<std::uint64_t> origin_in_region;  
    std::vector<std::uint64_t> box;               
    bool covered = true;
  };
 
  void validate_region(const std::vector<std::uint64_t>& origin,
                       const std::vector<std::uint64_t>& shape, std::size_t size,
                       const char* what) const {
    const std::size_t rank = meta_.shape.size();
    if (origin.size() != rank || shape.size() != rank) {
      throw error(std::string(what) + ": origin/shape rank must be " + std::to_string(rank));
    }
    for (std::size_t d = 0; d < rank; ++d) {
      if (shape[d] > meta_.shape[d] || origin[d] > meta_.shape[d] - shape[d]) {
        throw error(std::string(what) + ": region [" + std::to_string(origin[d]) + ", " +
                    std::to_string(origin[d]) + "+" + std::to_string(shape[d]) +
                    ") exceeds dimension " + std::to_string(d) + " (extent " +
                    std::to_string(meta_.shape[d]) + ")");
      }
    }
    const std::uint64_t bytes = detail::checked_product(shape, what) * meta_.dtype.itemsize;
    if (size != detail::checked_size(bytes, what)) {
      throw error(std::string(what) + ": buffer is " + std::to_string(size) +
                  " bytes, region needs " + std::to_string(bytes));
    }
  }
 
  template <typename Fn>
  void for_each_region_chunk(const std::vector<std::uint64_t>& origin,
                             const std::vector<std::uint64_t>& shape, const Fn& fn) const {
    const std::size_t rank = meta_.shape.size();
    std::vector<std::uint64_t> first(rank, 0);
    std::vector<std::uint64_t> last(rank, 0);
    for (std::size_t d = 0; d < rank; ++d) {
      first[d] = origin[d] / meta_.chunk_shape[d];
      last[d] = (origin[d] + shape[d] - 1) / meta_.chunk_shape[d];
    }
 
    RegionChunk rc;
    rc.origin_in_chunk.assign(rank, 0);
    rc.origin_in_region.assign(rank, 0);
    rc.box.assign(rank, 0);
    visit_shard_major(0, first, last, [&](const std::vector<std::uint64_t>& index) {
      rc.index = index;
      rc.covered = true;
      for (std::size_t d = 0; d < rank; ++d) {
        const std::uint64_t chunk_start = rc.index[d] * meta_.chunk_shape[d];
        const std::uint64_t valid_end =
            std::min(chunk_start + meta_.chunk_shape[d], meta_.shape[d]);
        const std::uint64_t begin = std::max(chunk_start, origin[d]);
        const std::uint64_t end = std::min(valid_end, origin[d] + shape[d]);
        rc.origin_in_chunk[d] = begin - chunk_start;
        rc.origin_in_region[d] = begin - origin[d];
        rc.box[d] = end - begin;
        rc.covered = rc.covered && begin == chunk_start && end == valid_end;
      }
      fn(rc);
    });
  }
 
  template <typename Leaf>
  void visit_shard_major(std::size_t level, const std::vector<std::uint64_t>& lo,
                         const std::vector<std::uint64_t>& hi, const Leaf& leaf) const {
    const std::size_t rank = lo.size();
    if (level < meta_.shard_levels.size()) {
      // Level-`level` shard extents in chunk-grid units; shard_shape is a
      // validated multiple of chunk_shape.
      std::vector<std::uint64_t> per(rank, 1);
      std::vector<std::uint64_t> shard_lo(rank, 0);
      std::vector<std::uint64_t> shard_hi(rank, 0);
      for (std::size_t d = 0; d < rank; ++d) {
        per[d] = meta_.shard_levels[level].shard_shape[d] / meta_.chunk_shape[d];
        shard_lo[d] = lo[d] / per[d];
        shard_hi[d] = hi[d] / per[d];
      }
      std::vector<std::uint64_t> shard = shard_lo;
      std::vector<std::uint64_t> sub_lo(rank, 0);
      std::vector<std::uint64_t> sub_hi(rank, 0);
      while (true) {
        for (std::size_t d = 0; d < rank; ++d) {
          sub_lo[d] = std::max(lo[d], shard[d] * per[d]);
          sub_hi[d] = std::min(hi[d], shard[d] * per[d] + per[d] - 1);
        }
        visit_shard_major(level + 1, sub_lo, sub_hi, leaf);
        if (!detail::next_index_box(shard, shard_lo, shard_hi)) {
          return;
        }
      }
    }
    std::vector<std::uint64_t> index = lo;
    while (true) {
      leaf(index);
      if (!detail::next_index_box(index, lo, hi)) {
        return;
      }
    }
  }
 
  std::shared_ptr<Store> store_;
  std::string path_;
  ArrayMeta meta_;
  CodecPipeline pipeline_;
  std::shared_ptr<Store> chunk_store_;
};
 
}  // namespace zarr
 
#endif  // LIBZARR_ARRAY_HPP