feat: Subqueries in SELECT for hierarchical data (includes)

.changeset/includes-aggregates.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+fix: support aggregates (e.g. count) in child/includes subqueries with per-parent scoping

.changeset/includes-parent-referencing-filters.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+feat: support parent-referencing WHERE filters in includes child queries

.changeset/includes-subqueries.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': minor
+---
+
+feat: support for subqueries for including hierarchical data in live queries

.changeset/includes-to-array.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': patch
+---
+
+feat: add `toArray()` wrapper for includes subqueries to materialize child results as plain arrays instead of live Collections

packages/db/src/query/builder/functions.ts

@@ -2,7 +2,8 @@ import { Aggregate, Func } from '../ir'
 import { toExpression } from './ref-proxy.js'
 import type { BasicExpression } from '../ir'
 import type { RefProxy } from './ref-proxy.js'
-import type { RefLeaf } from './types.js'
+import type { Context, GetResult, RefLeaf } from './types.js'
+import type { QueryBuilder } from './index.js'
 
 type StringRef =
   | RefLeaf<string>
@@ -376,3 +377,14 @@ export const operators = [
 ] as const
 
 export type OperatorName = (typeof operators)[number]
+
+export class ToArrayWrapper<T = any> {
+  declare readonly _type: T
+  constructor(public readonly query: QueryBuilder<any>) {}
+}
+
+export function toArray<TContext extends Context>(
+  query: QueryBuilder<TContext>,
+): ToArrayWrapper<GetResult<TContext>> {
+  return new ToArrayWrapper(query)
+}

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

