Add grouped parallel I/O for graph placements of unpaired reads

include/vg/io/alignment_io.hpp

@@ -40,6 +40,16 @@ size_t paired_for_each_parallel_after_wait(function<bool(T&, T&)> get_pair_if_av
                                            function<bool(void)> single_threaded_until_true,
                                            uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE);
 
+/// Group consecutive records sharing the same key (produced by get_key) into "runs",
+/// processes each run with the provided lambda function
+/// batch_size, is the number of runs per dispatched task
+/// returns the no. of runs processed
+template<typename T>
+size_t grouped_for_each_parallel(function<bool(T&)> get_record_if_available,
+                                 function<string(const T&)> get_key,
+                                 function<void(vector<T>&)> lambda,
+                                 uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE);
+
 // Opens an htsFile, reads GAF header lines, and closes the file.
 // Does nothing if the file refers to stdin ("-"), as we probably can't rewind it.
 // Returns the header lines without the trailing newline characters.
@@ -77,6 +87,20 @@ size_t gaf_paired_interleaved_for_each_parallel_after_wait(const HandleGraph& gr
                                                            function<void(Alignment&, Alignment&)> lambda,
                                                            function<bool(void)> single_threaded_until_true,
                                                            uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE);
+
+// single
+// grouped (same read name) iteration, for callers that need all placements
+// of a read (e.g. primary + secondaries) delivered together
+size_t gam_grouped_for_each_parallel(std::istream& in,
+                                     function<void(vector<Alignment>&)> lambda,
+                                     uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE);
+size_t gaf_grouped_for_each_parallel(function<size_t(nid_t)> node_to_length, function<string(nid_t, bool)> node_to_sequence, const string& filename,
+                                     function<void(vector<Alignment>&)> lambda,
+                                     uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE);
+size_t gaf_grouped_for_each_parallel(const HandleGraph& graph, const string& filename,
+                                     function<void(vector<Alignment>&)> lambda,
+                                     uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE);
+
 // gaf conversion
 
 /// Convert an alignment to GAF. The alignment must be in node ID space.
@@ -307,6 +331,55 @@ inline size_t paired_for_each_parallel_after_wait(function<bool(T&, T&)> get_pai
     return nLines;
 }
 
+template<typename T>
+inline size_t grouped_for_each_parallel(function<bool(T&)> get_record_if_available,
+                                        function<string(const T&)> get_key,
+                                        function<void(vector<T>&)> lambda,
+                                        uint64_t batch_size) {
+
+    // State for the run currently being assembled from the record source.
+    // Only ever touched serially (from within unpaired_for_each_parallel's
+    // single-threaded batch-filling loop), so no synchronization is needed.
+    vector<T> current_run;
+    string current_key;
+    bool source_exhausted = false;
+
+    // Adapts the flat record source into a source of same-key runs, so grouped
+    // iteration can reuse unpaired_for_each_parallel's bounded task backpressure
+    function<bool(vector<T>&)> get_run_if_available = [&](vector<T>& out_run) -> bool {
+        if (source_exhausted && current_run.empty()) {
+            return false;
+        }
+        T record;
+        while (get_record_if_available(record)) {
+            string key = get_key(record);
+            if (current_run.empty()) {
+                current_key = key;
+                current_run.emplace_back(std::move(record));
+            } else if (key == current_key) {
+                current_run.emplace_back(std::move(record));
+            } else {
+                // Found the start of the next run: hand back the finished one
+                // and stash this record as the start of the next.
+                out_run = std::move(current_run);
+                current_run.clear();
+                current_key = key;
+                current_run.emplace_back(std::move(record));
+                return true;
+            }
+        }
+        source_exhausted = true;
+        if (!current_run.empty()) {
+            out_run = std::move(current_run);
+            current_run.clear();
+            return true;
+        }
+        return false;
+    };
+
+    return unpaired_for_each_parallel<vector<T>>(get_run_if_available, lambda, batch_size);
+}
+
 }
 }
 #endif

src/alignment_io.cpp

