implement multi column sort

Author: severo

src/lib/tableProvider.ts

@@ -1,15 +1,41 @@
-import { DataFrame, ResolvableRow, resolvableRow } from 'hightable'
+import { DataFrame, OrderBy, ResolvableRow, resolvableRow } from 'hightable'
 import { FileMetaData, parquetSchema } from 'hyparquet'
-import { parquetQueryWorker, parquetSortIndexWorker } from './workers/parquetWorkerClient.js'
+import { parquetColumnRanksWorker, parquetQueryWorker } from './workers/parquetWorkerClient.js'
 import type { AsyncBufferFrom } from './workers/types.d.ts'
 
+/*
+ * sortIndex[0] gives the index of the first row in the sorted table
+ */
+export function computeSortIndex(orderByRanks: { direction: 'ascending' | 'descending', ranks: number[] }[]): number[] {
+  if (!(0 in orderByRanks)) {
+    throw new Error('orderByRanks should have at least one element')
+  }
+  const numRows = orderByRanks[0].ranks.length
+  return Array
+    .from({ length: numRows }, (_, i) => i)
+    .sort((a, b) => {
+      for (const { direction, ranks } of orderByRanks) {
+        const rankA = ranks[a]
+        const rankB = ranks[b]
+        if (rankA === undefined || rankB === undefined) {
+          throw new Error('Invalid ranks')
+        }
+        const value = direction === 'ascending' ? 1 : -1
+        if (rankA < rankB) return -value
+        if (rankA > rankB) return value
+      }
+      return 0
+    })
+}
+
 /**
  * Convert a parquet file into a dataframe.
  */
 export function parquetDataFrame(from: AsyncBufferFrom, metadata: FileMetaData): DataFrame {
   const { children } = parquetSchema(metadata)
   const header = children.map(child => child.element.name)
   const sortCache = new Map<string, Promise<number[]>>()
+  const columnRanksCache = new Map<string, Promise<number[]>>()
   const data = new Array<ResolvableRow | undefined>(Number(metadata.num_rows))
   const groups = new Array(metadata.row_groups.length).fill(false)
   let groupStart = 0
@@ -34,7 +60,8 @@ export function parquetDataFrame(from: AsyncBufferFrom, metadata: FileMetaData):
               throw new Error(`Missing data row for index ${i}`)
             }
             dataRow.index.resolve(i)
-            const row = groupData[i - rowStart]
+            const j = i - rowStart
+            const row = groupData[j]
             if (row === undefined) {
               throw new Error(`Missing row in groupData for index: ${i - rowStart}`)
             }
@@ -54,20 +81,33 @@ export function parquetDataFrame(from: AsyncBufferFrom, metadata: FileMetaData):
     }
   }
 