@@ -3,6 +3,7 @@ import {
   Aggregate as AggregateExpr,
   CollectionRef,
   Func as FuncExpr,
+  IncludesSubquery,
   PropRef,
   QueryRef,
   Value as ValueExpr,
@@ -23,6 +24,7 @@ import {
   isRefProxy,
   toExpression,
 } from './ref-proxy.js'
+import { ToArrayWrapper } from './functions.js'
 import type { NamespacedRow, SingleResult } from '../../types.js'
 import type {
   Aggregate,
@@ -31,6 +33,7 @@ import type {
   OrderBy,
   OrderByDirection,
   QueryIR,
+  Where,
 } from '../ir.js'
 import type {
   CompareOptions,
@@ -491,7 +494,7 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
     const aliases = this._getCurrentAliases()
     const refProxy = createRefProxy(aliases) as RefsForContext<TContext>
     const selectObject = callback(refProxy)
-    const select = buildNestedSelect(selectObject)
+    const select = buildNestedSelect(selectObject, aliases)
 
     return new BaseQueryBuilder({
       ...this.query,
@@ -867,7 +870,7 @@ function isPlainObject(value: any): value is Record<string, any> {
   )
 }
 
-function buildNestedSelect(obj: any): any {
+function buildNestedSelect(obj: any, parentAliases: Array<string> = []): any {
   if (!isPlainObject(obj)) return toExpr(obj)
   const out: Record<string, any> = {}
   for (const [k, v] of Object.entries(obj)) {
@@ -876,11 +879,267 @@ function buildNestedSelect(obj: any): any {
       out[k] = v
       continue
     }
-    out[k] = buildNestedSelect(v)
+    if (v instanceof BaseQueryBuilder) {
+      out[k] = buildIncludesSubquery(v, k, parentAliases, false)
+      continue
+    }
+    if (v instanceof ToArrayWrapper) {
+      if (!(v.query instanceof BaseQueryBuilder)) {
+        throw new Error(`toArray() must wrap a subquery builder`)
+      }
+      out[k] = buildIncludesSubquery(v.query, k, parentAliases, true)
+      continue
+    }
+    out[k] = buildNestedSelect(v, parentAliases)
   }
   return out
 }
 
+/**
+ * Recursively collects all PropRef nodes from an expression tree.
+ */
+function collectRefsFromExpression(expr: BasicExpression): Array<PropRef> {
+  const refs: Array<PropRef> = []
+  switch (expr.type) {
+    case `ref`:
+      refs.push(expr)
+      break
+    case `func`:
+      for (const arg of (expr as any).args ?? []) {
+        refs.push(...collectRefsFromExpression(arg))
+      }
+      break
+    default:
+      break
+  }
+  return refs
+}
+
+/**
+ * Checks whether a WHERE clause references any parent alias.
+ */
+function referencesParent(where: Where, parentAliases: Array<string>): boolean {
+  const expr =
+    typeof where === `object` && `expression` in where
+      ? where.expression
+      : where
+  return collectRefsFromExpression(expr).some(
+    (ref) => ref.path[0] != null && parentAliases.includes(ref.path[0]),
+  )
+}
+
+/**
+ * Builds an IncludesSubquery IR node from a child query builder.
+ * Extracts the correlation condition from the child's WHERE clauses by finding
+ * an eq() predicate that references both a parent alias and a child alias.
+ */
+function buildIncludesSubquery(
+  childBuilder: BaseQueryBuilder,
+  fieldName: string,
+  parentAliases: Array<string>,
+  materializeAsArray: boolean,
+): IncludesSubquery {
+  const childQuery = childBuilder._getQuery()
+
+  // Collect child's own aliases
+  const childAliases: Array<string> = [childQuery.from.alias]
+  if (childQuery.join) {
+    for (const j of childQuery.join) {
+      childAliases.push(j.from.alias)
+    }
+  }
+
+  // Walk child's WHERE clauses to find the correlation condition.
+  // The correlation eq() may be a standalone WHERE or nested inside a top-level and().
+  let parentRef: PropRef | undefined
+  let childRef: PropRef | undefined
+  let correlationWhereIndex = -1
+  let correlationAndArgIndex = -1 // >= 0 when found inside an and()
+
+  if (childQuery.where) {
+    for (let i = 0; i < childQuery.where.length; i++) {
+      const where = childQuery.where[i]!
+      const expr =
+        typeof where === `object` && `expression` in where
+          ? where.expression
+          : where
+
+      // Try standalone eq()
+      if (
+        expr.type === `func` &&
+        expr.name === `eq` &&
+        expr.args.length === 2
+      ) {
+        const result = extractCorrelation(
+          expr.args[0]!,
+          expr.args[1]!,
+          parentAliases,
+          childAliases,
+        )
+        if (result) {
+          parentRef = result.parentRef
+          childRef = result.childRef
+          correlationWhereIndex = i
+          break
+        }
+      }
+
+      // Try inside top-level and()
+      if (
+        expr.type === `func` &&
+        expr.name === `and` &&
+        expr.args.length >= 2
+      ) {
+        for (let j = 0; j < expr.args.length; j++) {
+          const arg = expr.args[j]!
+          if (
+            arg.type === `func` &&
+            arg.name === `eq` &&
+            arg.args.length === 2
+          ) {
+            const result = extractCorrelation(
+              arg.args[0]!,
+              arg.args[1]!,
+              parentAliases,
+              childAliases,
+            )
+            if (result) {
+              parentRef = result.parentRef
+              childRef = result.childRef
+              correlationWhereIndex = i
+              correlationAndArgIndex = j
+              break
+            }
+          }
+        }
+        if (parentRef) break
+      }
+    }
+  }
+
+  if (!parentRef || !childRef || correlationWhereIndex === -1) {
+    throw new Error(
+      `Includes subquery for "${fieldName}" must have a WHERE clause with an eq() condition ` +
+        `that correlates a parent field with a child field. ` +
+        `Example: .where(({child}) => eq(child.parentId, parent.id))`,
+    )
+  }
+
+  // Remove the correlation eq() from the child query's WHERE clauses.
+  // If it was inside an and(), remove just that arg (collapsing the and() if needed).
+  const modifiedWhere = [...childQuery.where!]
+  if (correlationAndArgIndex >= 0) {
+    const where = modifiedWhere[correlationWhereIndex]!
+    const expr =
+      typeof where === `object` && `expression` in where
+        ? where.expression
+        : where
+    const remainingArgs = (expr as any).args.filter(
+      (_: any, idx: number) => idx !== correlationAndArgIndex,
+    )
+    if (remainingArgs.length === 1) {
+      // Collapse and() with single remaining arg to just that expression
+      const isResidual =
+        typeof where === `object` && `expression` in where && where.residual
+      modifiedWhere[correlationWhereIndex] = isResidual
+        ? { expression: remainingArgs[0], residual: true }
+        : remainingArgs[0]
+    } else {
+      // Rebuild and() without the extracted arg
+      const newAnd = new FuncExpr(`and`, remainingArgs)
+      const isResidual =
+        typeof where === `object` && `expression` in where && where.residual
+      modifiedWhere[correlationWhereIndex] = isResidual
+        ? { expression: newAnd, residual: true }
+        : newAnd
+    }
+  } else {
+    modifiedWhere.splice(correlationWhereIndex, 1)
+  }
+
+  // Separate remaining WHEREs into pure-child vs parent-referencing
+  const pureChildWhere: Array<Where> = []
+  const parentFilters: Array<Where> = []
+  for (const w of modifiedWhere) {
+    if (referencesParent(w, parentAliases)) {
+      parentFilters.push(w)
+    } else {
+      pureChildWhere.push(w)
+    }
+  }
+
+  // Collect distinct parent PropRefs from parent-referencing filters
+  let parentProjection: Array<PropRef> | undefined
+  if (parentFilters.length > 0) {
+    const seen = new Set<string>()
+    parentProjection = []
+    for (const w of parentFilters) {
+      const expr = typeof w === `object` && `expression` in w ? w.expression : w
+      for (const ref of collectRefsFromExpression(expr)) {
+        if (
+          ref.path[0] != null &&
+          parentAliases.includes(ref.path[0]) &&
+          !seen.has(ref.path.join(`.`))
+        ) {
+          seen.add(ref.path.join(`.`))
+          parentProjection.push(ref)
+        }
+      }
+    }
+  }
+
+  const modifiedQuery: QueryIR = {
+    ...childQuery,
+    where: pureChildWhere.length > 0 ? pureChildWhere : undefined,
+  }
+
+  return new IncludesSubquery(
+    modifiedQuery,
+    parentRef,
+    childRef,
+    fieldName,
+    parentFilters.length > 0 ? parentFilters : undefined,
+    parentProjection,
+    materializeAsArray,
+  )
+}
+
+/**
+ * Checks if two eq() arguments form a parent-child correlation.
+ * Returns the parent and child PropRefs if found, undefined otherwise.
+ */
+function extractCorrelation(
+  argA: BasicExpression,
+  argB: BasicExpression,
+  parentAliases: Array<string>,
+  childAliases: Array<string>,
+): { parentRef: PropRef; childRef: PropRef } | undefined {
+  if (argA.type === `ref` && argB.type === `ref`) {
+    const aAlias = argA.path[0]
+    const bAlias = argB.path[0]
+
+    if (
+      aAlias &&
+      bAlias &&
+      parentAliases.includes(aAlias) &&
+      childAliases.includes(bAlias)
+    ) {
+      return { parentRef: argA, childRef: argB }
+    }
+
+    if (
+      aAlias &&
+      bAlias &&
+      parentAliases.includes(bAlias) &&
+      childAliases.includes(aAlias)
+    ) {
+      return { parentRef: argB, childRef: argA }
+    }
+  }
+
+  return undefined
+}
+
 // Internal function to build a query from a callback
 // used by liveQueryCollectionOptions.query
 export function buildQuery<TContext extends Context>(