Inject dataset reading via Hdf5Reader.

CMakeLists.txt

@@ -96,6 +96,7 @@ set(SONATA_SRC
     src/edge_index.cpp
     src/edges.cpp
     src/hdf5_mutex.cpp
+    src/hdf5_reader.cpp
     src/node_sets.cpp
     src/nodes.cpp
     src/population.cpp

include/bbp/sonata/config.h

@@ -182,8 +182,8 @@ class SONATA_API CircuitConfig
      *
      * \throws SonataError if the given population does not exist in any node network.
      */
+    NodePopulation getNodePopulation(const std::string& name, const Hdf5Reader& hdf5_reader) const;
     NodePopulation getNodePopulation(const std::string& name) const;
-
     /**
      * Returns a set with all available population names across all the edge networks.
      */
@@ -195,6 +195,7 @@ class SONATA_API CircuitConfig
      *
      * \throws SonataError if the given population does not exist in any edge network.
      */
+    EdgePopulation getEdgePopulation(const std::string& name, const Hdf5Reader& hdf5_reader) const;
     EdgePopulation getEdgePopulation(const std::string& name) const;
 
     /**

include/bbp/sonata/edges.h

@@ -15,6 +15,8 @@
 #include <string>
 #include <vector>
 
+#include "hdf5_reader.h"
+
 
 namespace bbp {
 namespace sonata {
@@ -30,6 +32,11 @@ class SONATA_API EdgePopulation: public Population
                    const std::string& csvFilePath,
                    const std::string& name);
 
+    EdgePopulation(const std::string& h5FilePath,
+                   const std::string& csvFilePath,
+                   const std::string& name,
+                   const Hdf5Reader& hdf5_reader);
+
     /**
      * Name of source population extracted from 'source_node_id' dataset
      */

include/bbp/sonata/hdf5_reader.h

@@ -0,0 +1,183 @@
+#pragma once
+
+#include <tuple>
+#include <vector>
+
+#include <bbp/sonata/selection.h>
+#include <highfive/H5File.hpp>
+
+namespace bbp {
+namespace sonata {
+
+/// Interface for implementing `readSelection<T>(dset, selection)`.
+template <class T>
+class Hdf5PluginRead1DInterface
+{
+  public:
+    virtual ~Hdf5PluginRead1DInterface() = default;
+
+    /// Read the selected subset of the one-dimensional array.
+    ///
+    /// The selection is canonical, i.e. sorted and non-overlapping. The dataset
+    /// is obtained from a `HighFive::File` opened via `this->openFile`.
+    virtual std::vector<T> readSelection(const HighFive::DataSet& dset,
+                                         const Selection& selection) const = 0;
+};
+
+template <class T>
+class Hdf5PluginRead2DInterface
+{
+  public:
+    virtual ~Hdf5PluginRead2DInterface() = default;
+
+    /// Read the Cartesian product of the two selections.
+    ///
+    /// Both selections are canonical, i.e. sorted and non-overlapping. The dataset
+    /// is obtained from a `HighFive::File` opened via `this->openFile`.
+    virtual std::vector<std::array<uint64_t, 2>> readSelection(const HighFive::DataSet& dset,
+                                                               const Selection& xsel,
+                                                               const Selection& ysel) const = 0;
+};
+
+template <class T, class U>
+class Hdf5PluginInterface;
+
+/// Interface of Plugins for reading HDF5 datasets.
+///
+/// All method must be called in an MPI-collective manner. Each method is free
+/// to break any MPI collective requirements.
+template <class... Ts, class... Us>
+class Hdf5PluginInterface<std::tuple<Ts...>, std::tuple<Us...>>
+    : virtual public Hdf5PluginRead1DInterface<Ts>...,
+      virtual public Hdf5PluginRead2DInterface<Us>...
+{
+  public:
+    /// Open the HDF5 file.
+    ///
+    /// This allows setting File Access Properties.
+    virtual HighFive::File openFile(const std::string& path) const = 0;
+};
+
+/// Abstraction for reading HDF5 datasets.
+///
+/// The Hdf5Reader provides an interface for reading canonical selections from
+/// datasets. Selections are canonical if they are sorted and don't overlap.
+/// This allows implementing different optimization strategies, such as
+/// minimizing bytes read, aggregating nearby reads or using MPI collective I/O.
+///
+/// The design uses virtual inheritance, which enables users to inject their own
+/// reader if needed. This class is the interface used within libsonata. It
+/// simply delegates to a "plugin", that satisfies the interface
+/// `Hdf5PluginInterface`.
+///
+/// To enable MPI collective I/O, `libsonata` must call all methods in an
+/// MPI-collective manner. This implies that the number of times any function in
+/// `libsonata` calls any of the `Hdf5Reader` methods must not depend on the
+/// arguments to the function.
+///
+/// Examples:
+///
+///    void wrong(Selection selection) {
+///      // Wrong because some MPI ranks might return without
+///      // calling `readSelection`.
+///      if(selection.empty()) {
+///        return;
+///      }
+///      hdf5_reader.readSelection(dset, selection);
+///    }
+///
+///    void also_wrong(Selection selection) {
+///      // Wrong because `hdf5_reader` is called `selection.ranges().size()`
+///      // number of times. Which could be different on each MPI rank.
+///      for(auto range : selection.ranges()) {
+///        hdf5_reader.readSelection(dset, Selection(std::vector{range}));
+///      }
+///    }
+///
+///    void subtle(Selection selection, bool flag) {
+///      // If the flag can differ between MPI ranks, this is wrong because
+///      // `readSelection` is called with different `dset`s. If the `flag` must
+///      // be the same on all MPI ranks, this is correct. If this happens in
+///      // the libsonata API, then passing the same `flag` on all MPI ranks becomes
+///      // a requirement for the users, when using a collective reader. Example:
+///      //     pop.get_attribute(attr_name, selection)
+///      if(flag) {
+///        hdf5_reader.readSelection(dset1, selection);
+///      } else {
+///        hdf5_reader.readSelection(dset2, selection);
+///      }
+///    }
+///
+///    void correct(Selection selection) {
+///      // Correct because no matter which branch is taken
+///      // `hdf5_reader.readSelection` is called exactly once.
+///      if(selection.size % 2 == 0) {
+///        hdf5_reader.readSelection(dset, selection);
+///      } else {
+///        hdf5_reader.readSelection(dset, {});
+///      }
+///    }
+///
+class SONATA_API Hdf5Reader
+{
+  public:
+    // The issue here is that on a mac `size_t` is different from
+    // `{,u}int{8,16,32,64}_t` but not on the other two OSes.
+    using supported_1D_types = std::tuple<uint8_t,
+                                          uint16_t,
+                                          uint32_t,
+                                          uint64_t,
+                                          int8_t,
+                                          int16_t,
+                                          int32_t,
+                                          int64_t,
+                                          float,
+                                          double,
+#ifdef __APPLE__
+                                          size_t,
+#endif
+                                          std::string>;
+
+    using supported_2D_types = std::tuple<std::array<uint64_t, 2>>;
+
+    /// Create a valid Hdf5Reader with the default plugin.
+    Hdf5Reader();
+
+    /// Create an Hdf5Reader with a user supplied plugin.
+    Hdf5Reader(std::shared_ptr<Hdf5PluginInterface<supported_1D_types, supported_2D_types>> impl);
+
+    /// Read the selected subset of the one-dimensional array.
+    ///
+    /// Both selections are canonical, i.e. sorted and non-overlapping. The dataset
+    /// is obtained from a `HighFive::File` opened via `this->openFile`.
+    template <class T>
+    std::vector<T> readSelection(const HighFive::DataSet& dset, const Selection& selection) const {
+        return static_cast<const Hdf5PluginRead1DInterface<T>&>(*impl).readSelection(dset,
+                                                                                     selection);
+    }
+
+    /// Open the HDF5.
+    ///
+    /// The dataset passed to `readSelection` must be obtained from a file open
+    /// via this method.
+    HighFive::File openFile(const std::string& filename) const;
+
+    /// Read the Cartesian product of the two selections.
+    ///
+    /// Both selections are canonical, i.e. sorted and non-overlapping. The dataset
+    /// is obtained from a `HighFive::File` opened via `this->openFile`.
+    template <class T>
+    std::vector<T> readSelection(const HighFive::DataSet& dset,
+                                 const Selection& xsel,
+                                 const Selection& ysel) const {
+        return static_cast<const Hdf5PluginRead2DInterface<T>&>(*impl).readSelection(dset,
+                                                                                     xsel,
+                                                                                     ysel);
+    }
+
+  private:
+    std::shared_ptr<Hdf5PluginInterface<supported_1D_types, supported_2D_types>> impl;
+};
+
+}  // namespace sonata
+}  // namespace bbp

include/bbp/sonata/nodes.h

@@ -30,6 +30,11 @@ class SONATA_API NodePopulation: public Population
                    const std::string& csvFilePath,
                    const std::string& name);
 
