Batch aggregation (two-level batching)

metagraph/src/annotation/binary_matrix/row_vector/vector_row_binmat.cpp

@@ -156,8 +156,9 @@ double VectorRowBinMat<RowType>::density() const {
     return static_cast<double>(num_relations()) / num_columns() / num_rows();
 }
 
-template class VectorRowBinMat<SmallVector<uint32_t>>;
 template class VectorRowBinMat<Vector<uint64_t>>;
+template class VectorRowBinMat<Vector<uint32_t>>;
+template class VectorRowBinMat<SmallVector<uint32_t>>;
 
 } // namespace matrix
 } // namespace annot

metagraph/src/cli/config/config.cpp

@@ -169,6 +169,8 @@ Config::Config(int argc, char *argv[]) {
             num_columns_cached = atoi(get_value(i++));
         } else if (!strcmp(argv[i], "--batch-size")) {
             query_batch_size = atoll(get_value(i++));
+        } else if (!strcmp(argv[i], "--batch-min-matches")) {
+            batch_min_matches = std::stod(get_value(i++));
         } else if (!strcmp(argv[i], "-p") || !strcmp(argv[i], "--parallel")) {
             set_num_threads(atoi(get_value(i++)));
         } else if (!strcmp(argv[i], "--parallel-nodes")) {
@@ -1325,6 +1327,7 @@ if (advanced) {
             fprintf(stderr, "\t-p --parallel [INT] \tuse multiple threads for computation [1]\n");
             // fprintf(stderr, "\t   --cache-size [INT] \tnumber of uncompressed rows to store in the cache [0]\n");
             fprintf(stderr, "\t   --batch-size [INT] \tquery batch size in bp (0 to disable batch query) [100'000'000]\n");
+            fprintf(stderr, "\t   --batch-min-matches [FLOAT] \taggregate batches unless this they have this ratio of k-mer matches [0.0]\n");
 if (advanced) {
             fprintf(stderr, "\t   --threads-each [INT]\tnumber of parallel batches [1]\n");
             fprintf(stderr, "\t   --RA-ivbuff-size [INT] \tsize (in bytes) of int_vector_buffer used in random access mode (e.g. by row disk annotator) [16384]\n");

metagraph/src/cli/config/config.hpp

@@ -144,6 +144,7 @@ class Config {
     double min_fraction = 0.0;
     double max_fraction = 1.0;
     double cleaning_threshold_percentile = 0.001;
+    double batch_min_matches = 0.0;
     std::vector<double> count_slice_quantiles;
     std::vector<double> count_quantiles;
 

metagraph/src/cli/load/load_annotated_graph.cpp

@@ -20,20 +20,39 @@ using namespace mtg::graph;
 using mtg::common::logger;
 
 
-std::unique_ptr<AnnotatedDBG> initialize_annotated_dbg(std::shared_ptr<DeBruijnGraph> graph,
-                                                       const Config &config,
-                                                       size_t max_chunks_open) {
+std::unique_ptr<annot::MultiLabelAnnotation<std::string>>
+load_annotation(std::shared_ptr<DeBruijnGraph> graph,
+                const Config &config,
+                size_t max_chunks_open) {
+    std::shared_future<std::shared_ptr<DeBruijnGraph>> graph_future = std::async(std::launch::async, [graph]{ return graph; });
+    return load_annotation(graph_future, config, max_chunks_open);
+}
+
+std::unique_ptr<annot::MultiLabelAnnotation<std::string>>
+load_annotation(std::shared_future<std::shared_ptr<DeBruijnGraph>> graph_future,
+                const Config &config,
+                size_t max_chunks_open) {
+
+    std::unique_ptr<annot::MultiLabelAnnotation<std::string>> annotation_temp;
+    if (config.infbase_annotators.size())
+        annotation_temp = initialize_annotation(config.infbase_annotators.at(0), config, 0, max_chunks_open);
+
+    std::shared_ptr<DeBruijnGraph> graph = graph_future.get();
     uint64_t max_index = graph->max_index();
 
-    auto base_graph = graph;
+    const auto *base_graph_ptr = graph.get();
     if (graph->get_mode() == DeBruijnGraph::PRIMARY) {
         graph = std::make_shared<CanonicalDBG>(graph);
         logger->trace("Primary graph wrapped into canonical");
+    } else if (const auto *canonical = dynamic_cast<const CanonicalDBG *>(graph.get())) {
+        base_graph_ptr = &canonical->get_graph();
+        max_index = base_graph_ptr->max_index();
     }
 
-    auto annotation_temp = config.infbase_annotators.size()
-            ? initialize_annotation(config.infbase_annotators.at(0), config, 0, max_chunks_open)
-            : initialize_annotation(config.anno_type, config, max_index, max_chunks_open);
+    if (!config.infbase_annotators.size())
+        annotation_temp = initialize_annotation(config.anno_type, config, max_index, max_chunks_open);
+
+    assert(annotation_temp);
 
     if (config.infbase_annotators.size()) {
         bool loaded = false;
@@ -56,7 +75,7 @@ std::unique_ptr<AnnotatedDBG> initialize_annotated_dbg(std::shared_ptr<DeBruijnG
         using namespace annot::matrix;
         BinaryMatrix &matrix = const_cast<BinaryMatrix &>(annotation_temp->get_matrix());
         if (IRowDiff *row_diff = dynamic_cast<IRowDiff*>(&matrix)) {
-            row_diff->set_graph(base_graph.get());
+            row_diff->set_graph(base_graph_ptr);
 
             if (auto *row_diff_column = dynamic_cast<RowDiff<ColumnMajor> *>(&matrix)) {
                 row_diff_column->load_anchor(config.infbase + kRowDiffAnchorExt);
@@ -65,9 +84,21 @@ std::unique_ptr<AnnotatedDBG> initialize_annotated_dbg(std::shared_ptr<DeBruijnG
         }
     }
 
-    // load graph
+    return annotation_temp;
+}
+
+std::unique_ptr<AnnotatedDBG> initialize_annotated_dbg(std::shared_ptr<DeBruijnGraph> graph,
+                                                       const Config &config,
+                                                       size_t max_chunks_open) {
+    if (graph->get_mode() == DeBruijnGraph::PRIMARY) {
+        graph = std::make_shared<CanonicalDBG>(graph);
+        logger->trace("Primary graph wrapped into canonical");
+    }
+
+    auto annotation = load_annotation(graph, config, max_chunks_open);
+
     auto anno_graph
-            = std::make_unique<AnnotatedDBG>(std::move(graph), std::move(annotation_temp));
+            = std::make_unique<AnnotatedDBG>(std::move(graph), std::move(annotation));
 
     if (!anno_graph->check_compatibility()) {
         logger->error("Graph and annotation are not compatible");

metagraph/src/cli/load/load_annotated_graph.hpp

@@ -1,18 +1,21 @@
 #ifndef __LOAD_ANNOTATED_GRAPH_HPP__
 #define __LOAD_ANNOTATED_GRAPH_HPP__
 
-
+#include <future>
 #include <memory>
 
 namespace mtg {
 
 namespace graph {
-
 class DeBruijnGraph;
 class AnnotatedDBG;
-
 } // namespace graph
 
+namespace annot {
+template <typename LabelType>
+class MultiLabelAnnotation;
+} // namespace annot
+
 namespace cli {
 
 class Config;
@@ -22,6 +25,16 @@ initialize_annotated_dbg(std::shared_ptr<graph::DeBruijnGraph> graph,
                          const Config &config,
                          size_t max_chunks_open = 2000);
 
+std::unique_ptr<annot::MultiLabelAnnotation<std::string>>
+load_annotation(std::shared_ptr<graph::DeBruijnGraph> graph,
+                const Config &config,
+                size_t max_chunks_open = 2000);
+
+std::unique_ptr<annot::MultiLabelAnnotation<std::string>>
+load_annotation(std::shared_future<std::shared_ptr<graph::DeBruijnGraph>> graph,
+                const Config &config,
+                size_t max_chunks_open = 2000);
+
 std::unique_ptr<graph::AnnotatedDBG> initialize_annotated_dbg(const Config &config);
 
 } // namespace cli