API for optimized access order in Collection (#453)

The point of this API is to enable collections to load morphologies
in any order they see fit. For HDF5 containers tests have shown that the
access pattern matters a lot in terms of performance.

The idea is a generic API that allows optimizing the iteration order for
reorderable loops, e.g.,

    for k, morph_name in enumerate(morphology_names):
        morph = collection.load(morph_name)
        f(k, morph)

can be replaced with

    for k, morph in collection.load_unordered(morphology_names):
        assert collection.load(morphology_names[k]) == morph
        f(k, morph)

Author: 1uc

3rdparty/HighFive

@@ -323,11 +323,87 @@ void bind_misc(py::module& m) {
             "options"_a = morphio::enums::Option::NO_MODIFIER,
             "mutable"_a = false,
             "Load the morphology named 'morph_name' form the collection.")
+        .def(
+            "load_unordered",
+            [](morphio::Collection* collection,
+               std::vector<std::string> morphology_names,
+               unsigned int options,
+               bool is_mutable) -> py::object {
+                if (is_mutable) {
+                    return py::cast(
+                        collection->load_unordered<morphio::mut::Morphology>(morphology_names,
+                                                                             options));
+                } else {
+                    return py::cast(
+                        collection->load_unordered<morphio::Morphology>(morphology_names, options));
+                }
+            },
+            "morphology_names"_a,
+            "options"_a = morphio::enums::Option::NO_MODIFIER,
+            "mutable"_a = false,
+            R"(Create an iterable of loop index and morphology.
+
+When reading from containers, the order in which morphologies are read can
+have a large impact on the overall time to load those morphologies.
+
+This iterator provides means of reordering loops to optimize the access
+pattern. Loops such as the following
+
+    for k, morph_name in enumerate(morphology_names):
+        morph = collection.load(morphology_names[k])
+        f(k, morph)
+
+can be replaced with
+
+    for k, morph in collection.load_unordered(morphology_names):
+      assert collection.load(morphology_names[k]) == morph
+      f(k, morph)
+
+The order in which the morphologies are returned in unspecified, but the
+loop index `k` can be used to retrieve the correct state corresponding to
+iteration `k` of the original loop.
+
+The iterable returned by `Collection.load_unordered` should only be used while
+`collection` is valid, e.g. within its context or before calling
+`Collection.close`.
+
+Note: This API is 'experimental', meaning it might change in the future.
+)")
+
+        .def("argsort",
+             &morphio::Collection::argsort,
+             "morphology_names"_a,
+             R"(Argsort `morphology_names` by optimal access order.
+
+Note: This API is 'experimental', meaning it might change in the future.
+)")
         .def("__enter__", [](morphio::Collection* collection) { return collection; })
         .def("__exit__",
              [](morphio::Collection* collection,
                 const py::object&,
                 const py::object&,
                 const py::object&) { collection->close(); })
         .def("close", &morphio::Collection::close);
+
+    py::class_<morphio::LoadUnordered<morphio::Morphology>>(
+        m, "LoadImmutableUnordered", "An iterable of immutable morphologies.")
+        .def(
+            "__iter__",
+            [](const morphio::LoadUnordered<morphio::Morphology>& iterable) {
+                return py::make_iterator(iterable.begin(), iterable.end());
+            },
+            // Bind the lifetime of the `morphio::LoadUnordered` (1) to the
+            // lifetime of the returned iterator (0).
+            py::keep_alive<0, 1>());
+
+    py::class_<morphio::LoadUnordered<morphio::mut::Morphology>>(
+        m, "LoadMutableUnordered", "An iterable of mutable morphologies.")
+        .def(
+            "__iter__",
+            [](const morphio::LoadUnordered<morphio::mut::Morphology>& iterable) {
+                return py::make_iterator(iterable.begin(), iterable.end());
+            },
+            // Bind the lifetime of the `morphio::LoadUnordered` (1) to the
+            // lifetime of the returned iterator (0).
+            py::keep_alive<0, 1>());
 }

binds/python/bind_misc.cpp