+    NodePopulation(const std::string& h5FilePath,
+                   const std::string& csvFilePath,
+                   const std::string& name,
+                   const Hdf5Reader& hdf5_reader);
+
     /**
      * Return selection of where attribute values match value
      *

include/bbp/sonata/population.h

@@ -19,6 +19,7 @@
 #include <utility>  // std::move
 #include <vector>
 
+#include <bbp/sonata/hdf5_reader.h>
 #include <bbp/sonata/selection.h>
 
 namespace bbp {
@@ -169,7 +170,8 @@ class SONATA_API Population
     Population(const std::string& h5FilePath,
                const std::string& csvFilePath,
                const std::string& name,
-               const std::string& prefix);
+               const std::string& prefix,
+               const Hdf5Reader& hdf5_reader);
 
     Population(const Population&) = delete;
 
@@ -194,7 +196,12 @@ template <typename Population>
 class SONATA_API PopulationStorage
 {
   public:
-    PopulationStorage(const std::string& h5FilePath, const std::string& csvFilePath = "");
+    PopulationStorage(const std::string& h5FilePath);
+    PopulationStorage(const std::string& h5FilePath, const std::string& csvFilePath);
+    PopulationStorage(const std::string& h5FilePath, const Hdf5Reader& hdf5_reader);
+    PopulationStorage(const std::string& h5FilePath,
+                      const std::string& csvFilePath,
+                      const Hdf5Reader& hdf5_reader);
 
     PopulationStorage(const PopulationStorage&) = delete;
 

python/bindings.cpp

@@ -305,11 +305,14 @@ py::class_<Storage> bindStorageClass(py::module& m, const char* clsName, const c
     };
     return py::class_<Storage>(
                m, clsName, imbuePopulationClassName(DOC(bbp, sonata, PopulationStorage)).c_str())
-        .def(py::init([](py::object h5_filepath, py::object csv_filepath) {
-                 return Storage(py::str(h5_filepath), py::str(csv_filepath));
+        .def(py::init([](py::object h5_filepath, py::object csv_filepath, Hdf5Reader hdf5_reader) {
+                 return Storage(py::str(h5_filepath),
+                                py::str(csv_filepath),
+                                std::move(hdf5_reader));
              }),
              "h5_filepath"_a,
-             "csv_filepath"_a = "")
+             "csv_filepath"_a = "",
+             "hdf5_reader"_a = Hdf5Reader())
         .def_property_readonly("population_names",
                                &Storage::populationNames,
                                imbuePopulationClassName(DOC_POP_STOR(populationNames)).c_str())
@@ -403,6 +406,8 @@ void bindReportReader(py::module& m, const std::string& prefix) {
 
 
 PYBIND11_MODULE(_libsonata, m) {
+    py::class_<Hdf5Reader>(m, "Hdf5Reader").def(py::init([]() { return Hdf5Reader(); }));
+
     py::class_<Selection>(m,
                           "Selection",
                           "ID sequence in the form convenient for querying attributes")
@@ -591,9 +596,19 @@ PYBIND11_MODULE(_libsonata, m) {
         .def_property_readonly("config_status", &CircuitConfig::getCircuitConfigStatus)
         .def_property_readonly("node_sets_path", &CircuitConfig::getNodeSetsPath)
         .def_property_readonly("node_populations", &CircuitConfig::listNodePopulations)
-        .def("node_population", &CircuitConfig::getNodePopulation)
+        .def("node_population",
+             [](const CircuitConfig& config, const std::string& name) {
+                 return config.getNodePopulation(name);
+             })
         .def_property_readonly("edge_populations", &CircuitConfig::listEdgePopulations)
-        .def("edge_population", &CircuitConfig::getEdgePopulation)
+        .def("edge_population",
+             [](const CircuitConfig& config, const std::string& name) {
+                 return config.getEdgePopulation(name);
+             })
+        .def("edge_population",
+             [](const CircuitConfig& config, const std::string& name, Hdf5Reader hdf5_reader) {
+                 return config.getEdgePopulation(name, hdf5_reader);
+             })
         .def("node_population_properties", &CircuitConfig::getNodePopulationProperties, "name"_a)
         .def("edge_population_properties", &CircuitConfig::getEdgePopulationProperties, "name"_a)
         .def_property_readonly("expanded_json", &CircuitConfig::getExpandedJSON);

python/libsonata/__init__.py

@@ -23,6 +23,7 @@
     SpikePopulation,
     SpikeReader,
     version,
+    Hdf5Reader,
 )
 
 __all__ = [
@@ -44,4 +45,5 @@
     "SpikePopulation",
     "SpikeReader",
     "version",
+    "Hdf5Reader",
 ]

src/config.cpp

@@ -202,9 +202,13 @@ PopulationType getPopulationProperties(
 
 template <typename PopulationType, typename PopulationPropertiesT>
 PopulationType getPopulation(const std::string& populationName,
-                             const std::unordered_map<std::string, PopulationPropertiesT>& src) {
+                             const std::unordered_map<std::string, PopulationPropertiesT>& src,
+                             const Hdf5Reader& hdf5_reader) {
     const auto properties = getPopulationProperties(populationName, src);
-    return PopulationType(properties.elementsPath, properties.typesPath, populationName);
+    return PopulationType(properties.elementsPath,
+                          properties.typesPath,
+                          populationName,
+                          hdf5_reader);
 }
 
 std::map<std::string, std::string> replaceVariables(std::map<std::string, std::string> variables) {
@@ -918,15 +922,25 @@ std::set<std::string> CircuitConfig::listNodePopulations() const {
 }
 
 NodePopulation CircuitConfig::getNodePopulation(const std::string& name) const {
-    return getPopulation<NodePopulation>(name, _nodePopulationProperties);
+    return getNodePopulation(name, Hdf5Reader());
+}
+
+NodePopulation CircuitConfig::getNodePopulation(const std::string& name,
+                                                const Hdf5Reader& hdf5_reader) const {
+    return getPopulation<NodePopulation>(name, _nodePopulationProperties, hdf5_reader);
 }
 
 std::set<std::string> CircuitConfig::listEdgePopulations() const {
     return getMapKeys(_edgePopulationProperties);
 }
 
 EdgePopulation CircuitConfig::getEdgePopulation(const std::string& name) const {
-    return getPopulation<EdgePopulation>(name, _edgePopulationProperties);
+    return getPopulation<EdgePopulation>(name, _edgePopulationProperties, Hdf5Reader());
+}
+
+EdgePopulation CircuitConfig::getEdgePopulation(const std::string& name,
+                                                const Hdf5Reader& hdf5_reader) const {
+    return getPopulation<EdgePopulation>(name, _edgePopulationProperties, hdf5_reader);
 }
 
 NodePopulationProperties CircuitConfig::getNodePopulationProperties(const std::string& name) const {