sharding.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: MIT
 
#ifndef LIBZARR_SHARDING_HPP
#define LIBZARR_SHARDING_HPP
 
#include <cstdint>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <vector>
 
#include "libzarr/codecs.hpp"
#include "libzarr/detail/common.hpp"
#include "libzarr/metadata.hpp"
#include "libzarr/store.hpp"
#include "libzarr/types.hpp"
#include "libzarr/v2.hpp"
#include "libzarr/v3.hpp"
 
 
namespace zarr::detail_shard {
 
inline constexpr std::uint64_t kSentinel = std::numeric_limits<std::uint64_t>::max();
 
struct IndexEntry {
  std::uint64_t offset = kSentinel;
  std::uint64_t nbytes = kSentinel;
  [[nodiscard]] bool missing() const { return offset == kSentinel && nbytes == kSentinel; }
};
 
inline std::uint64_t index_encoded_size(const std::vector<CodecSpec>& index_codecs,
                                        std::uint64_t entry_count, const std::string& ctx) {
  std::uint64_t size = entry_count * 2 * 8;
  bool have_bytes = false;
  for (const CodecSpec& codec : index_codecs) {
    if (codec.name == "bytes") {
      have_bytes = true;
    } else if (codec.name == "crc32c") {
      size += 4;
    } else {
      throw error(ctx +
                  ": index_codecs may only contain 'bytes' and 'crc32c' (the index must "
                  "have a fixed encoded size), got '" +
                  codec.name + "'");
    }
  }
  if (!have_bytes) {
    throw error(ctx + ": index_codecs must contain the 'bytes' codec");
  }
  return size;
}
 
struct ShardParams {
  std::string chunk_prefix;
  ChunkKeyKind key_encoding = ChunkKeyKind::v3_default;
  char separator = '/';
  std::vector<std::uint64_t> per_shard;
  std::vector<std::uint64_t> inner_grid;
  std::vector<CodecSpec> index_codecs;
  bool index_at_end = true;
};
 
inline ShardParams params_for_level(const ArrayMeta& meta, std::size_t level,
                                    const std::string& prefix) {
  const ShardLevel& lvl = meta.shard_levels[level];
  const std::vector<std::uint64_t>& inner_shape = level + 1 < meta.shard_levels.size()
                                                      ? meta.shard_levels[level + 1].shard_shape
                                                      : meta.chunk_shape;
  ShardParams params;
  params.chunk_prefix = prefix;
  params.key_encoding = meta.key_encoding;
  params.separator = meta.dimension_separator;
  params.index_codecs = lvl.index_codecs;
  params.index_at_end = lvl.index_at_end;
  params.per_shard.resize(inner_shape.size());
  params.inner_grid.resize(inner_shape.size());
  for (std::size_t d = 0; d < inner_shape.size(); ++d) {
    params.per_shard[d] = lvl.shard_shape[d] / inner_shape[d];
    params.inner_grid[d] = detail::ceil_div(meta.shape[d], inner_shape[d]);
  }
  return params;
}
 
inline void locate_index(const ShardParams& params, const std::vector<std::uint64_t>& index,
                         std::string& shard_key, std::uint64_t& slot) {
  const std::size_t rank = params.per_shard.size();
  std::vector<std::uint64_t> outer(rank, 0);
  slot = 0;
  for (std::size_t d = 0; d < rank; ++d) {
    outer[d] = index[d] / params.per_shard[d];
    slot = slot * params.per_shard[d] + index[d] % params.per_shard[d];
  }
  shard_key = params.chunk_prefix + (params.key_encoding == ChunkKeyKind::v3_default
                                         ? v3::chunk_key(outer, params.separator)
                                         : v2::chunk_key(outer, params.separator));
}
 
inline ArrayMeta index_array_meta(std::uint64_t entry_count,
                                  const std::vector<CodecSpec>& index_codecs) {
  ArrayMeta meta;
  meta.shape = {entry_count * 2};
  meta.chunk_shape = {entry_count * 2};
  meta.dtype = DataType::of(DType::uint64);
  meta.codecs = index_codecs;
  return meta;
}
 
inline std::vector<IndexEntry> decode_index(const CodecPipeline& pipeline,
                                            std::uint64_t entry_count, Bytes stored,
                                            const std::string& ctx) {
  const Bytes decoded = pipeline.decode(std::move(stored));
  if (decoded.size() != entry_count * 16) {
    throw error(ctx + ": corrupt shard index");
  }
  std::vector<IndexEntry> entries(static_cast<std::size_t>(entry_count));
  for (std::size_t i = 0; i < entries.size(); ++i) {
    std::memcpy(&entries[i].offset, decoded.data() + i * 16, 8);
    std::memcpy(&entries[i].nbytes, decoded.data() + i * 16 + 8, 8);
    const bool sentinel_mismatch =
        (entries[i].offset == kSentinel) != (entries[i].nbytes == kSentinel);
    const bool overflow =
        !entries[i].missing() && entries[i].nbytes > kSentinel - entries[i].offset;
    if (sentinel_mismatch || overflow) {
      throw error(ctx + ": corrupt shard index");
    }
  }
  return entries;
}
 
inline Bytes assemble_shard(const std::vector<std::optional<Bytes>>& entries,
                            const std::vector<CodecSpec>& index_codecs, bool index_at_end,
                            const std::string& ctx) {
  bool any = false;
  for (const auto& entry : entries) {
    any = any || entry.has_value();
  }
  if (!any) {
    return {};  // all-fill shard: nothing to store
  }
  const std::uint64_t entry_count = entries.size();
  const std::uint64_t index_size = index_encoded_size(index_codecs, entry_count, ctx);
  const CodecPipeline index_pipeline =
      CodecPipeline::resolve(index_array_meta(entry_count, index_codecs));
 
  Bytes body;
  std::vector<std::uint64_t> raw_index(entries.size() * 2, kSentinel);
  const std::uint64_t base = index_at_end ? 0 : index_size;  // chunks follow a leading index
  for (std::size_t i = 0; i < entries.size(); ++i) {
    const std::optional<Bytes>& entry = entries[i];
    if (!entry) {
      continue;
    }
    raw_index[i * 2] = base + body.size();
    raw_index[i * 2 + 1] = entry->size();
    body.insert(body.end(), entry->begin(), entry->end());
  }
  Bytes index_bytes(raw_index.size() * 8);
  std::memcpy(index_bytes.data(), raw_index.data(), index_bytes.size());
  index_bytes = index_pipeline.encode(std::move(index_bytes));
 
  Bytes shard;
  shard.reserve(body.size() + index_bytes.size());
  if (index_at_end) {
    shard = std::move(body);
    shard.insert(shard.end(), index_bytes.begin(), index_bytes.end());
  } else {
    shard = std::move(index_bytes);
    shard.insert(shard.end(), body.begin(), body.end());
  }
  return shard;
}
 
class ShardStore final : public Store {
 public:
  ShardStore(std::shared_ptr<Store> source, ShardParams params)
      : source_(std::move(source)),
        params_(std::move(params)),
        entry_count_(detail::checked_product(params_.per_shard, "shard grid")),
        index_size_(detail_shard::index_encoded_size(params_.index_codecs, entry_count_,
                                                     "sharding_indexed")),
        index_pipeline_(CodecPipeline::resolve(index_meta())) {
    if (!source_) {
      throw error("ShardStore: null store");
    }
  }
 
