H5Utils.h

Go to the documentation of this file.

// -*- C++ -*-
//
// This file is part of YODA -- Yet more Objects for Data Analysis
// Copyright (C) 2008-2023 The YODA collaboration (see AUTHORS for details)
//
#ifndef YODA_H5UTILS_H
#define YODA_H5UTILS_H
 
#include <YODA/Utils/BinningUtils.h>
#include <YODA/Config/BuildConfig.h>
#ifdef WITH_HIGHFIVE
#include <YODA/highfive/H5File.hpp>
#else
#include <highfive/H5File.hpp>
#define YODA_H5 HighFive
#endif
 
#include <map>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>
 
using std::map;
using std::string;
using std::vector;
 
namespace YODA {
 
  template <typename T>
  YODA_H5::DataSet H5DataSet(YODA_H5::File& h5file, const string& label,
                              vector<T>&& data, bool compress) {
 
    YODA_H5::DataSetCreateProps props;
    if (compress && data.size()) {
      props.add(YODA_H5::Chunking(vector<hsize_t>{data.size()}));
      props.add(YODA_H5::Deflate(9));
    }
    return h5file.createDataSet(label, std::move(data), props);
  }
 
 
  template <typename T>
  YODA_H5::DataSet H5DataSet(YODA_H5::File& h5file, const string& label,
                              size_t nrows, size_t chunksize, bool compress) {
 
    YODA_H5::DataSetCreateProps props;
    if (compress) { // enable compression
      props.add(YODA_H5::Chunking(vector<hsize_t>{1,chunksize}));
      props.add(YODA_H5::Deflate(9));
    }
    // create the dataset
    return h5file.createDataSet(label, YODA_H5::DataSpace({nrows,chunksize}),
                                YODA_H5::create_datatype<T>(), props);
  }
 
  class H5DataSetReader {
    public:
 
    H5DataSetReader() { }
 
    // @brief Constructor from existing DataSet
    H5DataSetReader(const YODA_H5::File& h5file, const string& label)
          : _nrows(0), _thisrow(0), _ncols(0), _thiscol(0), _ds(nullptr) {
      _ds = std::make_unique<YODA_H5::DataSet>(h5file.getDataSet(label));
      const auto& dims = _ds->getDimensions();
      _nrows = dims.at(0);
      if (dims.size() > 1)  _ncols = dims.at(1);
    }
 
    void skip(size_t len) noexcept { _thisrow += len; }
 
    template <typename T = size_t>
    T next() noexcept {
      T item;
      _ds->select({_thisrow++}, {1}).read(item);
      return item;
    }
 
    template <typename T>
    vector<T> read() noexcept {
      if (_nrows == 0)  return vector<T>{};
      vector<T> data; data.reserve(_nrows);
      _ds->select({0}, {_nrows}).read(data);
      return data;
    }
 
    template <typename T>
    vector<T> read(size_t len) noexcept {
      if (len == 0)  return {};
      vector<T> data; data.reserve(len);
      _ds->select({_thisrow}, {len}).read(data);
      _thisrow += len;
      return data;
    }
 
    template <typename T>
    vector<T> readSlice(size_t len) noexcept {
      if (len == 0)  return {};
      vector<T> data; data.reserve(len);
      while (len) {
        vector<vector<T>> tmp;
        size_t ncols = std::min(len, _ncols - _thiscol);
        size_t nrows = 1+(ncols-1)/_ncols; // C-style ceil
        _ds->select({_thisrow,_thiscol}, {nrows,ncols}).read(tmp);
        for (size_t i=0; i < tmp.size(); ++i) {
          data.insert(data.end(), std::make_move_iterator(std::begin(tmp[i])),
                                  std::make_move_iterator(std::end(tmp[i])));
        }
        if ((_thiscol + ncols) == _ncols)  ++_thisrow;
        _thiscol = (_thiscol + ncols) % _ncols;
        len -= ncols;
      }
      return data;
    }
 
