diff --git a/include/vg/io/alignment_io.hpp b/include/vg/io/alignment_io.hpp index f91edaa..9147495 100644 --- a/include/vg/io/alignment_io.hpp +++ b/include/vg/io/alignment_io.hpp @@ -34,12 +34,37 @@ size_t unpaired_for_each_parallel(function get_read_if_available, function lambda, uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); +template +size_t unpaired_for_each_parallel_after_wait(function get_read_if_available, + function lambda, + function single_threaded_until_true, + uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); + template size_t paired_for_each_parallel_after_wait(function get_pair_if_available, function lambda, 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); + +/// Group consecutive records sharing the same key. Process groups serially +/// until single_threaded_until_true returns true, and in parallel thereafter. +template +size_t grouped_for_each_parallel_after_wait(function get_record_if_available, + function get_key, + function&)> lambda, + function single_threaded_until_true, + 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 +102,46 @@ 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); + +/// Read interleaved GAM pairs and group consecutive alternative placements +/// belonging to the same fragment. Process groups serially until +/// single_threaded_until_true returns true, and in parallel thereafter. +size_t gam_paired_grouped_for_each_parallel_after_wait( + istream& in, + function>&)> lambda, + function single_threaded_until_true, + uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); + +/// GAF equivalent using explicit graph accessors. +size_t gaf_paired_grouped_for_each_parallel_after_wait( + function node_to_length, + function node_to_sequence, + const string& filename, + function>&)> lambda, + function single_threaded_until_true, + uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); + +/// GAF convenience overload using a HandleGraph. +size_t gaf_paired_grouped_for_each_parallel_after_wait( + const HandleGraph& graph, + const string& filename, + function>&)> lambda, + function single_threaded_until_true, + uint64_t batch_size = DEFAULT_PARALLEL_BATCHSIZE); + // gaf conversion /// Convert an alignment to GAF. The alignment must be in node ID space. @@ -127,15 +192,16 @@ void alignment_quality_short_to_char(Alignment& alignment); // implementation template -inline size_t unpaired_for_each_parallel(function get_read_if_available, - function lambda, - uint64_t batch_size) { +inline size_t unpaired_for_each_parallel_after_wait(function get_read_if_available, + function lambda, + function single_threaded_until_true, + uint64_t batch_size) { assert(batch_size % 2 == 0); size_t nLines = 0; vector *batch = nullptr; // number of batches currently being processed uint64_t batches_outstanding = 0; -#pragma omp parallel default(none) shared(batches_outstanding, batch, nLines, get_read_if_available, lambda, batch_size) +#pragma omp parallel default(none) shared(batches_outstanding, batch, nLines, get_read_if_available, lambda, single_threaded_until_true, batch_size) #pragma omp single { @@ -175,8 +241,9 @@ inline size_t unpaired_for_each_parallel(function get_read_if_availabl uint64_t current_batches_outstanding; #pragma omp atomic capture current_batches_outstanding = ++batches_outstanding; - - if (current_batches_outstanding >= max_batches_outstanding) { + + bool do_single_threaded = !single_threaded_until_true(); + if (current_batches_outstanding >= max_batches_outstanding || do_single_threaded) { // do this batch in the current thread because we've spawned the maximum number of // concurrent batch tasks for (auto& aln : *batch) { @@ -187,7 +254,8 @@ inline size_t unpaired_for_each_parallel(function get_read_if_availabl current_batches_outstanding = --batches_outstanding; if (4 * current_batches_outstanding / 3 < max_batches_outstanding - && max_batches_outstanding < max_max_batches_outstanding) { + && max_batches_outstanding < max_max_batches_outstanding + && !do_single_threaded) { // we went through at least 1/4 of the batch buffer while we were doing this thread's batch // this looks risky, since we want the batch buffer to stay populated the entire time we're // occupying this thread on compute, so let's increase the batch buffer size @@ -213,6 +281,14 @@ inline size_t unpaired_for_each_parallel(function get_read_if_availabl return nLines; } +template +inline size_t unpaired_for_each_parallel(function get_read_if_available, + function lambda, + uint64_t batch_size) { + return unpaired_for_each_parallel_after_wait(get_read_if_available, lambda, + []() { return true; }, batch_size); +} + template inline size_t paired_for_each_parallel_after_wait(function get_pair_if_available, function lambda, @@ -307,6 +383,66 @@ inline size_t paired_for_each_parallel_after_wait(function get_pai return nLines; } +template +inline size_t grouped_for_each_parallel_after_wait(function get_record_if_available, + function get_key, + function&)> lambda, + function single_threaded_until_true, + 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_after_wait>(get_run_if_available, lambda, + single_threaded_until_true, batch_size); +} + +template +inline size_t grouped_for_each_parallel(function get_record_if_available, + function get_key, + function&)> lambda, + uint64_t batch_size) { + return grouped_for_each_parallel_after_wait(get_record_if_available, get_key, lambda, + []() { return true; }, batch_size); +} + } } #endif diff --git a/src/alignment_io.cpp b/src/alignment_io.cpp index c755e1b..e21a2a5 100644 --- a/src/alignment_io.cpp +++ b/src/alignment_io.cpp @@ -1,11 +1,13 @@ #include "vg/io/alignment_io.hpp" #include "vg/io/gafkluge.hpp" #include "vg/io/edit.hpp" +#include "vg/io/protobuf_iterator.hpp" #include #include #include #include +#include //#define debug_translation @@ -13,6 +15,13 @@ namespace vg { namespace io { +static string fragment_key(const string& first_name, const string& second_name) { + if (first_name <= second_name) { + return first_name + '\n' + second_name; + } + return second_name + '\n' + first_name; +} + std::vector read_gaf_header_lines(const std::string& filename) { std::vector header_lines; if (filename == "-") { @@ -220,6 +229,189 @@ 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 gam_paired_grouped_for_each_parallel_after_wait( + istream& in, + function>&)> lambda, + function single_threaded_until_true, + uint64_t batch_size) { + + using AlignmentPair = pair; + ProtobufIterator it(in); + bool unmatched_record = false; + + function get_pair = [&](AlignmentPair& alignment_pair) { + if (!it.has_current()) { + return false; + } + Alignment first = it.take(); + if (!it.has_current()) { + unmatched_record = true; + return false; + } + Alignment second = it.take(); + alignment_pair = make_pair(std::move(first), std::move(second)); + return true; + }; + function get_key = [](const AlignmentPair& alignment_pair) { + return fragment_key(alignment_pair.first.name(), alignment_pair.second.name()); + }; + + size_t count = grouped_for_each_parallel_after_wait( + get_pair, get_key, lambda, single_threaded_until_true, batch_size); + if (unmatched_record) { + throw runtime_error("interleaved GAM input contains an unmatched alignment"); + } + return count; +} + + +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); +} + +size_t gaf_paired_grouped_for_each_parallel_after_wait( + function node_to_length, + function node_to_sequence, + const string& filename, + function>&)> lambda, + function single_threaded_until_true, + uint64_t batch_size) { + + using GafPair = pair; + + htsFile* in = hts_open(filename.c_str(), "r"); + if (in == nullptr) { + cerr << "error: [vg::io::alignment_io.cpp] couldn't open " << filename << endl; + exit(EXIT_FAILURE); + } + + kstring_t s_buffer = KS_INITIALIZE; + bool unmatched_record = false; + function get_pair = [&](GafPair& gaf_pair) { + gafkluge::GafRecord first; + if (!get_next_record_from_gaf(node_to_length, node_to_sequence, in, s_buffer, first)) { + return false; + } + gafkluge::GafRecord second; + if (!get_next_record_from_gaf(node_to_length, node_to_sequence, in, s_buffer, second)) { + unmatched_record = true; + return false; + } + gaf_pair = make_pair(std::move(first), std::move(second)); + return true; + }; + function get_key = [](const GafPair& gaf_pair) { + return fragment_key(gaf_pair.first.query_name, gaf_pair.second.query_name); + }; + function&)> convert_and_call = [&](vector& gaf_pairs) { + vector> alignment_pairs; + alignment_pairs.reserve(gaf_pairs.size()); + for (auto& gaf_pair : gaf_pairs) { + Alignment first; + Alignment second; + gaf_to_alignment(node_to_length, node_to_sequence, gaf_pair.first, first); + gaf_to_alignment(node_to_length, node_to_sequence, gaf_pair.second, second); + alignment_pairs.emplace_back(std::move(first), std::move(second)); + } + lambda(alignment_pairs); + }; + + size_t count = grouped_for_each_parallel_after_wait( + get_pair, get_key, convert_and_call, single_threaded_until_true, batch_size); + hts_close(in); + if (unmatched_record) { + throw runtime_error("interleaved GAF input contains an unmatched alignment"); + } + return count; +} + +size_t gaf_paired_grouped_for_each_parallel_after_wait( + const HandleGraph& graph, + const string& filename, + function>&)> lambda, + function single_threaded_until_true, + 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_paired_grouped_for_each_parallel_after_wait( + node_to_length, node_to_sequence, filename, lambda, + single_threaded_until_true, batch_size); +} + + string supplementary_tag_value(const Alignment& primary) { stringstream strm;