@gmod/bam 7.3.4 → 7.5.0

package/src/bamFile.ts

@@ -1,5 +1,4 @@
 import { unzip, unzipChunkSlice } from '@gmod/bgzf-filehandle'
-import QuickLRU from '@jbrowse/quick-lru'
 import crc32 from 'crc/calculators/crc32'
 import { LocalFile, RemoteFile } from 'generic-filehandle2'
 
@@ -8,16 +7,11 @@ import CSI from './csi.ts'
 import NullFilehandle from './nullFilehandle.ts'
 import BAMFeature from './record.ts'
 import { parseHeaderText } from './sam.ts'
-import {
-  appendInRange,
-  applyFilters,
-  filterCacheKey,
-  parseRefSeqs,
-} from './util.ts'
+import { appendInRange, parseRefSeqs } from './util.ts'
 
 import type Chunk from './chunk.ts'
 import type { Bytes } from './record.ts'
-import type { BamOpts, BaseOpts, FilterBy } from './util.ts'
+import type { BamOpts, BaseOpts } from './util.ts'
 import type { GenericFilehandle } from 'generic-filehandle2'
 
 export interface BamRecordLike {
@@ -54,14 +48,81 @@ function resolveFilehandle(
 }
 
 interface ChunkEntry<T> {
-  minBlock: number
-  maxBlock: number
+  // decompressed size of the chunk these features are views into
+  bytes: number
   features: T[]
 }
 
-function chunkCacheKey(chunk: Chunk, filterBy?: FilterBy) {
+function chunkCacheKey(chunk: Chunk) {
   const { minv, maxv } = chunk
-  return `${minv.blockPosition}:${minv.dataPosition}-${maxv.blockPosition}:${maxv.dataPosition}${filterCacheKey(filterBy)}`
+  return `${minv.blockPosition}:${minv.dataPosition}-${maxv.blockPosition}:${maxv.dataPosition}`
+}
+
+// Every record in an entry is a view into its chunk's decompressed buffer, so
+// caching one entry pins that whole buffer — 8MB apiece on the nanopore and
+// 2kb-read test files, and optimizeChunks merges spans up to 5MB *compressed*,
+// so tens of MB is possible. A count-based LRU therefore gives no bound on
+// memory at all, which is why this budgets by decompressed bytes instead. It is
+// the *only* bound: a query keeps every chunk it parsed, since a chunk dropped
+// here costs a re-download and re-decompress the next time the view moves.
+export const DEFAULT_MAX_CACHE_BYTES = 100 * 1024 * 1024
+
+class ChunkFeatureCache<T> {
+  public maxBytes: number
+  private entries = new Map<string, ChunkEntry<T>>()
+  private bytes = 0
+
+  constructor(maxBytes: number) {
+    this.maxBytes = maxBytes
+  }
+
+  get size() {
+    return this.entries.size
+  }
+
+  get byteSize() {
+    return this.bytes
+  }
+
+  get(key: string) {
+    const entry = this.entries.get(key)
+    if (entry) {
+      // re-insert so Map iteration order stays least-recently-used first
+      this.entries.delete(key)
+      this.entries.set(key, entry)
+    }
+    return entry
+  }
+
+  set(key: string, entry: ChunkEntry<T>) {
+    this.delete(key)
+    this.entries.set(key, entry)
+    this.bytes += entry.bytes
+    // 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.
+    const lru = this.entries.keys()
+    while (this.bytes > this.maxBytes && this.entries.size > 1) {
+      this.delete(lru.next().value!)
+    }
+  }
+
+  delete(key: string) {
+    const entry = this.entries.get(key)
+    if (entry) {
+      this.entries.delete(key)
+      this.bytes -= entry.bytes
+    }
+  }
+
+  clear() {
+    this.entries.clear()
+    this.bytes = 0
+  }
+
+  [Symbol.iterator]() {
+    return this.entries[Symbol.iterator]()
+  }
 }
 
 export default class BamFile<T extends BamRecordLike = BAMFeature> {
@@ -74,11 +135,8 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
   public htsget = false
   public headerP?: ReturnType<BamFile<T>['getHeaderPre']>
 
-  // Cache for parsed features by chunk
-  // When a new chunk overlaps a cached chunk, we evict the cached one
-  public chunkFeatureCache = new QuickLRU<string, ChunkEntry<T>>({
-    maxSize: 100,
-  })
+  // Cache for parsed features by chunk, bounded by decompressed bytes
+  public chunkFeatureCache: ChunkFeatureCache<T>
 
   private RecordClass: BamRecordClass<T>
 
@@ -95,6 +153,7 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
     htsget,
     renameRefSeqs = n => n,
     recordClass,
+    maxCacheBytes = DEFAULT_MAX_CACHE_BYTES,
   }: {
     bamFilehandle?: GenericFilehandle
     bamPath?: string
@@ -108,9 +167,12 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
     renameRefSeqs?: (a: string) => string
     htsget?: boolean
     recordClass?: BamRecordClass<T>
+    /** budget for the parsed-chunk cache, in decompressed bytes */
+    maxCacheBytes?: number
   }) {
     this.renameRefSeq = renameRefSeqs
     this.RecordClass = (recordClass ?? BAMFeature) as BamRecordClass<T>
+    this.chunkFeatureCache = new ChunkFeatureCache<T>(maxCacheBytes)
 
     const bamFh = resolveFilehandle(bamFilehandle, bamPath, bamUrl)
     if (bamFh) {
@@ -162,7 +224,11 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
         ? await this.bam.readFile()
         : await this.bam.read(readLen, 0)
     let uncba = await unzip(buffer)
-    const dataView = new DataView(uncba.buffer)
+    const dataView = new DataView(
+      uncba.buffer,
+      uncba.byteOffset,
+      uncba.byteLength,
+    )
 
     if (dataView.getInt32(0, true) !== BAM_MAGIC) {
       throw new Error('Not a BAM file')
@@ -221,26 +287,31 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
       return []
     }
     const chunks = await this.index.blocksForRange(chrId, min - 1, max, opts)
-    return this._fetchChunkFeaturesDirect(chunks, chrId, min, max, opts)
+    return this._fetchChunkFeatures(chunks, chrId, min, max, opts)
   }
 
-  // Evict any cached chunks whose block range overlaps [minBlock, maxBlock]
-  private evictOverlappingChunks(minBlock: number, maxBlock: number) {
-    for (const [key, entry] of this.chunkFeatureCache) {
-      if (minBlock <= entry.maxBlock && maxBlock >= entry.minBlock) {
-        this.chunkFeatureCache.delete(key)
-      }
+  // Parsed records for a chunk, reading and decompressing it only on a miss.
+  // Every path that wants a chunk's features goes through here — mate lookups
+  // included, since a viewAsPairs query revisits the same mate chunks each time
+  // the view moves.
+  private async _cachedChunkFeatures(chunk: Chunk, opts: BaseOpts) {
+    const cacheKey = chunkCacheKey(chunk)
+    let entry = this.chunkFeatureCache.get(cacheKey)
+    if (!entry) {
+      entry = await this._readChunkFeatures(chunk, opts)
+      this.chunkFeatureCache.set(cacheKey, entry)
     }
+    return entry.features
   }
 
-  private async _fetchChunkFeaturesDirect(
+  private async _fetchChunkFeatures(
     chunks: Chunk[],
     chrId: number,
     min: number,
     max: number,
     opts: BamOpts = {},
   ) {
-    const { viewAsPairs, filterBy, onProgress } = opts
+    const { viewAsPairs, onProgress } = opts
     const result: T[] = []
 
     let totalBytes = 0
@@ -252,28 +323,11 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
 
     for (let ci = 0, cl = chunks.length; ci < cl; ci++) {
       const chunk = chunks[ci]!
-      const cacheKey = chunkCacheKey(chunk, filterBy)
-      const minBlock = chunk.minv.blockPosition
-      const maxBlock = chunk.maxv.blockPosition
-
-      let records: T[]
-      const cached = this.chunkFeatureCache.get(cacheKey)
-      if (cached) {
-        records = cached.features
-      } else {
-        this.evictOverlappingChunks(minBlock, maxBlock)
-        const allRecords = await this._readChunkFeatures(chunk, opts)
-        records = filterBy ? applyFilters(allRecords, filterBy) : allRecords
-        this.chunkFeatureCache.set(cacheKey, {
-          minBlock,
-          maxBlock,
-          features: records,
-        })
-      }
+      const features = await this._cachedChunkFeatures(chunk, opts)
 
       downloadedBytes += chunk.fetchedSize()
       onProgress?.(downloadedBytes, totalBytes)
-      appendInRange(records, chrId, min, max, result)
+      appendInRange(features, chrId, min, max, result)
     }
 
     if (viewAsPairs) {
@@ -288,13 +342,16 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
 
   async fetchPairs(chrId: number, records: T[], opts: BamOpts) {
     const { pairAcrossChr, maxInsertSize = 200000 } = opts
-    const readNameCounts: Record<string, number> = {}
+    // Map, not a plain object: read names come from the file, and on a plain
+    // object a read named "constructor" would read back Object.prototype's and
+    // make the count NaN.
+    const readNameCounts = new Map<string, number>()
     const readIds = new Set<number>()
 
     for (let i = 0, l = records.length; i < l; i++) {
       const r = records[i]!
       const name = r.name
-      readNameCounts[name] = (readNameCounts[name] ?? 0) + 1
+      readNameCounts.set(name, (readNameCounts.get(name) ?? 0) + 1)
       readIds.add(r.fileOffset)
     }
 
@@ -304,7 +361,7 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
       const name = f.name
       if (
         this.index &&
-        readNameCounts[name] === 1 &&
+        readNameCounts.get(name) === 1 &&
         (pairAcrossChr ||
           (f.next_refid === chrId &&
             Math.abs(f.start - f.next_pos) < maxInsertSize))
@@ -332,12 +389,12 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
 
     const mateFeatLists = await Promise.all(
       [...map.values()].map(async c => {
-        const features = await this._readChunkFeatures(c, opts)
+        const features = await this._cachedChunkFeatures(c, opts)
         const mateRecs = [] as T[]
         for (let i = 0, l = features.length; i < l; i++) {
           const feature = features[i]!
           if (
-            readNameCounts[feature.name] === 1 &&
+            readNameCounts.get(feature.name) === 1 &&
             !readIds.has(feature.fileOffset)
           ) {
             mateRecs.push(feature)
@@ -355,18 +412,25 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
     // filehandle's own streaming onProgress also fired this callback it would
     // report a different `total` (this chunk's size, not the whole query),
     // making the determinate bar jump around.
-    const buf = await this.bam.read(chunk.fetchedSize(), chunk.minv.blockPosition, {
-      signal: opts.signal,
-    })
+    const buf = await this.bam.read(
+      chunk.fetchedSize(),
+      chunk.minv.blockPosition,
+      {
+        signal: opts.signal,
+      },
+    )
     const {
       buffer: data,
       cpositions,
       dpositions,
     } = await unzipChunkSlice(buf, chunk)
-    return this.readBamFeatures(data, cpositions, dpositions, chunk)
+    return {
+      features: this.readBamFeatures(data, cpositions, dpositions, chunk),
+      bytes: data.byteLength,
+    }
   }
 
-  async readBamFeatures(
+  readBamFeatures(
     ba: Uint8Array,
     cpositions: number[],
     dpositions: number[],
@@ -376,7 +440,7 @@ export default class BamFile<T extends BamRecordLike = BAMFeature> {
     const sink = [] as T[]
     let pos = 0
 
-    const dataView = new DataView(ba.buffer)
+    const dataView = new DataView(ba.buffer, ba.byteOffset, ba.byteLength)
     const hasDpositions = dpositions.length > 0
     const hasCpositions = cpositions.length > 0
 

package/src/csi.ts

@@ -36,7 +36,11 @@ export default class CSI extends IndexFile {
   }
 
   parseAuxData(bytes: Uint8Array, offset: number) {
-    const dataView = new DataView(bytes.buffer)
+    const dataView = new DataView(
+      bytes.buffer,
+      bytes.byteOffset,
+      bytes.byteLength,
+    )
     const formatFlags = dataView.getUint32(offset, true)
     const coordinateType =
       formatFlags & 0x10000 ? 'zero-based-half-open' : '1-based-closed'
@@ -74,7 +78,11 @@ export default class CSI extends IndexFile {
     const buffer = await this.filehandle.readFile(opts)
     const bytes = await unzip(buffer)
 
-    const dataView = new DataView(bytes.buffer)
+    const dataView = new DataView(
+      bytes.buffer,
+      bytes.byteOffset,
+      bytes.byteLength,
+    )
     let csiVersion
     const magic = dataView.getUint32(0, true)
 

package/src/htsget.ts

@@ -78,7 +78,7 @@ export default class HtsgetFile<
     })
 
     const zero = new VirtualOffset(0, 0)
-    const allRecords = await this.readBamFeatures(
+    const allRecords = this.readBamFeatures(
       uncba,
       [],
       [],
@@ -95,7 +95,11 @@ export default class HtsgetFile<
     const uncba = await fetchAndConcat(data.htsget.urls, {
       signal: opts.signal,
     })
-    const dataView = new DataView(uncba.buffer)
+    const dataView = new DataView(
+      uncba.buffer,
+      uncba.byteOffset,
+      uncba.byteLength,
+    )
 
     if (dataView.getInt32(0, true) !== BAM_MAGIC) {
       throw new Error('Not a BAM file')

package/src/index.ts

@@ -1,9 +1,8 @@
 export { default as BAI } from './bai.ts'
-export { default as BamFile } from './bamFile.ts'
+export { DEFAULT_MAX_CACHE_BYTES, default as BamFile } from './bamFile.ts'
 export { default as CSI } from './csi.ts'
 export { default as BamRecord } from './record.ts'
 export { default as HtsgetFile } from './htsget.ts'
 
 export type { Bytes } from './record.ts'
-export type { FilterBy, TagFilter } from './util.ts'
 export type { BamRecordClass, BamRecordLike } from './bamFile.ts'