@tanstack/db 0.0.14 → 0.0.16

package/src/query/compiler/group-by.ts

@@ -0,0 +1,428 @@
+import { filter, groupBy, groupByOperators, map } from "@electric-sql/d2mini"
+import { Func, PropRef } from "../ir.js"
+import { compileExpression } from "./evaluators.js"
+import type {
+  Aggregate,
+  BasicExpression,
+  GroupBy,
+  Having,
+  Select,
+} from "../ir.js"
+import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
+
+const { sum, count, avg, min, max } = groupByOperators
+
+/**
+ * Interface for caching the mapping between GROUP BY expressions and SELECT expressions
+ */
+interface GroupBySelectMapping {
+  selectToGroupByIndex: Map<string, number> // Maps SELECT alias to GROUP BY expression index
+  groupByExpressions: Array<any> // The GROUP BY expressions for reference
+}
+
+/**
+ * Validates that all non-aggregate expressions in SELECT are present in GROUP BY
+ * and creates a cached mapping for efficient lookup during processing
+ */
+function validateAndCreateMapping(
+  groupByClause: GroupBy,
+  selectClause?: Select
+): GroupBySelectMapping {
+  const selectToGroupByIndex = new Map<string, number>()
+  const groupByExpressions = [...groupByClause]
+
+  if (!selectClause) {
+    return { selectToGroupByIndex, groupByExpressions }
+  }
+
+  // Validate each SELECT expression
+  for (const [alias, expr] of Object.entries(selectClause)) {
+    if (expr.type === `agg`) {
+      // Aggregate expressions are allowed and don't need to be in GROUP BY
+      continue
+    }
+
+    // Non-aggregate expression must be in GROUP BY
+    const groupIndex = groupByExpressions.findIndex((groupExpr) =>
+      expressionsEqual(expr, groupExpr)
+    )
+
+    if (groupIndex === -1) {
+      throw new Error(
+        `Non-aggregate expression '${alias}' in SELECT must also appear in GROUP BY clause`
+      )
+    }
+
+    // Cache the mapping
+    selectToGroupByIndex.set(alias, groupIndex)
+  }
+
+  return { selectToGroupByIndex, groupByExpressions }
+}
+
+/**
+ * Processes the GROUP BY clause with optional HAVING and SELECT
+ * Works with the new __select_results structure from early SELECT processing
+ */
+export function processGroupBy(
+  pipeline: NamespacedAndKeyedStream,
+  groupByClause: GroupBy,
+  havingClauses?: Array<Having>,
+  selectClause?: Select,
+  fnHavingClauses?: Array<(row: any) => any>
+): NamespacedAndKeyedStream {
+  // Handle empty GROUP BY (single-group aggregation)
+  if (groupByClause.length === 0) {
+    // For single-group aggregation, create a single group with all data
+    const aggregates: Record<string, any> = {}
+
+    if (selectClause) {
+      // Scan the SELECT clause for aggregate functions
+      for (const [alias, expr] of Object.entries(selectClause)) {
+        if (expr.type === `agg`) {
+          const aggExpr = expr
+          aggregates[alias] = getAggregateFunction(aggExpr)
+        }
+      }
+    }
+
+    // Use a constant key for single group
+    const keyExtractor = () => ({ __singleGroup: true })
+
+    // Apply the groupBy operator with single group
+    pipeline = pipeline.pipe(
+      groupBy(keyExtractor, aggregates)
+    ) as NamespacedAndKeyedStream
+
+    // Update __select_results to include aggregate values
+    pipeline = pipeline.pipe(
+      map(([, aggregatedRow]) => {
+        // Start with the existing __select_results from early SELECT processing
+        const selectResults = (aggregatedRow as any).__select_results || {}
+        const finalResults: Record<string, any> = { ...selectResults }
+
+        if (selectClause) {
+          // Update with aggregate results
+          for (const [alias, expr] of Object.entries(selectClause)) {
+            if (expr.type === `agg`) {
+              finalResults[alias] = aggregatedRow[alias]
+            }
+            // Non-aggregates keep their original values from early SELECT processing
+          }
+        }
+
+        // Use a single key for the result and update __select_results
+        return [
+          `single_group`,
+          {
+            ...aggregatedRow,
+            __select_results: finalResults,
+          },
+        ] as [unknown, Record<string, any>]
+      })
+    )
+
+    // Apply HAVING clauses if present
+    if (havingClauses && havingClauses.length > 0) {
+      for (const havingClause of havingClauses) {
+        const transformedHavingClause = transformHavingClause(
+          havingClause,
+          selectClause || {}
+        )
+        const compiledHaving = compileExpression(transformedHavingClause)
+
+        pipeline = pipeline.pipe(
+          filter(([, row]) => {
+            // Create a namespaced row structure for HAVING evaluation
+            const namespacedRow = { result: (row as any).__select_results }
+            return compiledHaving(namespacedRow)
+          })
+        )
+      }
+    }
+
+    // Apply functional HAVING clauses if present
+    if (fnHavingClauses && fnHavingClauses.length > 0) {
+      for (const fnHaving of fnHavingClauses) {
+        pipeline = pipeline.pipe(
+          filter(([, row]) => {
+            // Create a namespaced row structure for functional HAVING evaluation
+            const namespacedRow = { result: (row as any).__select_results }
+            return fnHaving(namespacedRow)
+          })
+        )
+      }
+    }
+
+    return pipeline
+  }
+
+  // Multi-group aggregation logic...
+  // Validate and create mapping for non-aggregate expressions in SELECT
+  const mapping = validateAndCreateMapping(groupByClause, selectClause)
+
+  // Pre-compile groupBy expressions
+  const compiledGroupByExpressions = groupByClause.map(compileExpression)
+
+  // Create a key extractor function using simple __key_X format
+  const keyExtractor = ([, row]: [
+    string,
+    NamespacedRow & { __select_results?: any },
+  ]) => {
+    // Use the original namespaced row for GROUP BY expressions, not __select_results
+    const namespacedRow = { ...row }
+    delete (namespacedRow as any).__select_results
+
+    const key: Record<string, unknown> = {}
+
+    // Use simple __key_X format for each groupBy expression
+    for (let i = 0; i < groupByClause.length; i++) {
+      const compiledExpr = compiledGroupByExpressions[i]!
+      const value = compiledExpr(namespacedRow)
+      key[`__key_${i}`] = value
+    }
+
+    return key
+  }
+
+  // Create aggregate functions for any aggregated columns in the SELECT clause
+  const aggregates: Record<string, any> = {}
+
+  if (selectClause) {
+    // Scan the SELECT clause for aggregate functions
+    for (const [alias, expr] of Object.entries(selectClause)) {
+      if (expr.type === `agg`) {
+        const aggExpr = expr
+        aggregates[alias] = getAggregateFunction(aggExpr)
+      }
+    }
+  }
+
+  // Apply the groupBy operator
+  pipeline = pipeline.pipe(groupBy(keyExtractor, aggregates))
+
+  // Update __select_results to handle GROUP BY results
+  pipeline = pipeline.pipe(
+    map(([, aggregatedRow]) => {
+      // Start with the existing __select_results from early SELECT processing
+      const selectResults = (aggregatedRow as any).__select_results || {}
+      const finalResults: Record<string, any> = {}
+
+      if (selectClause) {
+        // Process each SELECT expression
+        for (const [alias, expr] of Object.entries(selectClause)) {
+          if (expr.type !== `agg`) {
+            // Use cached mapping to get the corresponding __key_X for non-aggregates
+            const groupIndex = mapping.selectToGroupByIndex.get(alias)
+            if (groupIndex !== undefined) {
+              finalResults[alias] = aggregatedRow[`__key_${groupIndex}`]
+            } else {
+              // Fallback to original SELECT results
+              finalResults[alias] = selectResults[alias]
+            }
+          } else {
+            // Use aggregate results
+            finalResults[alias] = aggregatedRow[alias]
+          }
+        }
+      } else {
+        // No SELECT clause - just use the group keys
+        for (let i = 0; i < groupByClause.length; i++) {
+          finalResults[`__key_${i}`] = aggregatedRow[`__key_${i}`]
+        }
+      }
+
+      // Generate a simple key for the live collection using group values
+      let finalKey: unknown
+      if (groupByClause.length === 1) {
+        finalKey = aggregatedRow[`__key_0`]
+      } else {
+        const keyParts: Array<unknown> = []
+        for (let i = 0; i < groupByClause.length; i++) {
+          keyParts.push(aggregatedRow[`__key_${i}`])
+        }
+        finalKey = JSON.stringify(keyParts)
+      }
+
+      return [
+        finalKey,
+        {
+          ...aggregatedRow,
+          __select_results: finalResults,
+        },
+      ] as [unknown, Record<string, any>]
+    })
+  )
+
+  // Apply HAVING clauses if present
+  if (havingClauses && havingClauses.length > 0) {
+    for (const havingClause of havingClauses) {
+      const transformedHavingClause = transformHavingClause(
+        havingClause,
+        selectClause || {}
+      )
+      const compiledHaving = compileExpression(transformedHavingClause)
+
+      pipeline = pipeline.pipe(
+        filter(([, row]) => {
+          // Create a namespaced row structure for HAVING evaluation
+          const namespacedRow = { result: (row as any).__select_results }
+          return compiledHaving(namespacedRow)
+        })
+      )
+    }
+  }
+
+  // Apply functional HAVING clauses if present
+  if (fnHavingClauses && fnHavingClauses.length > 0) {
+    for (const fnHaving of fnHavingClauses) {
+      pipeline = pipeline.pipe(
+        filter(([, row]) => {
+          // Create a namespaced row structure for functional HAVING evaluation
+          const namespacedRow = { result: (row as any).__select_results }
+          return fnHaving(namespacedRow)
+        })
+      )
+    }
+  }
+
+  return pipeline
+}
+
+/**
+ * Helper function to check if two expressions are equal
+ */
+function expressionsEqual(expr1: any, expr2: any): boolean {
+  if (!expr1 || !expr2) return false
+  if (expr1.type !== expr2.type) return false
+
+  switch (expr1.type) {
+    case `ref`:
+      // Compare paths as arrays
+      if (!expr1.path || !expr2.path) return false
+      if (expr1.path.length !== expr2.path.length) return false
+      return expr1.path.every(
+        (segment: string, i: number) => segment === expr2.path[i]
+      )
+    case `val`:
+      return expr1.value === expr2.value
+    case `func`:
+      return (
+        expr1.name === expr2.name &&
+        expr1.args?.length === expr2.args?.length &&
+        (expr1.args || []).every((arg: any, i: number) =>
+          expressionsEqual(arg, expr2.args[i])
+        )
+      )
+    case `agg`:
+      return (
+        expr1.name === expr2.name &&
+        expr1.args?.length === expr2.args?.length &&
+        (expr1.args || []).every((arg: any, i: number) =>
+          expressionsEqual(arg, expr2.args[i])
+        )
+      )
+    default:
+      return false
+  }
+}
+
+/**
+ * Helper function to get an aggregate function based on the Agg expression
+ */
+function getAggregateFunction(aggExpr: Aggregate) {
+  // Pre-compile the value extractor expression
+  const compiledExpr = compileExpression(aggExpr.args[0]!)
+
+  // Create a value extractor function for the expression to aggregate
+  const valueExtractor = ([, namespacedRow]: [string, NamespacedRow]) => {
+    const value = compiledExpr(namespacedRow)
+    // Ensure we return a number for numeric aggregate functions
+    return typeof value === `number` ? value : value != null ? Number(value) : 0
+  }
+
+  // Return the appropriate aggregate function
+  switch (aggExpr.name.toLowerCase()) {
+    case `sum`:
+      return sum(valueExtractor)
+    case `count`:
+      return count() // count() doesn't need a value extractor
+    case `avg`:
+      return avg(valueExtractor)
+    case `min`:
+      return min(valueExtractor)
+    case `max`:
+      return max(valueExtractor)
+    default:
+      throw new Error(`Unsupported aggregate function: ${aggExpr.name}`)
+  }
+}
+
+/**
+ * Transforms a HAVING clause to replace Agg expressions with references to computed values
+ */
+function transformHavingClause(
+  havingExpr: BasicExpression | Aggregate,
+  selectClause: Select
+): BasicExpression {
+  switch (havingExpr.type) {
+    case `agg`: {
+      const aggExpr = havingExpr
+      // Find matching aggregate in SELECT clause
+      for (const [alias, selectExpr] of Object.entries(selectClause)) {
+        if (selectExpr.type === `agg` && aggregatesEqual(aggExpr, selectExpr)) {
+          // Replace with a reference to the computed aggregate
+          return new PropRef([`result`, alias])
+        }
+      }
+      // If no matching aggregate found in SELECT, throw error
+      throw new Error(
+        `Aggregate function in HAVING clause must also be in SELECT clause: ${aggExpr.name}`
+      )
+    }
+
+    case `func`: {
+      const funcExpr = havingExpr
+      // Transform function arguments recursively
+      const transformedArgs = funcExpr.args.map(
+        (arg: BasicExpression | Aggregate) =>
+          transformHavingClause(arg, selectClause)
+      )
+      return new Func(funcExpr.name, transformedArgs)
+    }
+
+    case `ref`: {
+      const refExpr = havingExpr
+      // Check if this is a direct reference to a SELECT alias
+      if (refExpr.path.length === 1) {
+        const alias = refExpr.path[0]!
+        if (selectClause[alias]) {
+          // This is a reference to a SELECT alias, convert to result.alias
+          return new PropRef([`result`, alias])
+        }
+      }
+      // Return as-is for other refs
+      return havingExpr as BasicExpression
+    }
+
+    case `val`:
+      // Return as-is
+      return havingExpr as BasicExpression
+
+    default:
+      throw new Error(
+        `Unknown expression type in HAVING clause: ${(havingExpr as any).type}`
+      )
+  }
+}
+
+/**
+ * Checks if two aggregate expressions are equal
+ */
+function aggregatesEqual(agg1: Aggregate, agg2: Aggregate): boolean {
+  return (
+    agg1.name === agg2.name &&
+    agg1.args.length === agg2.args.length &&
+    agg1.args.every((arg, i) => expressionsEqual(arg, agg2.args[i]))
+  )
+}

