From 3ebf0102a9d8acf915d03e473eeba11c88e84cd7 Mon Sep 17 00:00:00 2001 From: Sagorika Nag Date: Thu, 23 Jul 2026 22:12:09 -0700 Subject: [PATCH] Add grouped I/O for single-read graph placements --- include/vg/io/alignment_io.hpp | 73 +++++++++++++++++++++++++++++++++ src/alignment_io.cpp | 75 ++++++++++++++++++++++++++++++++++ 2 files changed, 148 insertions(+) diff --git a/include/vg/io/alignment_io.hpp b/include/vg/io/alignment_io.hpp index f91edaa..106bdac 100644 --- a/include/vg/io/alignment_io.hpp +++ b/include/vg/io/alignment_io.hpp @@ -40,6 +40,16 @@ size_t paired_for_each_parallel_after_wait(function get_pair_if_av function 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 +size_t grouped_for_each_parallel(function get_record_if_available, + function get_key, + function&)> 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 lambda, function 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&)> lambda, + uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); +size_t gaf_grouped_for_each_parallel(function node_to_length, function node_to_sequence, const string& filename, + function&)> lambda, + uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); +size_t gaf_grouped_for_each_parallel(const HandleGraph& graph, const string& filename, + function&)> 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 get_pai return nLines; } +template +inline size_t grouped_for_each_parallel(function get_record_if_available, + function get_key, + function&)> 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 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&)> get_run_if_available = [&](vector& 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>(get_run_if_available, lambda, batch_size); +} + } } #endif diff --git a/src/alignment_io.cpp b/src/alignment_io.cpp index a0335b2..7701274 100644 --- a/src/alignment_io.cpp +++ b/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 #include @@ -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&)> lambda, + uint64_t batch_size) { + + ProtobufIterator it(in); + + function get_record = [&](Alignment& aln) -> bool { + if (!it.has_current()) { + return false; + } + aln = it.take(); + return true; + }; + function get_key = [](const Alignment& aln) { + return aln.name(); + }; + + return grouped_for_each_parallel(get_record, get_key, lambda, batch_size); +} + + +size_t gaf_grouped_for_each_parallel(function node_to_length, function node_to_sequence, const string& filename, + function&)> 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 get_record = [&](gafkluge::GafRecord& gaf) -> bool { + return get_next_record_from_gaf(node_to_length, node_to_sequence, in, s_buffer, gaf); + }; + function get_key = [](const gafkluge::GafRecord& gaf) { + return gaf.query_name; + }; + function&)> convert_and_call = [&](vector& gaf_run) { + vector 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(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&)> lambda, + uint64_t batch_size) { + function node_to_length = [&graph](nid_t node_id) { + return graph.get_length(graph.get_handle(node_id)); + }; + function 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;