  ~ShardStore() override = default;
  ShardStore(const ShardStore&) = delete;
  ShardStore& operator=(const ShardStore&) = delete;
  ShardStore(ShardStore&&) = delete;
  ShardStore& operator=(ShardStore&&) = delete;
 
  [[nodiscard]] std::optional<Bytes> read(std::string_view key) override {
    return read_range(key, ByteRange::full());
  }
 
  [[nodiscard]] std::optional<Bytes> read_range(std::string_view key, ByteRange range) override {
    const Location loc = locate(key);
    if (assembly_ && assembly_->shard_key == loc.shard_key) {
      // Serve pending writes so read-modify-write sequences stay coherent.
      const auto& entry = assembly_->entries[loc.intra];
      if (!entry) {
        return std::nullopt;
      }
      return slice_value(*entry, range, key);
    }
    const auto* index = load_index(loc.shard_key);
    if (index == nullptr || (*index)[loc.intra].missing()) {
      return std::nullopt;
    }
    const detail_shard::IndexEntry entry = (*index)[loc.intra];
    std::uint64_t begin = 0;
    std::uint64_t count = entry.nbytes;
    resolve_range(range, entry.nbytes, begin, count, key);
    return source_->read_range(loc.shard_key, ByteRange::slice(entry.offset + begin, count));
  }
 