-  function getSortIndex(orderBy: string) {
-    let sortIndex = sortCache.get(orderBy)
+  function getColumnRanks(column: string): Promise<number[]> {
+    let columnRanks = columnRanksCache.get(column)
+    if (!columnRanks) {
+      columnRanks = parquetColumnRanksWorker({ from, metadata, column })
+      columnRanksCache.set(column, columnRanks)
+    }
+    return columnRanks
+  }
+
+  function getSortIndex(orderBy: OrderBy): Promise<number[]> {
+    const orderByKey = JSON.stringify(orderBy)
+    let sortIndex = sortCache.get(orderByKey)
     if (!sortIndex) {
-      sortIndex = parquetSortIndexWorker({ from, metadata, orderBy })
-      sortCache.set(orderBy, sortIndex)
+      const orderByRanksPromise = Promise.all(
+        orderBy.map(({ column, direction }) => getColumnRanks(column).then(ranks => ({ direction, ranks })))
+      )
+      sortIndex = orderByRanksPromise.then(orderByRanks => computeSortIndex(orderByRanks))
+      sortCache.set(orderByKey, sortIndex)
     }
     return sortIndex
   }
 
   return {
     header,
     numRows: Number(metadata.num_rows),
-    rows({ start, end, orderBy }: { start: number, end: number, orderBy?: string}) {
-      if (orderBy) {
+    rows({ start, end, orderBy }: { start: number, end: number, orderBy?: OrderBy}) {
+      if (orderBy?.length) {
         const numRows = end - start
         const wrapped = new Array(numRows).fill(null).map(() => resolvableRow(header))
 

src/lib/workers/parquetWorker.ts

@@ -1,7 +1,7 @@
 import { ColumnData, parquetQuery } from 'hyparquet'
 import { compressors } from 'hyparquet-compressors'
 import { asyncBufferFrom } from '../utils.js'
-import type { ChunkMessage, ErrorMessage, IndicesMessage, ParquetReadWorkerOptions, ResultMessage } from './types.js'
+import type { ChunkMessage, ClientMessage, ColumnRanksMessage, ErrorMessage, ResultMessage } from './types.js'
 
 function postChunkMessage ({ chunk, queryId }: ChunkMessage) {
   self.postMessage({ chunk, queryId })
@@ -12,35 +12,51 @@ function postResultMessage ({ result, queryId }: ResultMessage) {
 function postErrorMessage ({ error, queryId }: ErrorMessage) {
   self.postMessage({ error, queryId })
 }
-function postIndicesMessage ({ indices, queryId }: IndicesMessage) {
-  self.postMessage({ indices, queryId })
+function postColumnRanksMessage ({ columnRanks, queryId }: ColumnRanksMessage) {
+  self.postMessage({ columnRanks, queryId })
 }
 
-self.onmessage = async ({ data }: {
-  data: ParquetReadWorkerOptions & { queryId: number; chunks: boolean };
-}) => {
-  const { metadata, from, rowStart, rowEnd, orderBy, columns, queryId, chunks, sortIndex } = data
+self.onmessage = async ({ data }: { data: ClientMessage }) => {
+  const { metadata, from, kind, queryId } = data
   const file = await asyncBufferFrom(from)
-  if (sortIndex === undefined) {
-    const onChunk = chunks ? (chunk: ColumnData) => { postChunkMessage({ chunk, queryId }) } : undefined
+  if (kind === 'columnRanks') {
+    const { column } = data
+    // return the column ranks in ascending order
+    // we can get the descending order replacing the rank with numRows - rank - 1. It's not exactly the rank of
+    // 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 result = await parquetQuery({ metadata, file, rowStart, rowEnd, orderBy, columns, compressors, onChunk })
-      postResultMessage({ result, queryId })
+      const sortColumn = await parquetQuery({ metadata, file, columns: [column], compressors })
+      const valuesWithIndex = sortColumn.map((row, index) => ({ value: row[column] as unknown, 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) => {
+        const { lastValue, lastRank, ranks } = accumulator
+        const { value, index } = currentValue
+        if (value === lastValue) {
+          ranks[index] = lastRank
+          return { ranks, lastValue, lastRank }
+        } else {
+          ranks[index] = rank
+          return { ranks, lastValue: value, lastRank: rank }
+        }
+      }, {
+        ranks: Array(numRows).fill(-1) as number[],
+        lastValue: undefined as unknown,
+        lastRank: 0,
+      }).ranks
+      postColumnRanksMessage({ columnRanks: columnRanks, queryId })
     } catch (error) {
       postErrorMessage({ error: error as Error, queryId })
     }
   } else {
+    const { rowStart, rowEnd, chunks } = data
+    const onChunk = chunks ? (chunk: ColumnData) => { postChunkMessage({ chunk, queryId }) } : undefined
     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 })
+      const result = await parquetQuery({ metadata, file, rowStart, rowEnd, compressors, onChunk })
+      postResultMessage({ result, queryId })
     } catch (error) {
       postErrorMessage({ error: error as Error, queryId })
     }

src/lib/workers/parquetWorkerClient.ts

@@ -2,38 +2,40 @@ 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 { ColumnData } from 'hyparquet'
-import type { ParquetMessage, ParquetReadWorkerOptions, ParquetSortIndexOptions, Row, SortParquetMessage } from './types.js'
+import type { Cells, ColumnRanksClientMessage, ColumnRanksWorkerMessage, ColumnRanksWorkerOptions, QueryClientMessage, QueryWorkerMessage, QueryWorkerOptions } from './types.js'
 
 let worker: Worker | undefined
 let nextQueryId = 0
-interface SortQueryAgent {
-  kind: 'sortIndex';
-  resolve: (value: number[]) => void;
-  reject: (error: Error) => void;
-}
 interface RowsQueryAgent {
   kind: 'query';
-  resolve: (value: Row[]) => void;
+  resolve: (value: Cells[]) => void;
   reject: (error: Error) => void;
   onChunk?: (chunk: ColumnData) => void;
 }
-type QueryAgent = SortQueryAgent | RowsQueryAgent
+interface ColumnRanksQueryAgent {
+  kind: 'columnRanks';
+  resolve: (value: number[]) => void;
+  reject: (error: Error) => void;
+}
+type QueryAgent = RowsQueryAgent | ColumnRanksQueryAgent
 
 const pending = new Map<number, QueryAgent>()
 
 function getWorker() {
   if (!worker) {
     worker = new ParquetWorker()
-    worker.onmessage = ({ data }: { data: ParquetMessage | SortParquetMessage }) => {
+    worker.onmessage = ({ data }: { data: QueryWorkerMessage | ColumnRanksWorkerMessage }) => {
       const pendingQueryAgent = pending.get(data.queryId)
       if (!pendingQueryAgent) {
         console.warn(
           `Unexpected: no pending promise found for queryId: ${data.queryId.toString()}`
         )
         return
       }
+
       if (pendingQueryAgent.kind === 'query') {
-        const { resolve, reject, onChunk } = pendingQueryAgent
+        const { resolve, reject } = pendingQueryAgent
+        const { onChunk } = pendingQueryAgent
         if ('error' in data) {
           reject(data.error)
         } else if ('result' in data) {
@@ -43,15 +45,16 @@ function getWorker() {
         } else {
           reject(new Error('Unexpected message from worker'))
         }
+        return
+      }
+
+      const { resolve, reject } = pendingQueryAgent
+      if ('error' in data) {
+        reject(data.error)
+      } else if ('columnRanks' in data) {
+        resolve(data.columnRanks)
       } else {
-        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'))
-        }
+        reject(new Error('Unexpected message from worker'))
       }
     }
   }
@@ -64,25 +67,26 @@ function getWorker() {
  * Instead of taking an AsyncBuffer, it takes a AsyncBufferFrom, because it needs
  * to be serialized to the worker.
  */
-export function parquetQueryWorker({ metadata, from, rowStart, rowEnd, orderBy, onChunk }: ParquetReadWorkerOptions): Promise<Row[]> {
+export function parquetQueryWorker({ metadata, from, rowStart, rowEnd, onChunk }: QueryWorkerOptions): Promise<Cells[]> {
+  // TODO(SL) Support passing columns?
   return new Promise((resolve, reject) => {
     const queryId = nextQueryId++
     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
     const chunks = onChunk !== undefined
-    worker.postMessage({ queryId, metadata, from, rowStart, rowEnd, orderBy, chunks })
+    const message: QueryClientMessage = { queryId, metadata, from, rowStart, rowEnd, chunks, kind: 'query' }
+    worker.postMessage(message)
   })
 }
 
-export function parquetSortIndexWorker({ metadata, from, orderBy }: ParquetSortIndexOptions): Promise<number[]> {
+export function parquetColumnRanksWorker({ metadata, from, column }: ColumnRanksWorkerOptions): Promise<number[]> {
   return new Promise((resolve, reject) => {
     const queryId = nextQueryId++
-    pending.set(queryId, { kind: 'sortIndex', resolve, reject })
+    pending.set(queryId, { kind: 'columnRanks', resolve, reject })
     const worker = getWorker()
-    worker.postMessage({
-      queryId, metadata, from, orderBy, sortIndex: true,
-    })
+    const message: ColumnRanksClientMessage = { queryId, metadata, from, column, kind: 'columnRanks' }
+    worker.postMessage(message)
   })
 }

src/lib/workers/types.ts

@@ -1,4 +1,4 @@
-import { ColumnData, FileMetaData, ParquetReadOptions } from 'hyparquet'
+import { ColumnData, FileMetaData } from 'hyparquet'
 
 // Serializable constructors for AsyncBuffers
 interface AsyncBufferFromFile {
@@ -11,37 +11,48 @@ interface AsyncBufferFromUrl {
   requestInit?: RequestInit
 }
 export type AsyncBufferFrom = AsyncBufferFromFile | AsyncBufferFromUrl
+// Cells is defined in hightable, but uses any, not unknown
+export type Cells = Record<string, unknown> ;
 
-// Same as ParquetReadOptions, but AsyncBufferFrom instead of AsyncBuffer
-export interface ParquetReadWorkerOptions extends Omit<ParquetReadOptions, 'file'> {
+export interface CommonWorkerOptions {
+  metadata: FileMetaData,
   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 ErrorMessage extends Message {
+  error: Error
+}
+
+/* Query worker */
+export interface QueryWorkerOptions extends CommonWorkerOptions {
+  rowStart?: number,
+  rowEnd?: number,
+  onChunk?: (chunk: ColumnData) => void
+}
+export interface QueryClientMessage extends QueryWorkerOptions, Message {
+  kind: 'query',
+  chunks?: boolean
+}
 export interface ChunkMessage extends Message {
   chunk: ColumnData
 }
 export interface ResultMessage extends Message {
-  result: Row[]
+  result: Cells[]
 }
-export interface IndicesMessage extends Message {
-  indices: number[]
+export type QueryWorkerMessage = ChunkMessage | ResultMessage | ErrorMessage
+
+/* ColumnRanks worker */
+export interface ColumnRanksWorkerOptions extends CommonWorkerOptions {
+  column: string
 }
-export interface ErrorMessage extends Message {
-  error: Error
+export interface ColumnRanksClientMessage extends ColumnRanksWorkerOptions, Message {
+  kind: 'columnRanks'
 }
-
-export type ParquetMessage = ChunkMessage | ResultMessage | ErrorMessage
-export type SortParquetMessage = IndicesMessage | ErrorMessage
-
-export interface ParquetSortIndexOptions {
-  metadata: FileMetaData
-  from: AsyncBufferFrom
-  orderBy: string
+export interface ColumnRanksMessage extends Message {
+  columnRanks: number[]
 }
+export type ColumnRanksWorkerMessage = ColumnRanksMessage | ErrorMessage
+
+export type ClientMessage = QueryClientMessage | ColumnRanksClientMessage