@tanstack/db 0.1.12 → 0.2.0

package/src/query/compiler/select.ts

@@ -1,4 +1,5 @@
 import { map } from "@tanstack/db-ivm"
+import { PropRef, Value as ValClass, isExpressionLike } from "../ir.js"
 import { compileExpression } from "./evaluators.js"
 import type { Aggregate, BasicExpression, Select } from "../ir.js"
 import type {
@@ -8,139 +9,135 @@ import type {
 } from "../../types.js"
 
 /**
- * Processes the SELECT clause and places results in __select_results
- * while preserving the original namespaced row for ORDER BY access
+ * Type for operations array used in select processing
  */
-export function processSelectToResults(
-  pipeline: NamespacedAndKeyedStream,
-  select: Select,
-  _allInputs: Record<string, KeyedStream>
-): NamespacedAndKeyedStream {
-  // Pre-compile all select expressions
-  const compiledSelect: Array<{
-    alias: string
-    compiledExpression: (row: NamespacedRow) => any
-  }> = []
-  const spreadAliases: Array<string> = []
-
-  for (const [alias, expression] of Object.entries(select)) {
-    if (alias.startsWith(`__SPREAD_SENTINEL__`)) {
-      // Extract the table alias from the sentinel key
-      const tableAlias = alias.replace(`__SPREAD_SENTINEL__`, ``)
-      spreadAliases.push(tableAlias)
-    } else {
-      if (isAggregateExpression(expression)) {
-        // For aggregates, we'll store the expression info for GROUP BY processing
-        // but still compile a placeholder that will be replaced later
-        compiledSelect.push({
-          alias,
-          compiledExpression: () => null, // Placeholder - will be handled by GROUP BY
-        })
-      } else {
-        compiledSelect.push({
-          alias,
-          compiledExpression: compileExpression(expression as BasicExpression),
-        })
-      }
+type SelectOp =
+  | {
+      kind: `merge`
+      targetPath: Array<string>
+      source: (row: NamespacedRow) => any
     }
-  }
+  | { kind: `field`; alias: string; compiled: (row: NamespacedRow) => any }
 
-  return pipeline.pipe(
-    map(([key, namespacedRow]) => {
-      const selectResults: Record<string, any> = {}
-
-      // First pass: spread table data for any spread sentinels
-      for (const tableAlias of spreadAliases) {
-        const tableData = namespacedRow[tableAlias]
-        if (tableData && typeof tableData === `object`) {
-          // Spread the table data into the result, but don't overwrite explicit fields
-          for (const [fieldName, fieldValue] of Object.entries(tableData)) {
-            if (!(fieldName in selectResults)) {
-              selectResults[fieldName] = fieldValue
+/**
+ * Unwraps any Value expressions
+ */
+function unwrapVal(input: any): any {
+  if (input instanceof ValClass) return input.value
+  return input
+}
+
+/**
+ * Processes a merge operation by merging source values into the target path
+ */
+function processMerge(
+  op: Extract<SelectOp, { kind: `merge` }>,
+  namespacedRow: NamespacedRow,
+  selectResults: Record<string, any>
+): void {
+  const value = op.source(namespacedRow)
+  if (value && typeof value === `object`) {
+    // Ensure target object exists
+    let cursor: any = selectResults
+    const path = op.targetPath
+    if (path.length === 0) {
+      // Top-level merge
+      for (const [k, v] of Object.entries(value)) {
+        selectResults[k] = unwrapVal(v)
+      }
+    } else {
+      for (let i = 0; i < path.length; i++) {
+        const seg = path[i]!
+        if (i === path.length - 1) {
+          const dest = (cursor[seg] ??= {})
+          if (typeof dest === `object`) {
+            for (const [k, v] of Object.entries(value)) {
+              dest[k] = unwrapVal(v)
             }
           }
+        } else {
+          const next = cursor[seg]
+          if (next == null || typeof next !== `object`) {
+            cursor[seg] = {}
+          }
+          cursor = cursor[seg]
         }
       }
+    }
+  }
+}
 
-      // Second pass: evaluate all compiled select expressions (non-aggregates)
-      for (const { alias, compiledExpression } of compiledSelect) {
-        selectResults[alias] = compiledExpression(namespacedRow)
+/**
+ * Processes a non-merge operation by setting the field value at the specified alias path
+ */
+function processNonMergeOp(
+  op: Extract<SelectOp, { kind: `field` }>,
+  namespacedRow: NamespacedRow,
+  selectResults: Record<string, any>
+): void {
+  // Support nested alias paths like "meta.author.name"
+  const path = op.alias.split(`.`)
+  if (path.length === 1) {
+    selectResults[op.alias] = op.compiled(namespacedRow)
+  } else {
+    let cursor: any = selectResults
+    for (let i = 0; i < path.length - 1; i++) {
+      const seg = path[i]!
+      const next = cursor[seg]
+      if (next == null || typeof next !== `object`) {
+        cursor[seg] = {}
       }
+      cursor = cursor[seg]
+    }
+    cursor[path[path.length - 1]!] = unwrapVal(op.compiled(namespacedRow))
+  }
+}
 
-      // Return the namespaced row with __select_results added
-      return [
-        key,
-        {
-          ...namespacedRow,
-          __select_results: selectResults,
-        },
-      ] as [
-        string,
-        typeof namespacedRow & { __select_results: typeof selectResults },
-      ]
-    })
-  )
+/**
+ * Processes a single row to generate select results
+ */
+function processRow(
+  [key, namespacedRow]: [unknown, NamespacedRow],
+  ops: Array<SelectOp>
+): [unknown, typeof namespacedRow & { __select_results: any }] {
+  const selectResults: Record<string, any> = {}
+
+  for (const op of ops) {
+    if (op.kind === `merge`) {
+      processMerge(op, namespacedRow, selectResults)
+    } else {
+      processNonMergeOp(op, namespacedRow, selectResults)
+    }
+  }
+
+  // Return the namespaced row with __select_results added
+  return [
+    key,
+    {
+      ...namespacedRow,
+      __select_results: selectResults,
+    },
+  ] as [
+    unknown,
+    typeof namespacedRow & { __select_results: typeof selectResults },
+  ]
 }
 
 /**
- * Processes the SELECT clause (legacy function - kept for compatibility)
+ * Processes the SELECT clause and places results in __select_results
+ * while preserving the original namespaced row for ORDER BY access
  */
 export function processSelect(
   pipeline: NamespacedAndKeyedStream,
   select: Select,
   _allInputs: Record<string, KeyedStream>
-): KeyedStream {
-  // Pre-compile all select expressions
-  const compiledSelect: Array<{
-    alias: string
-    compiledExpression: (row: NamespacedRow) => any
-  }> = []
-  const spreadAliases: Array<string> = []
-
-  for (const [alias, expression] of Object.entries(select)) {
-    if (alias.startsWith(`__SPREAD_SENTINEL__`)) {
-      // Extract the table alias from the sentinel key
-      const tableAlias = alias.replace(`__SPREAD_SENTINEL__`, ``)
-      spreadAliases.push(tableAlias)
-    } else {
-      if (isAggregateExpression(expression)) {
-        // Aggregates should be handled by GROUP BY processing, not here
-        throw new Error(
-          `Aggregate expressions in SELECT clause should be handled by GROUP BY processing`
-        )
-      }
-      compiledSelect.push({
-        alias,
-        compiledExpression: compileExpression(expression as BasicExpression),
-      })
-    }
-  }
-
-  return pipeline.pipe(
-    map(([key, namespacedRow]) => {
-      const result: Record<string, any> = {}
-
-      // First pass: spread table data for any spread sentinels
-      for (const tableAlias of spreadAliases) {
-        const tableData = namespacedRow[tableAlias]
-        if (tableData && typeof tableData === `object`) {
-          // Spread the table data into the result, but don't overwrite explicit fields
-          for (const [fieldName, fieldValue] of Object.entries(tableData)) {
-            if (!(fieldName in result)) {
-              result[fieldName] = fieldValue
-            }
-          }
-        }
-      }
+): NamespacedAndKeyedStream {
+  // Build ordered operations to preserve authoring order (spreads and fields)
+  const ops: Array<SelectOp> = []
 
-      // Second pass: evaluate all compiled select expressions
-      for (const { alias, compiledExpression } of compiledSelect) {
-        result[alias] = compiledExpression(namespacedRow)
-      }
+  addFromObject([], select, ops)
 
-      return [key, result] as [string, typeof result]
-    })
-  )
+  return pipeline.pipe(map((row) => processRow(row, ops)))
 }
 
 /**
@@ -171,3 +168,99 @@ export function processArgument(
 
   return value
 }
+
+/**
+ * Helper function to check if an object is a nested select object
+ *
+ * .select({
+ *   id: users.id,
+ *   profile: { // <-- this is a nested select object
+ *     name: users.name,
+ *     email: users.email
+ *   }
+ * })
+ */
+function isNestedSelectObject(obj: any): boolean {
+  return obj && typeof obj === `object` && !isExpressionLike(obj)
+}
+
+/**
+ * Helper function to process select objects and build operations array
+ */
+function addFromObject(
+  prefixPath: Array<string>,
+  obj: any,
+  ops: Array<SelectOp>
+) {
+  for (const [key, value] of Object.entries(obj)) {
+    if (key.startsWith(`__SPREAD_SENTINEL__`)) {
+      const rest = key.slice(`__SPREAD_SENTINEL__`.length)
+      const splitIndex = rest.lastIndexOf(`__`)
+      const pathStr = splitIndex >= 0 ? rest.slice(0, splitIndex) : rest
+      const isRefExpr =
+        value &&
+        typeof value === `object` &&
+        `type` in (value as any) &&
+        (value as any).type === `ref`
+      if (pathStr.includes(`.`) || isRefExpr) {
+        // Merge into the current destination (prefixPath) from the referenced source path
+        const targetPath = [...prefixPath]
+        const expr = isRefExpr
+          ? (value as BasicExpression)
+          : (new PropRef(pathStr.split(`.`)) as BasicExpression)
+        const compiled = compileExpression(expr)
+        ops.push({ kind: `merge`, targetPath, source: compiled })
+      } else {
+        // Table-level: pathStr is the alias; merge from namespaced row at the current prefix
+        const tableAlias = pathStr
+        const targetPath = [...prefixPath]
+        ops.push({
+          kind: `merge`,
+          targetPath,
+          source: (row) => (row as any)[tableAlias],
+        })
+      }
+      continue
+    }
+
+    const expression = value as any
+    if (isNestedSelectObject(expression)) {
+      // Nested selection object
+      addFromObject([...prefixPath, key], expression, ops)
+      continue
+    }
+
+    if (isAggregateExpression(expression)) {
+      // Placeholder for group-by processing later
+      ops.push({
+        kind: `field`,
+        alias: [...prefixPath, key].join(`.`),
+        compiled: () => null,
+      })
+    } else {
+      if (expression === undefined || !isExpressionLike(expression)) {
+        ops.push({
+          kind: `field`,
+          alias: [...prefixPath, key].join(`.`),
+          compiled: () => expression,
+        })
+        continue
+      }
+      // If the expression is a Value wrapper, embed the literal to avoid re-compilation mishaps
+      if (expression instanceof ValClass) {
+        const val = expression.value
+        ops.push({
+          kind: `field`,
+          alias: [...prefixPath, key].join(`.`),
+          compiled: () => val,
+        })
+      } else {
+        ops.push({
+          kind: `field`,
+          alias: [...prefixPath, key].join(`.`),
+          compiled: compileExpression(expression as BasicExpression),
+        })
+      }
+    }
+  }
+}

package/src/query/ir.ts

@@ -27,7 +27,7 @@ export interface QueryIR {
 export type From = CollectionRef | QueryRef
 
 export type Select = {
-  [alias: string]: BasicExpression | Aggregate
+  [alias: string]: BasicExpression | Aggregate | Select
 }
 
 export type Join = Array<JoinClause>
@@ -131,6 +131,19 @@ export class Aggregate<T = any> extends BaseExpression<T> {
   }
 }
 
+/**
+ * Runtime helper to detect IR expression-like objects.
+ * Prefer this over ad-hoc local implementations to keep behavior consistent.
+ */
+export function isExpressionLike(value: any): boolean {
+  return (
+    value instanceof Aggregate ||
+    value instanceof Func ||
+    value instanceof PropRef ||
+    value instanceof Value
+  )
+}
+
 /**
  * Helper functions for working with Where clauses
  */

package/src/query/optimizer.ts

@@ -125,13 +125,14 @@ import { CannotCombineEmptyExpressionListError } from "../errors.js"
 import {
   CollectionRef as CollectionRefClass,
   Func,
+  PropRef,
   QueryRef as QueryRefClass,
   createResidualWhere,
   getWhereExpression,
   isResidualWhere,
 } from "./ir.js"
 import { isConvertibleToCollectionFilter } from "./compiler/expressions.js"
-import type { BasicExpression, From, QueryIR, Where } from "./ir.js"
+import type { BasicExpression, From, QueryIR, Select, Where } from "./ir.js"
 
 /**
  * Represents a WHERE clause after source analysis
@@ -660,6 +661,7 @@ function applyOptimizations(
     orderBy: query.orderBy ? [...query.orderBy] : undefined,
     limit: query.limit,
     offset: query.offset,
+    distinct: query.distinct,
     fnSelect: query.fnSelect,
     fnWhere: query.fnWhere ? [...query.fnWhere] : undefined,
     fnHaving: query.fnHaving ? [...query.fnHaving] : undefined,
@@ -763,7 +765,7 @@ function optimizeFromWithTracking(
   // SAFETY CHECK: Only check safety when pushing WHERE clauses into existing subqueries
   // We need to be careful about pushing WHERE clauses into subqueries that already have
   // aggregates, HAVING, or ORDER BY + LIMIT since that could change their semantics
-  if (!isSafeToPushIntoExistingSubquery(from.query)) {
+  if (!isSafeToPushIntoExistingSubquery(from.query, whereClause, from.alias)) {
     // Return a copy without optimization to maintain immutability
     // Do NOT mark as optimized since we didn't actually optimize it
     return new QueryRefClass(deepCopyQuery(from.query), from.alias)
@@ -780,73 +782,143 @@ function optimizeFromWithTracking(
   return new QueryRefClass(optimizedSubQuery, from.alias)
 }
 
+function unsafeSelect(
+  query: QueryIR,
+  whereClause: BasicExpression<boolean>,
+  outerAlias: string
+): boolean {
+  if (!query.select) return false
+
+  return (
+    selectHasAggregates(query.select) ||
+    whereReferencesComputedSelectFields(query.select, whereClause, outerAlias)
+  )
+}
+
+function unsafeGroupBy(query: QueryIR) {
+  return query.groupBy && query.groupBy.length > 0
+}
+
+function unsafeHaving(query: QueryIR) {
+  return query.having && query.having.length > 0
+}
+
+function unsafeOrderBy(query: QueryIR) {
+  return (
+    query.orderBy &&
+    query.orderBy.length > 0 &&
+    (query.limit !== undefined || query.offset !== undefined)
+  )
+}
+
+function unsafeFnSelect(query: QueryIR) {
+  return (
+    query.fnSelect ||
+    (query.fnWhere && query.fnWhere.length > 0) ||
+    (query.fnHaving && query.fnHaving.length > 0)
+  )
+}
+
+function isSafeToPushIntoExistingSubquery(
+  query: QueryIR,
+  whereClause: BasicExpression<boolean>,
+  outerAlias: string
+): boolean {
+  return !(
+    unsafeSelect(query, whereClause, outerAlias) ||
+    unsafeGroupBy(query) ||
+    unsafeHaving(query) ||
+    unsafeOrderBy(query) ||
+    unsafeFnSelect(query)
+  )
+}
+
 /**
- * Determines if it's safe to push WHERE clauses into an existing subquery.
- *
- * Pushing WHERE clauses into existing subqueries can break semantics in several cases:
- *
- * 1. **Aggregates**: Pushing predicates before GROUP BY changes what gets aggregated
- * 2. **ORDER BY + LIMIT/OFFSET**: Pushing predicates before sorting+limiting changes the result set
- * 3. **HAVING clauses**: These operate on aggregated data, predicates should not be pushed past them
- * 4. **Functional operations**: fnSelect, fnWhere, fnHaving could have side effects
- *
- * Note: This safety check only applies when pushing WHERE clauses into existing subqueries.
- * Creating new subqueries from collection references is always safe.
- *
- * @param query - The existing subquery to check for safety
- * @returns True if it's safe to push WHERE clauses into this subquery, false otherwise
+ * Detects whether a SELECT projection contains any aggregate expressions.
+ * Recursively traverses nested select objects.
  *
- * @example
- * ```typescript
- * // UNSAFE: has GROUP BY - pushing WHERE could change aggregation
- * { from: users, groupBy: [dept], select: { count: agg('count', '*') } }
- *
- * // UNSAFE: has ORDER BY + LIMIT - pushing WHERE could change "top 10"
- * { from: users, orderBy: [salary desc], limit: 10 }
- *
- * // SAFE: plain SELECT without aggregates/limits
- * { from: users, select: { id, name } }
- * ```
+ * @param select - The SELECT object from the IR
+ * @returns True if any field is an aggregate, false otherwise
  */
-function isSafeToPushIntoExistingSubquery(query: QueryIR): boolean {
-  // Check for aggregates in SELECT clause
-  if (query.select) {
-    const hasAggregates = Object.values(query.select).some(
-      (expr) => expr.type === `agg`
-    )
-    if (hasAggregates) {
-      return false
+function selectHasAggregates(select: Select): boolean {
+  for (const value of Object.values(select)) {
+    if (typeof value === `object`) {
+      const v: any = value
+      if (v.type === `agg`) return true
+      if (!(`type` in v)) {
+        if (selectHasAggregates(v as unknown as Select)) return true
+      }
     }
   }
+  return false
+}
 
-  // Check for GROUP BY clause
-  if (query.groupBy && query.groupBy.length > 0) {
-    return false
-  }
-
-  // Check for HAVING clause
-  if (query.having && query.having.length > 0) {
-    return false
+/**
+ * Recursively collects all PropRef references from an expression.
+ *
+ * @param expr - The expression to traverse
+ * @returns Array of PropRef references found in the expression
+ */
+function collectRefs(expr: any): Array<PropRef> {
+  const refs: Array<PropRef> = []
+
+  if (expr == null || typeof expr !== `object`) return refs
+
+  switch (expr.type) {
+    case `ref`:
+      refs.push(expr as PropRef)
+      break
+    case `func`:
+    case `agg`:
+      for (const arg of expr.args ?? []) {
+        refs.push(...collectRefs(arg))
+      }
+      break
+    default:
+      break
   }
 
-  // Check for ORDER BY with LIMIT or OFFSET (dangerous combination)
-  if (query.orderBy && query.orderBy.length > 0) {
-    if (query.limit !== undefined || query.offset !== undefined) {
-      return false
-    }
-  }
+  return refs
+}
 
-  // Check for functional variants that might have side effects
-  if (
-    query.fnSelect ||
-    (query.fnWhere && query.fnWhere.length > 0) ||
-    (query.fnHaving && query.fnHaving.length > 0)
-  ) {
-    return false
+/**
+ * Determines whether the provided WHERE clause references fields that are
+ * computed by a subquery SELECT rather than pass-through properties.
+ *
+ * If true, pushing the WHERE clause into the subquery could change semantics
+ * (since computed fields do not necessarily exist at the subquery input level),
+ * so predicate pushdown must be avoided.
+ *
+ * @param select - The subquery SELECT map
+ * @param whereClause - The WHERE expression to analyze
+ * @param outerAlias - The alias of the subquery in the outer query
+ * @returns True if WHERE references computed fields, otherwise false
+ */
+function whereReferencesComputedSelectFields(
+  select: Select,
+  whereClause: BasicExpression<boolean>,
+  outerAlias: string
+): boolean {
+  // Build a set of computed field names at the top-level of the subquery output
+  const computed = new Set<string>()
+  for (const [key, value] of Object.entries(select)) {
+    if (key.startsWith(`__SPREAD_SENTINEL__`)) continue
+    if (value instanceof PropRef) continue
+    // Nested object or non-PropRef expression counts as computed
+    computed.add(key)
+  }
+
+  const refs = collectRefs(whereClause)
+
+  for (const ref of refs) {
+    const path = (ref as any).path as Array<string>
+    if (!Array.isArray(path) || path.length < 2) continue
+    const alias = path[0]
+    const field = path[1] as string
+    if (alias !== outerAlias) continue
+    if (computed.has(field)) return true
   }
-
-  // If none of the unsafe conditions are present, it's safe to optimize
-  return true
+  return false
 }
 
 /**