  [[nodiscard]] std::optional<std::uint64_t> size(std::string_view key) override {
    const Location loc = locate(key);
    if (assembly_ && assembly_->shard_key == loc.shard_key) {
      const auto& entry = assembly_->entries[loc.intra];
      if (!entry) {
        return std::nullopt;
      }
      return entry->size();
    }
    const auto* index = load_index(loc.shard_key);
    if (index == nullptr || (*index)[loc.intra].missing()) {
      return std::nullopt;
    }
    return (*index)[loc.intra].nbytes;
  }
 
  [[nodiscard]] bool exists(std::string_view key) override { return size(key).has_value(); }
 
  void write(std::string_view key, Bytes value) override {
    put(locate(key), std::optional<Bytes>(std::move(value)));
  }
 
  void erase(std::string_view key) override { put(locate(key), std::nullopt); }
 
  void flush() override {
    flush_assembly();
    source_->flush();
  }
 
  [[nodiscard]] std::vector<std::string> list_prefix(std::string_view /*prefix*/) override {
    throw error("ShardStore is an internal adapter; listing is not supported");
  }
  [[nodiscard]] DirListing list_dir(std::string_view /*prefix*/) override {
    throw error("ShardStore is an internal adapter; listing is not supported");
  }
 
 private:
  struct Location {
    std::string shard_key;
    std::size_t intra = 0;  // C-order position within the shard
  };
 
  struct Assembly {
    std::string shard_key;
    std::vector<std::optional<Bytes>> entries;
    bool dirty = false;
  };
 
  [[nodiscard]] ArrayMeta index_meta() const {
    return index_array_meta(entry_count_, params_.index_codecs);
  }
 
  static void resolve_range(ByteRange range, std::uint64_t size, std::uint64_t& begin,
                            std::uint64_t& count, std::string_view key) {
    if (range.kind == ByteRange::Kind::slice) {
      if (range.length > size || range.offset > size - range.length) {
        throw error("read_range: slice out of bounds for inner chunk '" + std::string(key) + "' (" +
                    std::to_string(size) + " bytes)");
      }
      begin = range.offset;
      count = range.length;
    } else if (range.kind == ByteRange::Kind::suffix) {
      if (range.length > size) {
        throw error("read_range: suffix out of bounds for inner chunk '" + std::string(key) +
                    "' (" + std::to_string(size) + " bytes)");
      }
      begin = size - range.length;
      count = range.length;
    }
  }
 
  [[nodiscard]] static std::optional<Bytes> slice_value(const Bytes& value, ByteRange range,
                                                        std::string_view key) {
    std::uint64_t begin = 0;
    std::uint64_t count = value.size();
    resolve_range(range, value.size(), begin, count, key);
    const auto first = value.begin() + static_cast<std::ptrdiff_t>(begin);
    return Bytes(first, first + static_cast<std::ptrdiff_t>(count));
  }
 
  [[nodiscard]] Location locate(std::string_view key) const {
    const std::size_t rank = params_.per_shard.size();
    std::string_view rest = key;
    if (!detail::starts_with(rest, params_.chunk_prefix)) {
      throw error("ShardStore: key '" + std::string(key) + "' is outside the array's chunks");
    }
    rest = rest.substr(params_.chunk_prefix.size());
    if (params_.key_encoding == ChunkKeyKind::v3_default) {
      if (rest.empty() || rest[0] != 'c' || (rest.size() > 1 && rest[1] != params_.separator)) {
        throw error("ShardStore: malformed chunk key '" + std::string(key) + "'");
      }
      rest = rest.size() > 1 ? rest.substr(2) : rest.substr(1);
    }
    std::vector<std::uint64_t> index(rank, 0);
    std::size_t pos = 0;
    for (std::size_t d = 0; d < rank; ++d) {
      std::uint64_t value = 0;
      const std::size_t start = pos;
      while (pos < rest.size() && rest[pos] >= '0' && rest[pos] <= '9') {
        value = value * 10 + static_cast<std::uint64_t>(rest[pos] - '0');
        ++pos;
      }
      if (pos == start || value >= params_.inner_grid[d]) {
        throw error("ShardStore: malformed chunk key '" + std::string(key) + "'");
      }
      index[d] = value;
      if (d + 1 < rank) {
        if (pos >= rest.size() || rest[pos] != params_.separator) {
          throw error("ShardStore: malformed chunk key '" + std::string(key) + "'");
        }
        ++pos;
      }
    }
    if (pos != rest.size()) {
      throw error("ShardStore: malformed chunk key '" + std::string(key) + "'");
    }
 
    Location loc;
    std::uint64_t slot = 0;
    locate_index(params_, index, loc.shard_key, slot);
    loc.intra = detail::checked_size(slot, "shard entry");
    return loc;
  }
 
