feat: Add optional maxRows argument to loadTSV

Author: effigies

src/files/tsv.test.ts

@@ -62,6 +62,48 @@ Deno.test('TSV loading', async (t) => {
     }
   })
 
+  await t.step('maxRows limits the number of rows read', async () => {
+    const file = pathToFile('/long.tsv')
+    // Use 1500 to avoid overlap with default initial capacity
+    const text = 'a\tb\tc\n' + '1\t2\t3\n'.repeat(1500)
+    file.stream = streamFromString(text)
+
+    let map = await loadTSV(file, 0)
+    assertEquals(map.a, [])
+    assertEquals(map.b, [])
+    assertEquals(map.c, [])
+
+    // Clear memoization cache. We currently do not key on maxRows.
+    loadTSV.cache.clear()
+    file.stream = streamFromString(text)
+    map = await loadTSV(file, 1)
+    assertEquals(map.a, ['1'])
+    assertEquals(map.b, ['2'])
+    assertEquals(map.c, ['3'])
+
+    loadTSV.cache.clear()
+    file.stream = streamFromString(text)
+    map = await loadTSV(file, 2)
+    assertEquals(map.a, ['1', '1'])
+    assertEquals(map.b, ['2', '2'])
+    assertEquals(map.c, ['3', '3'])
+
+    loadTSV.cache.clear()
+    file.stream = streamFromString(text)
+    map = await loadTSV(file, -1)
+    assertEquals(map.a, Array(1500).fill('1'))
+    assertEquals(map.b, Array(1500).fill('2'))
+    assertEquals(map.c, Array(1500).fill('3'))
+
+    loadTSV.cache.clear()
+    // Check that maxRows does not truncate shorter files
+    file.stream = streamFromString('a\tb\tc\n1\t2\t3\n4\t5\t6\n7\t8\t9\n')
+    map = await loadTSV(file, 4)
+    assertEquals(map.a, ['1', '4', '7'])
+    assertEquals(map.b, ['2', '5', '8'])
+    assertEquals(map.c, ['3', '6', '9'])
+  })
+
   // Tests will have populated the memoization cache
   await loadTSV.cache.clear()
 })

src/files/tsv.ts

@@ -8,7 +8,7 @@ import type { BIDSFile } from '../types/filetree.ts'
 import { filememoizeAsync } from '../utils/memoize.ts'
 import { createUTF8Stream } from './streams.ts'
 
-async function _loadTSV(file: BIDSFile): Promise<ColumnsMap> {
+async function _loadTSV(file: BIDSFile, maxRows: number = -1): Promise<ColumnsMap> {
   const reader = file.stream
     .pipeThrough(createUTF8Stream())
     .pipeThrough(new TextLineStream())
@@ -19,11 +19,12 @@ async function _loadTSV(file: BIDSFile): Promise<ColumnsMap> {
     const headers = (headerRow.done || !headerRow.value) ? [] : headerRow.value.split('\t')
 
     // Initialize columns in array for construction efficiency
-    const initialCapacity = 1000
+    const initialCapacity = maxRows >= 0 ? maxRows : 1000
     const columns: string[][] = headers.map(() => new Array<string>(initialCapacity))
 
+    maxRows = maxRows >= 0 ? maxRows : Infinity
     let rowIndex = 0 // Keep in scope after loop
-    for (; ; rowIndex++) {
+    for (; rowIndex < maxRows; rowIndex++) {
       const { done, value } = await reader.read()
       if (done) break
 

src/utils/memoize.ts

@@ -15,18 +15,18 @@ export const memoize = <T>(
 }
 
 export function filememoizeAsync<F extends HasParent, T>(
-  fn: (file: F) => Promise<T>,
-): WithCache<(file: F) => Promise<T>> {
+  fn: (file: F, ...args: any[]) => Promise<T>,
+): WithCache<(file: F, ...args: any[]) => Promise<T>> {
   const cache = new Map<string, Map<F, T>>()
-  const cached = async function (this: any, file: F): Promise<T> {
+  const cached = async function (this: any, file: F, ...args: any[]): Promise<T> {
     let subcache = cache.get(file.parent.path)
     if (!subcache) {
       subcache = new Map()
       cache.set(file.parent.path, subcache)
     }
     let val = subcache.get(file)
     if (!val) {
-      val = await fn.call(this, file)
+      val = await fn.call(this, file, ...args)
       subcache.set(file, val)
     }
     return val