Skip to content

Fix prebuilt-index alignment, split-index output, and write failures #34

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DylanCosto wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix prebuilt-index alignment, split-index output, and write failures #34

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions bonefish-core/src/align/split.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,8 +154,9 @@ pub struct MergeConfig<'a> {
pub n_parts: usize,
pub k: usize,
pub w: usize,
pub query_path: &'a str,
pub query_paths: &'a [String],
pub sam_rg: Option<&'a str>,
pub input_rg: &'a [String],
pub copy_comment: bool,
}

Expand All @@ -171,7 +172,6 @@ pub fn merge_split_results(
let n_parts = config.n_parts;
let k = config.k;
let w = config.w;
let query_path = config.query_path;
let sam_rg = config.sam_rg;
let copy_comment = config.copy_comment;
// 1. Read all temp file headers, build global seqs + rid_shifts
Expand All @@ -196,29 +196,32 @@ pub fn merge_split_results(
let expanded = rg_line.replace("\\t", "\t");
writeln!(handle, "{}", expanded)?;
}
for line in config.input_rg { writeln!(handle, "{}", line)?; }
}

// 4. Re-open query files for output
let query_reader = crate::fasta::open(query_path)
.map_err(|e| anyhow::anyhow!("Error re-opening query for merge: {}", e))?;

let mut total_stats = AlignmentStats::default();
let mut record_iter = query_reader.records();
let mut output_buffer = String::new();

