A pushed-down WHERE no longer widens the parquet projection

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,

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}

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')

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<string, SqlPrimitive>[]} rows
+ * @returns {Promise<AsyncDataSource>}
+ */
+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 = []