Predicate Pushdown - #100
Conversation
|
Please feel free to ask me questions about this implementation! |
|
This is AWESOME @park-brian!! 🙌 The code looks really nice. Give me a few days to read carefully and run some benchmarks. |
|
I can also work on some benchmarks! The query function also has new offset, limit, and order direction clauses (this will allow us to support basic sql without too much effort - we can even use template functions to escape values). eg: |
|
I have an existing benchmark suite at hyparquet-perf. It might need tests added that leverage to predicate pushdown. Also I have been developing a SQL to Mongo function for that reason. It works pretty well so far, and predicate pushdown will make it even better! sqlToMongo(sql: string)import type { OrderBy } from 'hightable'
import type { ParquetQueryFilter } from 'hyparquet/src/types.js'
interface SqlToMongo {
collection: string // table name
columns: string[] // selected columns
filter: ParquetQueryFilter // query filter (where clause)
orderBy?: OrderBy // optional order by clause
limit?: number // optional limit for results
// TODO: parse offset
offset?: number // optional offset for results
}
type Direction = 'ascending' | 'descending'
interface ParseExpressionResult {
node: Record<string, any>
endIdx: number
}
/**
* Convert a small SQL subset into Mongo‑style
* collection, projection and filter information.
*/
export function sqlToMongo(sql: string): SqlToMongo {
sql = sql.trim()
// Combined regex to match:
// SELECT <fields> FROM <collection>
// This ensures we have both SELECT and FROM.
const combinedMatch = /^SELECT\s+(.*?)\s+FROM\s+([A-Za-z0-9_]+)/i.exec(sql)
if (!combinedMatch) {
// If we can't match SELECT ... FROM ..., let's see what is missing
if (!/^SELECT\s+/i.test(sql)) {
throw new Error('Invalid SQL: SELECT clause not found.')
} else if (!/FROM\s+/i.test(sql)) {
throw new Error('Invalid SQL: FROM clause not found.')
} else {
throw new Error('Invalid SQL: SELECT or FROM clause not found.')
}
}
const selectClause = combinedMatch[1].trim()
const collection = combinedMatch[2]
const columns = parseSelectFields(selectClause)
// Extract OFFSET clause if present
// Must be at the end, or just before LIMIT
const offsetRegex = /\bOFFSET\s+(\d+)(?=\s+LIMIT\b|\s*;?\s*$)/i
const offsetMatch = offsetRegex.exec(sql)
let offset: number | undefined
if (offsetMatch) {
offset = parseInt(offsetMatch[1], 10)
if (isNaN(offset) || offset < 0) {
throw new Error('Invalid SQL: OFFSET must be a non-negative integer.')
}
// Remove it from the original SQL
sql = sql.replace(offsetMatch[0], '').trim()
}
// Extract LIMIT clause if present
const limitRegex = /\bLIMIT\s+(\d+)(?:\s*;|\s*$)/i
const limitMatch = limitRegex.exec(sql)
let limit: number | undefined
if (limitMatch) {
limit = parseInt(limitMatch[1], 10)
if (isNaN(limit) || limit < 1) {
throw new Error('Invalid SQL: LIMIT must be a positive integer.')
}
// Trim it from the original SQL
sql = sql.slice(0, limitMatch.index).trim()
}
// Extract ORDER BY clause if present
const orderByRegex = /\bORDER\s+BY\s+([^;]+?)(?:\s*;|\s*$)/i
const orderByMatch = orderByRegex.exec(sql)
let orderBy: OrderBy | undefined
if (orderByMatch) {
orderBy = parseOrderByClause(orderByMatch[1].trim())
// Trim it from the original SQL
sql = sql.slice(0, orderByMatch.index).trim()
}
// Extract WHERE clause if present
const whereRegex = /\bWHERE\s+([\s\S]+)/i
const whereMatch = whereRegex.exec(sql)
let whereClause = ''
if (whereMatch) {
// Trim trailing semicolons, etc.
whereClause = whereMatch[1].trim().replace(/;$/, '')
}
let filter: ParquetQueryFilter = {}
if (whereClause.trim()) {
const tokens = tokenizeWhereClause(whereClause)
// Convert certain tokens to uppercase if they are known keywords/operators
const keywords = ['AND', 'OR', '(', ')', '!=', '=', '>=', '<=', '>', '<']
for (let j = 0; j < tokens.length; j++) {
const t = tokens[j]
if (
!(t.startsWith('\'') && t.endsWith('\'')) &&
!(t.startsWith('"') && t.endsWith('"')) &&
keywords.includes(t.toUpperCase())
) tokens[j] = t.toUpperCase()
}
filter = parseExpression(tokens).node
}
if (columns.length && columns[0] === '*') columns.length = 0
return { collection, filter, columns, orderBy, limit, offset }
}
// Recursive parser for the WHERE clause
function parseExpression(tokens: string[], startIdx = 0): ParseExpressionResult {
const conditions: Record<string, any>[] = []
let currentOp: 'AND' | 'OR' | undefined
let i = startIdx
function pushCondition(cond: Record<string, any>) {
if (currentOp === 'OR') {
// If last element is an $or array, push into it, else create
if (conditions.length === 1 && conditions[0].$or) {
conditions[0].$or.push(cond)
} else {
// Combine existing conditions into an $and if necessary
if (conditions.length && !conditions[0].$or) {
cond = { $or: [conditions.pop(), cond] }
} else {
cond = { $or: [cond] }
}
conditions.push(cond)
}
} else {
// currentOp === 'AND' or undefined
if (conditions.length === 1 && conditions[0].$and) {
conditions[0].$and.push(cond)
} else {
if (conditions.length && !conditions[0].$and) {
cond = { $and: [conditions.pop(), cond] }
} else if (conditions.length && conditions[0].$or) {
cond = { $and: [conditions.pop(), cond] }
}
conditions.push(cond)
}
}
}
while (i < tokens.length) {
const token = tokens[i]
if (token === '(') {
// Parse subexpression
const sub = parseExpression(tokens, i + 1)
pushCondition(sub.node)
i = sub.endIdx
} else if (token === ')') {
break
} else if (token === 'AND' || token === 'OR') {
currentOp = token as 'AND' | 'OR'
i++
} else {
// Expect a pattern: FIELD OP VALUE
const field = stripQuotes(token)
i++
if (i >= tokens.length) throw new Error('Invalid WHERE clause: operator expected.')
const opToken = tokens[i]
i++
if (i >= tokens.length) throw new Error('Invalid WHERE clause: value expected.')
const valToken = tokens[i]
if (['=', '!=', '>', '<', '>=', '<='].includes(opToken)) {
const op = mongoOperator(opToken)
if (!op) {
throw new Error('Invalid operator in WHERE clause: ' + opToken)
}
const val = parseValue(valToken)
pushCondition({ [field]: { [op]: val } })
i++
} else {
throw new Error('Invalid operator in WHERE clause: ' + opToken)
}
}
}
if (conditions.length > 1) {
return { node: { $and: conditions }, endIdx: i + 1 }
} else if (conditions.length === 1) {
return { node: conditions[0], endIdx: i + 1 }
} else {
return { node: {}, endIdx: i + 1 }
}
}
// Convert SQL operator to Mongo operator
function mongoOperator(op: string): '$eq' | '$ne' | '$gt' | '$lt' | '$gte' | '$lte' | undefined {
switch (op) {
case '=': return '$eq'
case '!=': return '$ne'
case '>': return '$gt'
case '<': return '$lt'
case '>=': return '$gte'
case '<=': return '$lte'
default: throw new Error('Invalid operator in WHERE clause: ' + op)
}
}
// Parse a value from a token
function parseValue(value: string): string | number | boolean {
// Check if it's a quoted string
if (
value.startsWith('\'') && value.endsWith('\'') ||
value.startsWith('"') && value.endsWith('"')
) {
return value.slice(1, -1)
}
// Check boolean
const lower = value.toLowerCase()
if (lower === 'true') return true
if (lower === 'false') return false
// Check number
const num = Number(value)
if (!isNaN(num)) {
return num
}
// Otherwise return as-is (string field name or other identifier)
return value
}
/**
* Parse the raw SELECT list into individual field names,
* rejecting anything that is *not* a single identifier
* (quoted or unquoted).
*
* Clause is text between SELECT and FROM.
*/
function parseSelectFields(clause: string): string[] {
if (clause === '*') return []
const fields: string[] = []
let buf = ''
let inQuote = false
let quoteChar = ''
/** push current buffer and reset */
function push() {
const token = buf.trim()
buf = ''
if (!token) throw new Error('Invalid SQL: empty field in SELECT list.')
// quoted identifier → always fine
const isQuoted = /^['"].*['"]$/.test(token)
// simple unquoted identifier
const isIdentifier = /^[A-Za-z_][A-Za-z0-9_]*$/.test(token)
if (!isQuoted && !isIdentifier) {
// anything with spaces (e.g. “name AS …”) falls in here
throw new Error(`Invalid SQL: unexpected token in SELECT list: ${token}`)
}
fields.push(stripQuotes(token))
}
for (const ch of clause) {
if (inQuote) {
buf += ch
if (ch === quoteChar) inQuote = false
continue
}
if (ch === '\'' || ch === '"') {
inQuote = true
quoteChar = ch
buf += ch
} else if (ch === ',') {
push()
} else {
buf += ch
}
}
push()
return fields
}
// Tokenize the WHERE clause respecting quotes
function tokenizeWhereClause(str: string): string[] {
const tokens: string[] = []
let current = ''
let inQuote = false
let quoteChar = ''
for (const c of str) {
if (inQuote) {
if (c === quoteChar) {
inQuote = false
tokens.push((current + c).trim())
current = ''
} else {
current += c
}
} else {
if (c === '\'' || c === '"') {
if (current.trim()) {
tokens.push(current.trim())
}
current = c
inQuote = true
quoteChar = c
} else if (/\s/.test(c)) {
if (current.trim()) {
tokens.push(current.trim())
current = ''
}
} else if (c === '(' || c === ')') {
if (current.trim()) {
tokens.push(current.trim())
}
tokens.push(c)
current = ''
} else {
current += c
}
}
}
if (current.trim()) {
tokens.push(current.trim())
}
return tokens
}
function parseOrderByClause(clause: string): OrderBy {
return clause.split(',').map(part => {
const bits = part.trim().split(/\s+/)
if (!bits[0]) throw new Error('Invalid SQL: empty column in ORDER BY.')
const column = stripQuotes(bits[0])
let dir: Direction = 'ascending'
if (bits[1]) {
const d = bits[1].toUpperCase()
if (d === 'DESC' || d === 'DESCENDING') dir = 'descending'
else if (d === 'ASC' || d === 'ASCENDING') dir = 'ascending'
else throw new Error(`Invalid ORDER BY direction: ${bits[1]}`)
}
return { column, direction: dir }
})
}
/**
* Remove matching leading & trailing quotes.
*/
function stripQuotes(str: string): string {
if (
str.startsWith('"') && str.endsWith('"') ||
str.startsWith('\'') && str.endsWith('\'')
) {
return str.slice(1, -1)
}
return str
} |
|
I did Haven't debugged it yet, but here's the test that failed: parquetQuery({ file, compressors, orderBy: 'l_extendedprice', rowEnd: 10 })Against tpch dataset. |
|
Thank you @platypii , I am writing some skippable local tests for large files - orderBy is kinda tricky to optimize, since we need to iterate through all possible page matches/indexes to find the actual min/max values (then we can apply offsets and limits). we can only begin assembling the results once we have ordered the entire dataset. First, I tried fetching all indexes and pages in parallel and then assembling the results. Unfortunately, this consumed way too much memory (which is why we are now reading each row group in a loop, then applying pushdown to fetch individual pages). However, this does come with a performance penalty. It's way faster to make all your batch requests for all pages at once. There is probably a good middle ground |
|
Yea the case with an
These each benefit certain types of queries over others. Ideally I would love if we had a full query planner... maybe we have a certain number of pre-defined query plans, generate plans for each, and automatically pick the one that will have the least network overhead? It's a hard problem, and I don't know how to reconcile it with query plans that are dynamic (plan 1 above: fetch rows and filter until row limit hit). There would be upper and lower bounds on the query plan? Starts to get messy. Anyway for this PR I don't want to increase the scope. But I would prefer if queries that used to work still do. It sounds like the problem is not the size of the data, but recursion limit? Or maybe we need to bring back special case for |
|
Oh, I agree completely. I don't want to get rid of pre-existing working code. For some reason, this test seems to work - just added directly to the hyparquet suite (BENCHMARK=1 && npm test). It's super-slow though, which is not entirely unexpected. Give me a couple of days to hack on this. import { createWriteStream, existsSync } from 'fs'
import { pipeline } from 'stream/promises'
import { describe, expect, it } from 'vitest'
import { parquetQuery } from '../src/query.js'
import { asyncBufferFromFile } from '../src/node.js'
describe.skipIf(!process.env.BENCHMARK) {
it('reads large files with an orderBy and rowEnd clause', async () => {
const url = 'https://s3.hyperparam.app/tpch-lineitem-v2.parquet'
const filename = 'test/tpch-lineitem-v2.parquet'
existsSync(filename) || await pipeline((await fetch(url)).body, createWriteStream(filename))
const file = await asyncBufferFromFile(filename)
const rows = await parquetQuery({ file, orderBy: 'l_extendedprice', rowEnd: 10 })
expect(rows).toHaveLength(10)
}, 30_000)
} |
|
Okay this works for now. The issue was a spread operator in the code (functions can't take more than a couple thousand args). I added page indexes to the tpch dataset (exported it to a csv first, and then used pyarrow). I'm going to spend some time to optimize the page index queries. The other code paths will revert to what's already in master for query.js. import pyarrow.csv as csv
import pyarrow.parquet as pq
import sys
pq.write_table(csv.read_csv(sys.argv[1]), sys.argv[2],
row_group_size=32*1024*1024,
write_page_index=True) |
2 participants
This is just a preliminary pull request that implements most of what we talked about last(!) year. Predicate pushdown has been implemented in plan.js and query.js. Here's a summary of the additions:
In plan.js, we have more utilities for query planning, which we use to construct predicates that are used in query.js. We use row group and page-level statistics here to skip chunks that don't match the predicates. I added some utilities in column.js for reading pages.