@@ -1,6 +1,7 @@
 #include "vg/io/alignment_io.hpp"
 #include "vg/io/gafkluge.hpp"
 #include "vg/io/edit.hpp"
+#include "vg/io/protobuf_iterator.hpp"
 
 #include <sstream>
 #include <regex>
@@ -220,6 +221,80 @@ size_t gaf_paired_interleaved_for_each_parallel_after_wait(const HandleGraph& gr
     return gaf_paired_interleaved_for_each_parallel_after_wait(node_to_length, node_to_sequence, filename, lambda, single_threaded_until_true, batch_size);
 }
 
+size_t gam_grouped_for_each_parallel(std::istream& in,
+                                     function<void(vector<Alignment>&)> lambda,
+                                     uint64_t batch_size) {
+
+    ProtobufIterator<Alignment> it(in);
+
+    function<bool(Alignment&)> get_record = [&](Alignment& aln) -> bool {
+        if (!it.has_current()) {
+            return false;
+        }
+        aln = it.take();
+        return true;
+    };
+    function<string(const Alignment&)> get_key = [](const Alignment& aln) {
+        return aln.name();
+    };
+
+    return grouped_for_each_parallel<Alignment>(get_record, get_key, lambda, batch_size);
+}
+
+
+size_t gaf_grouped_for_each_parallel(function<size_t(nid_t)> node_to_length, function<string(nid_t, bool)> node_to_sequence, const string& filename,
+                                     function<void(vector<Alignment>&)> lambda,
+                                     uint64_t batch_size) {
+
+    htsFile* in = hts_open(filename.c_str(), "r");
+    if (in == NULL) {
+        cerr << "error: [vg::io::alignment_io.cpp] couldn't open " << filename << endl; exit(1);
+    }
+
+    kstring_t s_buffer = KS_INITIALIZE;
+
+    // Only reads and parses the GAF line into a GafRecord (cheap: no CIGAR/cs
+    // decoding, no sequence reconstruction). The expensive gaf_to_alignment
+    // conversion happens per-group below, inside the dispatched task, so it
+    // stays parallelized across worker threads instead of running on the
+    // single fetch thread.
+    function<bool(gafkluge::GafRecord&)> get_record = [&](gafkluge::GafRecord& gaf) -> bool {
+        return get_next_record_from_gaf(node_to_length, node_to_sequence, in, s_buffer, gaf);
+    };
+    function<string(const gafkluge::GafRecord&)> get_key = [](const gafkluge::GafRecord& gaf) {
+        return gaf.query_name;
+    };
+    function<void(vector<gafkluge::GafRecord>&)> convert_and_call = [&](vector<gafkluge::GafRecord>& gaf_run) {
+        vector<Alignment> aln_run;
+        aln_run.reserve(gaf_run.size());
+        for (auto& gaf : gaf_run) {
+            Alignment aln;
+            gaf_to_alignment(node_to_length, node_to_sequence, gaf, aln);
+            aln_run.emplace_back(std::move(aln));
+        }
+        lambda(aln_run);
+    };
+
+    size_t nLines = grouped_for_each_parallel<gafkluge::GafRecord>(get_record, get_key, convert_and_call, batch_size);
+
+    hts_close(in);
+    return nLines;
+}
+
+
+size_t gaf_grouped_for_each_parallel(const HandleGraph& graph, const string& filename,
+                                     function<void(vector<Alignment>&)> lambda,
+                                     uint64_t batch_size) {
+    function<size_t(nid_t)> node_to_length = [&graph](nid_t node_id) {
+        return graph.get_length(graph.get_handle(node_id));
+    };
+    function<string(nid_t, bool)> node_to_sequence = [&graph](nid_t node_id, bool is_reversed) {
+        return graph.get_sequence(graph.get_handle(node_id, is_reversed));
+    };
+    return gaf_grouped_for_each_parallel(node_to_length, node_to_sequence, filename, lambda, batch_size);
+}
+
+
 string supplementary_tag_value(const Alignment& primary) {
 
     stringstream strm;