diff --git a/src/core/query/parquet-source.js b/src/core/query/parquet-source.js index b620d54e..2e0a877d 100644 --- a/src/core/query/parquet-source.js +++ b/src/core/query/parquet-source.js @@ -53,11 +53,18 @@ export function parquetDataSource(file, metadata) { // With a WHERE present the engine owns LIMIT/OFFSET so it applies // them to the *filtered* result rather than to raw row positions. const appliedLimitOffset = !hints.where - // When a filter is pushed down it may reference columns outside the - // engine's projection; read all columns so hyparquet can evaluate - // it, and let the engine project. Without a filter, honor the - // requested projection. - const readColumns = filter ? undefined : hints.columns + // Honor the requested projection whether or not a filter is pushed + // down. A filter may reference columns outside the projection, but + // hyparquet unions its own `columnsNeededForFilter` into the read + // plan and deletes the extras from the returned rows, so passing the + // narrow projection is safe. Reading all columns instead would decode + // every payload column on any filtered scan. Because the emitted row + // set is now exactly `columns`, a caller that wraps this scan and + // reports its own `appliedWhere: false` (letting the engine re-apply + // the predicate over rows this scan already returned, as `unionSources` + // does) must include the predicate's columns in `columns`, or the + // engine has nothing to re-filter on. + const readColumns = hints.columns return { appliedWhere, diff --git a/src/core/query/union-source.js b/src/core/query/union-source.js index 1d3eb01e..2db61f0c 100644 --- a/src/core/query/union-source.js +++ b/src/core/query/union-source.js @@ -26,7 +26,13 @@ import { normalizeScanColumn } from './scan-column.js' * reading the column as null. When a partition can't satisfy the predicate we * drop `where` for it and let the engine filter the concatenated stream (it * already owns the filter via `appliedWhere: false`). `columns` is always - * forwarded: projecting an absent column reads as null, never throws. + * forwarded: projecting an absent column reads as null, never throws. Because + * a sub-source now emits exactly the columns it is asked for (see + * `parquet-source.js`), forwarding `columns` also determines what the engine + * gets to re-filter on: it relies on squirreling folding the WHERE columns + * into the projection it hands to `scan()`, so the predicate's columns are + * already present even though `appliedWhere: false` never asks for them + * explicitly. * * @param {AsyncDataSource[]} sources * @returns {AsyncDataSource} diff --git a/test/core/parquet-source.test.js b/test/core/parquet-source.test.js index f112a9c6..5c8cbfe3 100644 --- a/test/core/parquet-source.test.js +++ b/test/core/parquet-source.test.js @@ -153,6 +153,60 @@ test('WHERE on a non-projected column still filters correctly', async () => { assert.deepEqual(rows.map((r) => r.name), ['carol', 'dave', 'eve']) }) +// A filtered scan must still honor the projection. The tests above only prove +// the rows are right, which stays true when the scan reads every column and +// the engine projects afterwards; on a table whose unselected columns hold +// message bodies, that difference is the whole read. Assert the narrow read +// directly, both in what the scan emits and in what it pulls off the file. +test('a pushed-down filter does not widen the projection', async () => { + const columnData = rowsToColumnSources(COLUMNS, ROWS) + const arrayBuffer = parquetWriteBuffer({ columnData, codec: 'SNAPPY', rowGroupSize: 2 }) + const bytes = new Uint8Array(arrayBuffer) + + /** + * Scan `columns` under a filter on `score`, reporting the bytes pulled from + * the file. `score` is deliberately absent from every projection so the read + * can only cover it because hyparquet folds filter columns into its own plan. + * + * @param {string[] | undefined} columns + * @returns {Promise<{ read: number, emitted: string[][] }>} + */ + async function scanWithFilter(columns) { + let read = 0 + const counting = asyncBufferFromBytes(bytes) + const file = { + byteLength: counting.byteLength, + /** + * @param {number} start + * @param {number} [end] + */ + slice(start, end) { + read += (end ?? bytes.byteLength) - start + return counting.slice(start, end) + }, + } + const source = parquetDataSource(file, await parquetMetadataAsync(file)) + const scan = source.scan({ columns, where: whereOf('SELECT name FROM t WHERE score > 3') }) + assert.equal(scan.appliedWhere, true) + /** @type {string[][]} */ + const emitted = [] + for await (const row of scan.rows()) emitted.push(row.columns) + return { read, emitted } + } + + const projected = await scanWithFilter(['name']) + const everything = await scanWithFilter(undefined) + + // The filter is honored on a column the scan never emits. + assert.equal(projected.emitted.length, 3) + for (const columns of projected.emitted) assert.deepEqual(columns, ['name']) + // ...and the unselected columns were never read off the file. + assert.ok( + projected.read < everything.read, + `projected scan read ${projected.read} bytes, unprojected read ${everything.read}` + ) +}) + test('range WHERE (AND) returns the inclusive window', async () => { const source = await makeSource() const rows = await run(source, 'SELECT id FROM t WHERE id >= 2 AND id <= 4') diff --git a/test/core/union-source.test.js b/test/core/union-source.test.js index 3a47d9cb..192e355d 100644 --- a/test/core/union-source.test.js +++ b/test/core/union-source.test.js @@ -3,12 +3,18 @@ import assert from 'node:assert/strict' import test from 'node:test' -import { asyncRow, parseSql } from 'squirreling' +import { parquetMetadataAsync } from 'hyparquet' +import { parquetWriteBuffer } from 'hyparquet-writer' +import { asyncRow, collect, executeSql, parseSql } from 'squirreling' import { unionSources, emptySource } from '../../src/core/query/union-source.js' import { normalizeScanColumn } from '../../src/core/query/scan-column.js' +import { parquetDataSource } from '../../src/core/query/parquet-source.js' +import { rowsToColumnSources } from '../../hypaware-core/plugins-workspace/format-parquet/src/columns.js' /** + * @import { AsyncBuffer } from 'hyparquet' * @import { AsyncDataSource, ExprNode, IdentifierNode, ScanOptions, SqlPrimitive } from 'squirreling/src/types.js' + * @import { ColumnSpec } from '../../hypaware-plugin-kernel-types.js' */ /** @@ -120,6 +126,74 @@ test('unionSources forwards where/columns to sub-sources that have the predicate } }) +/** @type {ColumnSpec[]} */ +const PARQUET_PARTITION_COLUMNS = [ + { name: 'id', type: 'INT64', nullable: false }, + { name: 'name', type: 'STRING', nullable: false }, + { name: 'score', type: 'DOUBLE', nullable: false }, +] + +/** + * @param {Uint8Array} bytes + * @returns {AsyncBuffer} + */ +function asyncBufferFromBytes(bytes) { + return { + byteLength: bytes.byteLength, + slice(start, end) { + const sliced = bytes.subarray(start, end) + const out = new ArrayBuffer(sliced.byteLength) + new Uint8Array(out).set(sliced) + return out + }, + } +} + +/** + * Build a real, on-disk-shaped parquet `AsyncDataSource` partition (same + * construction as `test/core/parquet-source.test.js`), so the union test + * below exercises actual hyparquet reads and pushdown, not a fake source. + * + * @param {Record[]} rows + * @returns {Promise} + */ +async function makeParquetPartition(rows) { + const columnData = rowsToColumnSources(PARQUET_PARTITION_COLUMNS, rows) + const arrayBuffer = parquetWriteBuffer({ columnData, codec: 'SNAPPY', rowGroupSize: 2 }) + const file = asyncBufferFromBytes(new Uint8Array(arrayBuffer)) + const metadata = await parquetMetadataAsync(file) + return parquetDataSource(file, metadata) +} + +// @ref LLP 0015#multi-partition-union [tests]: appliedWhere: false only stays correct end-to-end because +// squirreling folds WHERE columns into the projection it hands to scan(); pin that at the layer +// that depends on it, with two real parquet partitions, not a fake source. +test('unionSources over two real parquet partitions filters correctly through executeSql (WHERE column folded into projection)', async () => { + const partitionA = await makeParquetPartition([ + { id: 1, name: 'alice', score: 1.5 }, + { id: 2, name: 'bob', score: 2.5 }, + { id: 3, name: 'carol', score: 3.5 }, + ]) + const partitionB = await makeParquetPartition([ + { id: 4, name: 'dave', score: 4.5 }, + { id: 5, name: 'eve', score: 5.5 }, + ]) + const union = unionSources([partitionA, partitionB]) + + // union.scan() always reports appliedWhere: false, handing the predicate + // back to the engine to re-apply over the merged stream. That only returns + // the right rows here because squirreling's planner folds `score` (the + // WHERE column) into the projection it passes to scan(), even though the + // query only selects `name`; each parquet partition then emits `score` + // alongside `name`, and the engine can actually filter on it. If squirreling + // ever stopped folding, this would start throwing `ColumnNotFoundError` at + // query time, when the engine re-filters on a column the rows no longer + // carry, while every other test in this file (and in parquet-source.test.js) + // stayed green, since they exercise a single source, not the union path. + const rows = await collect(executeSql({ tables: { t: union }, query: 'SELECT name FROM t WHERE score > 3' })) + assert.deepEqual(rows, [{ name: 'carol' }, { name: 'dave' }, { name: 'eve' }]) +}) + test('unionSources drops where for a partition that lacks a predicate column but keeps it for one that has it', async () => { /** @type {ScanOptions[]} */ const seen = []