@gmod/tabix 3.5.0 → 3.5.2

package/src/tabixIndexedFile.ts

@@ -1,6 +1,5 @@
 import AbortablePromiseCache from '@gmod/abortable-promise-cache'
 import { unzip, unzipChunkSlice } from '@gmod/bgzf-filehandle'
-import LRU from '@jbrowse/quick-lru'
 import { LocalFile, RemoteFile } from 'generic-filehandle2'
 
 import CSI from './csi.ts'
@@ -17,6 +16,118 @@ const NEWLINE = 10
 const CARRIAGE_RETURN = 13
 const SEMICOLON = 59
 
+// Ceiling on how many chunk reads getLines keeps in flight ahead of the one it
+// is parsing. Six is the HTTP/1.1 per-host connection cap browsers enforce, so
+// going much above it buys nothing on the transport that matters.
+const MAX_READ_AHEAD_CHUNKS = 6
+
+// SYNC: ~/src/gmod/bam-js/src/bamFile.ts DEFAULT_MAX_CACHE_BYTES
+//
+// We fetch compressed and cache decompressed, and an entry is a whole chunk, so
+// entry count says nothing about memory. How little is easy to underestimate:
+// a dense VCF (test/data/1kg.chr1.subset.vcf.gz — 213MB over 600kb of chr1)
+// has single index bins of 17MB compressed, 120MB decompressed. Panning it
+// under the old 80-entry cache peaked at 2GB RSS.
+const DEFAULT_CHUNK_CACHE_BYTES = 100 * 2 ** 20
+
+// The entry type AbortablePromiseCache stores in its backing cache. Not exported
+// by the package, so it is recovered from the constructor signature rather than
+// restated (and left to drift) here.
+type CacheEntry = NonNullable<
+  ReturnType<
+    ConstructorParameters<
+      typeof AbortablePromiseCache<Chunk, ReadChunk>
+    >[0]['cache']['get']
+  >
+>
+
+interface CachedChunk {
+  entry: CacheEntry
+  /** 0 until the read settles and the decompressed size is known */
+  bytes: number
+}
+
+/**
+ * Backing store for the chunk cache, bounded by the decompressed size of the
+ * chunks it holds rather than by entry count.
+ *
+ * A chunk's size is only known once its read settles, so `set` records the
+ * entry immediately and charges it to the budget later. Unsettled entries are
+ * therefore free, which is what we want: they are reads a query is waiting on.
+ */
+class ByteBoundedChunkCache {
+  private entries = new Map<string, CachedChunk>()
+  private bytes = 0
+  private maxBytes: number
+
+  constructor(maxBytes: number) {
+    this.maxBytes = maxBytes
+  }
+
+  get byteSize() {
+    return this.bytes
+  }
+
+  get size() {
+    return this.entries.size
+  }
+
+  has(key: string) {
+    return this.entries.has(key)
+  }
+
+  get(key: string) {
+    const cached = this.entries.get(key)
+    if (cached) {
+      // re-insert so Map iteration order stays least-recently-used first
+      this.entries.delete(key)
+      this.entries.set(key, cached)
+    }
+    return cached?.entry
+  }
+
+  set(key: string, entry: CacheEntry) {
+    this.delete(key)
+    const cached = { entry, bytes: 0 }
+    this.entries.set(key, cached)
+    void entry.promise
+      .then(chunk => {
+        // a later set() may have replaced this key while the read was in
+        // flight; charging these bytes to it would then double-count
+        if (this.entries.get(key) === cached) {
+          cached.bytes = chunk.buffer.byteLength
+          this.bytes += cached.bytes
+          this.evict()
+        }
+      })
+      .catch(() => {
+        // a failed or aborted read caches nothing and costs nothing
+      })
+  }
+
+  delete(key: string) {
+    const cached = this.entries.get(key)
+    if (cached) {
+      this.entries.delete(key)
+      this.bytes -= cached.bytes
+    }
+  }
+
+  keys() {
+    return this.entries.keys()
+  }
+
+  // Evict from the least-recently-used end. The size > 1 guard means a single
+  // chunk larger than the whole budget is still kept: the caller needs it for
+  // the query in flight, and dropping it would only force a re-decompress.
+  private evict() {
+    const lru = this.entries.keys()
+    while (this.bytes > this.maxBytes && this.entries.size > 1) {
+      this.delete(lru.next().value!)
+    }
+  }
+}
+
 type GetLinesCallback = (
   line: string,
   fileOffset: number,
@@ -209,7 +320,7 @@ export default class TabixIndexedFile {
     csiPath,
     csiUrl,
     csiFilehandle,
-    chunkCacheSize = 5 * 2 ** 20,
+    chunkCacheSize = DEFAULT_CHUNK_CACHE_BYTES,
   }: {
     path?: string
     filehandle?: GenericFilehandle
@@ -220,6 +331,7 @@ export default class TabixIndexedFile {
     csiPath?: string
     csiUrl?: string
     csiFilehandle?: GenericFilehandle
+    /** budget for the decompressed chunk cache, in bytes */
     chunkCacheSize?: number
   }) {
     this.filehandle = resolveFilehandle(filehandle, path, url)
@@ -235,7 +347,7 @@ export default class TabixIndexedFile {
     })
 
     this.chunkCache = new AbortablePromiseCache<Chunk, ReadChunk>({
-      cache: new LRU({ maxSize: Math.floor(chunkCacheSize / (1 << 16)) }),
+      cache: new ByteBoundedChunkCache(chunkCacheSize),
       fill: (args: Chunk, signal?: AbortSignal) =>
         this.readChunk(args, { signal }),
     })
@@ -327,12 +439,39 @@ export default class TabixIndexedFile {
     let downloadedBytes = 0
     onProgress?.(0, totalBytes)
 
-    for (const c of chunks) {
-      const { buffer, cpositions, dpositions } = await this.chunkCache.get(
-        c.toString(),
-        c,
-        signal,
-      )
+    // Read ahead, but only as far as the scan has earned. Every chunk is its
+    // own range request, so a query spanning many of them pays a network round
+    // trip apiece, serially — for a remote file that dominates, well ahead of
+    // decompression (1kg.chr1 over a 1Mb window reads 22 chunks in a row).
+    //
+    // A fixed window would be wrong, though: blocksForRange offers a chunk per
+    // overlapping bin across every level, and on a sparse file the early return
+    // below stops the scan inside the first one, leaving the rest untouched
+    // (chr22_nanopore_subset offers 7 chunks and reads 1). Prefetching those
+    // would multiply the bytes such a query fetches for no gain.
+    //
+    // Finishing a chunk without hitting the early return proves the next chunk
+    // has to be examined, so the window starts at one and doubles per chunk
+    // consumed. A query that stops in its first chunk issues exactly the reads
+    // a sequential scan did; a long scan reaches full concurrency after three.
+    let readAhead = 1
+    const reads: Promise<ReadChunk>[] = []
+    const ensureReadsStarted = (count: number) => {
+      while (reads.length < Math.min(count, chunks.length)) {
+        const c = chunks[reads.length]!
+        const read = this.chunkCache.get(c.toString(), c, signal)
+        void read.catch(() => {
+          // a prefetch the early return skips is never awaited, so swallow its
+          // rejection here rather than let it surface unhandled
+        })
+        reads.push(read)
+      }
+    }
+    ensureReadsStarted(1)
+
+    for (let ci = 0, cl = chunks.length; ci < cl; ci++) {
+      const c = chunks[ci]!
+      const { buffer, cpositions, dpositions } = await reads[ci]!
       downloadedBytes += c.fetchedSize()
       onProgress?.(downloadedBytes, totalBytes)
       const minvDataPosition = c.minv.dataPosition
@@ -375,14 +514,14 @@ export default class TabixIndexedFile {
           blockStart = n + 1
           continue
         }
-        let refMatch = true
+        let isRefMatch = true
         for (let i = 0; i < refLen; i++) {
           if (buffer[refStart + i] !== regionRefNameBytes[i]) {
-            refMatch = false
+            isRefMatch = false
             break
           }
         }
-        if (!refMatch) {
+        if (!isRefMatch) {
           blockStart = n + 1
           continue
         }
@@ -436,6 +575,11 @@ export default class TabixIndexedFile {
         }
         blockStart = n + 1
       }
+
+      // every line in this chunk was still inside the query, so the next chunk
+      // has to be examined too - widen the window and start the reads for it
+      readAhead = Math.min(readAhead * 2, MAX_READ_AHEAD_CHUNKS)
+      ensureReadsStarted(ci + 1 + readAhead)
     }
   }
 

package/src/tbi.ts

@@ -4,8 +4,8 @@ import Chunk from './chunk.ts'
 import IndexFile from './indexFile.ts'
 import {
   clampChunkEnds,
-  findFirstData,
   memoizeByRefId,
+  minVirtualOffset,
   optimizeChunks,
   parseAuxData,
   parsePseudoBin,
@@ -78,8 +78,15 @@ export default class TabixIndex extends IndexFile {
     // nameSectionLength is at TBI offset 32; re-read to find where bin data starts
     const nameSectionLength = dataView.getInt32(32, true)
 
-    // SYNC: ~/src/gmod/bam-js/src/csi.ts _parse — two-pass structure
-    // First pass: record per-refId byte offsets and find firstDataLine
+    // SYNC: ~/src/gmod/bam-js/src/bai.ts _parse — two-pass structure
+    // First pass: record per-refId byte offsets and find firstDataLine.
+    //
+    // Only the linear index is consulted. Its entry for a window is the
+    // smallest virtual offset of any record overlapping that window, so the
+    // minimum over the linear index is already the minimum over the bin chunks
+    // — walking the chunks too only re-derives it. Checked against every .tbi
+    // in test/data: same answer on all 23, and no ref has bins without a
+    // linear index.
     let curr = 36 + nameSectionLength
     let firstDataLine: VirtualOffset | undefined
     const offsets: number[] = []
@@ -97,21 +104,13 @@ export default class TabixIndex extends IndexFile {
           throw new Error(
             'tabix index contains too many bins, please use a CSI index',
           )
-        } else if (bin === maxBinNumber + 1) {
-          curr += 16 * chunkCount
-        } else {
-          for (let k = 0; k < chunkCount; k++) {
-            firstDataLine = findFirstData(firstDataLine, fromBytes(bytes, curr))
-            curr += 16
-          }
         }
+        curr += 16 * chunkCount
       }
       const linearCount = dataView.getInt32(curr, true)
       curr += 4
-      for (let k = 0; k < linearCount; k++) {
-        firstDataLine = findFirstData(firstDataLine, fromBytes(bytes, curr))
-        curr += 8
-      }
+      firstDataLine = minVirtualOffset(bytes, curr, linearCount, firstDataLine)
+      curr += 8 * linearCount
     }
 
     function getIndices(refId: number): RefIndex | undefined {

package/src/util.ts

@@ -2,8 +2,7 @@ import LRU from '@jbrowse/quick-lru'
 
 import Chunk from './chunk.ts'
 import { longFromBytesToUnsigned } from './long.ts'
-
-import type VirtualOffset from './virtualOffset.ts'
+import VirtualOffset from './virtualOffset.ts'
 
 // SYNC: ~/src/gmod/bam-js/src/util.ts optimizeChunks
 export function optimizeChunks(chunks: Chunk[], lowest?: VirtualOffset) {
@@ -118,13 +117,41 @@ export function clampChunkEnds(
   }
 }
 
-export function findFirstData(
-  currentFdl: VirtualOffset | undefined,
-  virtualOffset: VirtualOffset,
+/**
+ * The smallest of `current` and the `count` packed virtual offsets starting at
+ * `offset`, allocating at most one VirtualOffset rather than one per entry.
+ *
+ * The index first pass exists only to find this minimum, and it visits every
+ * linear-index entry in the file to do it — 301k of them on
+ * test/data/failing_tabix.vcf.gz.tbi against 19k bin chunks. Building a
+ * VirtualOffset per entry to compare and discard it is the bulk of that pass.
+ */
+export function minVirtualOffset(
+  bytes: Uint8Array,
+  offset: number,
+  count: number,
+  current: VirtualOffset | undefined,
 ) {
-  return !currentFdl || currentFdl.compareTo(virtualOffset) > 0
-    ? virtualOffset
-    : currentFdl
+  let minBlock = current ? current.blockPosition : Infinity
+  let minData = current ? current.dataPosition : 0
+  let found = false
+  for (let i = 0; i < count; i++) {
+    const p = offset + i * 8
+    const block =
+      bytes[p + 7]! * 0x1_00_00_00_00_00 +
+      bytes[p + 6]! * 0x1_00_00_00_00 +
+      bytes[p + 5]! * 0x1_00_00_00 +
+      bytes[p + 4]! * 0x1_00_00 +
+      bytes[p + 3]! * 0x1_00 +
+      bytes[p + 2]!
+    const data = (bytes[p + 1]! << 8) | bytes[p]!
+    if (block < minBlock || (block === minBlock && data < minData)) {
+      minBlock = block
+      minData = data
+      found = true
+    }
+  }
+  return found ? new VirtualOffset(minBlock, minData) : current
 }
 
 export function parseNameBytes(namesBytes: Uint8Array) {