@gmod/tabix 3.5.3 → 3.5.4

package/src/tabixIndexedFile.ts

@@ -9,6 +9,7 @@ import { optimizeChunks } from './util.ts'
 import type Chunk from './chunk.ts'
 import type IndexFile from './indexFile.ts'
 import type { Options } from './indexFile.ts'
+import type { ChunkSlice } from '@gmod/bgzf-filehandle'
 import type { GenericFilehandle } from 'generic-filehandle2'
 
 const TAB = 9
@@ -148,11 +149,11 @@ interface GetLinesOpts {
   onProgress?: (bytesDownloaded: number, totalBytes?: number) => void
 }
 
-interface ReadChunk {
-  buffer: Uint8Array
-  cpositions: number[]
-  dpositions: number[]
-}
+// The decompressed chunk plus its block offsets, as @gmod/bgzf-filehandle
+// returns them: cpositions/dpositions are Float64Arrays, which is what the
+// wasm decompressor produces, and only indexed reads and `.length` are used
+// here. Taken from the package rather than restated so the two can't drift.
+type ReadChunk = ChunkSlice
 
 function resolveFilehandle(
   filehandle?: GenericFilehandle,
@@ -222,8 +223,8 @@ function resolveIndex({
 }
 
 function calculateFileOffset(
-  cpositions: number[],
-  dpositions: number[],
+  cpositions: ArrayLike<number>,
+  dpositions: ArrayLike<number>,
   pos: number,
   blockStart: number,
   minvDataPosition: number,
@@ -289,6 +290,50 @@ function getVcfEnd(
   return endCoordinate
 }
 
+const textDecoder = new TextDecoder()
+
+/**
+ * The leading run of meta-character lines — what `tabix -H` prints. Everything
+ * from the first line that doesn't begin with the meta character is dropped.
+ */
+function trimToMetaLines(bytes: Uint8Array, metaChar: string) {
+  let lastNewline = -1
+  const metaByte = metaChar.charCodeAt(0)
+
+  for (let i = 0, l = bytes.length; i < l; i++) {
+    const byte = bytes[i]
+    if (i === lastNewline + 1 && byte !== metaByte) {
+      break
+    }
+    if (byte === NEWLINE) {
+      lastNewline = i
+    }
+  }
+  return bytes.subarray(0, lastNewline + 1)
+}
+
+/**
+ * The first `count` lines. Scans for the count-th newline and decodes only that
+ * far, rather than decoding and splitting the whole buffer to keep its first
+ * few lines — the buffer runs to the first data line, which for a file with a
+ * long commented preamble above its counted rows can be megabytes.
+ */
+function firstLines(bytes: Uint8Array, count: number) {
+  let end = 0
+  for (let i = 0; i < count; i++) {
+    const n = bytes.indexOf(NEWLINE, end)
+    if (n === -1) {
+      end = bytes.length
+      break
+    }
+    end = n + 1
+  }
+  return textDecoder
+    .decode(bytes.subarray(0, end))
+    .split(/\r?\n/)
+    .slice(0, count)
+}
+
 function parseIntFromBytes(buffer: Uint8Array, start: number, end: number) {
   let val = 0
   for (let i = start; i < end; i++) {
@@ -309,6 +354,7 @@ export default class TabixIndexedFile {
   private filehandle: GenericFilehandle
   private index: IndexFile
   private chunkCache: AbortablePromiseCache<Chunk, ReadChunk>
+  private headerP?: Promise<{ header: string; skippedLines: string[] }>
 
   constructor({
     path,
@@ -588,34 +634,59 @@ export default class TabixIndexedFile {
     return this.index.getMetadata(opts)
   }
 
-  async getHeaderBuffer(opts: Options = {}) {
-    const { firstDataLine, metaChar, maxBlockSize } =
-      await this.getMetadata(opts)
+  /**
+   * The file's leading blocks, decompressed: everything from the start through
+   * the end of the block holding the first data line.
+   */
+  private async readHeaderBytes(opts: Options) {
+    const { firstDataLine, maxBlockSize } = await this.getMetadata(opts)
 
     const maxFetch = (firstDataLine?.blockPosition ?? 0) + maxBlockSize
     // TODO: what if we don't have a firstDataLine, and the header actually
     // takes up more than one block? this case is not covered here
 
     const buf = await this.filehandle.read(maxFetch, 0, opts)
-    const bytes = (await unzip(buf)) as Uint8Array
-
-    // trim off lines after the last meta line
-    if (metaChar) {
-      let lastNewline = -1
-      const metaByte = metaChar.charCodeAt(0)
+    return unzip(buf)
+  }
 
-      for (let i = 0, l = bytes.length; i < l; i++) {
-        const byte = bytes[i]
-        if (i === lastNewline + 1 && byte !== metaByte) {
-          break
-        }
-        if (byte === NEWLINE) {
-          lastNewline = i
-        }
-      }
-      return bytes.subarray(0, lastNewline + 1)
+  /**
+   * Both header forms, from one read of the same bytes: the commented block
+   * `tabix -H` prints, and the rows `tabix -S N` counted.
+   *
+   * Memoized, because asking for one and then the other is the normal way to
+   * find a header whichever way the file keeps it (see `getHeaderLines`), and
+   * that used to fetch and decompress the file's leading blocks twice. Only
+   * the parsed results are retained — the decompressed bytes are dropped,
+   * which matters for a VCF header that can run to megabytes.
+   */
+  private async parseHeader(opts: Options) {
+    const { metaChar, skipLines = 0 } = await this.getMetadata(opts)
+    const bytes = await this.readHeaderBytes(opts)
+    return {
+      header: textDecoder.decode(
+        metaChar ? trimToMetaLines(bytes, metaChar) : bytes,
+      ),
+      skippedLines: skipLines > 0 ? firstLines(bytes, skipLines) : [],
     }
-    return bytes
+  }
+
+  private async getParsedHeader(opts: Options = {}) {
+    this.headerP ??= this.parseHeader(opts).catch((error: unknown) => {
+      this.headerP = undefined
+      throw error
+    })
+    return this.headerP
+  }
+
+  /**
+   * The bytes of the commented header. Deliberately not memoized, unlike the
+   * string form: this hands back the buffer, and holding one for the lifetime
+   * of the file is the caller's decision to make.
+   */
+  async getHeaderBuffer(opts: Options = {}) {
+    const { metaChar } = await this.getMetadata(opts)
+    const bytes = await this.readHeaderBytes(opts)
+    return metaChar ? trimToMetaLines(bytes, metaChar) : bytes
   }
 
   /**
@@ -634,29 +705,34 @@ export default class TabixIndexedFile {
    * way.
    */
   async getSkippedLines(opts: Options = {}) {
-    const {
-      firstDataLine,
-      skipLines = 0,
-      maxBlockSize,
-    } = await this.getMetadata(opts)
+    const { skipLines = 0 } = await this.getMetadata(opts)
+    // the index already answers this without reading the file at all
     if (skipLines <= 0) {
       return []
     }
-
-    // same read getHeaderBuffer makes, and the same caveat: a header spanning
-    // more blocks than this is not covered
-    const buf = await this.filehandle.read(
-      (firstDataLine?.blockPosition ?? 0) + maxBlockSize,
-      0,
-      opts,
-    )
-    const bytes = (await unzip(buf)) as Uint8Array
-    return new TextDecoder().decode(bytes).split(/\r?\n/).slice(0, skipLines)
+    return (await this.getParsedHeader(opts)).skippedLines
   }
 
   async getHeader(opts: Options = {}) {
-    const bytes = await this.getHeaderBuffer(opts)
-    return new TextDecoder().decode(bytes)
+    return (await this.getParsedHeader(opts)).header
+  }
+
+  /**
+   * The file's header lines, however that file keeps them: the meta-character
+   * block when there is one, and otherwise the rows the index counted. Empty
+   * lines are dropped.
+   *
+   * The two halves answer different questions (see `getSkippedLines`), but
+   * "what are this file's header lines" is nearly always the question a caller
+   * actually has, and answering it from `getHeader` alone is wrong in a way
+   * that doesn't announce itself: a bare header row comes back as the empty
+   * string, indistinguishable from a file that has no header, so callers fall
+   * back to an assumed column layout and quietly mis-name columns. Deciding it
+   * here also means one read of the leading blocks instead of two.
+   */
+  async getHeaderLines(opts: Options = {}) {
+    const { header, skippedLines } = await this.getParsedHeader(opts)
+    return (header ? header.split(/\r?\n/) : skippedLines).filter(Boolean)
   }
 
   async getReferenceSequenceNames(opts: Options = {}) {

package/src/tbi.ts

@@ -1,12 +1,9 @@
-import { unzip } from '@gmod/bgzf-filehandle'
-
 import Chunk from './chunk.ts'
 import IndexFile from './indexFile.ts'
 import {
   clampChunkEnds,
   memoizeByRefId,
   minVirtualOffset,
-  optimizeChunks,
   parseAuxData,
   parsePseudoBin,
 } from './util.ts'
@@ -19,44 +16,57 @@ const TBI_MAGIC = 21_578_324 // TBI\1
 
 // TBI is always the fixed depth-5, 14-bit-leaf binning scheme
 const TBI_DEPTH = 5
-const TBI_MAX_REF_LENGTH = 2 ** (14 + TBI_DEPTH * 3)
-
-/**
- * calculate the list of bins that may overlap with region [beg,end)
- * (zero-based half-open)
- */
-function reg2bins(beg: number, end: number) {
-  // Clamp into the binning scheme's range. Past 2**31 the shifts below
-  // truncate to int32 and go negative, which would silently yield no bins.
-  const b = Math.min(Math.max(beg, 0), TBI_MAX_REF_LENGTH)
-  const e = Math.min(end, TBI_MAX_REF_LENGTH) - 1
-  const bins: (readonly [number, number])[] = []
-  if (b <= e) {
-    bins.push(
-      [0, 0],
-      [1 + (b >> 26), 1 + (e >> 26)],
-      [9 + (b >> 23), 9 + (e >> 23)],
-      [73 + (b >> 20), 73 + (e >> 20)],
-      [585 + (b >> 17), 585 + (e >> 17)],
-      [4681 + (b >> 14), 4681 + (e >> 14)],
-    )
-  }
-  return bins
+const TBI_MIN_SHIFT = 14
+const TBI_MAX_REF_LENGTH = 2 ** (TBI_MIN_SHIFT + TBI_DEPTH * 3)
+
+// Clamp a query coordinate into the binning scheme's range. Past 2**31 the
+// int32 shifts below truncate and go negative, which would silently yield no
+// bins.
+function clampToRefLength(pos: number) {
+  return Math.min(Math.max(pos, 0), TBI_MAX_REF_LENGTH)
 }
 
 export default class TabixIndex extends IndexFile {
+  /**
+   * The bins that may overlap [beg, end) (zero-based half-open). TBI's scheme
+   * is CSI's with minShift/depth fixed, so this is CSI's reg2bins with the
+   * levels unrolled and its float shifts replaced by int32 ones.
+   */
+  protected reg2bins(beg: number, end: number) {
+    if (beg > TBI_MAX_REF_LENGTH) {
+      console.warn('querying outside of possible tabix range')
+    }
+    const b = clampToRefLength(beg)
+    const e = Math.min(end, TBI_MAX_REF_LENGTH) - 1
+    const bins: (readonly [number, number])[] = []
+    if (b <= e) {
+      bins.push(
+        [0, 0],
+        [1 + (b >> 26), 1 + (e >> 26)],
+        [9 + (b >> 23), 9 + (e >> 23)],
+        [73 + (b >> 20), 73 + (e >> 20)],
+        [585 + (b >> 17), 585 + (e >> 17)],
+        [4681 + (b >> 14), 4681 + (e >> 14)],
+      )
+    }
+    return bins
+  }
+
+  /**
+   * The linear index is monotonically non-decreasing, so the minimum virtual
+   * offset for chunks that could overlap [beg, ...) is its entry for beg.
+   * SYNC: ~/src/gmod/bam-js/src/bai.ts getLowestChunk
+   */
+  protected lowestOffset(ref: RefIndex, beg: number) {
+    const linearIndex = ref.linearIndex
+    return linearIndex?.[
+      Math.min(clampToRefLength(beg) >> TBI_MIN_SHIFT, linearIndex.length - 1)
+    ]
+  }
+
   /** @internal */
   async _parse(opts: Options = {}) {
-    const buf = await this.filehandle.readFile({
-      signal: opts.signal,
-      onProgress: opts.onProgress,
-    })
-    const bytes = (await unzip(buf)) as Uint8Array
-    const dataView = new DataView(
-      bytes.buffer,
-      bytes.byteOffset,
-      bytes.byteLength,
-    )
+    const { bytes, dataView } = await this.readIndexBytes(opts)
 
     if (dataView.getUint32(0, true) !== TBI_MAGIC) {
       throw new Error('Not a TBI file')
@@ -169,6 +179,8 @@ export default class TabixIndex extends IndexFile {
     return {
       indices: memoizeByRefId(getIndices),
       metaChar,
+      minShift: TBI_MIN_SHIFT,
+      depth: TBI_DEPTH,
       maxBinNumber,
       maxRefLength: TBI_MAX_REF_LENGTH,
       skipLines,
@@ -181,50 +193,4 @@ export default class TabixIndex extends IndexFile {
       maxBlockSize: 1 << 16,
     }
   }
-
-  async blocksForRange(
-    refName: string,
-    min: number,
-    max: number,
-    opts: Options = {},
-  ) {
-    const indexData = await this.parse(opts)
-    const refId = indexData.refNameToId[refName]
-    if (refId === undefined) {
-      return []
-    }
-    const ba = indexData.indices(refId)
-    if (!ba) {
-      return []
-    }
-
-    if (min > TBI_MAX_REF_LENGTH) {
-      console.warn('querying outside of possible tabix range')
-    }
-    // clamp before the >> 14 below, which would truncate to int32 and go
-    // negative past 2**31
-    const beg = Math.min(Math.max(min, 0), TBI_MAX_REF_LENGTH)
-    const overlappingBins = reg2bins(min, max) // List of bin #s that overlap min, max
-    const chunks: Chunk[] = []
-
-    // Find chunks in overlapping bins.  Leaf bins (< 4681) are not pruned
-    for (const [start, end] of overlappingBins) {
-      for (let bin = start; bin <= end; bin++) {
-        const binChunks = ba.binIndex[bin]
-        if (binChunks) {
-          for (const c of binChunks) {
-            chunks.push(c)
-          }
-        }
-      }
-    }
-
-    // The linear index is monotonically non-decreasing, so the minimum virtual
-    // offset for chunks that could overlap [min, ...) is at index minLin.
-    // SYNC: ~/src/gmod/bam-js/src/bai.ts getLowestChunk
-    const linearIndex = ba.linearIndex
-    const lowest = linearIndex?.[Math.min(beg >> 14, linearIndex.length - 1)]
-
-    return optimizeChunks(chunks, lowest)
-  }
 }