  [[nodiscard]] const std::vector<detail_shard::IndexEntry>* load_index(
      const std::string& shard_key) {
    for (std::size_t i = 0; i < cache_.size(); ++i) {
      if (cache_[i].first == shard_key) {
        if (i != 0) {
          std::rotate(cache_.begin(), cache_.begin() + static_cast<std::ptrdiff_t>(i),
                      cache_.begin() + static_cast<std::ptrdiff_t>(i) + 1);
        }
        return &cache_.front().second;
      }
    }
    auto stored =
        source_->read_range(shard_key, params_.index_at_end ? ByteRange::suffix(index_size_)
                                                            : ByteRange::slice(0, index_size_));
    if (!stored) {
      return nullptr;
    }
    std::vector<detail_shard::IndexEntry> entries =
        decode_index(index_pipeline_, entry_count_, std::move(*stored), shard_key);
    cache_.insert(cache_.begin(), {shard_key, std::move(entries)});
    if (cache_.size() > kCacheCapacity) {
      cache_.pop_back();
    }
    return &cache_.front().second;
  }
 
  void put(const Location& loc, std::optional<Bytes> value) {
    if (assembly_ && assembly_->shard_key != loc.shard_key) {
      flush_assembly();  // early flush: writes moved on to another shard
    }
    if (!assembly_) {
      assembly_ = load_assembly(loc.shard_key);
    }
    assembly_->entries[loc.intra] = std::move(value);
    assembly_->dirty = true;
  }
 
  [[nodiscard]] Assembly load_assembly(const std::string& shard_key) {
    Assembly assembly;
    assembly.shard_key = shard_key;
    assembly.entries.assign(static_cast<std::size_t>(entry_count_), std::nullopt);
    const auto stored = source_->read(shard_key);
    if (stored) {
      const auto* index = load_index(shard_key);
      assert(index != nullptr);
      for (std::size_t i = 0; i < index->size(); ++i) {
        const detail_shard::IndexEntry entry = (*index)[i];
        if (entry.missing()) {
          continue;
        }
        if (entry.offset > stored->size() || entry.nbytes > stored->size() - entry.offset) {
          throw error(shard_key + ": shard index points outside the shard");
        }
        const auto first = stored->begin() + static_cast<std::ptrdiff_t>(entry.offset);
        assembly.entries[i] = Bytes(first, first + static_cast<std::ptrdiff_t>(entry.nbytes));
      }
    }
    return assembly;
  }
 
  void flush_assembly() {
    if (!assembly_) {
      return;
    }
    const Assembly assembly = *std::move(assembly_);
    assembly_.reset();
    if (!assembly.dirty) {
      return;
    }
    drop_cached(assembly.shard_key);
    Bytes shard = assemble_shard(assembly.entries, params_.index_codecs, params_.index_at_end,
                                 assembly.shard_key);
    if (shard.empty()) {
      source_->erase(assembly.shard_key);  // all-fill shards are not stored
    } else {
      source_->write(assembly.shard_key, std::move(shard));
    }
  }
 
  void drop_cached(const std::string& shard_key) {
    for (std::size_t i = 0; i < cache_.size(); ++i) {
      if (cache_[i].first == shard_key) {
        cache_.erase(cache_.begin() + static_cast<std::ptrdiff_t>(i));
        return;
      }
    }
  }
 
  static constexpr std::size_t kCacheCapacity = 16;
 