@@ -10,6 +10,21 @@ namespace morphio {
 
 class CollectionImpl;
 
+template <class M>
+class LoadUnordered;
+
+/**
+ * Enable if `T` is a immutable morphology.
+ */
+template <class T, class U = void>
+struct enable_if_immutable: public std::enable_if<std::is_same<T, Morphology>::value, U> {};
+
+/**
+ * Enable if `T` is a mutable morphology.
+ */
+template <class T, class U = void>
+struct enable_if_mutable: public std::enable_if<std::is_same<T, mut::Morphology>::value, U> {};
+
 class Collection
 {
   public:
@@ -34,15 +49,34 @@ class Collection
      * Load the morphology as an immutable morphology.
      */
     template <class M>
-    typename std::enable_if<std::is_same<M, Morphology>::value, M>::type load(
-        const std::string& morph_name, unsigned int options = NO_MODIFIER) const;
+    typename enable_if_mutable<M, M>::type load(const std::string& morph_name,
+                                                unsigned int options = NO_MODIFIER) const;
 
     /**
      * Load the morphology as a mutable morphology.
      */
     template <class M>
-    typename std::enable_if<std::is_same<M, mut::Morphology>::value, M>::type load(
-        const std::string& morph_name, unsigned int options = NO_MODIFIER) const;
+    typename enable_if_immutable<M, M>::type load(const std::string& morph_name,
+                                                  unsigned int options = NO_MODIFIER) const;
+
+    /**
+     * Returns an iterable of loop index, morphology pairs.
+     *
+     * See `LoadUnordered` for details.
+     */
+    template <class M>
+    LoadUnordered<M> load_unordered(std::vector<std::string> morphology_names,
+                                    unsigned int options = NO_MODIFIER) const;
+
+    /**
+     * Returns the reordered loop indices.
+     *
+     * This is the suggested order in which one should load the morphologies to
+     * minimize seeking within the file.
+     *
+     * Note: This API is 'experimental', meaning it might change in the future.
+     */
+    std::vector<size_t> argsort(const std::vector<std::string>& morphology_names) const;
 
     /**
      * Close the collection.
@@ -61,10 +95,99 @@ class Collection
     std::shared_ptr<CollectionImpl> _collection;
 };
 
-extern template mut::Morphology Collection::load<mut::Morphology>(const std::string& morph_name,
-                                                                  unsigned int options) const;
+class LoadUnorderedImpl;
+
+/**
+ * An iterable of loop index and morphologies.
+ *
+ * When reading from containers, the order in which morphologies are read can
+ * have a large impact on the overall time to load those morphologies.
+ *
+ * This iterator provides means of reordering loops to optimize the access
+ * pattern. Loops such as the following
+ *
+ *     for(size_t k = 0; k < morphology_names.size; ++k) {
+ *       auto morph = collection.load<M>(morphology_names[k]);
+ *       f(k, morph);
+ *     }
+ *
+ * can be replaced with
+ *
+ *     for(auto [k, morph] : collection.load_unordered<M>(morphology_names)) {
+ *       assert(collection.load<M>(morphology_names[k]) == morph);
+ *       f(k, morph);
+ *     }
+ *
+ * The order in which the morphologies are returned in unspecified, but the
+ * loop index `k` can be used to retrieve the correct state corresponding to
+ * iteration `k` of the original loop.
+ *
+ * Note, that it is safe for an `LoadUnordered` object to outlive its
+ * `collection`. Internally a shallow copy of the original `collection` is
+ * stored inside of and kept alive for the life time of the `LoadUnordered`
+ * object.
+ *
+ * Note: This API is 'experimental', meaning it might change in the future.
+ */
+template <class M>
+class LoadUnordered
+{
+  protected:
+    class Iterator
+    {
+      public:
+        Iterator(std::shared_ptr<LoadUnorderedImpl> load_unordered_impl, size_t k);
+
+        template <class U = M>
+        typename enable_if_immutable<U, std::pair<size_t, M>>::type operator*() const;
+
+        template <class U = M>
+        typename enable_if_mutable<U, std::pair<size_t, M>>::type operator*() const;
+
+        Iterator& operator++();
+        Iterator operator++(int);
+
+        bool operator==(const Iterator& other) const;
+        bool operator!=(const Iterator& other) const;
+
+      private:
+        size_t _k;
+        std::shared_ptr<LoadUnorderedImpl> _load_unordered_impl;
+    };
+
+  public:
+    LoadUnordered(std::shared_ptr<LoadUnorderedImpl> load_unordered_impl);
+
+    Iterator begin() const;
+    Iterator end() const;
+
+  protected:
+    std::shared_ptr<LoadUnorderedImpl> _load_unordered_impl;
+};
+
+extern template class LoadUnordered<Morphology>;
+extern template class LoadUnordered<mut::Morphology>;
+
+extern template class LoadUnordered<Morphology>::Iterator;
+extern template class LoadUnordered<mut::Morphology>::Iterator;
+
+extern template typename enable_if_immutable<Morphology, std::pair<size_t, Morphology>>::type
+    LoadUnordered<Morphology>::Iterator::operator*<Morphology>() const;
+
+extern template
+    typename enable_if_mutable<mut::Morphology, std::pair<size_t, mut::Morphology>>::type
+        LoadUnordered<mut::Morphology>::Iterator::operator*<mut::Morphology>() const;
+
+extern template typename enable_if_mutable<mut::Morphology, mut::Morphology>::type
+Collection::load<mut::Morphology>(const std::string& morph_name, unsigned int options) const;
+
+extern template typename enable_if_immutable<Morphology, Morphology>::type
+Collection::load<Morphology>(const std::string& morph_name, unsigned int options) const;
+
+extern template LoadUnordered<Morphology> Collection::load_unordered<Morphology>(
+    std::vector<std::string> morphology_names, unsigned int options) const;
 
-extern template Morphology Collection::load<Morphology>(const std::string& morph_name,
-                                                        unsigned int options) const;
+extern template LoadUnordered<mut::Morphology> Collection::load_unordered<mut::Morphology>(
+    std::vector<std::string> morphology_names, unsigned int options) const;
 
 }  // namespace morphio

include/morphio/collection.h

@@ -22,7 +22,7 @@ class Morphology
   public:
     virtual ~Morphology() = default;
 
-    Morphology(Morphology&) noexcept = default;
+    Morphology(const Morphology&) noexcept = default;
     Morphology& operator=(const Morphology&) noexcept = default;
     Morphology(Morphology&&) noexcept = default;
     Morphology& operator=(Morphology&&) noexcept = default;