@tanstack/db 0.4.8 → 0.4.10

package/src/query/compiler/joins.ts

@@ -1,18 +1,13 @@
-import {
-  consolidate,
-  filter,
-  join as joinOperator,
-  map,
-  tap,
-} from "@tanstack/db-ivm"
+import { filter, join as joinOperator, map, tap } from "@tanstack/db-ivm"
 import {
   CollectionInputNotFoundError,
   InvalidJoinCondition,
-  InvalidJoinConditionLeftTableError,
-  InvalidJoinConditionRightTableError,
-  InvalidJoinConditionSameTableError,
-  InvalidJoinConditionTableMismatchError,
+  InvalidJoinConditionLeftSourceError,
+  InvalidJoinConditionRightSourceError,
+  InvalidJoinConditionSameSourceError,
+  InvalidJoinConditionSourceMismatchError,
   JoinCollectionNotFoundError,
+  SubscriptionNotFoundError,
   UnsupportedJoinSourceTypeError,
   UnsupportedJoinTypeError,
 } from "../../errors.js"
@@ -36,34 +31,41 @@ import type {
   NamespacedAndKeyedStream,
   NamespacedRow,
 } from "../../types.js"
-import type { QueryCache, QueryMapping } from "./types.js"
+import type { QueryCache, QueryMapping, WindowOptions } from "./types.js"
 import type { CollectionSubscription } from "../../collection/subscription.js"
 
+/** Function type for loading specific keys into a lazy collection */
 export type LoadKeysFn = (key: Set<string | number>) => void