  std::shared_ptr<Store> source_;
  ShardParams params_;
  std::uint64_t entry_count_;
  std::uint64_t index_size_;
  CodecPipeline index_pipeline_;
  std::vector<std::pair<std::string, std::vector<detail_shard::IndexEntry>>> cache_;
  std::optional<Assembly> assembly_;
};
 
}  // namespace zarr::detail_shard
 
namespace zarr::shard {
 
struct Placement {
  std::string shard_key;
  std::uint64_t slot = 0;
  std::uint64_t index_size = 0;
  bool index_at_end = true;
};
 
struct Extent {
  std::uint64_t offset = 0;
  std::uint64_t nbytes = 0;
  bool missing = false;
};
 
[[nodiscard]] inline Placement place(const ArrayMeta& meta, const std::string& path,
                                     const std::vector<std::uint64_t>& inner_index,
                                     std::size_t level = 0) {
  if (level >= meta.shard_levels.size()) {
    throw error("zarr::shard::place: level " + std::to_string(level) + " is not a shard level (" +
                std::to_string(meta.shard_levels.size()) + " levels)");
  }
  const std::string prefix = path.empty() ? "" : path + "/";
  const detail_shard::ShardParams params = detail_shard::params_for_level(meta, level, prefix);
  if (inner_index.size() != params.per_shard.size()) {
    throw error("zarr::shard::place: inner_index rank " + std::to_string(inner_index.size()) +
                " != array rank " + std::to_string(params.per_shard.size()));
  }
  for (std::size_t d = 0; d < inner_index.size(); ++d) {
    if (inner_index[d] >= params.inner_grid[d]) {
      throw error("zarr::shard::place: inner_index[" + std::to_string(d) +
                  "] = " + std::to_string(inner_index[d]) + " out of range");
    }
  }
  Placement out;
  detail_shard::locate_index(params, inner_index, out.shard_key, out.slot);
  const std::uint64_t entry_count = detail::checked_product(params.per_shard, "shard grid");
  out.index_size =
      detail_shard::index_encoded_size(params.index_codecs, entry_count, "sharding_indexed");
  out.index_at_end = params.index_at_end;
  return out;
}
 
[[nodiscard]] inline Extent extent(const ArrayMeta& meta, const Bytes& index_bytes,
                                   std::uint64_t slot, std::size_t level = 0) {
  if (level >= meta.shard_levels.size()) {
    throw error("zarr::shard::extent: level " + std::to_string(level) + " is not a shard level");
  }
  const detail_shard::ShardParams params = detail_shard::params_for_level(meta, level, "");
  const std::uint64_t entry_count = detail::checked_product(params.per_shard, "shard grid");
  if (slot >= entry_count) {
    throw error("zarr::shard::extent: slot " + std::to_string(slot) + " out of range (" +
                std::to_string(entry_count) + " entries)");
  }
  const CodecPipeline pipeline =
      CodecPipeline::resolve(detail_shard::index_array_meta(entry_count, params.index_codecs));
  const std::vector<detail_shard::IndexEntry> entries =
      detail_shard::decode_index(pipeline, entry_count, index_bytes, "sharding_indexed");
  const detail_shard::IndexEntry& e = entries[static_cast<std::size_t>(slot)];
  Extent out;
  if (e.missing()) {
    out.missing = true;
  } else {
    out.offset = e.offset;
    out.nbytes = e.nbytes;
  }
  return out;
}
 
[[nodiscard]] inline Bytes pack(const ArrayMeta& meta,
                                const std::vector<std::optional<Bytes>>& entries,
                                std::size_t level = 0) {
  if (level >= meta.shard_levels.size()) {
    throw error("zarr::shard::pack: level " + std::to_string(level) + " is not a shard level");
  }
  const detail_shard::ShardParams params = detail_shard::params_for_level(meta, level, "");
  const std::uint64_t entry_count = detail::checked_product(params.per_shard, "shard grid");
  if (entries.size() != entry_count) {
    throw error("zarr::shard::pack: got " + std::to_string(entries.size()) +
                " entries but shard has " + std::to_string(entry_count) + " slots");
  }
  return detail_shard::assemble_shard(entries, params.index_codecs, params.index_at_end,
                                      "sharding_indexed");
}
 
}  // namespace zarr::shard
 
#endif  // LIBZARR_SHARDING_HPP