    template <typename T>
    vector<T> readAt(size_t row, size_t len) noexcept {
      _thisrow = row;
      return read<T>(len);
    }
 
    template <typename T = size_t>
    vector<T> readAttribute(const string& label) const noexcept {
      vector<T> data;
      _ds->getAttribute(label).read(data);
      return data;
    }
 
    private:
 
    size_t _nrows, _thisrow, _ncols, _thiscol;
 
    std::unique_ptr<YODA_H5::DataSet> _ds;
 
  };
 
 
  template <typename T>
  class H5DataSetWriter {
    public:
 
    // Constructor for empty DataSet
    H5DataSetWriter(YODA_H5::File& h5file, const string& label,
                    size_t datalen, size_t chunksize, bool compress)
          : _ncols(chunksize), _thiscol(0), _nrows(1+(datalen-1)/chunksize), _thisrow(0),
            _ds(H5DataSet<T>(h5file, label, _nrows, chunksize, compress)) {  }
 
    void writeSlice(vector<T>&& data) noexcept {
 
      if (data.empty())  return;
 
      size_t offset = 0, len = data.size();
      const auto itr = data.cbegin();
      while (len) {
        vector<T> tmp;
        size_t ncols = std::min(len, _ncols - _thiscol);
        auto first = itr + offset;
        auto last = first + ncols;
        _ds.select({_thisrow,_thiscol},
                   {1,ncols}).write(vector<vector<T>>{{std::make_move_iterator(first),
                                                       std::make_move_iterator(last)}});
        offset += ncols;
        _thiscol += ncols;
        if (_thiscol == _ncols) {
          _thiscol = 0; ++_thisrow;
        }
        len -= ncols;
      }
    }
 
    template <typename U>
    void createAttribute(const string& label, vector<U>&& data) noexcept {
      _ds.createAttribute(label, std::forward<vector<U>>(data));
    }
 
    size_t colPos() const noexcept { return _thiscol; }
 
    size_t rowPos() const noexcept { return _thisrow; }
 
    private:
 
    size_t _ncols, _thiscol, _nrows, _thisrow;
 
    YODA_H5::DataSet _ds;
 
  };
 
 
 
 
  class EdgeHandlerBase {
    public:
 
    EdgeHandlerBase() { }
 
    virtual ~EdgeHandlerBase() { }
 
    virtual void writeToFile(const string&, YODA_H5::File&, bool compress) = 0;
 
  };
 
 
 
 
  template<typename T>
  class EdgeHandler : public EdgeHandlerBase {
    public:
 
    using EdgeHandlerBase::EdgeHandlerBase;
 
    // Constructor from a vector of edges
    EdgeHandler(const vector<T>& edges)
        : _edges(edges) { }
 
    // Constructor from an rvalue vector of edges
    EdgeHandler(vector<T>&& edges)
        : _edges(std::move(edges)) { }
 
    // Method to extend the vector of edges
    void extend(vector<T>&& edges) {
      _edges.insert(_edges.end(), std::make_move_iterator(std::begin(edges)),
                                  std::make_move_iterator(std::end(edges)));
    }
 
    auto begin() const { return _edges.cbegin(); }
 
    // Method to commit vector of edges to H5 @a file using @a label
    void writeToFile(const string& label, YODA_H5::File& file, bool compress) {
      (void)H5DataSet(file, label, std::move(_edges), compress);
    }
 
    private:
 
    vector<T> _edges;
  };
 
  template <typename T>
  using EdgeHandlerPtr = typename std::shared_ptr<EdgeHandler<T>>;
 
  using EdgeHandlerBasePtr = typename std::shared_ptr<EdgeHandlerBase>;
 
 
 
 
  class H5FileManager {
 
    public:
 
    static constexpr size_t AO_META = 2;
 
