Convert parquet integration to TypeScript

Author: platypii

src/lib/parquet/parquetDataSource.js src/lib/parquet/parquetDataSource.tssrc/lib/parquet/parquetDataSource.js renamed to src/lib/parquet/parquetDataSource.ts

@@ -1,34 +1,24 @@
 import { parquetReadObjects, parquetSchema } from 'hyparquet'
+import type { AsyncBuffer, Compressors, FileMetaData } from 'hyparquet'
+import type { AsyncDataSource, AsyncRow, ScanOptions, SqlPrimitive } from 'squirreling'
 import { whereToParquetFilter } from './parquetFilter.js'
 
-/**
- * @import { AsyncBuffer, Compressors, FileMetaData, ParquetQueryFilter } from 'hyparquet'
- * @import { AsyncCells, AsyncDataSource, AsyncRow, SqlPrimitive } from 'squirreling'
- */
-
 /**
  * Creates a parquet data source for use with squirreling SQL engine.
- *
- * @param {AsyncBuffer} file
- * @param {FileMetaData} metadata
- * @param {Compressors} compressors
- * @returns {AsyncDataSource}
  */
-export function parquetDataSource(file, metadata, compressors) {
+export function parquetDataSource(file: AsyncBuffer, metadata: FileMetaData, compressors: Compressors): AsyncDataSource {
   const schema = parquetSchema(metadata)
   return {
     columns: schema.children.map(child => child.element.name),
-    scan({ columns, where, limit, offset, signal }) {
+    scan({ columns, where, limit, offset, signal }: ScanOptions) {
       // Convert WHERE AST to hyparquet filter format
       const whereFilter = where && whereToParquetFilter(where)
-      /** @type {ParquetQueryFilter | undefined} */
       const filter = where ? whereFilter : undefined
       const appliedWhere = Boolean(filter && whereFilter)
       const appliedLimitOffset = !where || appliedWhere
 
       // Ensure columns exist in metadata if provided
       if (columns) {
-        const schema = parquetSchema(metadata)
         for (const col of columns) {
           if (!schema.children.some(child => child.element.name === col)) {
             throw new Error(`Column "${col}" not found in parquet schema`)
@@ -61,7 +51,6 @@ export function parquetDataSource(file, metadata, compressors) {
             }
 
             // Read objects from this row group
-            // TODO: move to worker
             const data = await parquetReadObjects({
               file,
               metadata,
@@ -90,15 +79,8 @@ export function parquetDataSource(file, metadata, compressors) {
   }
 }
 
-/**
- * Creates an async row accessor that wraps a plain JavaScript object
- *
- * @param {Record<string, SqlPrimitive>} obj - the plain object
- * @returns {AsyncRow} a row accessor interface
- */
-function asyncRow(obj) {
-  /** @type {AsyncCells} */
-  const cells = {}
+function asyncRow(obj: Record<string, SqlPrimitive>): AsyncRow {
+  const cells: Record<string, () => Promise<SqlPrimitive>> = {}
   for (const [key, value] of Object.entries(obj)) {
     cells[key] = () => Promise.resolve(value)
   }

src/lib/parquet/parquetFilter.js

@@ -0,0 +1,121 @@
+import type { ParquetQueryFilter } from 'hyparquet'
+import type { BinaryNode, BinaryOp, ComparisonOp, ExprNode, InValuesNode, SqlPrimitive } from 'squirreling/src/ast.js'
+
+/**
+ * Converts a WHERE clause AST to hyparquet filter format.
+ * Returns undefined if the expression cannot be fully converted.
+ */
+export function whereToParquetFilter(where: ExprNode | undefined): ParquetQueryFilter | undefined {
+  if (!where) return undefined
+  return convertExpr(where, false)
+}
+
+function convertExpr(node: ExprNode, negate: boolean): ParquetQueryFilter | undefined {
+  if (node.type === 'unary' && node.op === 'NOT') {
+    return convertExpr(node.argument, !negate)
+  }
+  if (node.type === 'binary') {
+    return convertBinary(node, negate)
+  }
+  if (node.type === 'in valuelist') {
+    return convertInValues(node, negate)
+  }
+  if (node.type === 'cast') {
+    return convertExpr(node.expr, negate)
+  }
+  return undefined
+}
+
+function convertBinary({ op, left, right }: BinaryNode, negate: boolean): ParquetQueryFilter | undefined {
+  if (op === 'AND') {
+    const leftFilter = convertExpr(left, negate)
+    const rightFilter = convertExpr(right, negate)
+    if (!leftFilter || !rightFilter) return
+    return negate
+      ? { $or: [leftFilter, rightFilter] }
+      : { $and: [leftFilter, rightFilter] }
+  }
+  if (op === 'OR') {
+    const leftFilter = convertExpr(left, false)
+    const rightFilter = convertExpr(right, false)
+    if (!leftFilter || !rightFilter) return
+    return negate
+      ? { $nor: [leftFilter, rightFilter] }
+      : { $or: [leftFilter, rightFilter] }
+  }
+
+  if (op === 'LIKE') return
+
+  const { column, value, flipped } = extractColumnAndValue(left, right)
+  if (!column || value === undefined) return
+
+  const mongoOp = mapOperator(op, flipped, negate)
+  if (!mongoOp) return
+  return { [column]: { [mongoOp]: value } }
+}
+
+function extractColumnAndValue(left: ExprNode, right: ExprNode): { column: string | undefined; value: SqlPrimitive | undefined; flipped: boolean } {
+  if (left.type === 'identifier' && right.type === 'literal') {
+    return { column: left.name, value: right.value, flipped: false }
+  }
+  if (left.type === 'literal' && right.type === 'identifier') {
+    return { column: right.name, value: left.value, flipped: true }
+  }
+  return { column: undefined, value: undefined, flipped: false }
+}
+
+export function isBinaryOp(op: string): op is ComparisonOp {
+  return ['AND', 'OR', 'LIKE', '=', '!=', '<>', '<', '>', '<=', '>='].includes(op)
+}
+
+function mapOperator(op: BinaryOp, flipped: boolean, negate: boolean): string | undefined {
+  if (!isBinaryOp(op)) return
+
+  let mappedOp: ComparisonOp = op
+  if (negate) mappedOp = neg(mappedOp)
+  if (flipped) mappedOp = flip(mappedOp)
+  switch (mappedOp) {
+    case '=': return '$eq'
+    case '!=': case '<>': return '$ne'
+    case '<': return '$lt'
+    case '<=': return '$lte'
+    case '>': return '$gt'
+    case '>=': return '$gte'
+  }
+}
+
+function neg(op: ComparisonOp): ComparisonOp {
+  switch (op) {
+    case '<': return '>='
+    case '<=': return '>'
+    case '>': return '<='
+    case '>=': return '<'
+    case '=': return '!='
+    case '!=': return '='
+    case '<>': return '='
+  }
+}
+
+function flip(op: ComparisonOp): ComparisonOp {
+  switch (op) {
+    case '<': return '>'
+    case '<=': return '>='
+    case '>': return '<'
+    case '>=': return '<='
+    case '=': return '='
+    case '!=': return '!='
+    case '<>': return '='
+  }
+}
+
+function convertInValues(node: InValuesNode, negate: boolean): ParquetQueryFilter | undefined {
+  if (node.expr.type !== 'identifier') return
+
+  const values: SqlPrimitive[] = []
+  for (const val of node.values) {
+    if (val.type !== 'literal') return
+    values.push(val.value)
+  }
+
+  return { [node.expr.name]: { [negate ? '$nin' : '$in']: values } }
+}