Parallel output writer and optional LCP array

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
 

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<T: Int> Iterator for FileAccessIter<'_, T> {
@@ -141,22 +141,22 @@ impl<T: Int> 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::<T>();
-                let mut buffer: Vec<T> = 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::<T>()) as u64;
-
-                self.buffer = buffer;
                 self.buffer_pos = 0;
             }
 
@@ -168,7 +168,26 @@ impl<T: Int> 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<usize> {
+    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<usize> {
+    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<()> {

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::<u32>::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::<u64>::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::<u32>::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::<u32>::new(builder_args);

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<u8> {
+    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<usize>,
+    write_lcp: bool,
+) -> Result<(NamedTempFile, SufrFile<u32>)> {
+    let outfile = NamedTempFile::new()?;
+    let outpath = outfile.path().to_string_lossy().to_string();
+    SufrBuilder::<u32>::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<u32> = 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<u32>, Vec<u32>) {
+    let key = |pos: usize| -> Vec<Option<u8>> {
+        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<usize> = 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<u32> = 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<usize> = (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<u32> = 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(())
+}

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)?;