Also optimize orderBy for singleton array (#477)

Author: kevin-dp

.changeset/afraid-plums-clean.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+Fixed an optimization bug where orderBy clauses using a single-column array were not recognized as optimizable. Queries that order by a single column are now correctly optimized even when specified as an array.

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

@@ -14,7 +14,6 @@ import type {
   BasicExpression,
   JoinClause,
   OrderBy,
-  OrderByClause,
   OrderByDirection,
   QueryIR,
 } from "../ir.js"
@@ -501,17 +500,23 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
                 : undefined,
           }
 
-    // Create the new OrderBy structure with expression and direction
-    const orderByClause: OrderByClause = {
-      expression: toExpression(result),
-      compareOptions: opts,
+    const makeOrderByClause = (res: any) => {
+      return {
+        expression: toExpression(res),
+        compareOptions: opts,
+      }
     }
 
+    // Create the new OrderBy structure with expression and direction
+    const orderByClauses = Array.isArray(result)
+      ? result.map((r) => makeOrderByClause(r))
+      : [makeOrderByClause(result)]
+
     const existingOrderBy: OrderBy = this.query.orderBy || []
 
     return new BaseQueryBuilder({
       ...this.query,
-      orderBy: [...existingOrderBy, orderByClause],
+      orderBy: [...existingOrderBy, ...orderByClauses],
     }) as any
   }
 

packages/db/tests/query/order-by.test.ts

@@ -1241,6 +1241,92 @@ function createOrderByTests(autoIndex: `off` | `eager`): void {
       })
     })
 
+    describe(`OrderBy Optimization Tests`, () => {
+      const itWhenAutoIndex = autoIndex === `eager` ? it : it.skip
+
+      itWhenAutoIndex(
+        `optimizes single-column orderBy when passed as single value`,
+        async () => {
+          // Patch getConfig to expose the builder on the returned config for test access
+          const { CollectionConfigBuilder } = await import(
+            `../../src/query/live/collection-config-builder.js`
+          )
+          const originalGetConfig = CollectionConfigBuilder.prototype.getConfig
+
+          CollectionConfigBuilder.prototype.getConfig = function (this: any) {
+            const cfg = originalGetConfig.call(this)
+            ;(cfg as any).__builder = this
+            return cfg
+          }
+
+          try {
+            const collection = createLiveQueryCollection((q) =>
+              q
+                .from({ employees: employeesCollection })
+                .orderBy(({ employees }) => employees.salary, `desc`)
+                .limit(3)
+                .select(({ employees }) => ({
+                  id: employees.id,
+                  name: employees.name,
+                  salary: employees.salary,
+                }))
+            )
+
+            await collection.preload()
+
+            const builder = (collection as any).config.__builder
+            expect(builder).toBeTruthy()
+            expect(
+              Object.keys(builder.optimizableOrderByCollections)
+            ).toContain(employeesCollection.id)
+          } finally {
+            CollectionConfigBuilder.prototype.getConfig = originalGetConfig
+          }
+        }
+      )
+
+      itWhenAutoIndex(
+        `optimizes single-column orderBy when passed as array with single element`,
+        async () => {
+          // Patch getConfig to expose the builder on the returned config for test access
+          const { CollectionConfigBuilder } = await import(
+            `../../src/query/live/collection-config-builder.js`
+          )
+          const originalGetConfig = CollectionConfigBuilder.prototype.getConfig
+
+          CollectionConfigBuilder.prototype.getConfig = function (this: any) {
+            const cfg = originalGetConfig.call(this)
+            ;(cfg as any).__builder = this
+            return cfg
+          }
+
+          try {
+            const collection = createLiveQueryCollection((q) =>
+              q
+                .from({ employees: employeesCollection })
+                .orderBy(({ employees }) => [employees.salary], `desc`)
+                .limit(3)
+                .select(({ employees }) => ({
+                  id: employees.id,
+                  name: employees.name,
+                  salary: employees.salary,
+                }))
+            )
+
+            await collection.preload()
+
+            const builder = (collection as any).config.__builder
+            expect(builder).toBeTruthy()
+            expect(
+              Object.keys(builder.optimizableOrderByCollections)
+            ).toContain(employeesCollection.id)
+          } finally {
+            CollectionConfigBuilder.prototype.getConfig = originalGetConfig
+          }
+        }
+      )
+    })
+
     describe(`String Comparison Tests`, () => {
       it(`handles case differently in lexical vs locale string comparison`, async () => {
         const numericEmployees = [