From 75d19d357098013c09578177a6818f90cb7546c0 Mon Sep 17 00:00:00 2001 From: Travis Date: Tue, 6 Oct 2026 16:41:21 -0700 Subject: [PATCH] Write output with positioned parallel writes; make LCP optional Port of the parallel-writer branch from my attempts to revise main (PR #13: 4cd61e9, 93e5e32, 022761b) onto dev. The writer no longer streams the file through one BufWriter. open_output writes the header and the text; each partition's suffix array and LCP are written in parallel at precomputed offsets in 64 MiB chunks; finish_output appends the sequence names and patches the header. write_at sits next to read_exact_at in file_access.rs with the same cfg(unix)/cfg(windows) split. The LCP array is optional and off by default (SufrBuilderArgs.write_lcp, `sufr create --write-lcp`). lcp_pos = 0 in the header means no LCP array; SufrMetadata.has_lcp reports it, summarize prints it, and `list --show-lcp` exits with an error when it is absent. Without an LCP, the reader serves a max-query-len shorter than the build-time value by comparing text slices (care positions for a seed mask, the first mql bytes otherwise), streaming the suffix array in 4M-element blocks. FileAccessIter refills 16 MiB at a time instead of 1 GiB and reuses its buffer. Implemented with assistance from Claude (Fable 5.1) --- CHANGELOG.md | 14 ++ libsufr/src/file_access.rs | 33 ++++- libsufr/src/lib.rs | 8 + libsufr/src/subsample_tests.rs | 143 ++++++++++++++++++ libsufr/src/suffix_array.rs | 2 + libsufr/src/sufr_builder.rs | 262 +++++++++++++++++++++++---------- libsufr/src/sufr_file.rs | 151 ++++++++++++++++++- libsufr/src/types.rs | 8 + sufr/src/lib.rs | 10 ++ 9 files changed, 543 insertions(+), 88 deletions(-) create mode 100644 libsufr/src/subsample_tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index be6d405..22ef8ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 improving build time on some inputs - Implement an LCP cache, significantly improving build time on some inputs with very long duplications or tandem repeats +- Add `SufrBuilderArgs::write_lcp` and `sufr create --write-lcp`. The LCP + array is now optional and off by default. `search`, `count`, `locate`, + and `extract` do not read it; `list --show-lcp` fails on a file without + one. The header's LCP offset is 0 when there is no LCP array, and + `SufrMetadata::has_lcp` and `summarize` report whether one is present. +- Subsample the suffix array for a query-time max query length without an + LCP array by streaming the suffix array in blocks of four million + elements and comparing each suffix with its predecessor in the text ### Changed @@ -28,6 +36,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Reduce some unnecessary internal allocations - During build, the merge sort makes better use of the available threads with work-stealing via `rayon::join` +- Write the output file with positioned writes (`pwrite` on Unix, + `seek_write` on Windows) in parallel 64 MiB chunks instead of one + sequential `BufWriter`, and patch the header with the array offsets at the end +- `FileAccessIter` refills its buffer with 16 MiB (2^24 bytes) of whole + elements instead of 1 GiB, so iterating over an on-disk array smaller + than 1 GiB no longer reads it whole ### Removed diff --git a/libsufr/src/file_access.rs b/libsufr/src/file_access.rs index 7243f7b..4d4eee3 100644 --- a/libsufr/src/file_access.rs +++ b/libsufr/src/file_access.rs @@ -123,8 +123,8 @@ pub struct FileAccessIter<'a, T: Int> { } impl<'a, T: Int> FileAccessIter<'a, T> { - /// The maximum size in bytes of the buffer (currently 2^30) - const BUFFER_SIZE: usize = 2usize.pow(30); + /// The maximum size in bytes of the buffer (currently 2^24, 16 MiB) + const BUFFER_SIZE: usize = 2usize.pow(24); } impl Iterator for FileAccessIter<'_, T> { @@ -141,22 +141,22 @@ impl Iterator for FileAccessIter<'_, T> { return None; } + // Read whole elements, at most BUFFER_SIZE bytes, straight + // into the reused typed buffer let bytes_wanted = min( Self::BUFFER_SIZE, (self.file_access.end_position - self.current_position) as usize, ); let num_vals = bytes_wanted / mem::size_of::(); - let mut buffer: Vec = vec![T::from_usize(0); num_vals]; + self.buffer.resize(num_vals, T::from_usize(0)); read_exact_at( &self.file_access.file, - slice_int_to_slice_u8(&mut buffer), + slice_int_to_slice_u8(&mut self.buffer), self.current_position, ) .unwrap(); self.current_position += (num_vals * mem::size_of::()) as u64; - - self.buffer = buffer; self.buffer_pos = 0; } @@ -168,7 +168,26 @@ impl Iterator for FileAccessIter<'_, T> { } } -// Position-independent reads for both Unix and Windows +// Position-independent reads and writes for both Unix and Windows + +/// Write from `buf` at an absolute `offset` without using the file cursor, +/// so several threads may write disjoint regions of one file at once. +/// Returns the number of bytes written, which may be fewer than requested. +#[cfg(unix)] +pub(crate) fn write_at(file: &File, buf: &[u8], offset: u64) -> io::Result { + use std::os::unix::fs::FileExt; + + file.write_at(buf, offset) +} + +// Windows version of `write_at`; `seek_write` also takes an absolute +// offset per call, so concurrent calls place their data correctly. +#[cfg(windows)] +pub(crate) fn write_at(file: &File, buf: &[u8], offset: u64) -> io::Result { + use std::os::windows::fs::FileExt; + + file.seek_write(buf, offset) +} #[cfg(unix)] fn read_exact_at(file: &File, buf: &mut [u8], offset: u64) -> io::Result<()> { diff --git a/libsufr/src/lib.rs b/libsufr/src/lib.rs index 6c677fd..df3d752 100644 --- a/libsufr/src/lib.rs +++ b/libsufr/src/lib.rs @@ -23,6 +23,8 @@ mod file_access; mod lcp; +#[cfg(test)] +mod subsample_tests; pub mod suffix_array; pub mod sufr_builder; pub mod sufr_file; @@ -62,6 +64,7 @@ mod tests { sequence_names: seq_data.sequence_names, num_partitions: 2, seed_mask: None, + write_lcp: true, }; let res = SufrBuilder::::new(args); assert!(res.is_ok()); @@ -110,6 +113,7 @@ mod tests { sequence_names: seq_data.sequence_names, num_partitions: 2, seed_mask: None, + write_lcp: true, }; let res = SufrBuilder::::new(args); @@ -157,6 +161,7 @@ mod tests { sequence_names: seq_data.sequence_names, num_partitions: 2, seed_mask: None, + write_lcp: true, }; let res = SufrBuilder::::new(builder_args); assert!(res.is_ok()); @@ -235,6 +240,7 @@ mod tests { sequence_names: seq_data.sequence_names, num_partitions: 1, seed_mask: Some("101".to_string()), + write_lcp: true, }; // 7 $ @@ -279,6 +285,7 @@ mod tests { sequence_names: seq_data.sequence_names, num_partitions: 1, seed_mask: Some("11011".to_string()), + write_lcp: true, }; // 0 16 $ @@ -336,6 +343,7 @@ mod tests { sequence_names: seq_data.sequence_names, num_partitions: 1, seed_mask: Some("11000111".to_string()), + write_lcp: true, }; let res = SufrBuilder::::new(builder_args); diff --git a/libsufr/src/subsample_tests.rs b/libsufr/src/subsample_tests.rs new file mode 100644 index 0000000..4c847f5 --- /dev/null +++ b/libsufr/src/subsample_tests.rs @@ -0,0 +1,143 @@ +//! Tests that subsampling a suffix array for a shorter max query length +//! gives the same result with and without an LCP array on disk. + +use crate::{sufr_builder::SufrBuilder, sufr_file::SufrFile, types::SufrBuilderArgs}; +use anyhow::Result; +use pretty_assertions::assert_eq; +use tempfile::NamedTempFile; + +/// A small deterministic generator (xorshift64*) so the test needs no +/// random-number crate +struct Rng(u64); + +impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545F4914F6CDD1D) + } + + fn below(&mut self, n: usize) -> usize { + (self.next() % n as u64) as usize + } +} + +fn random_text(rng: &mut Rng) -> Vec { + let mut text = vec![]; + for rec in 0..3 { + if rec > 0 { + text.push(b'%'); + } + let len = 300 + rng.below(400); + text.extend((0..len).map(|_| b"ACGT"[rng.below(4)])); + // A few Ns so some suffixes are skipped in DNA mode + let at = rng.below(len - 5); + let start = text.len() - len + at; + for b in text[start..start + 4].iter_mut() { + *b = b'N'; + } + } + text.push(b'$'); + text +} + +/// Build with or without an LCP array and return the open file +fn build( + text: &[u8], + seed_mask: Option<&str>, + max_query_len: Option, + write_lcp: bool, +) -> Result<(NamedTempFile, SufrFile)> { + let outfile = NamedTempFile::new()?; + let outpath = outfile.path().to_string_lossy().to_string(); + SufrBuilder::::new(SufrBuilderArgs { + text: text.to_vec(), + path: Some(outpath.clone()), + low_memory: false, + max_query_len, + is_dna: true, + allow_ambiguity: false, + ignore_softmask: false, + sequence_starts: vec![0], + sequence_names: vec!["s".to_string()], + num_partitions: 3, + seed_mask: seed_mask.map(String::from), + write_lcp, + })?; + let sufr_file: SufrFile = SufrFile::read(&outpath, false)?; + Ok((outfile, sufr_file)) +} + +/// Independent definition: keep rank i when the first `mql` key symbols of +/// SA[i] differ from those of SA[i-1] (a symbol past the text end differs +/// from everything, including another past-the-end symbol). +fn expected_subsample( + text: &[u8], + sa: &[u32], + offsets: &[usize], + mql: usize, +) -> (Vec, Vec) { + let key = |pos: usize| -> Vec> { + offsets[..mql.min(offsets.len())] + .iter() + .map(|&off| text.get(pos + off).copied()) + .collect() + }; + let mut sub_sa = vec![]; + let mut rank = vec![]; + for i in 0..sa.len() { + let keep = i == 0 || { + let a = key(sa[i - 1] as usize); + let b = key(sa[i] as usize); + !a.iter().zip(&b).all(|(x, y)| x.is_some() && x == y) + }; + if keep { + sub_sa.push(sa[i]); + rank.push(i as u32); + } + } + (sub_sa, rank) +} + +#[test] +fn subsample_without_lcp_matches_definition_and_lcp_version() -> Result<()> { + let mut rng = Rng(11); + let text = random_text(&mut rng); + + // Seed-mask index, query-time MQL shorter than the weight + let mask = "11101101101111"; + let offsets: Vec = mask + .bytes() + .enumerate() + .filter(|(_, b)| *b == b'1') + .map(|(i, _)| i) + .collect(); + for mql in [1, 2, 4, 7, 10] { + let (_f1, mut with_lcp) = build(&text, Some(mask), None, true)?; + let (_f2, mut without_lcp) = build(&text, Some(mask), None, false)?; + assert!(with_lcp.has_lcp); + assert!(!without_lcp.has_lcp); + let sa: Vec = without_lcp.suffix_array_file.iter().collect(); + let expected = expected_subsample(&text, &sa, &offsets, mql); + let from_text = without_lcp.subsample_suffix_array(mql); + assert_eq!(from_text, expected, "mask, mql={mql}, vs definition"); + let from_lcp = with_lcp.subsample_suffix_array(mql); + assert_eq!(from_text, from_lcp, "mask, mql={mql}, vs LCP array"); + } + + // Max-query-len index, query-time MQL shorter than the build-time one + let built_mql = 12; + let offsets: Vec = (0..built_mql).collect(); + for mql in [1, 3, 6, 11] { + let (_f1, mut with_lcp) = build(&text, None, Some(built_mql), true)?; + let (_f2, mut without_lcp) = build(&text, None, Some(built_mql), false)?; + let sa: Vec = without_lcp.suffix_array_file.iter().collect(); + let expected = expected_subsample(&text, &sa, &offsets, mql); + let from_text = without_lcp.subsample_suffix_array(mql); + assert_eq!(from_text, expected, "mql index, mql={mql}, vs definition"); + let from_lcp = with_lcp.subsample_suffix_array(mql); + assert_eq!(from_text, from_lcp, "mql index, mql={mql}, vs LCP array"); + } + Ok(()) +} diff --git a/libsufr/src/suffix_array.rs b/libsufr/src/suffix_array.rs index 980883f..65538e6 100644 --- a/libsufr/src/suffix_array.rs +++ b/libsufr/src/suffix_array.rs @@ -117,6 +117,7 @@ impl SuffixArray { /// sequence_names: seq_data.sequence_names, /// num_partitions: 16, /// seed_mask: None, + /// write_lcp: true, /// }; /// /// let suffix_array = SuffixArray::new(builder_args)?; @@ -446,6 +447,7 @@ impl SuffixArray { /// sequence_names: seq_data.sequence_names, /// num_partitions: 16, /// seed_mask: None, + /// write_lcp: true, /// }; /// /// let outpath = SuffixArray::write(builder_args)?; diff --git a/libsufr/src/sufr_builder.rs b/libsufr/src/sufr_builder.rs index 32afa79..4727626 100644 --- a/libsufr/src/sufr_builder.rs +++ b/libsufr/src/sufr_builder.rs @@ -9,6 +9,7 @@ //! use crate::{ + file_access::write_at, lcp::{lcp, LcpCache}, types::{ Int, SeedMask, SuffixSortType, SufrBuilderArgs, OUTFILE_VERSION, @@ -24,7 +25,7 @@ use std::{ cell::RefCell, cmp::{max, min, Ordering}, fs::{self, File, OpenOptions}, - io::{BufWriter, Read, Seek, SeekFrom, Write}, + io::{Read, Write}, mem, ops::Range, path::PathBuf, @@ -85,6 +86,51 @@ pub struct SufrBuilder> { /// The name of the output file pub path: String, + + /// Whether to write the LCP array + pub write_lcp: bool, +} + +// -------------------------------------------------- +/// The output file after the header and text have been written. +/// The suffix and LCP arrays are written into it at known offsets, +/// possibly from several threads. +struct Output { + file: File, + + /// Where the header records the text/SA/LCP positions + locs_pos: u64, + + /// Byte position of the text + text_pos: usize, + + /// Byte position of the suffix array + sa_pos: usize, + + /// Byte position of the LCP array, or 0 when none is written + lcp_pos: usize, + + /// Byte position after the arrays, where the sequence names go + end_pos: usize, +} + +/// Write all of `buf` at `offset`, in parallel 64 MB chunks. +fn write_all_at(file: &File, buf: &[u8], offset: u64) -> Result<()> { + let chunk_len = 1 << 26; + buf.par_chunks(chunk_len) + .enumerate() + .try_for_each(|(i, mut chunk)| -> Result<()> { + let mut pos = offset + (i * chunk_len) as u64; + while !chunk.is_empty() { + let n = write_at(file, chunk, pos)?; + if n == 0 { + bail!("Failed to write to output file"); + } + chunk = &chunk[n..]; + pos += n as u64; + } + Ok(()) + }) } // -------------------------------------------------- @@ -122,6 +168,7 @@ impl + Send + Sync> SufrBuilder { /// sequence_names: seq_data.sequence_names, /// num_partitions: 1024, /// seed_mask: None, + /// write_lcp: true, /// }; /// /// if text_len < u32::MAX as u64 { @@ -207,6 +254,7 @@ impl + Send + Sync> SufrBuilder { partitions: vec![], n_ranges, path: args.path.unwrap_or("out.sufr".to_string()), + write_lcp: args.write_lcp, }; sa.sort(args.num_partitions, occupancy)?; sa.write()?; @@ -625,8 +673,13 @@ impl + Send + Sync> SufrBuilder { // Write to disk let mut sa_buffer = B::default(); sa_buffer.extend_from_slice(&part_sa)?; - let mut lcp_buffer = B::default(); - lcp_buffer.extend_from_slice(&lcp)?; + let lcp_buffer = if self.write_lcp { + let mut lcp_buffer = B::default(); + lcp_buffer.extend_from_slice(&lcp)?; + Some(lcp_buffer) + } else { + None + }; *partition = Some(Partition { order: partition_num, @@ -839,39 +892,100 @@ impl + Send + Sync> SufrBuilder { // -------------------------------------------------- /// Serialize contents of the sorted partitions to a _.sufr_ file. + /// The header and text are written first, then each partition's + /// suffix array (and LCP array, when requested) is written in + /// parallel at its precomputed offset, then the header is patched + /// with the array positions. /// Returns the number of bytes written to disk. - /// - /// Args: - /// * `filename`: the name of the output file. fn write(&mut self) -> Result { - let filename = &self.path; - let mut file = BufWriter::new( - File::create(filename).map_err(|e| anyhow!("{filename}: {e}"))?, - ); + let now = Instant::now(); + let out = self.open_output(self.num_suffixes.to_usize())?; + let int_size = mem::size_of::(); + + // Offset (in elements) of each partition in the arrays, and the + // last suffix of the preceding partition for fixing the LCP boundary + let mut offsets = Vec::with_capacity(self.partitions.len()); + let mut total = 0; + let mut prev_last_suffix = None; + for partition in &self.partitions { + offsets.push((total, prev_last_suffix)); + total += partition.len; + prev_last_suffix = Some(partition.last_suffix); + } + + // NB: Taking the partitions and consuming their buffers frees + // memory/disk space along the way + let partitions = mem::take(&mut self.partitions); + partitions.into_par_iter().zip(offsets).try_for_each( + |(partition, (offset, prev_last_suffix))| -> Result<()> { + let sa = partition.sa_buffer.consume()?; + write_all_at( + &out.file, + vec_to_slice_u8(&sa), + (out.sa_pos + offset * int_size) as u64, + )?; + drop(sa); + + if let Some(lcp_buffer) = partition.lcp_buffer { + let mut lcp = lcp_buffer.consume()?; + if let Some(prev_last_suffix) = prev_last_suffix { + // Fix LCP boundary + if let Some(val) = lcp.first_mut() { + *val = self.find_lcp( + prev_last_suffix, + partition.first_suffix, + self.text_len, + 0, // start at beginning + None, + ); + } + } + write_all_at( + &out.file, + vec_to_slice_u8(&lcp), + (out.lcp_pos + offset * int_size) as u64, + )?; + } + Ok(()) + }, + )?; - let mut bytes_out: usize = 0; + let bytes_out = self.finish_output(&out)?; + info!( + "Wrote suffix array to '{}' in {:?}", + self.path, + now.elapsed() + ); + Ok(bytes_out) + } - // TODO (throughout this method): write() can write less than the whole buffer + // -------------------------------------------------- + /// Create the output file and write the header and text. The suffix + /// (and optional LCP) arrays are written afterwards at the positions + /// recorded in the returned `Output`. + fn open_output(&self, num_suffixes: usize) -> Result { + let now = Instant::now(); + let filename = &self.path; + let file = File::create(filename).map_err(|e| anyhow!("{filename}: {e}"))?; // Various metadata let is_dna: u8 = if self.is_dna { 1 } else { 0 }; let allow_ambiguity: u8 = if self.allow_ambiguity { 1 } else { 0 }; let ignore_softmask: u8 = if self.ignore_softmask { 1 } else { 0 }; - bytes_out += - file.write(&[OUTFILE_VERSION, is_dna, allow_ambiguity, ignore_softmask])?; + let mut header: Vec = + vec![OUTFILE_VERSION, is_dna, allow_ambiguity, ignore_softmask]; // Text length - bytes_out += file.write(&self.text_len.to_usize().to_le_bytes())?; + header.extend(self.text_len.to_usize().to_le_bytes()); - // Locations of text, suffix array, and LCP - // Will be corrected at the end - let locs_pos = file.stream_position()?; - bytes_out += file.write(&0usize.to_le_bytes())?; - bytes_out += file.write(&0usize.to_le_bytes())?; - bytes_out += file.write(&0usize.to_le_bytes())?; + // Locations of text, suffix array, and LCP; filled in by finish_output + let locs_pos = header.len() as u64; + header.extend(0usize.to_le_bytes()); + header.extend(0usize.to_le_bytes()); + header.extend(0usize.to_le_bytes()); // Number of suffixes - bytes_out += file.write(&self.num_suffixes.to_usize().to_le_bytes())?; + header.extend(num_suffixes.to_le_bytes()); // Max query length let max_query_len = if let SuffixSortType::MaxQueryLen(val) = &self.sort_type { @@ -879,73 +993,65 @@ impl + Send + Sync> SufrBuilder { } else { 0 }; - bytes_out += file.write(&max_query_len.to_le_bytes())?; + header.extend(max_query_len.to_le_bytes()); // Number of sequences - bytes_out += file.write(&self.sequence_starts.len().to_le_bytes())?; + header.extend(self.sequence_starts.len().to_le_bytes()); // Sequence starts - bytes_out += file.write(vec_to_slice_u8(&self.sequence_starts))?; + header.extend_from_slice(vec_to_slice_u8(&self.sequence_starts)); // Seed mask match &self.sort_type { SuffixSortType::Mask(seed_mask) => { - bytes_out += file.write(&seed_mask.bytes.len().to_le_bytes())?; - file.write_all(&seed_mask.bytes)?; - bytes_out += seed_mask.bytes.len(); + header.extend(seed_mask.bytes.len().to_le_bytes()); + header.extend_from_slice(&seed_mask.bytes); } - _ => bytes_out += file.write(&0usize.to_le_bytes())?, - } - - // Text - let text_pos = bytes_out; - file.write_all(&self.text)?; - bytes_out += self.text.len(); - - // Stitch partitioned suffix files together - // NB: Using mem::take() and consume() to free memory/disk space along the way - let sa_pos = bytes_out; - for partition in &mut self.partitions { - let sa_buffer = mem::take(&mut partition.sa_buffer).consume()?; - let sa_bytes = vec_to_slice_u8(&sa_buffer); - bytes_out += sa_bytes.len(); - file.write_all(sa_bytes)?; + _ => header.extend(0usize.to_le_bytes()), } - let lcp_pos = bytes_out; - - // Stitch partitioned LCP files together - for i in 0..self.partitions.len() { - let mut lcp = mem::take(&mut self.partitions[i].lcp_buffer).consume()?; - - if i != 0 { - // Fix LCP boundary - if let Some(val) = lcp.first_mut() { - *val = self.find_lcp( - self.partitions[i - 1].last_suffix, - self.partitions[i].first_suffix, - self.text_len, - 0, // start at beginning - None, - ); - } - } + write_all_at(&file, &header, 0)?; - let lcp_bytes = vec_to_slice_u8(&lcp); - bytes_out += lcp_bytes.len(); - file.write_all(lcp_bytes)?; - } + // Text + let text_pos = header.len(); + write_all_at(&file, &self.text, text_pos as u64)?; + + let int_size = mem::size_of::(); + let sa_pos = text_pos + self.text.len(); + let sa_end = sa_pos + num_suffixes * int_size; + let (lcp_pos, end_pos) = if self.write_lcp { + (sa_end, sa_end + num_suffixes * int_size) + } else { + (0, sa_end) + }; + info!("Wrote header and text in {:?}", now.elapsed()); + + Ok(Output { + file, + locs_pos, + text_pos, + sa_pos, + lcp_pos, + end_pos, + }) + } + // -------------------------------------------------- + /// Write the sequence names after the arrays and record the array + /// positions in the header. Returns the total number of bytes in + /// the file. + fn finish_output(&self, out: &Output) -> Result { // Sequence names are variable in length so they are at the end - bytes_out += file.write(&bincode::serialize(&self.sequence_names)?)?; + let names = bincode::serialize(&self.sequence_names)?; + write_all_at(&out.file, &names, out.end_pos as u64)?; // Go back to header and record the locations - file.seek(SeekFrom::Start(locs_pos))?; - let _ = file.write(&text_pos.to_le_bytes())?; - let _ = file.write(&sa_pos.to_le_bytes())?; - let _ = file.write(&lcp_pos.to_le_bytes())?; + let mut locs = out.text_pos.to_le_bytes().to_vec(); + locs.extend(out.sa_pos.to_le_bytes()); + locs.extend(out.lcp_pos.to_le_bytes()); + write_all_at(&out.file, &locs, out.locs_pos)?; - Ok(bytes_out) + Ok(out.end_pos + names.len()) } } @@ -1054,8 +1160,8 @@ struct Partition { /// The buffer containing the suffix array. sa_buffer: B, - /// The buffer containing the LCP array. - lcp_buffer: B, + /// The buffer containing the LCP array, when one is written. + lcp_buffer: Option, } // -------------------------------------------------- @@ -1268,6 +1374,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: None, + write_lcp: true, }; let sufr = SufrBuilder::::new(args)?; @@ -1309,6 +1416,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: None, + write_lcp: true, }; let sufr = SufrBuilder::::new(args)?; @@ -1353,6 +1461,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: Some("101".to_string()), + write_lcp: true, }; let sufr: SufrBuilder = SufrBuilder::new(args)?; @@ -1398,6 +1507,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: None, + write_lcp: true, }; let sufr: SufrBuilder = SufrBuilder::new(args)?; @@ -1445,6 +1555,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: Some("1101".to_string()), + write_lcp: true, }; let sufr: SufrBuilder = SufrBuilder::new(args)?; @@ -1482,6 +1593,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: None, + write_lcp: true, }; let sufr: SufrBuilder = SufrBuilder::new(args)?; @@ -1516,6 +1628,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: None, + write_lcp: true, }; let sufr: SufrBuilder = SufrBuilder::new(args)?; @@ -1570,6 +1683,7 @@ mod test { sequence_names: vec!["1".to_string()], num_partitions: 2, seed_mask: Some("101".to_string()), + write_lcp: true, }; let sufr: SufrBuilder = SufrBuilder::new(args)?; diff --git a/libsufr/src/sufr_file.rs b/libsufr/src/sufr_file.rs index e6773c6..087ef2d 100644 --- a/libsufr/src/sufr_file.rs +++ b/libsufr/src/sufr_file.rs @@ -59,9 +59,13 @@ pub struct SufrFile { /// The byte position in the file where the suffix array begins. pub suffix_array_pos: usize, - /// The byte position in the file where the LCP array begins. + /// The byte position in the file where the LCP array begins, + /// or 0 when the file has no LCP array. pub lcp_pos: usize, + /// Whether or not the file has an on-disk LCP array + pub has_lcp: bool, + /// How the suffixes were sorted (fully, max query length, spaced seeds) pub sort_type: SuffixSortType, @@ -219,9 +223,13 @@ impl SufrFile { FileAccess::new(filename, suffix_array_pos as u64, len_suffixes)?; file.seek_relative(suffix_array_file.size as i64)?; - // LCP - let lcp_file: FileAccess = - FileAccess::new(filename, lcp_pos as u64, len_suffixes)?; + // LCP (lcp_pos == 0 means the file was written without an LCP array) + let has_lcp = lcp_pos > 0; + let lcp_file: FileAccess = if has_lcp { + FileAccess::new(filename, lcp_pos as u64, len_suffixes)? + } else { + FileAccess::new(filename, suffix_array_pos as u64, 0)? + }; file.seek_relative(lcp_file.size as i64)?; // Sequence names are variable in length so they are at the end @@ -246,6 +254,7 @@ impl SufrFile { text_pos, suffix_array_pos, lcp_pos, + has_lcp, text_len: T::from_usize(text_len), len_suffixes: T::from_usize(len_suffixes), sort_type, @@ -410,7 +419,9 @@ impl SufrFile { } // -------------------------------------------------- - /// Subsample a suffix array using a maximum query length + /// Subsample a suffix array using a maximum query length. + /// Uses the LCP array when the file has one, and otherwise compares + /// each suffix with its predecessor in the text. /// /// Args: /// * `max_query_len`: prefix length @@ -418,6 +429,10 @@ impl SufrFile { &mut self, max_query_len: usize, ) -> (Vec, Vec) { + if !self.has_lcp { + return self.subsample_suffix_array_by_text(max_query_len); + } + let max_query_len = T::from_usize(max_query_len); let max_len = self.len_suffixes.to_usize(); @@ -439,6 +454,109 @@ impl SufrFile { (suffix_array, rank) } + // -------------------------------------------------- + /// Subsample without an LCP array: keep a suffix when its first + /// `max_query_len` key symbols differ from its predecessor's, comparing + /// the text directly. In max-query-len mode the key is the first + /// `max_query_len` bytes (compared as slices, which lowers to memcmp); + /// in seed-mask mode it is the first `max_query_len` care positions. + /// The suffix array is streamed from disk in blocks, so memory is the + /// text plus the subsample plus one block. + /// + /// Args: + /// * `max_query_len`: prefix length + fn subsample_suffix_array_by_text( + &mut self, + max_query_len: usize, + ) -> (Vec, Vec) { + let now = Instant::now(); + let text_len = self.text_len.to_usize(); + let text_owned: Vec; + let text: &[u8] = if self.text.is_empty() { + // Reading the whole text cannot be out of range, so this only + // fails on an I/O error, like the iterator over the SA file + text_owned = self + .text_file + .get_range(0..text_len) + .expect("failed to read text"); + &text_owned + } else { + &self.text + }; + + let differs: Box bool + Send + Sync> = match &self + .sort_type + { + SuffixSortType::Mask(seed_mask) => { + let positions = + &seed_mask.positions[..min(max_query_len, seed_mask.weight)]; + Box::new(move |a: usize, b: usize| { + !positions.iter().all(|&offset| { + matches!( + (text.get(a + offset), text.get(b + offset)), + (Some(x), Some(y)) if x == y + ) + }) + }) + } + SuffixSortType::MaxQueryLen(_) => Box::new(move |a: usize, b: usize| { + let end_a = min(a + max_query_len, text_len); + let end_b = min(b + max_query_len, text_len); + // A suffix shorter than max_query_len always starts a new group + end_a - a < max_query_len + || end_b - b < max_query_len + || text[a..end_a] != text[b..end_b] + }), + }; + + // Stream the suffix array in blocks; compare within a block in + // parallel, carrying the last suffix of the previous block across. + let block_len = 1 << 22; + let mut block: Vec = Vec::with_capacity(block_len); + let mut suffix_array: Vec = vec![]; + let mut rank: Vec = vec![]; + let mut prev: Option = None; + let mut seen = 0usize; + let mut suffixes = self.suffix_array_file.iter(); + loop { + block.clear(); + block.extend(suffixes.by_ref().take(block_len)); + if block.is_empty() { + break; + } + let keep: Vec = block + .par_iter() + .enumerate() + .map(|(i, &suffix)| { + let before = if i > 0 { + Some(block[i - 1].to_usize()) + } else { + prev + }; + match before { + None => true, + Some(a) => differs(a, suffix.to_usize()), + } + }) + .collect(); + for (i, (&suffix, &kept)) in block.iter().zip(&keep).enumerate() { + if kept { + suffix_array.push(suffix); + rank.push(T::from_usize(seen + i)); + } + } + seen += block.len(); + prev = block.last().map(|v| v.to_usize()); + } + info!( + "Subsampled SA ({}/{seen}) from text (no LCP array) in {:?}", + suffix_array.len(), + now.elapsed() + ); + + (suffix_array, rank) + } + // -------------------------------------------------- /// Retrieve file metadata /// @@ -461,6 +579,7 @@ impl SufrFile { /// assert_eq!(meta.sequence_starts, vec![0]); /// assert_eq!(meta.sequence_names, vec!["1".to_string()]); /// assert_eq!(meta.sort_type, SuffixSortType::MaxQueryLen(0)); + /// assert_eq!(meta.has_lcp, true); /// /// Ok(()) /// } @@ -487,6 +606,7 @@ impl SufrFile { .collect::>(), sequence_names: self.sequence_names.clone(), sort_type: self.sort_type.clone(), + has_lcp: self.has_lcp, }) } @@ -959,6 +1079,13 @@ impl SufrFile { /// ``` /// pub fn list(&mut self, args: ListOptions) -> Result<()> { + if args.show_lcp && !self.has_lcp { + bail!( + "{} has no LCP array; rebuild with --write-lcp to show LCP values", + self.filename + ); + } + let width = self.text_len.to_string().len(); let text_len = self.text_len.to_usize(); let suffix_len = args.len.unwrap_or(text_len); @@ -1004,7 +1131,12 @@ impl SufrFile { let number = args.number.unwrap_or(0); if args.ranks.is_empty() { for (rank, suffix) in self.suffix_array_file.iter().enumerate() { - print(rank, suffix.to_usize(), self.lcp_file.get(rank).unwrap())?; + let lcp = if self.has_lcp { + self.lcp_file.get(rank).unwrap() + } else { + T::from_usize(0) + }; + print(rank, suffix.to_usize(), lcp)?; if number > 0 && rank == number - 1 { break; @@ -1013,7 +1145,12 @@ impl SufrFile { } else { for rank in args.ranks { if let Some(suffix) = self.suffix_array_file.get(rank) { - print(rank, suffix.to_usize(), self.lcp_file.get(rank).unwrap())?; + let lcp = if self.has_lcp { + self.lcp_file.get(rank).unwrap() + } else { + T::from_usize(0) + }; + print(rank, suffix.to_usize(), lcp)?; } else { eprintln!("Invalid rank: {rank}"); } diff --git a/libsufr/src/types.rs b/libsufr/src/types.rs index edecbc1..b2d4d2d 100644 --- a/libsufr/src/types.rs +++ b/libsufr/src/types.rs @@ -574,6 +574,11 @@ pub struct SufrBuilderArgs { /// An optional seed mask of 1/0 for care/don't-care positions, /// cf. `SeedMask`. pub seed_mask: Option, + + /// Whether to write the LCP array. Searching does not use it; + /// listing LCP values does, and query-time max query lengths shorter + /// than the build-time value are faster with it. + pub write_lcp: bool, } // -------------------------------------------------- @@ -618,6 +623,9 @@ pub struct SufrMetadata { /// Sort type pub sort_type: SuffixSortType, + + /// Whether or not the file has an on-disk LCP array + pub has_lcp: bool, } #[cfg(test)] diff --git a/sufr/src/lib.rs b/sufr/src/lib.rs index ba589c4..3c42e98 100644 --- a/sufr/src/lib.rs +++ b/sufr/src/lib.rs @@ -118,6 +118,11 @@ pub struct CreateArgs { /// Spaced seeds mask #[arg(short, long, value_name = "MASK")] pub seed_mask: Option, + + /// Also write the LCP array (needed by "list --show-lcp"; speeds up + /// query-time max query lengths shorter than the build-time value) + #[arg(long)] + pub write_lcp: bool, } #[derive(Debug, Parser)] @@ -347,6 +352,7 @@ pub fn create(args: &CreateArgs) -> Result<()> { sequence_names: seq_data.sequence_names, num_partitions: args.num_partitions, seed_mask: args.seed_mask.clone(), + write_lcp: args.write_lcp, }; let now = Instant::now(); @@ -604,6 +610,10 @@ pub fn summarize(args: &SummarizeArgs) -> Result<()> { "Len Suffixes".to_string(), num_fmt.format(",.0", meta.len_suffixes as f64), ]); + rows.push(vec![ + "LCP array".to_string(), + if meta.has_lcp { "yes" } else { "no" }.to_string(), + ]); match meta.sort_type { SuffixSortType::Mask(seed_mask) => {