query optimiser that pushes where clauses down to subqueries closer to the source (#256)

Author: samwillis

.changeset/add-query-optimiser.md

@@ -0,0 +1,7 @@
+---
+"@tanstack/db": patch
+---
+
+Add query optimizer with predicate pushdown
+
+Implements automatic query optimization that moves WHERE clauses closer to data sources, reducing intermediate result sizes and improving performance for queries with joins.

packages/db/src/query/compiler/index.ts

@@ -1,4 +1,5 @@
 import { distinct, filter, map } from "@electric-sql/d2mini"
+import { optimizeQuery } from "../optimizer.js"
 import { compileExpression } from "./evaluators.js"
 import { processJoins } from "./joins.js"
 import { processGroupBy } from "./group-by.js"
@@ -10,30 +11,35 @@ import type {
   NamespacedAndKeyedStream,
   ResultStream,
 } from "../../types.js"
-
-/**
- * Cache for compiled subqueries to avoid duplicate compilation
- */
-type QueryCache = WeakMap<QueryIR, ResultStream>
+import type { QueryCache, QueryMapping } from "./types.js"
 
 /**
  * Compiles a query2 IR into a D2 pipeline
- * @param query The query IR to compile
+ * @param rawQuery The query IR to compile
  * @param inputs Mapping of collection names to input streams
  * @param cache Optional cache for compiled subqueries (used internally for recursion)
+ * @param queryMapping Optional mapping from optimized queries to original queries
  * @returns A stream builder representing the compiled query
  */
 export function compileQuery(
-  query: QueryIR,
+  rawQuery: QueryIR,
   inputs: Record<string, KeyedStream>,
-  cache: QueryCache = new WeakMap()
+  cache: QueryCache = new WeakMap(),
+  queryMapping: QueryMapping = new WeakMap()
 ): ResultStream {
-  // Check if this query has already been compiled
-  const cachedResult = cache.get(query)
+  // Check if the original raw query has already been compiled
+  const cachedResult = cache.get(rawQuery)
   if (cachedResult) {
     return cachedResult
   }
 
+  // Optimize the query before compilation
+  const query = optimizeQuery(rawQuery)
+
+  // Create mapping from optimized query to original for caching
+  queryMapping.set(query, rawQuery)
+  mapNestedQueries(query, rawQuery, queryMapping)
+
   // Create a copy of the inputs map to avoid modifying the original
   const allInputs = { ...inputs }
 
@@ -44,7 +50,8 @@ export function compileQuery(
   const { alias: mainTableAlias, input: mainInput } = processFrom(
     query.from,
     allInputs,
-    cache
+    cache,
+    queryMapping
   )
   tables[mainTableAlias] = mainInput
 
@@ -68,7 +75,8 @@ export function compileQuery(
       tables,
       mainTableAlias,
       allInputs,
-      cache
+      cache,
+      queryMapping
     )
   }
 
@@ -218,8 +226,8 @@ export function compileQuery(
     )
 
     const result = resultPipeline
-    // Cache the result before returning
-    cache.set(query, result)
+    // Cache the result before returning (use original query as key)
+    cache.set(rawQuery, result)
     return result
   } else if (query.limit !== undefined || query.offset !== undefined) {
     // If there's a limit or offset without orderBy, throw an error
@@ -241,8 +249,8 @@ export function compileQuery(
   )
 
   const result = resultPipeline
-  // Cache the result before returning
-  cache.set(query, result)
+  // Cache the result before returning (use original query as key)
+  cache.set(rawQuery, result)
   return result
 }
 
@@ -252,7 +260,8 @@ export function compileQuery(
 function processFrom(
   from: CollectionRef | QueryRef,
   allInputs: Record<string, KeyedStream>,
-  cache: QueryCache
+  cache: QueryCache,
+  queryMapping: QueryMapping
 ): { alias: string; input: KeyedStream } {
   switch (from.type) {
     case `collectionRef`: {
@@ -265,8 +274,16 @@ function processFrom(
       return { alias: from.alias, input }
     }
     case `queryRef`: {
+      // Find the original query for caching purposes
+      const originalQuery = queryMapping.get(from.query) || from.query
+
       // Recursively compile the sub-query with cache
-      const subQueryInput = compileQuery(from.query, allInputs, cache)
+      const subQueryInput = compileQuery(
+        originalQuery,
+        allInputs,
+        cache,
+        queryMapping
+      )
 
       // Subqueries may return [key, [value, orderByIndex]] (with ORDER BY) or [key, value] (without ORDER BY)
       // We need to extract just the value for use in parent queries
@@ -283,3 +300,53 @@ function processFrom(
       throw new Error(`Unsupported FROM type: ${(from as any).type}`)
   }
 }
+
+/**
+ * Recursively maps optimized subqueries to their original queries for proper caching.
+ * This ensures that when we encounter the same QueryRef object in different contexts,
+ * we can find the original query to check the cache.
+ */
+function mapNestedQueries(
+  optimizedQuery: QueryIR,
+  originalQuery: QueryIR,
+  queryMapping: QueryMapping
+): void {
+  // Map the FROM clause if it's a QueryRef
+  if (
+    optimizedQuery.from.type === `queryRef` &&
+    originalQuery.from.type === `queryRef`
+  ) {
+    queryMapping.set(optimizedQuery.from.query, originalQuery.from.query)
+    // Recursively map nested queries
+    mapNestedQueries(
+      optimizedQuery.from.query,
+      originalQuery.from.query,
+      queryMapping
+    )
+  }
+
+  // Map JOIN clauses if they exist
+  if (optimizedQuery.join && originalQuery.join) {
+    for (
+      let i = 0;
+      i < optimizedQuery.join.length && i < originalQuery.join.length;
+      i++
+    ) {
+      const optimizedJoin = optimizedQuery.join[i]!
+      const originalJoin = originalQuery.join[i]!
+
+      if (
+        optimizedJoin.from.type === `queryRef` &&
+        originalJoin.from.type === `queryRef`
+      ) {
+        queryMapping.set(optimizedJoin.from.query, originalJoin.from.query)
+        // Recursively map nested queries in joins
+        mapNestedQueries(
+          optimizedJoin.from.query,
+          originalJoin.from.query,
+          queryMapping
+        )
+      }
+    }
+  }
+}

packages/db/src/query/compiler/joins.ts

@@ -7,18 +7,13 @@ import {
 import { compileExpression } from "./evaluators.js"
 import { compileQuery } from "./index.js"
 import type { IStreamBuilder, JoinType } from "@electric-sql/d2mini"
-import type { CollectionRef, JoinClause, QueryIR, QueryRef } from "../ir.js"
+import type { CollectionRef, JoinClause, QueryRef } from "../ir.js"
 import type {
   KeyedStream,
   NamespacedAndKeyedStream,
   NamespacedRow,
-  ResultStream,
 } from "../../types.js"
-
-/**
- * Cache for compiled subqueries to avoid duplicate compilation
- */
-type QueryCache = WeakMap<QueryIR, ResultStream>
+import type { QueryCache, QueryMapping } from "./types.js"
 
 /**
  * Processes all join clauses in a query
@@ -29,7 +24,8 @@ export function processJoins(
   tables: Record<string, KeyedStream>,
   mainTableAlias: string,
   allInputs: Record<string, KeyedStream>,
-  cache: QueryCache
+  cache: QueryCache,
+  queryMapping: QueryMapping
 ): NamespacedAndKeyedStream {
   let resultPipeline = pipeline
 
@@ -40,7 +36,8 @@ export function processJoins(
       tables,
       mainTableAlias,
       allInputs,
-      cache
+      cache,
+      queryMapping
     )
   }
 
@@ -56,13 +53,15 @@ function processJoin(
   tables: Record<string, KeyedStream>,
   mainTableAlias: string,
   allInputs: Record<string, KeyedStream>,
-  cache: QueryCache
+  cache: QueryCache,
+  queryMapping: QueryMapping
 ): NamespacedAndKeyedStream {
   // Get the joined table alias and input stream
   const { alias: joinedTableAlias, input: joinedInput } = processJoinSource(
     joinClause.from,
     allInputs,
-    cache
+    cache,
+    queryMapping
   )
 
   // Add the joined table to the tables map
@@ -128,7 +127,8 @@ function processJoin(
 function processJoinSource(
   from: CollectionRef | QueryRef,
   allInputs: Record<string, KeyedStream>,
-  cache: QueryCache
+  cache: QueryCache,
+  queryMapping: QueryMapping
 ): { alias: string; input: KeyedStream } {
   switch (from.type) {
     case `collectionRef`: {
@@ -141,8 +141,16 @@ function processJoinSource(
       return { alias: from.alias, input }
     }
     case `queryRef`: {
+      // Find the original query for caching purposes
+      const originalQuery = queryMapping.get(from.query) || from.query
+
       // Recursively compile the sub-query with cache
-      const subQueryInput = compileQuery(from.query, allInputs, cache)
+      const subQueryInput = compileQuery(
+        originalQuery,
+        allInputs,
+        cache,
+        queryMapping
+      )
 
       // Subqueries may return [key, [value, orderByIndex]] (with ORDER BY) or [key, value] (without ORDER BY)
       // We need to extract just the value for use in parent queries

packages/db/src/query/compiler/types.ts

@@ -0,0 +1,12 @@
+import type { QueryIR } from "../ir.js"
+import type { ResultStream } from "../../types.js"
+
+/**
+ * Cache for compiled subqueries to avoid duplicate compilation
+ */
+export type QueryCache = WeakMap<QueryIR, ResultStream>
+
+/**
+ * Mapping from optimized queries back to their original queries for caching
+ */
+export type QueryMapping = WeakMap<QueryIR, QueryIR>