    H5FileManager(const YODA_H5::File& file)
        : _index(-1), _cachepos(-1),
          _labelindex(0), _h5file(file),
          _aoinfo(file, "aoinfo"),
          _layout(file, "sizeinfo"),
          _content(file, "content"),
          _annos(file, "annotations"),
          _labels(H5DataSetReader(file, "labels").read<string>()) {
 
      _meta = _aoinfo.readAttribute("meta");
      if (_meta.size() < 2)
        throw ReadError("No file metadata found!");
 
      _annosizes = _layout.read<size_t>(_meta.at(1));
      _datasizes = _layout.read<size_t>(_meta.at(1));
      _labelsizes = _layout.read<size_t>(_layout.next());
    }
 
    size_t version() const {
      return _meta.at(0);
    }
 
    size_t size() const {
      return _meta.at(1);
    }
 
    bool next() {
      ++_index;
      if ((size_t)_index == size())  return false;
 
      _aodims = _aoinfo.readAt<string>(AO_META*_index, AO_META);
      if (_aodims.empty()) {
        throw ReadError("No AO information found!");
      }
 
      _cachepos = -1;
 
      return true;
    }
 
    const string& path() const {
      return _aodims.at(0);
    }
 
    const string& type() const {
      return _aodims.at(1);
    }
 
    vector<string> loadAnnotations() noexcept {
      return _annos.readSlice<string>(_annosizes[_index]);
    }
 
    vector<double> loadContent() noexcept {
      return _content.readSlice<double>(_datasizes[_index]);
    }
 
    void skipCommon() noexcept {
      _annos.skip(_annosizes[_index]);
      _content.skip(_datasizes[_index]);
    }
 
    vector<size_t> loadMasks() noexcept {
      return _layout.read<size_t>(_layout.next());
    }
 
    void skipMasks() noexcept {
      _layout.skip(_layout.next());
    }
 
    vector<string> loadSources() noexcept {
 
      if (_cachepos < 0) { // fill label cache
        _labelcache = _layout.read<size_t>(_labelsizes[_labelindex++]);
        _cachepos = 0;
      }
 
      size_t len = _labelcache[_cachepos++];
      if (len == 0) return {};
 
      vector<size_t> indices(_labelcache.begin()+_cachepos, _labelcache.begin()+_cachepos+len);
      vector<string> sources; sources.reserve(indices.size());
      for (size_t idx : indices) {
        sources.emplace_back(_labels[idx]);
      }
      _cachepos += len;
      return sources;
    }
 
    void skipSources() noexcept {
      _layout.skip(_labelsizes[_labelindex++]);
    }
 
    template <typename EdgeT>
    vector<EdgeT> loadEdges() noexcept {
 
      const string label = string("edges_") + TypeID<EdgeT>::name();
      auto itr = datasets.find(label);
      if (itr == datasets.end()) {
        // put H5::DataSet into DataSet cache
        datasets[label] = H5DataSetReader(_h5file, label);
        itr = datasets.find(label);
      }
      return itr->second.read<EdgeT>(_layout.next());
    }
 
    template <typename EdgeT>
    void skipEdges() noexcept {
      const string label = string("edges_") + TypeID<EdgeT>::name();
      auto itr = datasets.find(label);
      if (itr == datasets.end()) {
        // put H5::DataSet into DataSet cache
        datasets[label] = H5DataSetReader(_h5file, label);
        itr = datasets.find(label);
      }
      itr->second.skip(_layout.next());
    }
 
    private:
 
    ssize_t _index, _cachepos;
 
    size_t _labelindex;
 
    const YODA_H5::File _h5file;
 
    H5DataSetReader _aoinfo, _layout, _content, _annos;
 
    vector<string> _labels, _aodims;
 
    vector<size_t> _meta, _annosizes, _datasizes, _labelsizes, _labelcache;
 
    map<string,H5DataSetReader> datasets;
  };
 
 
}
 
#endif