{
for query_path in config.query_paths {
let query_reader = crate::fasta::open(query_path)
.map_err(|e| anyhow::anyhow!("Error re-opening query for merge: {}", e))?;
let mut record_iter = query_reader.records();
// Single-end merge
loop {
let rec = match record_iter.next() {
Some(Ok(r)) => r,
Some(Err(e)) => { eprintln!("Warning: Error reading record: {}", e); continue; },
Some(Err(e)) => anyhow::bail!("{}: {}", query_path, e),
None => break,
};

let qname = rec.name().to_string();
let qseq = rec.sequence().to_vec();
let qual = rec.quality().map(|qs| String::from_utf8_lossy(qs).to_string());
let comment = if copy_comment { rec.description().map(|s| s.to_string()) } else { None };
#[cfg(not(target_arch = "wasm32"))]
let comment = if copy_comment { rec.description().and_then(crate::fasta::bam::filter_comment_tags) } else { None };
#[cfg(target_arch = "wasm32")]
let comment = if copy_comment { rec.description().map(str::to_string) } else { None };
let qlen = qseq.len();

// Read results from all parts
Expand Down
90 changes: 70 additions & 20 deletions bonefish/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -851,8 +851,6 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
anyhow::bail!("Input file '{}' not found.", qf);
}
}
let query_path = query_files.first().map(|s| s.as_str()).unwrap_or("-");

// Check the reference and the output combination before -o creates (truncates) its file.
std::fs::File::open(target_path).map_err(|e| anyhow::anyhow!("cannot open reference '{}': {}", target_path, e))?;
if cli.sort {
Expand Down Expand Up @@ -920,7 +918,8 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
|| Index::sniff_is_index_file(target_path);
let mut n_parts = 0usize;
let mut sam_header_written = false;
let split_prefix = cli.split_prefix.clone();
let mut split_output = SplitOutput { prefix: cli.split_prefix.clone(), temp_dir: None };
split_output.prepare(false, &query_files)?;
let mut saved_k = k;
let mut saved_w = w;

Expand Down Expand Up @@ -961,6 +960,13 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
None => break,
};
n_parts += 1;
if n_parts == 1 {
let has_more = match &mapped_parts {
Some(parts) => !parts.is_empty(),
None => !std::io::BufRead::fill_buf(&mut idx_reader)?.is_empty(),
};
split_output.prepare(cli.output_sam && has_more, &query_files)?;
}
if let Some(dw) = dump_writer.as_mut() {
if cli.idx_no_seq {
let mut idx_copy = mi.clone();
Expand All @@ -983,18 +989,16 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
}
saved_k = mi.kmer_size;
saved_w = mi.window_size;
bonefish::api::finalize_options(&mut opt, mi.kmer_size);
if n_parts == 1 && (cli.k.is_some_and(|v| v != mi.kmer_size) || cli.w.is_some_and(|v| v != mi.window_size)) {
eprintln!("Warning: -k/-w are ignored with a prebuilt index (it has k={}, w={})", mi.kmer_size, mi.window_size);
}
if let Some(ref alt_path) = cli.alt {
read_alt_list(&mut mi, alt_path)?;
}

if let Some(ref prefix) = split_prefix {
if let Some(ref prefix) = split_output.prefix {
// Split mode: write results to temp files
if query_path == "-" && n_parts > 1 {
anyhow::bail!("Cannot use stdin for queries with multi-part index. Use a named file or increase -I.");
}
let mut split_writer = split::split_init(prefix, n_parts - 1, &mi)
.map_err(|e| anyhow::anyhow!("Failed to init split temp file: {}", e))?;
for qf in &query_files {
Expand Down Expand Up @@ -1077,13 +1081,22 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
}
}

for qf in &query_files {
let part_stats = map_one_part(
&mi, &mut opt, &run_cfg,
qf, 1, &mut sam_header_written,
&mut handle,
)?;
total_stats = total_stats + part_stats;
if let Some(ref prefix) = split_output.prefix {
let mut split_writer = split::split_init(prefix, 0, &mi)?;
for qf in &query_files {
total_stats = total_stats + map_one_part_split(
&mi, &mut opt, &run_cfg, qf, 1, &mut split_writer,
)?;
}
} else {
for qf in &query_files {
let part_stats = map_one_part(
&mi, &mut opt, &run_cfg,
qf, 1, &mut sam_header_written,
&mut handle,
)?;
total_stats = total_stats + part_stats;
}
}
} else {
// ─── Multi-part from mmap: split all_seqs into batches ───
Expand All @@ -1101,6 +1114,7 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
n_parts: &mut usize, opt: &mut MapOptions, dump_writer: &mut Option<std::io::BufWriter<std::fs::File>>,
sam_header_written: &mut bool, total_stats: &mut AlignmentStats| -> anyhow::Result<()>
{
split_output.prepare(cli.output_sam, &query_files)?;
let seqs = std::mem::take(batch_seqs);
*batch_bases = 0;
*n_parts += 1;
Expand Down Expand Up @@ -1128,7 +1142,7 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
}
}

if let Some(ref prefix) = split_prefix {
if let Some(ref prefix) = split_output.prefix {
let mut split_writer = split::split_init(prefix, *n_parts - 1, &mi)
.map_err(|e| anyhow::anyhow!("Failed to init split temp file: {}", e))?;
for qf in &query_files {
Expand Down Expand Up @@ -1183,6 +1197,7 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
let (batch_seqs, is_eof) = reader.read_batch(batch_size)
.map_err(|e| anyhow::anyhow!("cannot read reference '{}': {}", target_path, e))?;
if batch_seqs.is_empty() { break; }
split_output.prepare(cli.output_sam && !is_eof, &query_files)?;

n_parts += 1;
let t0 = Instant::now();
Expand Down Expand Up @@ -1210,7 +1225,7 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
}
}

if let Some(ref prefix) = split_prefix {
if let Some(ref prefix) = split_output.prefix {
let mut split_writer = split::split_init(prefix, n_parts - 1, &mi)
.map_err(|e| anyhow::anyhow!("Failed to init split temp file: {}", e))?;
for qf in &query_files {
Expand All @@ -1237,20 +1252,21 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
}
}

if n_parts > 1 && split_prefix.is_none() {
if n_parts > 1 && split_output.prefix.is_none() {
eprintln!("[WARNING] Index has {} parts. Without --split-prefix, MAPQ and secondary flags may be inaccurate.", n_parts);
}

// Merge split results if --split-prefix was used
if let Some(ref prefix) = split_prefix && n_parts >= 1 {
if let Some(ref prefix) = split_output.prefix && n_parts >= 1 {
eprintln!("[*] Merging {} split parts...", n_parts);
let merge_cfg = split::MergeConfig {
prefix,
n_parts,
k: saved_k,
w: saved_w,
query_path,
query_paths: &query_files,
sam_rg: cli.sam_rg.as_deref(),
input_rg: &run_cfg.input_rg,
copy_comment,
};
let merge_stats = split::merge_split_results(
Expand All @@ -1259,6 +1275,7 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
total_stats = total_stats + merge_stats;
}

handle.flush()?;
drop(handle);
if let Some(w) = sorted_out {
let t_sort = Instant::now();
Expand Down Expand Up @@ -1299,6 +1316,36 @@ fn run(mut cli: AlignArgs) -> anyhow::Result<()> {
Ok(())
}

/// Automatically merge multi-part SAM/BAM output so its header contains every reference.
/// The user can still choose a temporary-file location with --split-prefix.
struct SplitOutput {
prefix: Option<String>,
temp_dir: Option<std::path::PathBuf>,
}

impl SplitOutput {
fn prepare(&mut self, multi_part_sam: bool, queries: &[String]) -> anyhow::Result<()> {
if self.prefix.is_none() && !multi_part_sam { return Ok(()); }
if queries.iter().any(|q| q == "-" || std::fs::metadata(q).is_ok_and(|m| !m.is_file())) {
anyhow::bail!("Split-index merging requires regular query files, not stdin or pipes. Increase -I to map in one part.");
}
if self.prefix.is_none() {
let stamp = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?.as_nanos();
let dir = std::env::temp_dir().join(format!("bonefish-split-{}-{}", std::process::id(), stamp));
std::fs::create_dir(&dir)?;
self.prefix = Some(dir.join("part").to_string_lossy().into_owned());
self.temp_dir = Some(dir);
}
Ok(())
}
}

impl Drop for SplitOutput {
fn drop(&mut self) {
if let Some(dir) = &self.temp_dir { let _ = std::fs::remove_dir_all(dir); }
}
}

/// Run-level configuration fields that stay constant across index parts.
struct RunConfig<'a> {
mid_occ_frac: f32,
Expand Down Expand Up @@ -1343,7 +1390,7 @@ fn map_one_part(
}
load_junctions(opt, mi, cli)?;

// SAM header: write once for single-part, or per-part for multi-part without split-prefix
// Single-part SAM header; multi-part SAM/BAM headers are written by the split merge.
if cli.output_sam && !*sam_header_written {
writeln!(handle, "@HD\tVN:1.6\tSO:unsorted\tGO:query")?;
for s in &mi.seqs {
Expand Down Expand Up @@ -1547,6 +1594,9 @@ fn map_one_part(
}
}
});
// A failed writer stops the coordinator from draining the input queue.
// Disconnect it before the scope waits for a reader blocked in send().
drop(rx);
drop(tx_out);
writer_handle.join().expect("writer thread panicked")
})?;
Expand Down
80 changes: 80 additions & 0 deletions bonefish/tests/regression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,86 @@ fn raw_index_gives_same_output() {
let direct = run(&["-a", "-t", "2", reference, reads]);
let mapped = run(&["-a", "-t", "2", raw_s, reads]);
assert!(direct == mapped, "raw index output differs from FASTA-built index");
// A loaded index supplies k for chaining as well as seeding. The preset's
// default k (19) must not change mappings from this k=15 index.
run(&["-k15", "-w10", "-d", raw_s, "--idx-raw", reference]);
let direct = run(&["-a", "-t2", "-k15", "-w10", reference, reads]);
let mapped = run(&["-a", "-t2", raw_s, reads]);
assert_eq!(direct, mapped, "loaded index used the preset's chaining penalty");
std::fs::remove_dir_all(&dir).ok();
}

#[test]
fn split_index_preserves_inputs_and_reference_headers() {
let dir = tmp_dir("split");
let reference = dir.join("ref.fa");
let gz = dir.join("ref.fa.gz");
let (q1, q2) = (dir.join("q1.fa"), dir.join("q2.fa"));
let mut state = 372u64;
let seq: String = (0..12000).map(|_| {
state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
b"ACGT"[(state >> 32) as usize & 3] as char
}).collect();
let text = format!(">a\n{}\n>b\n{}\n", &seq[..6000], &seq[6000..]);
std::fs::write(&reference, &text).unwrap();
let mut encoder = flate2::write::GzEncoder::new(std::fs::File::create(&gz).unwrap(), flate2::Compression::fast());
encoder.write_all(text.as_bytes()).unwrap();
encoder.finish().unwrap();
std::fs::write(&q1, format!(">read_a\n{}\n", &seq[1000..3000])).unwrap();
std::fs::write(&q2, format!(">read_b\n{}\n", &seq[7000..9000])).unwrap();
let (reference, q1, q2) = (reference.to_str().unwrap(), q1.to_str().unwrap(), q2.to_str().unwrap());
let expected = run(&["-a", "-t2", reference, q1, q2]);
let prefix = dir.join("split");
let prefix = prefix.to_str().unwrap();
// The single-part FASTA fast path must create the same split files as other paths.
assert_eq!(run(&["-a", "-t2", "--split-prefix", prefix, reference, q1, q2]), expected);
for target in [reference, gz.to_str().unwrap()] {
assert_eq!(run(&["-a", "-t2", "-I5k", "--split-prefix", prefix, target, q1, q2]), expected);
assert_eq!(run(&["-a", "-t2", "-I5k", target, q1, q2]), expected);
}
for (file, raw) in [("ref.idx", false), ("ref.raw", true)] {
let index = dir.join(file);
let mut args = vec!["-t2", "-I5k", "-d", index.to_str().unwrap()];
if raw { args.push("--idx-raw"); }
args.push(reference);
run(&args);
assert_eq!(run(&["-a", "-t2", index.to_str().unwrap(), q1, q2]), expected);
}
let bam = dir.join("out.bam");
run(&["-t2", "-I5k", "-o", bam.to_str().unwrap(), reference, q1, q2]);
let mut bytes = Vec::new();
flate2::read::MultiGzDecoder::new(std::fs::File::open(bam).unwrap()).read_to_end(&mut bytes).unwrap();
let header_len = u32::from_le_bytes(bytes[4..8].try_into().unwrap()) as usize;
assert_eq!(u32::from_le_bytes(bytes[8 + header_len..12 + header_len].try_into().unwrap()), 2);
std::fs::remove_dir_all(&dir).ok();
}

#[cfg(target_os = "linux")]
#[test]
fn output_errors_fail_without_hanging() {
use std::process::Stdio;
use std::time::{Duration, Instant};
let dir = tmp_dir("write-error");
let reads = dir.join("reads.fa");
for count in [1, 100] {
std::fs::write(&reads, (0..count).map(|i| format!(">r{i}\n{}\n", "A".repeat(2000))).collect::<String>()).unwrap();
let mut child = Command::new(env!("CARGO_BIN_EXE_bonefish"))
.args(["-a", "-t2", "-K1k", "-o", "/dev/full"])
.arg(data("ref.fa.gz")).arg(&reads)
.stdout(Stdio::null()).stderr(Stdio::piped()).spawn().unwrap();
let start = Instant::now();
while child.try_wait().unwrap().is_none() {
if start.elapsed() > Duration::from_secs(5) {
child.kill().unwrap();
child.wait().unwrap();
panic!("output failure left the reader blocked ({count} reads)");
}
std::thread::sleep(Duration::from_millis(20));
}
let out = child.wait_with_output().unwrap();
assert!(!out.status.success(), "output failure was reported as success");
assert!(String::from_utf8_lossy(&out.stderr).contains("No space left on device"));
}
std::fs::remove_dir_all(&dir).ok();
}

Expand Down
5 changes: 5 additions & 0 deletions docs/usage.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,11 @@ bonefish -x tx -d tx.raw --idx-raw transcripts.fa
Without `--idx-raw`, `-d` writes a smaller index that takes a few seconds to load. Indexes are
bonefish's own format. minimap2 `.mmi` files can be read but not written.

SAM/BAM output automatically merges results when a reference has multiple index parts
(`-I` controls the number of reference bases per part). The merge rereads the queries, so
use regular query files rather than stdin or pipes. `--split-prefix PATH` chooses where
the intermediate files are written; otherwise the system temporary directory is used.

## Input

- **Formats:** FASTQ, FASTA or BAM, plain or gzipped, from files, stdin (`-`) or a pipe. Several
Expand Down
Loading