only fetch the values for the column, without transposing to objects

Author: severo

src/lib/getParquetColumn.ts

@@ -0,0 +1,65 @@
+import { ColumnData, ParquetReadOptions, parquetRead } from 'hyparquet'
+
+type GetColumnOptions = Omit<ParquetReadOptions, 'columns' | 'rowStart' | 'rowEnd' | 'onChunk' | 'onComplete'> & {column: string}
+
+export async function getParquetColumn({ metadata, file, column, compressors }: GetColumnOptions): Promise<unknown[]> {
+  const numRows = Number(metadata?.num_rows)
+  if (isNaN(numRows)) {
+    throw new Error('metadata.num_rows is undefined')
+  }
+  if (numRows === 0) {
+    return []
+  }
+  const lastError: {error?: Error} = {}
+  const values: unknown[] = Array(numRows).fill(undefined)
+  const ranges: [number, number][] = []
+  function onChunk({ columnName, columnData, rowStart, rowEnd }: ColumnData) {
+    if (columnName !== column) {
+      lastError.error = new Error(`unexpected column name ${columnName}`)
+    }
+    for (let i = rowStart; i < rowEnd; i++) {
+      values[i] = columnData[i - rowStart]
+    }
+    ranges.push([rowStart, rowEnd])
+  }
+
+  // this awaits all the promises. When it returns, all the data should have already been sent using onChunk
+  await parquetRead({ metadata, file, columns: [column], compressors, onChunk })
+
+  // Do some checks before returning the data
+
+  // check for errors
+  if (lastError.error !== undefined) {
+    throw lastError.error
+  }
+
+  // check for missing data (should be faster than checking for undefined values in the array)
+  const sortedRanges = ranges.sort((a, b) => a[0] - b[0])
+  for (let i = 0; i < sortedRanges.length - 1; i++) {
+    const range = sortedRanges[i]
+    const nextRange = sortedRanges[i + 1]
+    if (!range || !nextRange) {
+      throw new Error('The ranges should not be undefined')
+    }
+    if (range[1] !== nextRange[0]) {
+      throw new Error(`missing data between rows ${range[1]} and ${nextRange[0]}`)
+    }
+  }
+  const firstRange = sortedRanges[0]
+  if (!firstRange) {
+    throw new Error('The first range should not be undefined')
+  }
+  if (firstRange[0] !== 0) {
+    throw new Error(`missing data before row ${firstRange[0]}`)
+  }
+  const lastRange = sortedRanges[sortedRanges.length - 1]
+  if (!lastRange) {
+    throw new Error('The last range should not be undefined')
+  }
+  if (lastRange[1] !== numRows) {
+    throw new Error(`missing data after row ${lastRange[1]}`)
+  }
+
+  // return the values
+  return values
+}

src/lib/workers/parquetWorker.ts

@@ -1,5 +1,6 @@
 import { ColumnData, parquetQuery } from 'hyparquet'
 import { compressors } from 'hyparquet-compressors'
+import { getParquetColumn } from '../getParquetColumn.js'
 import { asyncBufferFrom } from '../utils.js'
 import type { ChunkMessage, ClientMessage, ColumnRanksMessage, ErrorMessage, ResultMessage } from './types.js'
 
@@ -26,10 +27,9 @@ self.onmessage = async ({ data }: { data: ClientMessage }) => {
     // the descending order, because the rank is the first, not the last, of the ties. But it's enough for the
     // purpose of sorting.
 
-    // TODO(SL): ensure only the expected column is fetched
     try {
-      const sortColumn = await parquetQuery({ metadata, file, columns: [column], compressors })
-      const valuesWithIndex = sortColumn.map((row, index) => ({ value: row[column] as unknown, index }))
+      const sortColumn = await getParquetColumn({ metadata, file, column, compressors })
+      const valuesWithIndex = sortColumn.map((value, index) => ({ value, index }))
       const sortedValuesWithIndex = Array.from(valuesWithIndex).sort(({ value: a }, { value: b }) => compare<unknown>(a, b))
       const numRows = sortedValuesWithIndex.length
       const columnRanks = sortedValuesWithIndex.reduce((accumulator, currentValue, rank) => {

src/lib/workers/parquetWorkerClient.ts

@@ -34,8 +34,7 @@ function getWorker() {
       }
 
       if (pendingQueryAgent.kind === 'query') {
-        const { resolve, reject } = pendingQueryAgent
-        const { onChunk } = pendingQueryAgent
+        const { onChunk, resolve, reject } = pendingQueryAgent
         if ('error' in data) {
           reject(data.error)
         } else if ('result' in data) {