@gmod/gbz-base 1.0.0 → 1.1.0

package/tools/haplotype-index/src/main.rs

@@ -7,22 +7,39 @@ use simple_sds::serialize;
 
 use std::env;
 use std::process;
+use std::thread;
+use std::time::Instant;
 
-const USAGE: &str = "Usage: gbz-haplotype-index [--interval BP] [--forward-only] graph.gbz graph.gbz.db
-       gbz-haplotype-index [--interval BP] [--forward-only] --from-db graph.gbz.db
+const USAGE: &str = "Usage: gbz-haplotype-index [options] graph.gbz graph.gbz.db
+       gbz-haplotype-index [options] --output index.db graph.gbz [graph.gbz.db]
+       gbz-haplotype-index [options] --from-db graph.gbz.db
 
 Walks every path in both orientations and writes a sample every --interval bp
-(default 4096) into table HaplotypeSamples of the database, plus the path
-lengths into HaplotypeLengths. The path start and end are always sampled.
-Existing tables are replaced. With --from-db the walk reads node records from
-the database itself, so the GBZ is not needed.
+(default 4096) into table HaplotypeSamples, plus the path lengths into
+HaplotypeLengths. The path start and end are always sampled.
+
+By default the tables are written into the database itself, replacing any
+existing ones. With --output FILE they are written into FILE as a standalone
+companion database that the reader opens beside the graph database; the
+companion records the graph's path count so a mismatch is caught at open.
+
+With --from-db the walk reads node records from the database itself, so the
+GBZ is not needed. Walking a GBZ uses --threads (default: all cores).
+
+Options:
+  --interval BP     bp between samples along a path (default 4096)
+  --forward-only    sample only the forward orientation of each path
+  --output FILE     write a companion database instead of augmenting graph.gbz.db
+  --threads N       walker threads for the GBZ route
 ";
 
 struct Args {
     gbz: Option<String>,
-    db: String,
+    db: Option<String>,
+    output: Option<String>,
     interval: usize,
     forward_only: bool,
+    threads: usize,
 }
 
 fn parse_args() -> Args {
@@ -30,6 +47,8 @@ fn parse_args() -> Args {
     let mut interval = 4096;
     let mut forward_only = false;
     let mut from_db = false;
+    let mut output = None;
+    let mut threads = thread::available_parallelism().map(|n| n.get()).unwrap_or(1);
     let mut iter = env::args().skip(1);
     while let Some(arg) = iter.next() {
         match arg.as_str() {
@@ -40,6 +59,14 @@ fn parse_args() -> Args {
                     process::exit(1);
                 });
             }
+            "--threads" => {
+                let value = iter.next().unwrap_or_default();
+                threads = value.parse().unwrap_or_else(|_| {
+                    eprintln!("Invalid --threads: {}", value);
+                    process::exit(1);
+                });
+            }
+            "--output" => output = Some(iter.next().unwrap_or_default()),
             "--forward-only" => forward_only = true,
             "--from-db" => from_db = true,
             "-h" | "--help" => {
@@ -49,46 +76,52 @@ fn parse_args() -> Args {
             _ => positional.push(arg),
         }
     }
-    let expected = if from_db { 1 } else { 2 };
-    if positional.len() != expected || interval == 0 {
+    let valid = if from_db {
+        positional.len() == 1
+    } else if output.is_some() {
+        positional.len() == 1 || positional.len() == 2
+    } else {
+        positional.len() == 2
+    };
+    if !valid || interval == 0 || threads == 0 {
         eprint!("{}", USAGE);
         process::exit(1);
     }
-    let db = positional.pop().unwrap();
-    let gbz = positional.pop();
-    Args { gbz, db, interval, forward_only }
+    let (gbz, db) = if from_db {
+        (None, positional.pop())
+    } else {
+        let gbz = positional.remove(0);
+        (Some(gbz), positional.pop())
+    };
+    Args { gbz, db, output, interval, forward_only, threads }
 }
 
+#[derive(Clone, Copy)]
 struct Sample {
-    node_handle: usize,
-    node_offset: usize,
-    path_handle: usize,
-    orientation: usize,
-    path_offset: usize,
+    node_handle: u32,
+    node_offset: u32,
+    path_handle: u32,
+    orientation: u8,
+    path_offset: u32,
 }
 
 trait PathSource {
-    fn path_count(&mut self) -> usize;
-    fn start(&mut self, path_handle: usize, orientation: Orientation) -> Option<Pos>;
-    fn step(&mut self, pos: Pos) -> (usize, Option<Pos>);
-    fn node_len(&mut self, handle: usize) -> usize;
+    fn start(&self, path_handle: usize, orientation: Orientation) -> Option<Pos>;
+    fn step(&self, pos: Pos) -> (usize, Option<Pos>);
+    fn node_len(&self, handle: usize) -> usize;
 }
 
-struct GbzSource {
-    graph: GBZ,
+struct GbzSource<'a> {
+    graph: &'a GBZ,
 }
 
-impl PathSource for GbzSource {
-    fn path_count(&mut self) -> usize {
-        self.graph.metadata().map(|m| m.paths()).unwrap_or(0)
-    }
-
-    fn start(&mut self, path_handle: usize, orientation: Orientation) -> Option<Pos> {
+impl PathSource for GbzSource<'_> {
+    fn start(&self, path_handle: usize, orientation: Orientation) -> Option<Pos> {
         let index: &GBWT = self.graph.as_ref();
         index.start(support::encode_path(path_handle, orientation))
     }
 
-    fn step(&mut self, pos: Pos) -> (usize, Option<Pos>) {
+    fn step(&self, pos: Pos) -> (usize, Option<Pos>) {
         let node_len = self.graph.sequence_len(support::node_id(pos.node)).unwrap();
         let index: &GBWT = self.graph.as_ref();
         let bwt: &BWT = index.as_ref();
@@ -96,42 +129,37 @@ impl PathSource for GbzSource {
         (node_len, record.lf(pos.offset))
     }
 
-    fn node_len(&mut self, handle: usize) -> usize {
+    fn node_len(&self, handle: usize) -> usize {
         self.graph.sequence_len(support::node_id(handle)).unwrap()
     }
 }
 
 struct DbSource<'a> {
-    interface: GraphInterface<'a>,
-    paths: usize,
+    interface: std::cell::RefCell<GraphInterface<'a>>,
 }
 
-impl<'a> PathSource for DbSource<'a> {
-    fn path_count(&mut self) -> usize {
-        self.paths
-    }
-
-    fn start(&mut self, path_handle: usize, orientation: Orientation) -> Option<Pos> {
-        let path = self.interface.get_path(path_handle).unwrap()?;
+impl PathSource for DbSource<'_> {
+    fn start(&self, path_handle: usize, orientation: Orientation) -> Option<Pos> {
+        let path = self.interface.borrow_mut().get_path(path_handle).unwrap()?;
         Some(if orientation == Orientation::Forward { path.fw_start } else { path.rev_start })
     }
 
-    fn step(&mut self, pos: Pos) -> (usize, Option<Pos>) {
-        let record = self.interface.get_record(pos.node).unwrap().unwrap();
+    fn step(&self, pos: Pos) -> (usize, Option<Pos>) {
+        let record = self.interface.borrow_mut().get_record(pos.node).unwrap().unwrap();
         (record.sequence_len(), record.to_gbwt_record().lf(pos.offset))
     }
 
-    fn node_len(&mut self, handle: usize) -> usize {
-        self.interface.get_record(handle).unwrap().unwrap().sequence_len()
+    fn node_len(&self, handle: usize) -> usize {
+        self.interface.borrow_mut().get_record(handle).unwrap().unwrap().sequence_len()
     }
 }
 
-fn walk(source: &mut dyn PathSource, path_handle: usize, orientation: Orientation, interval: usize) -> (Vec<Sample>, usize) {
-    let mut samples = Vec::new();
+fn walk(source: &dyn PathSource, path_handle: usize, orientation: Orientation, interval: usize, samples: &mut Vec<Sample>) -> usize {
     let mut pos = source.start(path_handle, orientation);
     let mut offset = 0;
     let mut next_sample = 0;
     let mut last: Option<(Pos, usize)> = None;
+    let first = samples.len();
     while let Some(current) = pos {
         if current.node == ENDMARKER {
             break;
@@ -139,11 +167,11 @@ fn walk(source: &mut dyn PathSource, path_handle: usize, orientation: Orientatio
         let (node_len, next) = source.step(current);
         if offset >= next_sample {
             samples.push(Sample {
-                node_handle: current.node,
-                node_offset: current.offset,
-                path_handle,
-                orientation: orientation as usize,
-                path_offset: offset,
+                node_handle: current.node as u32,
+                node_offset: current.offset as u32,
+                path_handle: path_handle as u32,
+                orientation: orientation as u8,
+                path_offset: offset as u32,
             });
             next_sample = offset + interval;
             last = None;
@@ -155,74 +183,69 @@ fn walk(source: &mut dyn PathSource, path_handle: usize, orientation: Orientatio
     }
     if let Some((end, end_offset)) = last {
         samples.push(Sample {
-            node_handle: end.node,
-            node_offset: end.offset,
-            path_handle,
-            orientation: orientation as usize,
-            path_offset: end_offset,
+            node_handle: end.node as u32,
+            node_offset: end.offset as u32,
+            path_handle: path_handle as u32,
+            orientation: orientation as u8,
+            path_offset: end_offset as u32,
         });
     }
     let length = offset;
     if orientation == Orientation::Reverse {
-        for sample in samples.iter_mut() {
-            let node_len = source.node_len(sample.node_handle);
-            sample.path_offset = length - sample.path_offset - node_len;
+        for sample in samples[first..].iter_mut() {
+            let node_len = source.node_len(sample.node_handle as usize);
+            sample.path_offset = (length - sample.path_offset as usize - node_len) as u32;
         }
     }
-    (samples, length)
+    length
 }
 
-fn run(source: &mut dyn PathSource, connection: &mut Connection, args: &Args) {
-    let orientations: Vec<Orientation> = if args.forward_only {
+fn orientations(args: &Args) -> Vec<Orientation> {
+    if args.forward_only {
         vec![Orientation::Forward]
     } else {
         vec![Orientation::Forward, Orientation::Reverse]
-    };
-    let paths = source.path_count();
-    let mut inserted = 0;
-    let batch = 64;
-    let mut handle = 0;
-    while handle < paths {
-        let transaction = connection.transaction().unwrap();
-        {
-            let mut insert_sample = transaction
-                .prepare("INSERT INTO HaplotypeSamples(node_handle, node_offset, path_handle, orientation, path_offset) VALUES (?1, ?2, ?3, ?4, ?5)")
-                .unwrap();
-            let mut insert_length = transaction.prepare("INSERT INTO HaplotypeLengths(path_handle, length) VALUES (?1, ?2)").unwrap();
-            for path_handle in handle..(handle + batch).min(paths) {
-                let mut length = 0;
-                for &orientation in orientations.iter() {
-                    let (samples, walked) = walk(source, path_handle, orientation, args.interval);
-                    length = walked;
-                    for sample in samples {
-                        insert_sample
-                            .execute(params![
-                                sample.node_handle as i64,
-                                sample.node_offset as i64,
-                                sample.path_handle as i64,
-                                sample.orientation as i64,
-                                sample.path_offset as i64
-                            ])
-                            .unwrap();
-                        inserted += 1;
-                    }
-                }
-                insert_length.execute(params![path_handle as i64, length as i64]).unwrap();
-            }
+    }
+}
+
+fn walk_paths(source: &dyn PathSource, handles: std::ops::Range<usize>, args: &Args) -> (Vec<Sample>, Vec<(usize, usize)>) {
+    let mut samples = Vec::new();
+    let mut lengths = Vec::new();
+    for path_handle in handles {
+        let mut length = 0;
+        for &orientation in orientations(args).iter() {
+            length = walk(source, path_handle, orientation, args.interval, &mut samples);
         }
-        transaction.commit().unwrap();
-        handle += batch;
-        eprintln!("{} / {} paths, {} samples", handle.min(paths), paths, inserted);
+        lengths.push((path_handle, length));
     }
-    eprintln!("Inserted {} samples for {} paths", inserted, paths);
+    (samples, lengths)
 }
 
-fn path_count_from_db(db: &str) -> usize {
-    let connection = Connection::open(db).unwrap();
-    let value: String = connection
-        .query_row("SELECT value FROM Tags WHERE key = 'paths'", [], |row| row.get(0))
-        .unwrap_or_else(|_| "0".to_string());
-    value.parse().unwrap_or(0)
+fn walk_gbz(graph: &GBZ, paths: usize, args: &Args) -> (Vec<Sample>, Vec<(usize, usize)>) {
+    let chunk = (paths + args.threads - 1) / args.threads;
+    let started = Instant::now();
+    let results: Vec<(Vec<Sample>, Vec<(usize, usize)>)> = thread::scope(|scope| {
+        let workers: Vec<_> = (0..args.threads)
+            .map(|t| {
+                let range = (t * chunk).min(paths)..((t + 1) * chunk).min(paths);
+                scope.spawn(move || {
+                    let source = GbzSource { graph };
+                    let result = walk_paths(&source, range.clone(), args);
+                    eprintln!("thread {} walked paths {}..{} ({} samples) in {:.0} s", t, range.start, range.end, result.0.len(), started.elapsed().as_secs_f64());
+                    result
+                })
+            })
+            .collect();
+        workers.into_iter().map(|w| w.join().unwrap()).collect()
+    });
+    let mut samples = Vec::new();
+    let mut lengths = Vec::new();
+    for (s, l) in results {
+        samples.extend(s);
+        lengths.extend(l);
+    }
+    lengths.sort_unstable();
+    (samples, lengths)
 }
 
 const SCHEMA: &str = "CREATE TABLE HaplotypeSamples (
@@ -238,82 +261,93 @@ CREATE TABLE HaplotypeLengths (
     length INTEGER NOT NULL
 ) STRICT;";
 
-fn merge(db: &str, tmp: &str, args: &Args) {
-    let scratch = Connection::open(tmp).unwrap();
-    let mut connection = Connection::open(db).unwrap();
-    connection
-        .execute_batch(&format!("DROP TABLE IF EXISTS HaplotypeSamples; DROP TABLE IF EXISTS HaplotypeLengths; {}", SCHEMA))
-        .unwrap();
+fn write(target: &str, standalone: bool, mut samples: Vec<Sample>, lengths: &[(usize, usize)], paths: usize, args: &Args) {
+    let started = Instant::now();
+    samples.sort_unstable_by_key(|s| (s.node_handle, s.node_offset));
+    eprintln!("Sorted {} samples in {:.0} s", samples.len(), started.elapsed().as_secs_f64());
+    let mut connection = Connection::open(target).unwrap_or_else(|e| {
+        eprintln!("Cannot open {}: {}", target, e);
+        process::exit(1);
+    });
+    connection.execute_batch("PRAGMA journal_mode = OFF; PRAGMA synchronous = OFF;").unwrap();
+    let mut setup = String::from("DROP TABLE IF EXISTS HaplotypeSamples; DROP TABLE IF EXISTS HaplotypeLengths; ");
+    if standalone {
+        setup.push_str("CREATE TABLE IF NOT EXISTS Tags (key TEXT PRIMARY KEY, value TEXT NOT NULL) STRICT; ");
+    }
+    setup.push_str(SCHEMA);
+    connection.execute_batch(&setup).unwrap();
     let transaction = connection.transaction().unwrap();
     {
-        let mut read_samples = scratch
-            .prepare("SELECT node_handle, node_offset, path_handle, orientation, path_offset FROM HaplotypeSamples ORDER BY node_handle, node_offset")
-            .unwrap();
         let mut write_sample = transaction
             .prepare("INSERT INTO HaplotypeSamples(node_handle, node_offset, path_handle, orientation, path_offset) VALUES (?1, ?2, ?3, ?4, ?5)")
             .unwrap();
-        let mut rows = read_samples.query([]).unwrap();
-        while let Some(row) = rows.next().unwrap() {
-            let values: [i64; 5] = [row.get(0).unwrap(), row.get(1).unwrap(), row.get(2).unwrap(), row.get(3).unwrap(), row.get(4).unwrap()];
-            write_sample.execute(params![values[0], values[1], values[2], values[3], values[4]]).unwrap();
+        for s in samples.iter() {
+            write_sample
+                .execute(params![s.node_handle as i64, s.node_offset as i64, s.path_handle as i64, s.orientation as i64, s.path_offset as i64])
+                .unwrap();
         }
-        let mut read_lengths = scratch.prepare("SELECT path_handle, length FROM HaplotypeLengths ORDER BY path_handle").unwrap();
         let mut write_length = transaction.prepare("INSERT INTO HaplotypeLengths(path_handle, length) VALUES (?1, ?2)").unwrap();
-        let mut rows = read_lengths.query([]).unwrap();
-        while let Some(row) = rows.next().unwrap() {
-            let handle: i64 = row.get(0).unwrap();
-            let length: i64 = row.get(1).unwrap();
-            write_length.execute(params![handle, length]).unwrap();
+        for &(handle, length) in lengths {
+            write_length.execute(params![handle as i64, length as i64]).unwrap();
         }
-        transaction
-            .execute("INSERT OR REPLACE INTO Tags(key, value) VALUES ('haplotype_index_interval', ?1)", params![args.interval.to_string()])
-            .unwrap();
-        transaction
-            .execute(
-                "INSERT OR REPLACE INTO Tags(key, value) VALUES ('haplotype_index_orientations', ?1)",
-                params![if args.forward_only { "forward" } else { "both" }],
-            )
-            .unwrap();
+        let mut write_tag = transaction.prepare("INSERT OR REPLACE INTO Tags(key, value) VALUES (?1, ?2)").unwrap();
+        write_tag.execute(params!["haplotype_index_interval", args.interval.to_string()]).unwrap();
+        write_tag.execute(params!["haplotype_index_orientations", if args.forward_only { "forward" } else { "both" }]).unwrap();
+        write_tag.execute(params!["haplotype_index_paths", paths.to_string()]).unwrap();
     }
     transaction.commit().unwrap();
+    eprintln!("Wrote {} samples for {} paths to {} in {:.0} s", samples.len(), paths, target, started.elapsed().as_secs_f64());
+}
+
+fn path_count_from_db(db: &str) -> usize {
+    let connection = Connection::open(db).unwrap();
+    let value: String = connection
+        .query_row("SELECT value FROM Tags WHERE key = 'paths'", [], |row| row.get(0))
+        .unwrap_or_else(|_| "0".to_string());
+    value.parse().unwrap_or(0)
 }
 
 fn main() {
     let args = parse_args();
-    let tmp = format!("{}.haplotype-index.tmp", args.db);
-    let _ = std::fs::remove_file(&tmp);
-    {
-        let mut scratch = Connection::open(&tmp).unwrap_or_else(|e| {
-            eprintln!("Cannot create {}: {}", tmp, e);
-            process::exit(1);
-        });
-        scratch.execute_batch(SCHEMA).unwrap();
-        match &args.gbz {
-            Some(gbz) => {
-                let graph: GBZ = serialize::load_from(gbz).unwrap_or_else(|e| {
-                    eprintln!("Cannot load {}: {}", gbz, e);
-                    process::exit(1);
-                });
-                if graph.metadata().is_none() {
-                    eprintln!("The GBZ has no path metadata");
-                    process::exit(1);
-                }
-                let mut source = GbzSource { graph };
-                run(&mut source, &mut scratch, &args);
+    let (samples, lengths, paths) = match &args.gbz {
+        Some(gbz) => {
+            let started = Instant::now();
+            let graph: GBZ = serialize::load_from(gbz).unwrap_or_else(|e| {
+                eprintln!("Cannot load {}: {}", gbz, e);
+                process::exit(1);
+            });
+            let paths = graph.metadata().map(|m| m.paths()).unwrap_or(0);
+            if paths == 0 {
+                eprintln!("The GBZ has no path metadata");
+                process::exit(1);
             }
-            None => {
-                let paths = path_count_from_db(&args.db);
-                let database = GBZBase::open(&args.db).unwrap_or_else(|e| {
-                    eprintln!("Cannot open {} as a GBZ-base: {}", args.db, e);
+            eprintln!("Loaded {} with {} paths in {:.0} s", gbz, paths, started.elapsed().as_secs_f64());
+            if let Some(db) = &args.db {
+                let db_paths = path_count_from_db(db);
+                if db_paths != paths {
+                    eprintln!("{} has {} paths but {} has {}", gbz, paths, db, db_paths);
                     process::exit(1);
-                });
-                let interface = GraphInterface::new(&database).unwrap();
-                let mut source = DbSource { interface, paths };
-                run(&mut source, &mut scratch, &args);
+                }
             }
+            let (samples, lengths) = walk_gbz(&graph, paths, &args);
+            (samples, lengths, paths)
         }
+        None => {
+            let db = args.db.as_ref().unwrap();
+            let paths = path_count_from_db(db);
+            let database = GBZBase::open(db).unwrap_or_else(|e| {
+                eprintln!("Cannot open {} as a GBZ-base: {}", db, e);
+                process::exit(1);
+            });
+            let interface = GraphInterface::new(&database).unwrap();
+            let source = DbSource { interface: std::cell::RefCell::new(interface) };
+            let (samples, lengths) = walk_paths(&source, 0..paths, &args);
+            (samples, lengths, paths)
+        }
+    };
+    match (&args.output, &args.db) {
+        (Some(output), _) => write(output, true, samples, &lengths, paths, &args),
+        (None, Some(db)) => write(db, false, samples, &lengths, paths, &args),
+        (None, None) => unreachable!(),
     }
-    merge(&args.db, &tmp, &args);
-    let _ = std::fs::remove_file(&tmp);
-    eprintln!("Merged into {}", args.db);
 }