fix: close 5 unbounded memory violations in operator pipeline

- AggregateOperator: incrementally merge partials instead of accumulating
  all in an array — O(groups) memory instead of O(batches × groups)
- mergeQueryResults: use k-way merge for sorted results without limit,
  limit-aware concat with early termination for bounded results
- LocalExecutor: always pass FsSpillBackend + memory budget to pipeline
  options and HashJoinOperator so spill is available for joins
- ProjectOperator: create new projected row objects instead of deleting
  keys (avoids V8 hidden class deoptimization)
- HashJoinOperator: sample first batch to estimate row size and go
  straight to partitioned spill if right side likely exceeds budget

Author: teamchong

src/local-executor.ts

@@ -16,7 +16,7 @@ import { instantiateWasm, type WasmEngine } from "./wasm-engine.js";
 import { VipCache } from "./vip-cache.js";
 import { QueryModeError } from "./errors.js";
 import { parseLanceV2Columns, lanceV2ToColumnMeta, computeLanceV2Stats } from "./lance-v2.js";
-import { buildPipeline, drainPipeline, type FragmentSource, type PipelineOptions } from "./operators.js";
+import { buildPipeline, drainPipeline, DEFAULT_MEMORY_BUDGET, type FragmentSource, type PipelineOptions } from "./operators.js";
 
 /**
  * Executor for local mode (Node/Bun).
@@ -426,8 +426,11 @@ export class LocalExecutor implements QueryExecutor {
     // Step 3: Build streaming pipeline
     const wasm = await this.getWasm();
     const fragment = await this.makeFragmentSource(query.table, projectedColumns, isUrl);
-    const pipeOpts: PipelineOptions | undefined = this.memoryBudgetBytes
-      ? { memoryBudgetBytes: this.memoryBudgetBytes } : undefined;
+    const { FsSpillBackend } = await import("./operators.js");
+    const pipeOpts: PipelineOptions = {
+      memoryBudgetBytes: this.memoryBudgetBytes ?? DEFAULT_MEMORY_BUDGET,
+      spill: new FsSpillBackend(),
+    };
 
     // If join is specified, build left + right pipelines and combine with HashJoinOperator
     if (query.join) {
@@ -581,13 +584,16 @@ export class LocalExecutor implements QueryExecutor {
     }
 
     // Hash join: right is build side, left is probe side
-    const { HashJoinOperator } = await import("./operators.js");
+    const { HashJoinOperator, FsSpillBackend: FsSpillJoin } = await import("./operators.js");
+    const joinSpill = pipeOpts?.spill ?? new FsSpillJoin();
+    const joinBudget = pipeOpts?.memoryBudgetBytes ?? DEFAULT_MEMORY_BUDGET;
     let pipeline: import("./operators.js").Operator = new HashJoinOperator(
       leftPipeline, rightPipeline, join.leftKey, join.rightKey, join.type ?? "inner",
+      joinBudget, joinSpill,
     );
 
     // Apply post-join sort/limit/aggregates
-    const { buildPipeline: _, TopKOperator, InMemorySortOperator, LimitOperator, ProjectOperator, AggregateOperator, ExternalSortOperator, DEFAULT_MEMORY_BUDGET } = await import("./operators.js");
+    const { buildPipeline: _, TopKOperator, InMemorySortOperator, LimitOperator, ProjectOperator, AggregateOperator, ExternalSortOperator } = await import("./operators.js");
     const hasAgg = query.aggregates && query.aggregates.length > 0;
 
     if (hasAgg) {
@@ -780,8 +786,11 @@ export class LocalExecutor implements QueryExecutor {
       ? query.projections
       : (dataset.fragmentMetas.values().next().value?.columns.map((c: ColumnMeta) => c.name) ?? []);
 
-    const pipeOpts: PipelineOptions | undefined = this.memoryBudgetBytes
-      ? { memoryBudgetBytes: this.memoryBudgetBytes } : undefined;
+    const { FsSpillBackend: FsSpill } = await import("./operators.js");
+    const pipeOpts: PipelineOptions = {
+      memoryBudgetBytes: this.memoryBudgetBytes ?? DEFAULT_MEMORY_BUDGET,
+      spill: new FsSpill(),
+    };
     const { pipeline, scan, wasmAgg } = buildPipeline(fragments, query, wasm, outputColumns, pipeOpts);
     const rows = await drainPipeline(pipeline);
 

src/merge.ts

@@ -140,13 +140,13 @@ export function mergeQueryResults(
     };
   }
 
-  // Sort + limit path: k-way merge
-  if (query.sortColumn && query.limit) {
+  // Sort path: k-way merge (works with or without limit)
+  if (query.sortColumn) {
     const rows = kWayMerge(
       partials.map((p) => p.rows),
       query.sortColumn,
       query.sortDirection ?? "asc",
-      query.limit,
+      query.limit ?? Infinity,
     );
     return {
       rows,
@@ -158,10 +158,19 @@ export function mergeQueryResults(
     };
   }
 
-  // Default: concat and apply limit
-  let rows = partials.flatMap((p) => p.rows);
-  if (query.limit && rows.length > query.limit) {
-    rows = rows.slice(0, query.limit);
+  // Unsorted: limit-aware concat with early termination
+  let rows: Row[];
+  if (query.limit) {
+    rows = [];
+    for (const p of partials) {
+      for (const row of p.rows) {
+        rows.push(row);
+        if (rows.length >= query.limit) break;
+      }
+      if (rows.length >= query.limit) break;
+    }
+  } else {
+    rows = partials.flatMap((p) => p.rows);
   }
 
   return {

src/operators.ts

@@ -267,12 +267,12 @@ export class ProjectOperator implements Operator {
     const batch = await this.upstream.next();
     if (!batch) return null;
 
-    for (const row of batch) {
-      for (const key of Object.keys(row)) {
-        if (!this.keep.has(key)) delete row[key];
-      }
-    }
-    return batch;
+    // Create new projected row objects — no delete, no V8 hidden class deopt
+    return batch.map(row => {
+      const out: Row = {};
+      for (const k of this.keep) if (k in row) out[k] = row[k];
+      return out;
+    });
   }
 
   async close(): Promise<void> {
@@ -298,19 +298,18 @@ export class AggregateOperator implements Operator {
     if (this.consumed) return null;
     this.consumed = true;
 
-    // Stream all batches through partial aggregation
-    const partials: PartialAgg[] = [];
+    // Incrementally merge partials — O(groups) memory, not O(batches × groups)
+    let merged: PartialAgg | null = null;
     while (true) {
       const batch = await this.upstream.next();
       if (!batch) break;
-      partials.push(computePartialAgg(batch, this.query));
+      const partial = computePartialAgg(batch, this.query);
+      merged = merged ? mergePartialAggs([merged, partial]) : partial;
     }
 
-    if (partials.length === 0) {
+    if (!merged) {
       return finalizePartialAgg({ states: [] }, this.query);
     }
-
-    const merged = partials.length === 1 ? partials[0] : mergePartialAggs(partials);
     return finalizePartialAgg(merged, this.query);
   }
 
@@ -673,89 +672,144 @@ export class HashJoinOperator implements Operator {
   private async buildOrPartition(): Promise<void> {
     if (this.hashMap || this.partitionCount > 0) return;
 
-    const inMemoryRows: Row[] = [];
-    let exceeds = false;
+    // Sample first batch to estimate whether right side fits in memory.
+    // If estimated total exceeds budget, go straight to partitioned path
+    // instead of accumulating rows we'll have to re-spill.
+    const firstBatch = await this.right.next();
+    if (!firstBatch) {
+      // Empty right side — build empty hash map
+      this.hashMap = new Map<string, Row[]>();
+      return;
+    }
 
-    // Consume right side, tracking memory
-    while (true) {
-      const batch = await this.right.next();
-      if (!batch) break;
-      for (const row of batch) {
-        const rowSize = estimateRowSize(row);
-        this.buildSizeBytes += rowSize;
-        inMemoryRows.push(row);
+    let batchSizeBytes = 0;
+    for (const row of firstBatch) batchSizeBytes += estimateRowSize(row);
+    const avgRowSize = batchSizeBytes / firstBatch.length;
 
-        if (this.spill && this.buildSizeBytes > this.memoryBudget) {
-          exceeds = true;
-          break;
+    // Heuristic: if first batch already exceeds half the budget, go partitioned
+    const goPartitioned = this.spill && batchSizeBytes > this.memoryBudget / 2;
+
+    if (!goPartitioned) {
+      // Optimistic in-memory path: consume right side, tracking memory
+      const inMemoryRows: Row[] = [...firstBatch];
+      this.buildSizeBytes = batchSizeBytes;
+      let exceeds = false;
+
+      while (true) {
+        const batch = await this.right.next();
+        if (!batch) break;
+        for (const row of batch) {
+          const rowSize = estimateRowSize(row);
+          this.buildSizeBytes += rowSize;
+          inMemoryRows.push(row);
+
+          if (this.spill && this.buildSizeBytes > this.memoryBudget) {
+            exceeds = true;
+            break;
+          }
         }
+        if (exceeds) break;
       }
-      if (exceeds) break;
-    }
 
-    if (!exceeds) {
-      // Fits in memory — build hash map directly
-      this.hashMap = new Map<string, Row[]>();
-      for (const row of inMemoryRows) {
-        const key = this.toJoinKey(row[this.rightKey]);
-        const bucket = this.hashMap.get(key);
-        if (bucket) bucket.push(row);
-        else this.hashMap.set(key, [row]);
+      if (!exceeds) {
+        // Fits in memory — build hash map directly
+        this.hashMap = new Map<string, Row[]>();
+        for (const row of inMemoryRows) {
+          const key = this.toJoinKey(row[this.rightKey]);
+          const bucket = this.hashMap.get(key);
+          if (bucket) bucket.push(row);
+          else this.hashMap.set(key, [row]);
+        }
+        return;
       }
+
+      // Fell through: exceeded budget mid-stream, partition what we have
+      this.buildExceeded = true;
+      this.partitionCount = Math.max(4, Math.ceil(this.buildSizeBytes / (this.memoryBudget / 2)));
+      await this.partitionRightRows(inMemoryRows);
+      inMemoryRows.length = 0;
+
+      // Continue consuming remaining right-side rows
+      await this.consumeRemainingRight();
       return;
     }
 
-    // Switch to Grace hash join: partition both sides with bounded memory.
-    // Use partitionCount that ensures each partition fits in memory budget.
+    // Proactive partition path: first batch signals right side is large
     this.buildExceeded = true;
-    this.partitionCount = Math.max(4, Math.ceil(this.buildSizeBytes / (this.memoryBudget / 2)));
-    const bucketBudget = Math.floor(this.memoryBudget / this.partitionCount);
+    this.buildSizeBytes = batchSizeBytes;
+    // Estimate partition count from first batch size × expected batch count
+    // Use conservative estimate: assume at least 4× more data coming
+    const estimatedTotal = batchSizeBytes * 4;
+    this.partitionCount = Math.max(4, Math.ceil(estimatedTotal / (this.memoryBudget / 2)));
+
+    await this.partitionRightRows(firstBatch);
+    await this.consumeRemainingRight();
+  }
 
-    // Spill right-side rows we already consumed, flushing per-bucket when full
+  /** Partition an array of right-side rows into spill buckets. */
+  private async partitionRightRows(rows: Row[]): Promise<void> {
+    const bucketBudget = Math.floor(this.memoryBudget / this.partitionCount);
     const rightBuckets: Row[][] = Array.from({ length: this.partitionCount }, () => []);
     const rightBucketBytes: number[] = new Array(this.partitionCount).fill(0);
-    this.rightPartitionIds = new Array(this.partitionCount).fill("");
+    if (!this.rightPartitionIds.length) {
+      this.rightPartitionIds = new Array(this.partitionCount).fill("");
+    }
 
-    const flushRightBucket = async (bi: number): Promise<void> => {
+    const flushBucket = async (bi: number): Promise<void> => {
       if (rightBuckets[bi].length === 0) return;
       const spillId = await this.spill!.writeRun(rightBuckets[bi]);
-      // Append to existing partition by tracking multiple spill IDs per partition
       this.rightPartitionIds[bi] = this.rightPartitionIds[bi]
         ? this.rightPartitionIds[bi] + "|" + spillId
         : spillId;
       rightBuckets[bi] = [];
       rightBucketBytes[bi] = 0;
     };
 
-    for (const row of inMemoryRows) {
+    for (const row of rows) {
       const bi = this.hashPartition(row[this.rightKey], this.partitionCount);
       const rowSize = estimateRowSize(row);
       if (rightBucketBytes[bi] + rowSize > bucketBudget && rightBuckets[bi].length > 0) {
-        await flushRightBucket(bi);
+        await flushBucket(bi);
       }
       rightBuckets[bi].push(row);
       rightBucketBytes[bi] += rowSize;
     }
-    // Free inMemoryRows — no longer needed
-    inMemoryRows.length = 0;
 
-    // Continue consuming remaining right-side rows
+    for (let bi = 0; bi < this.partitionCount; bi++) await flushBucket(bi);
+  }
+
+  /** Consume remaining right-side batches into partitions. */
+  private async consumeRemainingRight(): Promise<void> {
+    const bucketBudget = Math.floor(this.memoryBudget / this.partitionCount);
+    const rightBuckets: Row[][] = Array.from({ length: this.partitionCount }, () => []);
+    const rightBucketBytes: number[] = new Array(this.partitionCount).fill(0);
+
+    const flushBucket = async (bi: number): Promise<void> => {
+      if (rightBuckets[bi].length === 0) return;
+      const spillId = await this.spill!.writeRun(rightBuckets[bi]);
+      this.rightPartitionIds[bi] = this.rightPartitionIds[bi]
+        ? this.rightPartitionIds[bi] + "|" + spillId
+        : spillId;
+      rightBuckets[bi] = [];
+      rightBucketBytes[bi] = 0;
+    };
+
     while (true) {
       const batch = await this.right.next();
       if (!batch) break;
       for (const row of batch) {
         const bi = this.hashPartition(row[this.rightKey], this.partitionCount);
         const rowSize = estimateRowSize(row);
+        this.buildSizeBytes += rowSize;
         if (rightBucketBytes[bi] + rowSize > bucketBudget && rightBuckets[bi].length > 0) {
-          await flushRightBucket(bi);
+          await flushBucket(bi);
         }
         rightBuckets[bi].push(row);
         rightBucketBytes[bi] += rowSize;
       }
     }
 
-    // Flush remaining right buckets
-    for (let bi = 0; bi < this.partitionCount; bi++) await flushRightBucket(bi);
+    for (let bi = 0; bi < this.partitionCount; bi++) await flushBucket(bi);
 
     // Consume and partition left side with same bounded approach
     const leftBuckets: Row[][] = Array.from({ length: this.partitionCount }, () => []);