package/src/query/compiler/index.ts

@@ -0,0 +1,276 @@
+import { filter, map } from "@electric-sql/d2mini"
+import { compileExpression } from "./evaluators.js"
+import { processJoins } from "./joins.js"
+import { processGroupBy } from "./group-by.js"
+import { processOrderBy } from "./order-by.js"
+import { processSelectToResults } from "./select.js"
+import type { CollectionRef, QueryIR, QueryRef } from "../ir.js"
+import type {
+  KeyedStream,
+  NamespacedAndKeyedStream,
+  ResultStream,
+} from "../../types.js"
+
+/**
+ * Cache for compiled subqueries to avoid duplicate compilation
+ */
+type QueryCache = WeakMap<QueryIR, ResultStream>
+
+/**
+ * Compiles a query2 IR into a D2 pipeline
+ * @param query 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)
+ * @returns A stream builder representing the compiled query
+ */
+export function compileQuery(
+  query: QueryIR,
+  inputs: Record<string, KeyedStream>,
+  cache: QueryCache = new WeakMap()
+): ResultStream {
+  // Check if this query has already been compiled
+  const cachedResult = cache.get(query)
+  if (cachedResult) {
+    return cachedResult
+  }
+
+  // Create a copy of the inputs map to avoid modifying the original
+  const allInputs = { ...inputs }
+
+  // Create a map of table aliases to inputs
+  const tables: Record<string, KeyedStream> = {}
+
+  // Process the FROM clause to get the main table
+  const { alias: mainTableAlias, input: mainInput } = processFrom(
+    query.from,
+    allInputs,
+    cache
+  )
+  tables[mainTableAlias] = mainInput
+
+  // Prepare the initial pipeline with the main table wrapped in its alias
+  let pipeline: NamespacedAndKeyedStream = mainInput.pipe(
+    map(([key, row]) => {
+      // Initialize the record with a nested structure
+      const ret = [key, { [mainTableAlias]: row }] as [
+        string,
+        Record<string, typeof row>,
+      ]
+      return ret
+    })
+  )
+
+  // Process JOIN clauses if they exist
+  if (query.join && query.join.length > 0) {
+    pipeline = processJoins(
+      pipeline,
+      query.join,
+      tables,
+      mainTableAlias,
+      allInputs,
+      cache
+    )
+  }
+
+  // Process the WHERE clause if it exists
+  if (query.where && query.where.length > 0) {
+    // Compile all WHERE expressions
+    const compiledWheres = query.where.map((where) => compileExpression(where))
+
+    // Apply each WHERE condition as a filter (they are ANDed together)
+    for (const compiledWhere of compiledWheres) {
+      pipeline = pipeline.pipe(
+        filter(([_key, namespacedRow]) => {
+          return compiledWhere(namespacedRow)
+        })
+      )
+    }
+  }
+
+  // Process functional WHERE clauses if they exist
+  if (query.fnWhere && query.fnWhere.length > 0) {
+    for (const fnWhere of query.fnWhere) {
+      pipeline = pipeline.pipe(
+        filter(([_key, namespacedRow]) => {
+          return fnWhere(namespacedRow)
+        })
+      )
+    }
+  }
+
+  // Process the SELECT clause early - always create __select_results
+  // This eliminates duplication and allows for future DISTINCT implementation
+  if (query.fnSelect) {
+    // Handle functional select - apply the function to transform the row
+    pipeline = pipeline.pipe(
+      map(([key, namespacedRow]) => {
+        const selectResults = query.fnSelect!(namespacedRow)
+        return [
+          key,
+          {
+            ...namespacedRow,
+            __select_results: selectResults,
+          },
+        ] as [string, typeof namespacedRow & { __select_results: any }]
+      })
+    )
+  } else if (query.select) {
+    pipeline = processSelectToResults(pipeline, query.select, allInputs)
+  } else {
+    // If no SELECT clause, create __select_results with the main table data
+    pipeline = pipeline.pipe(
+      map(([key, namespacedRow]) => {
+        const selectResults =
+          !query.join && !query.groupBy
+            ? namespacedRow[mainTableAlias]
+            : namespacedRow
+
+        return [
+          key,
+          {
+            ...namespacedRow,
+            __select_results: selectResults,
+          },
+        ] as [string, typeof namespacedRow & { __select_results: any }]
+      })
+    )
+  }
+
+  // Process the GROUP BY clause if it exists
+  if (query.groupBy && query.groupBy.length > 0) {
+    pipeline = processGroupBy(
+      pipeline,
+      query.groupBy,
+      query.having,
+      query.select,
+      query.fnHaving
+    )
+  } else if (query.select) {
+    // Check if SELECT contains aggregates but no GROUP BY (implicit single-group aggregation)
+    const hasAggregates = Object.values(query.select).some(
+      (expr) => expr.type === `agg`
+    )
+    if (hasAggregates) {
+      // Handle implicit single-group aggregation
+      pipeline = processGroupBy(
+        pipeline,
+        [], // Empty group by means single group
+        query.having,
+        query.select,
+        query.fnHaving
+      )
+    }
+  }
+
+  // Process the HAVING clause if it exists (only applies after GROUP BY)
+  if (query.having && (!query.groupBy || query.groupBy.length === 0)) {
+    // Check if we have aggregates in SELECT that would trigger implicit grouping
+    const hasAggregates = query.select
+      ? Object.values(query.select).some((expr) => expr.type === `agg`)
+      : false
+
+    if (!hasAggregates) {
+      throw new Error(`HAVING clause requires GROUP BY clause`)
+    }
+  }
+
+  // Process functional HAVING clauses outside of GROUP BY (treat as additional WHERE filters)
+  if (
+    query.fnHaving &&
+    query.fnHaving.length > 0 &&
+    (!query.groupBy || query.groupBy.length === 0)
+  ) {
+    // If there's no GROUP BY but there are fnHaving clauses, apply them as filters
+    for (const fnHaving of query.fnHaving) {
+      pipeline = pipeline.pipe(
+        filter(([_key, namespacedRow]) => {
+          return fnHaving(namespacedRow)
+        })
+      )
+    }
+  }
+
+  // Process orderBy parameter if it exists
+  if (query.orderBy && query.orderBy.length > 0) {
+    const orderedPipeline = processOrderBy(
+      pipeline,
+      query.orderBy,
+      query.limit,
+      query.offset
+    )
+
+    // Final step: extract the __select_results and include orderBy index
+    const resultPipeline = orderedPipeline.pipe(
+      map(([key, [row, orderByIndex]]) => {
+        // Extract the final results from __select_results and include orderBy index
+        const finalResults = (row as any).__select_results
+        return [key, [finalResults, orderByIndex]] as [unknown, [any, string]]
+      })
+    )
+
+    const result = resultPipeline
+    // Cache the result before returning
+    cache.set(query, result)
+    return result
+  } else if (query.limit !== undefined || query.offset !== undefined) {
+    // If there's a limit or offset without orderBy, throw an error
+    throw new Error(
+      `LIMIT and OFFSET require an ORDER BY clause to ensure deterministic results`
+    )
+  }
+
+  // Final step: extract the __select_results and return tuple format (no orderBy)
+  const resultPipeline: ResultStream = pipeline.pipe(
+    map(([key, row]) => {
+      // Extract the final results from __select_results and return [key, [results, undefined]]
+      const finalResults = (row as any).__select_results
+      return [key, [finalResults, undefined]] as [
+        unknown,
+        [any, string | undefined],
+      ]
+    })
+  )
+
+  const result = resultPipeline
+  // Cache the result before returning
+  cache.set(query, result)
+  return result
+}
+
+/**
+ * Processes the FROM clause to extract the main table alias and input stream
+ */
+function processFrom(
+  from: CollectionRef | QueryRef,
+  allInputs: Record<string, KeyedStream>,
+  cache: QueryCache
+): { alias: string; input: KeyedStream } {
+  switch (from.type) {
+    case `collectionRef`: {
+      const input = allInputs[from.collection.id]
+      if (!input) {
+        throw new Error(
+          `Input for collection "${from.collection.id}" not found in inputs map`
+        )
+      }
+      return { alias: from.alias, input }
+    }
+    case `queryRef`: {
+      // Recursively compile the sub-query with cache
+      const subQueryInput = compileQuery(from.query, allInputs, cache)
+
+      // 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
+      const extractedInput = subQueryInput.pipe(
+        map((data: any) => {
+          const [key, [value, _orderByIndex]] = data
+          return [key, value] as [unknown, any]
+        })
+      )
+
+      return { alias: from.alias, input: extractedInput }
+    }
+    default:
+      throw new Error(`Unsupported FROM type: ${(from as any).type}`)
+  }
+}