+
+/** Callbacks for managing lazy-loaded collections in optimized joins */
 export type LazyCollectionCallbacks = {
   loadKeys: LoadKeysFn
   loadInitialState: () => void
 }
 
 /**
- * Processes all join clauses in a query
+ * Processes all join clauses, applying lazy loading optimizations and maintaining
+ * alias tracking for per-alias subscriptions (enables self-joins).
  */
 export function processJoins(
   pipeline: NamespacedAndKeyedStream,
   joinClauses: Array<JoinClause>,
-  tables: Record<string, KeyedStream>,
-  mainTableId: string,
-  mainTableAlias: string,
+  sources: Record<string, KeyedStream>,
+  mainCollectionId: string,
+  mainSource: string,
   allInputs: Record<string, KeyedStream>,
   cache: QueryCache,
   queryMapping: QueryMapping,
   collections: Record<string, Collection>,
   subscriptions: Record<string, CollectionSubscription>,
   callbacks: Record<string, LazyCollectionCallbacks>,
-  lazyCollections: Set<string>,
+  lazySources: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
+  setWindowFn: (windowFn: (options: WindowOptions) => void) => void,
   rawQuery: QueryIR,
-  onCompileSubquery: CompileQueryFn
+  onCompileSubquery: CompileQueryFn,
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): NamespacedAndKeyedStream {
   let resultPipeline = pipeline
 
@@ -71,19 +73,22 @@ export function processJoins(
     resultPipeline = processJoin(
       resultPipeline,
       joinClause,
-      tables,
-      mainTableId,
-      mainTableAlias,
+      sources,
+      mainCollectionId,
+      mainSource,
       allInputs,
       cache,
       queryMapping,
       collections,
       subscriptions,
       callbacks,
-      lazyCollections,
+      lazySources,
       optimizableOrderByCollections,
+      setWindowFn,
       rawQuery,
-      onCompileSubquery
+      onCompileSubquery,
+      aliasToCollectionId,
+      aliasRemapping
     )
   }
 
@@ -91,28 +96,34 @@ export function processJoins(
 }
 
 /**
- * Processes a single join clause
+ * Processes a single join clause with lazy loading optimization.
+ * For LEFT/RIGHT/INNER joins, marks one side as "lazy" (loads on-demand based on join keys).
  */
 function processJoin(
   pipeline: NamespacedAndKeyedStream,
   joinClause: JoinClause,
-  tables: Record<string, KeyedStream>,
-  mainTableId: string,
-  mainTableAlias: string,
+  sources: Record<string, KeyedStream>,
+  mainCollectionId: string,
+  mainSource: string,
   allInputs: Record<string, KeyedStream>,
   cache: QueryCache,
   queryMapping: QueryMapping,
   collections: Record<string, Collection>,
   subscriptions: Record<string, CollectionSubscription>,
   callbacks: Record<string, LazyCollectionCallbacks>,
-  lazyCollections: Set<string>,
+  lazySources: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
+  setWindowFn: (windowFn: (options: WindowOptions) => void) => void,
   rawQuery: QueryIR,
-  onCompileSubquery: CompileQueryFn
+  onCompileSubquery: CompileQueryFn,
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): NamespacedAndKeyedStream {
-  // Get the joined table alias and input stream
+  const isCollectionRef = joinClause.from.type === `collectionRef`
+
+  // Get the joined source alias and input stream
   const {
-    alias: joinedTableAlias,
+    alias: joinedSource,
     input: joinedInput,
     collectionId: joinedCollectionId,
   } = processJoinSource(
@@ -121,40 +132,48 @@ function processJoin(
     collections,
     subscriptions,
     callbacks,
-    lazyCollections,
+    lazySources,
     optimizableOrderByCollections,
+    setWindowFn,
     cache,
     queryMapping,
-    onCompileSubquery
+    onCompileSubquery,
+    aliasToCollectionId,
+    aliasRemapping
   )
 
-  // Add the joined table to the tables map
-  tables[joinedTableAlias] = joinedInput
+  // Add the joined source to the sources map
+  sources[joinedSource] = joinedInput
+  if (isCollectionRef) {
+    // Only direct collection references form new alias bindings. Subquery
+    // aliases reuse the mapping returned from the recursive compilation above.
+    aliasToCollectionId[joinedSource] = joinedCollectionId
+  }
 
-  const mainCollection = collections[mainTableId]
+  const mainCollection = collections[mainCollectionId]
   const joinedCollection = collections[joinedCollectionId]
 
   if (!mainCollection) {
-    throw new JoinCollectionNotFoundError(mainTableId)
+    throw new JoinCollectionNotFoundError(mainCollectionId)
   }
 
   if (!joinedCollection) {
     throw new JoinCollectionNotFoundError(joinedCollectionId)
   }
 
-  const { activeCollection, lazyCollection } = getActiveAndLazyCollections(
+  const { activeSource, lazySource } = getActiveAndLazySources(
     joinClause.type,
     mainCollection,
     joinedCollection
   )
 
-  // Analyze which table each expression refers to and swap if necessary
-  const availableTableAliases = Object.keys(tables)
+  // Analyze which source each expression refers to and swap if necessary
+  const availableSources = Object.keys(sources)
   const { mainExpr, joinedExpr } = analyzeJoinExpressions(
     joinClause.left,
     joinClause.right,
-    availableTableAliases,
-    joinedTableAlias
+    availableSources,
+    joinedSource
   )
 
   // Pre-compile the join expressions
@@ -164,7 +183,7 @@ function processJoin(
   // Prepare the main pipeline for joining
   let mainPipeline = pipeline.pipe(
     map(([currentKey, namespacedRow]) => {
-      // Extract the join key from the main table expression
+      // Extract the join key from the main source expression
       const mainKey = compiledMainExpr(namespacedRow)
 
       // Return [joinKey, [originalKey, namespacedRow]]
@@ -179,9 +198,9 @@ function processJoin(
   let joinedPipeline = joinedInput.pipe(
     map(([currentKey, row]) => {
       // Wrap the row in a namespaced structure
-      const namespacedRow: NamespacedRow = { [joinedTableAlias]: row }
+      const namespacedRow: NamespacedRow = { [joinedSource]: row }
 
-      // Extract the join key from the joined table expression
+      // Extract the join key from the joined source expression
       const joinedKey = compiledJoinedExpr(namespacedRow)
 
       // Return [joinKey, [originalKey, namespacedRow]]
@@ -197,13 +216,12 @@ function processJoin(
     throw new UnsupportedJoinTypeError(joinClause.type)
   }
 
-  if (activeCollection) {
+  if (activeSource) {
     // If the lazy collection comes from a subquery that has a limit and/or an offset clause
     // then we need to deoptimize the join because we don't know which rows are in the result set
     // since we simply lookup matching keys in the index but the index contains all rows
     // (not just the ones that pass the limit and offset clauses)
-    const lazyFrom =
-      activeCollection === `main` ? joinClause.from : rawQuery.from
+    const lazyFrom = activeSource === `main` ? joinClause.from : rawQuery.from
     const limitedSubquery =
       lazyFrom.type === `queryRef` &&
       (lazyFrom.query.limit || lazyFrom.query.offset)
@@ -219,24 +237,25 @@ function processJoin(
       // based on the value of the joinKey and by looking up
       // matching rows in the index of the lazy collection
 
-      // Mark the lazy collection as lazy
+      // Mark the lazy source alias as lazy
       // this Set is passed by the liveQueryCollection to the compiler
       // such that the liveQueryCollection can check it after compilation
-      // to know which collections are lazy collections
-      lazyCollections.add(lazyCollection.id)
+      // to know which source aliases should load data lazily (not initially)
+      const lazyAlias = activeSource === `main` ? joinedSource : mainSource
+      lazySources.add(lazyAlias)
 
       const activePipeline =
-        activeCollection === `main` ? mainPipeline : joinedPipeline
+        activeSource === `main` ? mainPipeline : joinedPipeline
 
-      const lazyCollectionJoinExpr =
-        activeCollection === `main`
+      const lazySourceJoinExpr =
+        activeSource === `main`
           ? (joinedExpr as PropRef)
           : (mainExpr as PropRef)
 
       const followRefResult = followRef(
         rawQuery,
-        lazyCollectionJoinExpr,
-        lazyCollection
+        lazySourceJoinExpr,
+        lazySource
       )!
       const followRefCollection = followRefResult.collection
 
@@ -249,38 +268,51 @@ function processJoin(
         )
       }
 
+      // Set up lazy loading: intercept active side's stream and dynamically load
+      // matching rows from lazy side based on join keys.
       const activePipelineWithLoading: IStreamBuilder<
         [key: unknown, [originalKey: string, namespacedRow: NamespacedRow]]
       > = activePipeline.pipe(
         tap((data) => {
-          const lazyCollectionSubscription = subscriptions[lazyCollection.id]
-
-          if (!lazyCollectionSubscription) {
-            throw new Error(
-              `Internal error: subscription for collection is missing in join pipeline. Make sure the live query collection sets the subscription before running the pipeline.`
+          // Find the subscription for lazy loading.
+          // Subscriptions are keyed by the innermost alias (where the collection subscription
+          // was actually created). For subqueries, the join alias may differ from the inner alias.
+          // aliasRemapping provides a flattened one-hop lookup from outer → innermost alias.
+          // Example: .join({ activeUser: subquery }) where subquery uses .from({ user: collection })
+          // → aliasRemapping['activeUser'] = 'user' (always maps directly to innermost, never recursive)
+          const resolvedAlias = aliasRemapping[lazyAlias] || lazyAlias
+          const lazySourceSubscription = subscriptions[resolvedAlias]
+
+          if (!lazySourceSubscription) {
+            throw new SubscriptionNotFoundError(
+              resolvedAlias,
+              lazyAlias,
+              lazySource.id,
+              Object.keys(subscriptions)
             )
           }
 
-          if (lazyCollectionSubscription.hasLoadedInitialState()) {
+          if (lazySourceSubscription.hasLoadedInitialState()) {
             // Entire state was already loaded because we deoptimized the join
             return
           }
 
+          // Request filtered snapshot from lazy collection for matching join keys
           const joinKeys = data.getInner().map(([[joinKey]]) => joinKey)
           const lazyJoinRef = new PropRef(followRefResult.path)
-          const loaded = lazyCollectionSubscription.requestSnapshot({
+          const loaded = lazySourceSubscription.requestSnapshot({
             where: inArray(lazyJoinRef, joinKeys),
             optimizedOnly: true,
           })
 
           if (!loaded) {
             // Snapshot wasn't sent because it could not be loaded from the indexes
-            lazyCollectionSubscription.requestSnapshot()
+            lazySourceSubscription.requestSnapshot()
           }
         })
       )
 
-      if (activeCollection === `main`) {
+      if (activeSource === `main`) {
         mainPipeline = activePipelineWithLoading
       } else {
         joinedPipeline = activePipelineWithLoading
@@ -290,68 +322,66 @@ function processJoin(
 
   return mainPipeline.pipe(
     joinOperator(joinedPipeline, joinClause.type as JoinType),
-    consolidate(),
     processJoinResults(joinClause.type)
   )
 }
 
 /**
- * Analyzes join expressions to determine which refers to which table
- * and returns them in the correct order (available table expression first, joined table expression second)
+ * Analyzes join expressions to determine which refers to which source
+ * and returns them in the correct order (available source expression first, joined source expression second)
  */
 function analyzeJoinExpressions(
   left: BasicExpression,
   right: BasicExpression,
-  allAvailableTableAliases: Array<string>,
-  joinedTableAlias: string
+  allAvailableSourceAliases: Array<string>,
+  joinedSource: string
 ): { mainExpr: BasicExpression; joinedExpr: BasicExpression } {
-  // Filter out the joined table alias from the available table aliases
-  const availableTableAliases = allAvailableTableAliases.filter(
-    (alias) => alias !== joinedTableAlias
+  // Filter out the joined source alias from the available source aliases
+  const availableSources = allAvailableSourceAliases.filter(
+    (alias) => alias !== joinedSource
   )
 
-  const leftTableAlias = getTableAliasFromExpression(left)
-  const rightTableAlias = getTableAliasFromExpression(right)
+  const leftSourceAlias = getSourceAliasFromExpression(left)
+  const rightSourceAlias = getSourceAliasFromExpression(right)
 
-  // If left expression refers to an available table and right refers to joined table, keep as is
+  // If left expression refers to an available source and right refers to joined source, keep as is
   if (
-    leftTableAlias &&
-    availableTableAliases.includes(leftTableAlias) &&
-    rightTableAlias === joinedTableAlias
+    leftSourceAlias &&
+    availableSources.includes(leftSourceAlias) &&
+    rightSourceAlias === joinedSource
   ) {
     return { mainExpr: left, joinedExpr: right }
   }
 
-  // If left expression refers to joined table and right refers to an available table, swap them
+  // If left expression refers to joined source and right refers to an available source, swap them
   if (
-    leftTableAlias === joinedTableAlias &&
-    rightTableAlias &&
-    availableTableAliases.includes(rightTableAlias)
+    leftSourceAlias === joinedSource &&
+    rightSourceAlias &&
+    availableSources.includes(rightSourceAlias)
   ) {
     return { mainExpr: right, joinedExpr: left }
   }
 
-  // If one expression doesn't refer to any table, this is an invalid join
-  if (!leftTableAlias || !rightTableAlias) {
-    // For backward compatibility, use the first available table alias in error message
-    throw new InvalidJoinConditionTableMismatchError()
+  // If one expression doesn't refer to any source, this is an invalid join
+  if (!leftSourceAlias || !rightSourceAlias) {
+    throw new InvalidJoinConditionSourceMismatchError()
   }
 
   // If both expressions refer to the same alias, this is an invalid join
-  if (leftTableAlias === rightTableAlias) {
-    throw new InvalidJoinConditionSameTableError(leftTableAlias)
+  if (leftSourceAlias === rightSourceAlias) {
+    throw new InvalidJoinConditionSameSourceError(leftSourceAlias)
   }
 
-  // Left side must refer to an available table
+  // Left side must refer to an available source
   // This cannot happen with the query builder as there is no way to build a ref
-  // to an unavailable table, but just in case, but could happen with the IR
-  if (!availableTableAliases.includes(leftTableAlias)) {
-    throw new InvalidJoinConditionLeftTableError(leftTableAlias)
+  // to an unavailable source, but just in case, but could happen with the IR
+  if (!availableSources.includes(leftSourceAlias)) {
+    throw new InvalidJoinConditionLeftSourceError(leftSourceAlias)
   }
 
-  // Right side must refer to the joined table
-  if (rightTableAlias !== joinedTableAlias) {
-    throw new InvalidJoinConditionRightTableError(joinedTableAlias)
+  // Right side must refer to the joined source
+  if (rightSourceAlias !== joinedSource) {
+    throw new InvalidJoinConditionRightSourceError(joinedSource)
   }
 
   // This should not be reachable given the logic above, but just in case
@@ -359,27 +389,27 @@ function analyzeJoinExpressions(
 }
 
 /**
- * Extracts the table alias from a join expression
+ * Extracts the source alias from a join expression
  */
-function getTableAliasFromExpression(expr: BasicExpression): string | null {
+function getSourceAliasFromExpression(expr: BasicExpression): string | null {
   switch (expr.type) {
     case `ref`:
-      // PropRef path has the table alias as the first element
+      // PropRef path has the source alias as the first element
       return expr.path[0] || null
     case `func`: {
-      // For function expressions, we need to check if all arguments refer to the same table
-      const tableAliases = new Set<string>()
+      // For function expressions, we need to check if all arguments refer to the same source
+      const sourceAliases = new Set<string>()
       for (const arg of expr.args) {
-        const alias = getTableAliasFromExpression(arg)
+        const alias = getSourceAliasFromExpression(arg)
         if (alias) {
-          tableAliases.add(alias)
+          sourceAliases.add(alias)
         }
       }
-      // If all arguments refer to the same table, return that table alias
-      return tableAliases.size === 1 ? Array.from(tableAliases)[0]! : null
+      // If all arguments refer to the same source, return that source alias
+      return sourceAliases.size === 1 ? Array.from(sourceAliases)[0]! : null
     }
     default:
-      // Values (type='val') don't reference any table
+      // Values (type='val') don't reference any source
       return null
   }
 }
@@ -393,18 +423,26 @@ function processJoinSource(
   collections: Record<string, Collection>,
   subscriptions: Record<string, CollectionSubscription>,
   callbacks: Record<string, LazyCollectionCallbacks>,
-  lazyCollections: Set<string>,
+  lazySources: Set<string>,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
+  setWindowFn: (windowFn: (options: WindowOptions) => void) => void,
   cache: QueryCache,
   queryMapping: QueryMapping,
-  onCompileSubquery: CompileQueryFn
+  onCompileSubquery: CompileQueryFn,
+  aliasToCollectionId: Record<string, string>,
+  aliasRemapping: Record<string, string>
 ): { alias: string; input: KeyedStream; collectionId: string } {
   switch (from.type) {
     case `collectionRef`: {
-      const input = allInputs[from.collection.id]
+      const input = allInputs[from.alias]
       if (!input) {
-        throw new CollectionInputNotFoundError(from.collection.id)
+        throw new CollectionInputNotFoundError(
+          from.alias,
+          from.collection.id,
+          Object.keys(allInputs)
+        )
       }
+      aliasToCollectionId[from.alias] = from.collection.id
       return { alias: from.alias, input, collectionId: from.collection.id }
     }
     case `queryRef`: {
@@ -418,12 +456,39 @@ function processJoinSource(
         collections,
         subscriptions,
         callbacks,
-        lazyCollections,
+        lazySources,
         optimizableOrderByCollections,
+        setWindowFn,
         cache,
         queryMapping
       )
 
+      // Pull up alias mappings from subquery to parent scope.
+      // This includes both the innermost alias-to-collection mappings AND
+      // any existing remappings from nested subquery levels.
+      Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId)
+      Object.assign(aliasRemapping, subQueryResult.aliasRemapping)
+
+      // Create a flattened remapping from outer alias to innermost alias.
+      // For nested subqueries, this ensures one-hop lookups (not recursive chains).
+      //
+      // Example with 3-level nesting:
+      //   Inner:  .from({ user: usersCollection })
+      //   Middle: .from({ activeUser: innerSubquery })     → creates: activeUser → user
+      //   Outer:  .join({ author: middleSubquery }, ...)   → creates: author → user (not author → activeUser)
+      //
+      // We search through the PULLED-UP aliasToCollectionId (which contains the
+      // innermost 'user' alias), so we always map directly to the deepest level.
+      // This means aliasRemapping[lazyAlias] is always a single lookup, never recursive.
+      const innerAlias = Object.keys(subQueryResult.aliasToCollectionId).find(
+        (alias) =>
+          subQueryResult.aliasToCollectionId[alias] ===
+          subQueryResult.collectionId
+      )
+      if (innerAlias && innerAlias !== from.alias) {
+        aliasRemapping[from.alias] = innerAlias
+      }
+
       // Extract the pipeline from the compilation result
       const subQueryInput = subQueryResult.pipeline
 
@@ -517,41 +582,35 @@ function processJoinResults(joinType: string) {
 /**
  * Returns the active and lazy collections for a join clause.
  * The active collection is the one that we need to fully iterate over
- * and it can be the main table (i.e. left collection) or the joined table (i.e. right collection).
+ * and it can be the main source (i.e. left collection) or the joined source (i.e. right collection).
  * The lazy collection is the one that we should join-in lazily based on matches in the active collection.
  * @param joinClause - The join clause to analyze
  * @param leftCollection - The left collection
  * @param rightCollection - The right collection
  * @returns The active and lazy collections. They are undefined if we need to loop over both collections (i.e. both are active)
  */
-function getActiveAndLazyCollections(
+function getActiveAndLazySources(
   joinType: JoinClause[`type`],
   leftCollection: Collection,
   rightCollection: Collection
 ):
-  | { activeCollection: `main` | `joined`; lazyCollection: Collection }
-  | { activeCollection: undefined; lazyCollection: undefined } {
-  if (leftCollection.id === rightCollection.id) {
-    // We can't apply this optimization if there's only one collection
-    // because `liveQueryCollection` will detect that the collection is lazy
-    // and treat it lazily (because the collection is shared)
-    // and thus it will not load any keys because both sides of the join
-    // will be handled lazily
-    return { activeCollection: undefined, lazyCollection: undefined }
-  }
+  | { activeSource: `main` | `joined`; lazySource: Collection }
+  | { activeSource: undefined; lazySource: undefined } {
+  // Self-joins can now be optimized since we track lazy loading by source alias
+  // rather than collection ID. Each alias has its own subscription and lazy state.
 
   switch (joinType) {
     case `left`:
-      return { activeCollection: `main`, lazyCollection: rightCollection }
+      return { activeSource: `main`, lazySource: rightCollection }
     case `right`:
-      return { activeCollection: `joined`, lazyCollection: leftCollection }
+      return { activeSource: `joined`, lazySource: leftCollection }
     case `inner`:
       // The smallest collection should be the active collection
       // and the biggest collection should be lazy
       return leftCollection.size < rightCollection.size
-        ? { activeCollection: `main`, lazyCollection: rightCollection }
-        : { activeCollection: `joined`, lazyCollection: leftCollection }
+        ? { activeSource: `main`, lazySource: rightCollection }
+        : { activeSource: `joined`, lazySource: leftCollection }
     default:
-      return { activeCollection: undefined, lazyCollection: undefined }
+      return { activeSource: undefined, lazySource: undefined }
   }
 }

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

@@ -5,6 +5,7 @@ import { ensureIndexForField } from "../../indexes/auto-index.js"
 import { findIndexForField } from "../../utils/index-optimization.js"
 import { compileExpression } from "./evaluators.js"
 import { replaceAggregatesByRefs } from "./group-by.js"
+import type { WindowOptions } from "./types.js"
 import type { CompiledSingleRowExpression } from "./evaluators.js"
 import type { OrderBy, OrderByClause, QueryIR, Select } from "../ir.js"
 import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
@@ -13,6 +14,7 @@ import type { IndexInterface } from "../../indexes/base-index.js"
 import type { Collection } from "../../collection/index.js"
 
 export type OrderByOptimizationInfo = {
+  alias: string
   orderBy: OrderBy
   offset: number
   limit: number
@@ -37,6 +39,7 @@ export function processOrderBy(
   selectClause: Select,
   collection: Collection,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
+  setWindowFn: (windowFn: (options: WindowOptions) => void) => void,
   limit?: number,
   offset?: number
 ): IStreamBuilder<KeyValue<unknown, [NamespacedRow, string]>> {
@@ -106,6 +109,8 @@ export function processOrderBy(
 
   let setSizeCallback: ((getSize: () => number) => void) | undefined
 
+  let orderByOptimizationInfo: OrderByOptimizationInfo | undefined
+
   // Optimize the orderBy operator to lazily load elements
   // by using the range index of the collection.
   // Only for orderBy clause on a single column for now (no composite ordering)
@@ -155,7 +160,13 @@ export function processOrderBy(
 
       if (index && index.supports(`gt`)) {
         // We found an index that we can use to lazily load ordered data
-        const orderByOptimizationInfo = {
+        const orderByAlias =
+          orderByExpression.path.length > 1
+            ? String(orderByExpression.path[0])
+            : rawQuery.from.alias
+
+        orderByOptimizationInfo = {
+          alias: orderByAlias,
           offset: offset ?? 0,
           limit,
           comparator,
@@ -172,7 +183,7 @@ export function processOrderBy(
             ...optimizableOrderByCollections[followRefCollection.id]!,
             dataNeeded: () => {
               const size = getSize()
-              return Math.max(0, limit - size)
+              return Math.max(0, orderByOptimizationInfo!.limit - size)
             },
           }
         }
@@ -187,6 +198,23 @@ export function processOrderBy(
       offset,
       comparator: compare,
       setSizeCallback,
+      setWindowFn: (
+        windowFn: (options: { offset?: number; limit?: number }) => void
+      ) => {
+        setWindowFn(
+          // We wrap the move function such that we update the orderByOptimizationInfo
+          // because that is used by the `dataNeeded` callback to determine if we need to load more data
+          (options) => {
+            windowFn(options)
+            if (orderByOptimizationInfo) {
+              orderByOptimizationInfo.offset =
+                options.offset ?? orderByOptimizationInfo.offset
+              orderByOptimizationInfo.limit =
+                options.limit ?? orderByOptimizationInfo.limit
+            }
+          }
+        )
+      },
     })
     // orderByWithFractionalIndex returns [key, [value, index]] - we keep this format
   )