Parquet aware table provider

Author: platypii

public/build/app.min.js

@@ -24,7 +24,6 @@ if (arg === 'chat') {
     } else if (stat.isFile()) {
       const parent = path.split('/').slice(0, -1).join('/')
       const key = path.split('/').pop()
-      console.log('WTF1', parent, key)
       serve(parent, key)
     }
   }).catch(() => {

public/build/app.min.js.map

@@ -3,7 +3,7 @@ import React, { useCallback, useEffect, useState } from 'react'
 import { parquetDataFrame } from '../../tableProvider.js'
 import { Spinner } from '../Layout.js'
 import ContentHeader from './ContentHeader.js'
-import { asyncBufferFromUrl, FileMetaData, parquetMetadataAsync } from 'hyparquet'
+import { asyncBufferFromUrl, parquetMetadataAsync } from 'hyparquet'
 
 enum LoadingState {
   NotLoaded,
@@ -28,7 +28,6 @@ interface Content {
 export default function ParquetView({ file, setProgress, setError }: ViewerProps) {
   const [loading, setLoading] = useState<LoadingState>(LoadingState.NotLoaded)
   const [content, setContent] = useState<Content>()
-  const [metadata, setMetadata] = useState<FileMetaData>()
 
   const isUrl = file.startsWith('http://') || file.startsWith('https://')
   const url = isUrl ? file : '/api/store/get?key=' + file
@@ -41,8 +40,7 @@ export default function ParquetView({ file, setProgress, setError }: ViewerProps
         const from = { url, byteLength: asyncBuffer.byteLength }
         setProgress(0.66)
         const metadata = await parquetMetadataAsync(asyncBuffer)
-        setMetadata(metadata)
-        let dataframe = await parquetDataFrame(from, metadata)
+        let dataframe = parquetDataFrame(from, metadata)
         dataframe = rowCache(dataframe)
         const fileSize = asyncBuffer.byteLength
         setContent({ dataframe, fileSize })

public/build/worker.min.js

@@ -1,17 +1,82 @@
-import type { DataFrame } from 'hightable'
+import { DataFrame, resolvablePromise } from 'hightable'
 import { FileMetaData, parquetSchema } from 'hyparquet'
-import { AsyncBufferFrom, parquetQueryWorker } from './workers/parquetWorkerClient.js'
+import { AsyncBufferFrom, parquetQueryWorker, parquetSortIndexWorker } from './workers/parquetWorkerClient.js'
 
 /**
  * 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 data = new Array(Number(metadata.num_rows))
+  const groups = new Array(metadata.row_groups.length).fill(false)
+  let groupStart = 0
+  const groupEnds = metadata.row_groups.map(group => groupStart += Number(group.num_rows))
+
+  function fetchRowGroup(groupIndex: number) {
+    if (!groups[groupIndex]) {
+      const rowStart = groupEnds[groupIndex - 1] || 0
+      const rowEnd = groupEnds[groupIndex]
+      // Initialize with resolvable promises
+      for (let i = rowStart; i < rowEnd; i++) {
+        data[i] = Object.fromEntries(header.map(key => [key, resolvablePromise()]))
+      }
+      parquetQueryWorker({ from, metadata, rowStart, rowEnd }).then(groupData => {
+        for (let i = rowStart; i < rowEnd; i++) {
+          for (const [key, value] of Object.entries(groupData[i - rowStart])) {
+            data[i][key].resolve(value)
+          }
+        }
+      })
+      groups[groupIndex] = true
+    }
+  }
+
+  function getSortIndex(orderBy: string) {
+    let sortIndex = sortCache.get(orderBy)
+    if (!sortIndex) {
+      sortIndex = parquetSortIndexWorker({ from, metadata, orderBy })
+      sortCache.set(orderBy, sortIndex)
+    }
+    return sortIndex
+  }
+
   return {
-    header: children.map(child => child.element.name),
+    header,
     numRows: Number(metadata.num_rows),
     rows(rowStart: number, rowEnd: number, orderBy?: string) {
-      return parquetQueryWorker({ asyncBuffer: from, rowStart, rowEnd, orderBy })
+      if (orderBy) {
+        const numRows = rowEnd - rowStart
+        const wrapped = new Array(numRows).fill(null)
+          .map(() => Object.fromEntries(header.map(key => [key, resolvablePromise()])))
+
+        let sortIndex = getSortIndex(orderBy)
+        sortIndex.then(indices => {
+          // Compute row groups to fetch
+          for (const index of indices.slice(rowStart, rowEnd)) {
+            const groupIndex = groupEnds.findIndex(end => index < end)
+            fetchRowGroup(groupIndex)
+          }
+
+          // Re-assemble data in sorted order into wrapped
+          for (let i = rowStart; i < rowEnd; i++) {
+            for (const key of header) {
+              data[indices[i]][key].then((value: any) => wrapped[i - rowStart][key].resolve(value))
+            }
+          }
+        })
+
+        return wrapped
+      } else {
+        for (let i = 0; i < groups.length; i++) {
+          const groupStart = groupEnds[i - 1] || 0
+          if (rowStart < groupEnds[i] && rowEnd > groupStart) {
+            fetchRowGroup(i)
+          }
+        }
+        return data.slice(rowStart, rowEnd)
+      }
     },
     sortable: true,
   }

public/build/worker.min.js.map

@@ -1,16 +1,28 @@
 import { ColumnData, parquetQuery } from 'hyparquet'
 import { compressors } from 'hyparquet-compressors'
-import { asyncBufferFrom } from './parquetWorkerClient.js'
+import { asyncBufferFrom, compare } from './parquetWorkerClient.js'
 
 self.onmessage = async ({ data }) => {
-  const { metadata, asyncBuffer, rowStart, rowEnd, orderBy, queryId, chunks } = data
-  const file = await asyncBufferFrom(asyncBuffer)
+  const { metadata, from, rowStart, rowEnd, orderBy, columns, queryId, chunks, sortIndex } = data
+  const file = await asyncBufferFrom(from)
   const onChunk = chunks ? (chunk: ColumnData) => self.postMessage({ chunk, queryId }) : undefined
   try {
-    const result = await parquetQuery({
-      metadata, file, rowStart, rowEnd, orderBy, compressors, onChunk
-    })
-    self.postMessage({ result, queryId })
+    if (sortIndex) {
+      // Special case for sorted index
+      if (orderBy === undefined) throw new Error('sortIndex requires orderBy')
+      if (rowStart !== undefined || rowEnd !== undefined) throw new Error('sortIndex requires all rows')
+      const sortColumn = await parquetQuery({
+        metadata, file, columns: [orderBy], compressors
+      })
+      const result = Array.from(sortColumn, (_, index) => index)
+        .sort((a, b) => compare(sortColumn[a][orderBy], sortColumn[b][orderBy]))
+      self.postMessage({ result, queryId })
+    } else {
+      const result = await parquetQuery({
+        metadata, file, rowStart, rowEnd, orderBy, columns, compressors, onChunk
+      })
+      self.postMessage({ result, queryId })
+    }
   } catch (error) {
     self.postMessage({ error, queryId })
   }

src/cli.js

@@ -1,4 +1,4 @@
-import { asyncBufferFromUrl, cachedAsyncBuffer, AsyncBuffer, ParquetReadOptions } from 'hyparquet'
+import { asyncBufferFromUrl, cachedAsyncBuffer, AsyncBuffer, ParquetReadOptions, FileMetaData } from 'hyparquet'
 
 // Serializable constructor for AsyncBuffers
 export interface AsyncBufferFrom {
@@ -8,13 +8,38 @@ export interface AsyncBufferFrom {
 
 // Same as ParquetReadOptions, but AsyncBufferFrom instead of AsyncBuffer
 interface ParquetReadWorkerOptions extends Omit<ParquetReadOptions, 'file'> {
-  asyncBuffer: AsyncBufferFrom
+  from: AsyncBufferFrom
   orderBy?: string
+  sortIndex?: boolean
 }
 
 let worker: Worker | undefined
 let nextQueryId = 0
-const pending = new Map<number, { resolve: (value: any) => void, reject: (error: any) => void }>()
+interface QueryAgent {
+  resolve: (value: any) => void
+  reject: (error: any) => void
+  onChunk?: (chunk: any) => void
+}
+const pending = new Map<number, QueryAgent>()
+
+function getWorker() {
+  if (!worker) {
+    worker = new Worker(new URL('worker.min.js', import.meta.url))
+    worker.onmessage = ({ data }) => {
+      const { resolve, reject, onChunk } = pending.get(data.queryId)!
+      if (data.error) {
+        reject(data.error)
+      } else if (data.result) {
+        resolve(data.result)
+      } else if (data.chunk) {
+        onChunk?.(data.chunk)
+      } else {
+        reject(new Error('Unexpected message from worker'))
+      }
+    }
+  }
+  return worker
+}
 
 /**
  * Presents almost the same interface as parquetRead, but runs in a worker.
@@ -23,32 +48,35 @@ const pending = new Map<number, { resolve: (value: any) => void, reject: (error:
  * to be serialized to the worker.
  */
 export function parquetQueryWorker(
-  { metadata, asyncBuffer, rowStart, rowEnd, orderBy, onChunk }: ParquetReadWorkerOptions
+  { metadata, from, rowStart, rowEnd, orderBy, onChunk }: ParquetReadWorkerOptions
 ): Promise<Record<string, any>[]> {
   return new Promise((resolve, reject) => {
     const queryId = nextQueryId++
-    pending.set(queryId, { resolve, reject })
-    // Create a worker
-    if (!worker) {
-      worker = new Worker(new URL('worker.min.js', import.meta.url))
-      worker.onmessage = ({ data }) => {
-        const { resolve, reject } = pending.get(data.queryId)!
-        // Convert postmessage data to callbacks
-        if (data.error) {
-          reject(data.error)
-        } else if (data.result) {
-          resolve(data.result)
-        } else if (data.chunk) {
-          onChunk?.(data.chunk)
-        } else {
-          reject(new Error('Unexpected message from worker'))
-        }
-      }
-    }
+    pending.set(queryId, { 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, asyncBuffer, rowStart, rowEnd, orderBy, chunks
+      queryId, metadata, from, rowStart, rowEnd, orderBy, chunks
+    })
+  })
+}
+
+interface ParquetSortIndexOptions {
+  metadata: FileMetaData
+  from: AsyncBufferFrom
+  orderBy: string
+}
+
+export function parquetSortIndexWorker({ metadata, from, orderBy }: ParquetSortIndexOptions): Promise<number[]> {
+  return new Promise((resolve, reject) => {
+    const queryId = nextQueryId++
+    pending.set(queryId, { resolve, reject })
+    const worker = getWorker()
+
+    worker.postMessage({
+      queryId, metadata, from, orderBy, sortIndex: true
     })
   })
 }
@@ -65,3 +93,9 @@ export async function asyncBufferFrom(from: AsyncBufferFrom): Promise<AsyncBuffe
   return asyncBuffer
 }
 const cache = new Map<string, Promise<AsyncBuffer>>()
+
+export function compare(a: any, b: any): number {
+  if (a < b) return -1
+  if (a > b) return 1
+  return 1 // TODO: how to handle nulls?
+}