@gmod/gbz-base 1.1.0 → 2.0.0

package/src/query.ts

@@ -1,29 +1,18 @@
-import { GENERIC_SAMPLE } from './db.ts'
+import { pathNameFor } from './pathName.ts'
 import { Subgraph } from './subgraph.ts'
 
-import type { GBZBase, PathName } from './db.ts'
+import type { GBZBase } from './db.ts'
+import type { PathQuery } from './pathName.ts'
 import type { HaplotypeOutput, SnarlOutput } from './subgraph.ts'
 
-export interface QueryOptions {
-  context?: number
-  haplotypes?: HaplotypeOutput
-  snarls?: SnarlOutput
-  limit?: number
-}
-
-export interface PathQuery {
-  sample?: string
-  contig: string
-  haplotype?: number
-}
+export type { PathQuery } from './pathName.ts'
 
-function pathName(query: PathQuery, fragment: number): PathName {
-  return {
-    sample: query.sample ?? GENERIC_SAMPLE,
-    contig: query.contig,
-    haplotype: query.haplotype ?? 0,
-    fragment,
-  }
+export interface QueryOptions {
+  context?: number | undefined
+  haplotypes?: HaplotypeOutput | undefined
+  snarls?: SnarlOutput | undefined
+  limit?: number | undefined
+  signal?: AbortSignal | undefined
 }
 
 export async function subgraphAtOffset(
@@ -32,9 +21,8 @@ export async function subgraphAtOffset(
   offset: number,
   opts: QueryOptions = {},
 ) {
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
-  const reference = await subgraph.pathPosition(pathName(query, offset))
+  const subgraph = new Subgraph(db, opts)
+  const reference = await subgraph.pathPosition(pathNameFor(query, offset))
   await subgraph.aroundPosition(
     reference.position.handle,
     reference.position.nodeOffset,
@@ -52,9 +40,8 @@ export async function subgraphInInterval(
   end: number,
   opts: QueryOptions = {},
 ) {
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
-  const reference = await subgraph.pathPosition(pathName(query, start))
+  const subgraph = new Subgraph(db, opts)
+  const reference = await subgraph.pathPosition(pathNameFor(query, start))
   await subgraph.aroundInterval(
     reference.position,
     end - start,
@@ -80,8 +67,7 @@ export async function subgraphAroundNodes(
       'Overlapping snarls cannot be extracted for a node-based query with multiple nodes',
     )
   }
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
+  const subgraph = new Subgraph(db, opts)
   await subgraph.aroundNodes(nodes, opts.context ?? 100)
   await subgraph.extractSnarls(snarls)
   subgraph.extractPaths(undefined, haplotypes)
@@ -92,14 +78,13 @@ export async function subgraphBetween(
   db: GBZBase,
   start: number,
   end: number,
-  opts: Pick<QueryOptions, 'haplotypes' | 'limit'> = {},
+  opts: Pick<QueryOptions, 'haplotypes' | 'limit' | 'signal'> = {},
 ) {
   const haplotypes = opts.haplotypes ?? 'all'
   if (haplotypes === 'reference-only') {
     throw new Error('Cannot output a reference path in a node-based query')
   }
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
+  const subgraph = new Subgraph(db, opts)
   await subgraph.betweenNodes(start, end)
   subgraph.extractPaths(undefined, haplotypes)
   return subgraph

package/src/sqlite/btree.ts

@@ -15,8 +15,20 @@ interface PageHeader {
   cellPointers: number
 }
 
+function readUint16(page: Uint8Array, offset: number) {
+  return page[offset]! * 256 + page[offset + 1]!
+}
+
+function readUint32(page: Uint8Array, offset: number) {
+  return (
+    page[offset]! * 16777216 +
+    page[offset + 1]! * 65536 +
+    page[offset + 2]! * 256 +
+    page[offset + 3]!
+  )
+}
+
 function readHeader(page: Uint8Array, start: number): PageHeader {
-  const view = new DataView(page.buffer, page.byteOffset, page.byteLength)
   const type = page[start]
   if (
     type !== INTERIOR_INDEX &&
@@ -29,25 +41,27 @@ function readHeader(page: Uint8Array, start: number): PageHeader {
   const interior = type === INTERIOR_INDEX || type === INTERIOR_TABLE
   return {
     type,
-    cellCount: view.getUint16(start + 3),
-    rightChild: interior ? view.getUint32(start + 8) : 0,
+    cellCount: readUint16(page, start + 3),
+    rightChild: interior ? readUint32(page, start + 8) : 0,
     cellPointers: start + (interior ? 12 : 8),
   }
 }
 
 function cellOffset(page: Uint8Array, header: PageHeader, index: number) {
-  const view = new DataView(page.buffer, page.byteOffset, page.byteLength)
-  return view.getUint16(header.cellPointers + 2 * index)
+  return readUint16(page, header.cellPointers + 2 * index)
 }
 
-function readUint32(page: Uint8Array, offset: number) {
-  const view = new DataView(page.buffer, page.byteOffset, page.byteLength)
-  return view.getUint32(offset)
+interface IndexCell {
+  values: SqlValue[]
+  leftChild: number
 }
 
+const MAX_DECODED_INDEX_PAGES = 4096
+
 export class BTree {
   private readonly usable: number
   private pager: Pager
+  private decodedIndexPages = new Map<number, (IndexCell | undefined)[]>()
 
   constructor(pager: Pager, reservedBytes: number) {
     this.pager = pager
@@ -179,12 +193,42 @@ export class BTree {
     }
   }
 
-  private async indexCell(page: Uint8Array, header: PageHeader, index: number) {
-    const offset = cellOffset(page, header, index)
-    const interior = header.type === INTERIOR_INDEX
-    const [size, afterSize] = readVarint(page, interior ? offset + 4 : offset)
-    const values = decodeRecord(await this.payload(page, afterSize, size, true))
-    return { values, leftChild: interior ? readUint32(page, offset) : 0 }
+  private decodedCells(pageNumber: number, header: PageHeader) {
+    let cells = this.decodedIndexPages.get(pageNumber)
+    if (cells) {
+      this.decodedIndexPages.delete(pageNumber)
+    } else {
+      cells = new Array<IndexCell | undefined>(header.cellCount)
+      if (this.decodedIndexPages.size >= MAX_DECODED_INDEX_PAGES) {
+        const oldest = this.decodedIndexPages.keys().next().value
+        if (oldest !== undefined) {
+          this.decodedIndexPages.delete(oldest)
+        }
+      }
+    }
+    this.decodedIndexPages.set(pageNumber, cells)
+    return cells
+  }
+
+  private async indexCell(
+    pageNumber: number,
+    page: Uint8Array,
+    header: PageHeader,
+    index: number,
+  ): Promise<IndexCell> {
+    const cells = this.decodedCells(pageNumber, header)
+    let cell = cells[index]
+    if (!cell) {
+      const offset = cellOffset(page, header, index)
+      const interior = header.type === INTERIOR_INDEX
+      const [size, afterSize] = readVarint(page, interior ? offset + 4 : offset)
+      cell = {
+        values: decodeRecord(await this.payload(page, afterSize, size, true)),
+        leftChild: interior ? readUint32(page, offset) : 0,
+      }
+      cells[index] = cell
+    }
+    return cell
   }
 
   async *indexScanFrom(
@@ -196,7 +240,7 @@ export class BTree {
     let high = header.cellCount
     while (first < high) {
       const mid = (first + high) >> 1
-      const cell = await this.indexCell(page, header, mid)
+      const cell = await this.indexCell(root, page, header, mid)
       if (compareKey(cell.values, low) < 0) {
         first = mid + 1
       } else {
@@ -204,7 +248,7 @@ export class BTree {
       }
     }
     for (let i = first; i < header.cellCount; i++) {
-      const cell = await this.indexCell(page, header, i)
+      const cell = await this.indexCell(root, page, header, i)
       if (header.type === INTERIOR_INDEX) {
         yield* this.indexScanFrom(cell.leftChild, low)
       }
@@ -228,7 +272,7 @@ export class BTree {
       let child = 0
       while (low < high) {
         const mid = (low + high) >> 1
-        const cell = await this.indexCell(page, header, mid)
+        const cell = await this.indexCell(pageNumber, page, header, mid)
         if (compareKey(cell.values, key) <= 0) {
           best = cell.values
           low = mid + 1

package/src/subgraph.ts

@@ -1,4 +1,3 @@
-import { formatPathName } from './db.ts'
 import {
   ENDMARKER,
   edgeIsCanonical,
@@ -16,16 +15,12 @@ import {
 } from './gbwt/node.ts'
 import { gfaHeaderLines, sha256Hex, subgraphName } from './graphName.ts'
 import { weightedLcs } from './lcs.ts'
+import { formatPathName } from './pathName.ts'
 
-import type {
-  GBZBase,
-  GbzPath,
-  GbzRecord,
-  HaplotypeSample,
-  PathName,
-} from './db.ts'
+import type { GBZBase, GbzPath, GbzRecord, HaplotypeSample } from './db.ts'
 import type { NodeSide, Orientation } from './gbwt/node.ts'
 import type { Pos } from './gbwt/record.ts'
+import type { PathName } from './pathName.ts'
 
 export type HaplotypeOutput = 'all' | 'distinct' | 'reference-only' | 'none'
 
@@ -86,12 +81,8 @@ export interface SubgraphJson {
   paths: SubgraphPath[]
 }
 
-export interface HaplotypeAlignment {
-  pathHandle: number | undefined
-  name: PathName | undefined
+export interface AlignmentSpan {
   strand: '+' | '-'
-  hapStart: number | undefined
-  hapEnd: number | undefined
   refStart: number
   refEnd: number
   cigar: string
@@ -100,8 +91,27 @@ export interface HaplotypeAlignment {
   start: Pos
 }
 
-export interface ToJsonOptions {
-  names?: 'anonymous' | 'resolved'
+export type HaplotypeAlignment = AlignmentSpan &
+  (
+    | {
+        resolved: true
+        name: PathName
+        label: string
+        pathHandle: number
+        hapStart: number
+        hapEnd: number
+      }
+    | { resolved: false }
+  )
+
+export interface SubgraphOutputOptions {
+  cigar?: boolean | undefined
+  names?: 'anonymous' | 'resolved' | undefined
+}
+
+export interface SubgraphOptions {
+  limit?: number | undefined
+  signal?: AbortSignal | undefined
 }
 
 function sideBefore(
@@ -183,7 +193,6 @@ export class Subgraph {
   private refInterval: [number, number] | undefined
   private refIndexCache: Map<number, number[]> | undefined
   private refPrefixCache: number[] | undefined
-  limit: number | undefined
   readonly stats = {
     orderedAlignments: 0,
     lcsAlignments: 0,
@@ -192,9 +201,13 @@ export class Subgraph {
   }
 
   private db: GBZBase
+  private readonly limit: number | undefined
+  private readonly signal: AbortSignal | undefined
 
-  constructor(db: GBZBase) {
+  constructor(db: GBZBase, opts: SubgraphOptions = {}) {
     this.db = db
+    this.limit = opts.limit
+    this.signal = opts.signal
   }
 
   get nodeCount() {
@@ -236,6 +249,7 @@ export class Subgraph {
   }
 
   private async addNode(id: number) {
+    this.signal?.throwIfAborted()
     if (this.limit !== undefined && this.nodeCount >= this.limit) {
       throw new Error(`Subgraph size limit of ${this.limit} nodes exceeded`)
     }
@@ -800,6 +814,7 @@ export class Subgraph {
       let current: number | undefined = start
       let pos: Pos | undefined
       while (anchor === undefined) {
+        this.signal?.throwIfAborted()
         if (current !== undefined) {
           if (visited.has(current)) {
             break
@@ -1094,24 +1109,17 @@ export class Subgraph {
         last -= 1
       }
       const identity = info.identity
-      const strand =
+      const walkStrand =
         info.path.some(handle => isReverse(handle)) &&
         !info.path.some(handle => !isReverse(handle))
           ? '-'
           : '+'
-      result.push({
-        pathHandle: identity?.pathHandle,
-        name: identity?.name,
+      const span: AlignmentSpan = {
         strand: identity
           ? identity.orientation === 'forward'
             ? '+'
             : '-'
-          : strand,
-        hapStart: identity
-          ? identity.name.fragment + identity.hapStart
-          : undefined,
-        hapEnd: identity ? identity.name.fragment + identity.hapEnd : undefined,
-        start: info.positions[0]!,
+          : walkStrand,
         refStart: reference.start + leading,
         refEnd: reference.start + refTotal - trailing,
         cigar: edits
@@ -1120,7 +1128,26 @@ export class Subgraph {
           .join(''),
         weight: info.weight,
         path: info.path,
-      })
+        start: info.positions[0]!,
+      }
+      if (identity) {
+        const hapStart = identity.name.fragment + identity.hapStart
+        const hapEnd = identity.name.fragment + identity.hapEnd
+        result.push({
+          ...span,
+          resolved: true,
+          name: identity.name,
+          label: formatPathName(
+            { ...identity.name, fragment: hapStart },
+            hapEnd,
+          ),
+          pathHandle: identity.pathHandle,
+          hapStart,
+          hapEnd,
+        })
+      } else {
+        result.push({ ...span, resolved: false })
+      }
     })
     return result
   }
@@ -1165,7 +1192,8 @@ export class Subgraph {
     return sha256Hex(chunks)
   }
 
-  async toGFA(cigar: boolean, opts: ToJsonOptions = {}) {
+  async toGFA(opts: SubgraphOutputOptions = {}) {
+    const cigar = opts.cigar ?? false
     const lines = [
       this.refPath ? `H\tVN:Z:1.1\tRS:Z:${this.refPath.sample}` : 'H\tVN:Z:1.1',
       ...gfaHeaderLines(
@@ -1244,7 +1272,8 @@ export class Subgraph {
     return `${lines.join('\n')}\n`
   }
 
-  toJSON(cigar: boolean, opts: ToJsonOptions = {}): SubgraphJson {
+  toSubgraphJson(opts: SubgraphOutputOptions = {}): SubgraphJson {
+    const cigar = opts.cigar ?? false
     const handles = this.sortedHandles()
     const nodes = handles
       .filter(handle => !isReverse(handle))

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

@@ -208,15 +208,22 @@ fn orientations(args: &Args) -> Vec<Orientation> {
     }
 }
 
-fn walk_paths(source: &dyn PathSource, handles: std::ops::Range<usize>, args: &Args) -> (Vec<Sample>, Vec<(usize, usize)>) {
+fn walk_paths(source: &dyn PathSource, handles: std::ops::Range<usize>, args: &Args, label: &str) -> (Vec<Sample>, Vec<(usize, usize)>) {
     let mut samples = Vec::new();
     let mut lengths = Vec::new();
-    for path_handle in handles {
+    let started = Instant::now();
+    let total = handles.len();
+    let mut walked_bp: usize = 0;
+    for (done, path_handle) in handles.enumerate() {
         let mut length = 0;
         for &orientation in orientations(args).iter() {
             length = walk(source, path_handle, orientation, args.interval, &mut samples);
         }
+        walked_bp += length;
         lengths.push((path_handle, length));
+        if (done + 1) % 500 == 0 || done + 1 == total {
+            eprintln!("{}: {} / {} paths, {:.2} Gbp, {} samples, {:.0} s", label, done + 1, total, walked_bp as f64 / 1e9, samples.len(), started.elapsed().as_secs_f64());
+        }
     }
     (samples, lengths)
 }
@@ -230,8 +237,9 @@ fn walk_gbz(graph: &GBZ, paths: usize, args: &Args) -> (Vec<Sample>, Vec<(usize,
                 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());
+                    let label = format!("thread {} (paths {}..{})", t, range.start, range.end);
+                    let result = walk_paths(&source, range.clone(), args, &label);
+                    eprintln!("{} done in {:.0} s", label, started.elapsed().as_secs_f64());
                     result
                 })
             })
@@ -341,7 +349,7 @@ fn main() {
             });
             let interface = GraphInterface::new(&database).unwrap();
             let source = DbSource { interface: std::cell::RefCell::new(interface) };
-            let (samples, lengths) = walk_paths(&source, 0..paths, &args);
+            let (samples, lengths) = walk_paths(&source, 0..paths, &args, "database walk");
             (samples, lengths, paths)
         }
     };