merge the workers back together

Author: severo

packages/components/src/lib/tableProvider.ts

@@ -1,8 +1,6 @@
 import { DataFrame, ResolvablePromise, resolvablePromise } from 'hightable'
 import { FileMetaData, parquetSchema } from 'hyparquet'
-import { parquetQueryWorker } from '../workers/parquetWorkerClient.ts'
-import { parquetSortIndexWorker } from '../workers/sortParquetWorkerClient.ts'
-import type { AsyncBufferFrom } from '../workers/types.ts'
+import { AsyncBufferFrom, parquetQueryWorker, parquetSortIndexWorker } from '../workers/parquetWorkerClient.ts'
 
 type ResolvableRow = Record<string, ResolvablePromise<unknown>>;
 

packages/components/src/workers/parquetWorker.ts

@@ -1,12 +1,14 @@
 import { ColumnData, parquetQuery } from 'hyparquet'
 import { compressors } from 'hyparquet-compressors'
-import { asyncBufferFrom } from './parquetWorkerClient.ts'
-import type {
+import {
   ChunkMessage,
   ErrorMessage,
+  IndicesMessage,
   ParquetReadWorkerOptions,
   ResultMessage,
-} from './types.ts'
+  asyncBufferFrom,
+  compare,
+} from './parquetWorkerClient.ts'
 
 function postChunkMessage ({ chunk, queryId }: ChunkMessage) {
   self.postMessage({ chunk, queryId })
@@ -17,6 +19,9 @@ function postResultMessage ({ result, queryId }: ResultMessage) {
 function postErrorMessage ({ error, queryId }: ErrorMessage) {
   self.postMessage({ error, queryId })
 }
+function postIndicesMessage ({ indices, queryId }: IndicesMessage) {
+  self.postMessage({ indices, queryId })
+}
 
 self.onmessage = async ({
   data,
@@ -32,26 +37,49 @@ self.onmessage = async ({
     columns,
     queryId,
     chunks,
+    sortIndex,
   } = data
   const file = await asyncBufferFrom(from)
-  const onChunk = chunks
-    ? (chunk: ColumnData) => {
-      postChunkMessage({ chunk, queryId })
+  if (sortIndex === undefined) {
+    const onChunk = chunks
+      ? (chunk: ColumnData) => {
+        postChunkMessage({ chunk, queryId })
+      }
+      : undefined
+    try {
+      const result = await parquetQuery({
+        metadata,
+        file,
+        rowStart,
+        rowEnd,
+        orderBy,
+        columns,
+        compressors,
+        onChunk,
+      })
+      postResultMessage({ result, queryId })
+    } catch (error) {
+      postErrorMessage({ error: error as Error, queryId })
+    }
+  } else {
+    try {
+      // Special case for sorted index
+      if (orderBy === undefined)
+        throw new Error('sortParquetWorker requires orderBy')
+      if (rowStart !== undefined || rowEnd !== undefined)
+        throw new Error('sortIndex requires all rows')
+      const sortColumn = await parquetQuery({
+        metadata,
+        file,
+        columns: [orderBy],
+        compressors,
+      })
+      const indices = Array.from(sortColumn, (_, index) => index).sort((a, b) =>
+        compare<unknown>(sortColumn[a][orderBy], sortColumn[b][orderBy]),
+      )
+      postIndicesMessage({ indices, queryId })
+    } catch (error) {
+      postErrorMessage({ error: error as Error, queryId })
     }
-    : undefined
-  try {
-    const result = await parquetQuery({
-      metadata,
-      file,
-      rowStart,
-      rowEnd,
-      orderBy,
-      columns,
-      compressors,
-      onChunk,
-    })
-    postResultMessage({ result, queryId })
-  } catch (error) {
-    postErrorMessage({ error: error as Error, queryId })
   }
 }

packages/components/src/workers/parquetWorkerClient.ts

@@ -1,52 +1,105 @@
 import ParquetWorker from './parquetWorker?worker&inline'
 /// ^ the worker is bundled with the main thread code (inline) which is easier for users to import
 /// (no need to copy the worker file to the right place)
-import { AsyncBuffer, ColumnData } from 'hyparquet'
+import { AsyncBuffer, ColumnData, FileMetaData, ParquetReadOptions } from 'hyparquet'
 import { asyncBufferFromUrl } from '../lib/utils.ts'
-import type {
-  AsyncBufferFrom,
-  ParquetMessage,
-  ParquetReadWorkerOptions,
-  Row,
-} from './types.ts'
-// import { asyncBufferFromUrl, cachedAsyncBuffer, AsyncBuffer } from 'hyparquet'
+
+// Serializable constructor for AsyncBuffers
+export interface AsyncBufferFrom {
+  url: string
+  byteLength: number
+  headers?: Record<string, string>
+}
+// Same as ParquetReadOptions, but AsyncBufferFrom instead of AsyncBuffer
+export interface ParquetReadWorkerOptions extends Omit<ParquetReadOptions, 'file'> {
+  from: AsyncBufferFrom
+  orderBy?: string
+  sortIndex?: boolean
+}
+// Row is defined in hightable, but not exported + we change any to unknown
+export type Row = Record<string, unknown>;
+
+interface Message {
+  queryId: number
+}
+export interface ChunkMessage extends Message {
+  chunk: ColumnData
+}
+export interface ResultMessage extends Message {
+  result: Row[]
+}
+export interface IndicesMessage extends Message {
+  indices: number[]
+}
+export interface ErrorMessage extends Message {
+  error: Error
+}
+
+export type ParquetMessage = ChunkMessage | ResultMessage | ErrorMessage
+export type SortParquetMessage = IndicesMessage | ErrorMessage
+
+export interface ParquetSortIndexOptions {
+  metadata: FileMetaData
+  from: AsyncBufferFrom
+  orderBy: string
+}
 
 
 let worker: Worker | undefined
 let nextQueryId = 0
-interface QueryAgent {
+interface SortQueryAgent {
+  kind: 'sortIndex';
+  resolve: (value: number[]) => void;
+  reject: (error: Error) => void;
+}
+interface RowsQueryAgent {
+  kind: 'query';
   resolve: (value: Row[]) => void;
   reject: (error: Error) => void;
   onChunk?: (chunk: ColumnData) => void;
 }
+type QueryAgent = SortQueryAgent | RowsQueryAgent
+
 const pending = new Map<number, QueryAgent>()
 
 function getWorker() {
   if (!worker) {
     worker = new ParquetWorker()
-    worker.onmessage = ({ data }: { data: ParquetMessage }) => {
+    worker.onmessage = ({ data }: { data: ParquetMessage | SortParquetMessage }) => {
       const pendingQueryAgent = pending.get(data.queryId)
       if (!pendingQueryAgent) {
         throw new Error(
           `Unexpected: no pending promise found for queryId: ${data.queryId.toString()}`,
         )
         // TODO(SL): should never happen. But if it does, I'm not sure if throwing an error here helps.
       }
-      const { resolve, reject, onChunk } = pendingQueryAgent
-      if ('error' in data) {
-        reject(data.error)
-      } else if ('result' in data) {
-        resolve(data.result)
-      } else if ('chunk' in data) {
-        onChunk?.(data.chunk)
+      if (pendingQueryAgent.kind === 'query') {
+        const { resolve, reject, onChunk } = pendingQueryAgent
+        if ('error' in data) {
+          reject(data.error)
+        } else if ('result' in data) {
+          resolve(data.result)
+        } else if ('chunk' in data) {
+          onChunk?.(data.chunk)
+        } else {
+          reject(new Error('Unexpected message from worker'))
+        }
       } else {
-        reject(new Error('Unexpected message from worker'))
+        const { resolve, reject } = pendingQueryAgent
+        if ('error' in data) {
+          reject(data.error)
+        } else if ('indices' in data) {
+          resolve(data.indices)
+        } else {
+          reject(new Error('Unexpected message from worker'))
+        }
       }
     }
   }
   return worker
 }
 
+
 /**
  * Presents almost the same interface as parquetRead, but runs in a worker.
  * This is useful for reading large parquet files without blocking the main thread.
@@ -63,7 +116,7 @@ export function parquetQueryWorker({
 }: ParquetReadWorkerOptions): Promise<Row[]> {
   return new Promise((resolve, reject) => {
     const queryId = nextQueryId++
-    pending.set(queryId, { resolve, reject, onChunk })
+    pending.set(queryId, { kind: 'query', resolve, reject, onChunk })
     const worker = getWorker()
 
     // If caller provided an onChunk callback, worker will send chunks as they are parsed
@@ -80,6 +133,17 @@ export function parquetQueryWorker({
   })
 }
 
+export function parquetSortIndexWorker({ metadata, from, orderBy }: ParquetSortIndexOptions): Promise<number[]> {
+  return new Promise((resolve, reject) => {
+    const queryId = nextQueryId++
+    pending.set(queryId, { kind: 'sortIndex', resolve, reject })
+    const worker = getWorker()
+    worker.postMessage({
+      queryId, metadata, from, orderBy, sortIndex: true,
+    })
+  })
+}
+
 /**
  * Convert AsyncBufferFrom to AsyncBuffer and cache results.
  */