Support ordering grouped results based on aggregates (#481)

Author: kevin-dp

.changeset/lemon-queens-pump.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+Add support for queries to order results based on aggregated values

packages/db/src/query/compiler/group-by.ts

@@ -130,7 +130,7 @@ export function processGroupBy(
     if (havingClauses && havingClauses.length > 0) {
       for (const havingClause of havingClauses) {
         const havingExpression = getHavingExpression(havingClause)
-        const transformedHavingClause = transformHavingClause(
+        const transformedHavingClause = replaceAggregatesByRefs(
           havingExpression,
           selectClause || {}
         )
@@ -265,7 +265,7 @@ export function processGroupBy(
   if (havingClauses && havingClauses.length > 0) {
     for (const havingClause of havingClauses) {
       const havingExpression = getHavingExpression(havingClause)
-      const transformedHavingClause = transformHavingClause(
+      const transformedHavingClause = replaceAggregatesByRefs(
         havingExpression,
         selectClause || {}
       )
@@ -367,11 +367,12 @@ function getAggregateFunction(aggExpr: Aggregate) {
 }
 
 /**
- * Transforms a HAVING clause to replace Agg expressions with references to computed values
+ * Transforms basic expressions and aggregates to replace Agg expressions with references to computed values
  */
-function transformHavingClause(
+export function replaceAggregatesByRefs(
   havingExpr: BasicExpression | Aggregate,
-  selectClause: Select
+  selectClause: Select,
+  resultAlias: string = `result`
 ): BasicExpression {
   switch (havingExpr.type) {
     case `agg`: {
@@ -380,7 +381,7 @@ function transformHavingClause(
       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])
+          return new PropRef([resultAlias, alias])
         }
       }
       // If no matching aggregate found in SELECT, throw error
@@ -392,7 +393,7 @@ function transformHavingClause(
       // Transform function arguments recursively
       const transformedArgs = funcExpr.args.map(
         (arg: BasicExpression | Aggregate) =>
-          transformHavingClause(arg, selectClause)
+          replaceAggregatesByRefs(arg, selectClause)
       )
       return new Func(funcExpr.name, transformedArgs)
     }
@@ -404,7 +405,7 @@ function transformHavingClause(
         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 new PropRef([resultAlias, alias])
         }
       }
       // Return as-is for other refs

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

@@ -259,6 +259,7 @@ export function compileQuery(
       rawQuery,
       pipeline,
       query.orderBy,
+      query.select || {},
       collections[mainCollectionId]!,
       optimizableOrderByCollections,
       query.limit,

packages/db/src/query/compiler/order-by.ts

@@ -4,9 +4,10 @@ import { PropRef } from "../ir.js"
 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 { followRef } from "./index.js"
 import type { CompiledSingleRowExpression } from "./evaluators.js"
-import type { OrderByClause, QueryIR } from "../ir.js"
+import type { OrderByClause, QueryIR, Select } from "../ir.js"
 import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
 import type { IStreamBuilder, KeyValue } from "@tanstack/db-ivm"
 import type { BaseIndex } from "../../indexes/base-index.js"
@@ -33,16 +34,24 @@ export function processOrderBy(
   rawQuery: QueryIR,
   pipeline: NamespacedAndKeyedStream,
   orderByClause: Array<OrderByClause>,
+  selectClause: Select,
   collection: Collection,
   optimizableOrderByCollections: Record<string, OrderByOptimizationInfo>,
   limit?: number,
   offset?: number
 ): IStreamBuilder<KeyValue<unknown, [NamespacedRow, string]>> {
   // Pre-compile all order by expressions
-  const compiledOrderBy = orderByClause.map((clause) => ({
-    compiledExpression: compileExpression(clause.expression),
-    compareOptions: clause.compareOptions,
-  }))
+  const compiledOrderBy = orderByClause.map((clause) => {
+    const clauseWithoutAggregates = replaceAggregatesByRefs(
+      clause.expression,
+      selectClause,
+      `__select_results`
+    )
+    return {
+      compiledExpression: compileExpression(clauseWithoutAggregates),
+      compareOptions: clause.compareOptions,
+    }
+  })
 
   // Create a value extractor function for the orderBy operator
   const valueExtractor = (row: NamespacedRow & { __select_results?: any }) => {

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

@@ -2,7 +2,7 @@ import { beforeEach, describe, expect, it } from "vitest"
 import { createCollection } from "../../src/collection.js"
 import { mockSyncCollectionOptions } from "../utls.js"
 import { createLiveQueryCollection } from "../../src/query/live-query-collection.js"
-import { eq, gt } from "../../src/query/builder/functions.js"
+import { eq, gt, max } from "../../src/query/builder/functions.js"
 
 type Person = {
   id: string
@@ -711,6 +711,126 @@ function createOrderByTests(autoIndex: `off` | `eager`): void {
       })
     })
 
+    describe(`OrderBy with GroupBy`, () => {
+      it(`should order grouped results correctly`, async () => {
+        type VehicleDocument = {
+          id: number
+          vin: string
+          updatedAt: number
+        }
+
+        const vehicleDocumentsData = [
+          { id: 1, vin: `1`, updatedAt: new Date(`2023-01-01`).getTime() },
+          { id: 2, vin: `2`, updatedAt: new Date(`2023-01-02`).getTime() },
+          { id: 3, vin: `1`, updatedAt: new Date(`2023-01-05`).getTime() },
+        ]
+
+        const vehicleDocumentCollection = createCollection(
+          mockSyncCollectionOptions<VehicleDocument>({
+            id: `vehicle-document-collection`,
+            getKey: (doc) => doc.id,
+            autoIndex: `eager`,
+            initialData: vehicleDocumentsData,
+          })
+        )
+
+        const liveQuery = createLiveQueryCollection({
+          query: (q) =>
+            q
+              .from({ vehicleDocuments: vehicleDocumentCollection })
+              .groupBy((q) => q.vehicleDocuments.vin)
+              .orderBy((q) => q.vehicleDocuments.vin, `asc`)
+              .select((q) => ({
+                vin: q.vehicleDocuments.vin,
+              })),
+          startSync: true,
+        })
+
+        await liveQuery.stateWhenReady()
+        expect(liveQuery.toArray).toEqual([{ vin: `1` }, { vin: `2` }])
+
+        // Insert a vehicle document
+        vehicleDocumentCollection.utils.begin()
+        vehicleDocumentCollection.utils.write({
+          type: `insert`,
+          value: {
+            id: 4,
+            vin: `3`,
+            updatedAt: new Date(`2023-01-03`).getTime(),
+          },
+        })
+        vehicleDocumentCollection.utils.commit()
+
+        expect(liveQuery.toArray).toEqual([
+          { vin: `1` },
+          { vin: `2` },
+          { vin: `3` },
+        ])
+      })
+
+      it(`should order groups based on aggregates correctly`, async () => {
+        type VehicleDocument = {
+          id: number
+          vin: string
+          updatedAt: number
+        }
+
+        const vehicleDocumentsData = [
+          { id: 1, vin: `1`, updatedAt: new Date(`2023-01-01`).getTime() },
+          { id: 2, vin: `2`, updatedAt: new Date(`2023-01-02`).getTime() },
+          { id: 3, vin: `1`, updatedAt: new Date(`2023-01-05`).getTime() },
+        ]
+
+        const vehicleDocumentCollection = createCollection(
+          mockSyncCollectionOptions<VehicleDocument>({
+            id: `vehicle-document-collection`,
+            getKey: (doc) => doc.id,
+            autoIndex: `eager`,
+            initialData: vehicleDocumentsData,
+          })
+        )
+
+        const liveQuery = createLiveQueryCollection({
+          query: (q) =>
+            q
+              .from({ vehicleDocuments: vehicleDocumentCollection })
+              .groupBy((q) => q.vehicleDocuments.vin)
+              .orderBy((q) => max(q.vehicleDocuments.updatedAt), `desc`)
+              .select((q) => ({
+                vin: q.vehicleDocuments.vin,
+                updatedAt: max(q.vehicleDocuments.updatedAt),
+              }))
+              .offset(0)
+              .limit(10),
+          startSync: true,
+        })
+
+        await liveQuery.stateWhenReady()
+        expect(liveQuery.toArray).toEqual([
+          { vin: `1`, updatedAt: new Date(`2023-01-05`).getTime() },
+          { vin: `2`, updatedAt: new Date(`2023-01-02`).getTime() },
+        ])
+
+        // Insert a vehicle document
+        vehicleDocumentCollection.utils.begin()
+        vehicleDocumentCollection.utils.write({
+          type: `insert`,
+          value: {
+            id: 4,
+            vin: `3`,
+            updatedAt: new Date(`2023-01-03`).getTime(),
+          },
+        })
+        vehicleDocumentCollection.utils.commit()
+
+        expect(liveQuery.toArray).toEqual([
+          { vin: `1`, updatedAt: new Date(`2023-01-05`).getTime() },
+          { vin: `3`, updatedAt: new Date(`2023-01-03`).getTime() },
+          { vin: `2`, updatedAt: new Date(`2023-01-02`).getTime() },
+        ])
+      })
+    })
+
     describe(`OrderBy with Where Clauses`, () => {
       it(`orders filtered results correctly`, async () => {
         const collection = createLiveQueryCollection((q) =>