Configuration options for sortBy (#314)

Author: kevin-dp

.changeset/late-icons-fly.md

@@ -0,0 +1,5 @@
+---
+"@tanstack/db": patch
+---
+
+Add option to configure how orderBy compares values. This includes ascending/descending order, ordering of null values, and lexical vs locale comparison for strings.

packages/db/src/indexes/btree-index.ts

@@ -1,5 +1,5 @@
-import { ascComparator } from "../utils/comparison.js"
 import { BTree } from "../utils/btree.js"
+import { defaultComparator } from "../utils/comparison.js"
 import { BaseIndex } from "./base-index.js"
 import type { BasicExpression } from "../query/ir.js"
 import type { IndexOperation } from "./base-index.js"
@@ -43,7 +43,7 @@ export class BTreeIndex<
   private orderedEntries: BTree<any, undefined> // we don't associate values with the keys of the B+ tree (the keys are indexed values)
   private valueMap = new Map<any, Set<TKey>>() // instead we store a mapping of indexed values to a set of PKs
   private indexedKeys = new Set<TKey>()
-  private compareFn: (a: any, b: any) => number = ascComparator
+  private compareFn: (a: any, b: any) => number = defaultComparator
 
   constructor(
     id: number,
@@ -52,7 +52,7 @@ export class BTreeIndex<
     options?: any
   ) {
     super(id, expression, name, options)
-    this.compareFn = options?.compareFn ?? ascComparator
+    this.compareFn = options?.compareFn ?? defaultComparator
     this.orderedEntries = new BTree(this.compareFn)
   }
 

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

@@ -19,12 +19,14 @@ import type {
   QueryIR,
 } from "../ir.js"
 import type {
+  CompareOptions,
   Context,
   GroupByCallback,
   JoinOnCallback,
   MergeContext,
   MergeContextWithJoinType,
   OrderByCallback,
+  OrderByOptions,
   RefProxyForContext,
   ResultTypeFromSelect,
   SchemaFromSource,
@@ -478,16 +480,31 @@ export class BaseQueryBuilder<TContext extends Context = Context> {
    */
   orderBy(
     callback: OrderByCallback<TContext>,
-    direction: OrderByDirection = `asc`
+    options: OrderByDirection | OrderByOptions = `asc`
   ): QueryBuilder<TContext> {
     const aliases = this._getCurrentAliases()
     const refProxy = createRefProxy(aliases) as RefProxyForContext<TContext>
     const result = callback(refProxy)
 
+    const opts: CompareOptions =
+      typeof options === `string`
+        ? { direction: options, nulls: `first`, stringSort: `locale` }
+        : {
+            direction: options.direction ?? `asc`,
+            nulls: options.nulls ?? `first`,
+            stringSort: options.stringSort ?? `locale`,
+            locale:
+              options.stringSort === `locale` ? options.locale : undefined,
+            localeOptions:
+              options.stringSort === `locale`
+                ? options.localeOptions
+                : undefined,
+          }
+
     // Create the new OrderBy structure with expression and direction
     const orderByClause: OrderByClause = {
       expression: toExpression(result),
-      direction,
+      compareOptions: opts,
     }
 
     const existingOrderBy: OrderBy = this.query.orderBy || []

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

@@ -1,5 +1,5 @@
 import type { CollectionImpl } from "../../collection.js"
-import type { Aggregate, BasicExpression } from "../ir.js"
+import type { Aggregate, BasicExpression, OrderByDirection } from "../ir.js"
 import type { QueryBuilder } from "./index.js"
 
 export interface Context {
@@ -75,6 +75,29 @@ export type OrderByCallback<TContext extends Context> = (
   refs: RefProxyForContext<TContext>
 ) => any
 
+export type OrderByOptions = {
+  direction?: OrderByDirection
+  nulls?: `first` | `last`
+} & StringSortOpts
+
+export type StringSortOpts =
+  | {
+      stringSort?: `lexical`
+    }
+  | {
+      stringSort?: `locale`
+      locale?: string
+      localeOptions?: object
+    }
+
+export type CompareOptions = {
+  direction: OrderByDirection
+  nulls: `first` | `last`
+  stringSort: `lexical` | `locale`
+  locale?: string
+  localeOptions?: object
+}
+
 // Callback type for groupBy clauses
 export type GroupByCallback<TContext extends Context> = (
   refs: RefProxyForContext<TContext>

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

@@ -1,5 +1,5 @@
 import { orderByWithFractionalIndex } from "@tanstack/db-ivm"
-import { ascComparator, descComparator } from "../../utils/comparison.js"
+import { defaultComparator, makeComparator } from "../../utils/comparison.js"
 import { compileExpression } from "./evaluators.js"
 import type { OrderByClause } from "../ir.js"
 import type { NamespacedAndKeyedStream, NamespacedRow } from "../../types.js"
@@ -19,7 +19,7 @@ export function processOrderBy(
   // Pre-compile all order by expressions
   const compiledOrderBy = orderByClause.map((clause) => ({
     compiledExpression: compileExpression(clause.expression),
-    direction: clause.direction,
+    compareOptions: clause.compareOptions,
   }))
 
   // Create a value extractor function for the orderBy operator
@@ -53,35 +53,31 @@ export function processOrderBy(
   }
 
   // Create a multi-property comparator that respects the order and direction of each property
-  const makeComparator = () => {
-    return (a: unknown, b: unknown) => {
-      // If we're comparing arrays (multiple properties), compare each property in order
-      if (orderByClause.length > 1) {
-        const arrayA = a as Array<unknown>
-        const arrayB = b as Array<unknown>
-        for (let i = 0; i < orderByClause.length; i++) {
-          const direction = orderByClause[i]!.direction
-          const compareFn =
-            direction === `desc` ? descComparator : ascComparator
-          const result = compareFn(arrayA[i], arrayB[i])
-          if (result !== 0) {
-            return result
-          }
+  const comparator = (a: unknown, b: unknown) => {
+    // If we're comparing arrays (multiple properties), compare each property in order
+    if (orderByClause.length > 1) {
+      const arrayA = a as Array<unknown>
+      const arrayB = b as Array<unknown>
+      for (let i = 0; i < orderByClause.length; i++) {
+        const clause = orderByClause[i]!
+        const compareFn = makeComparator(clause.compareOptions)
+        const result = compareFn(arrayA[i], arrayB[i])
+        if (result !== 0) {
+          return result
         }
-        return arrayA.length - arrayB.length
-      }
-
-      // Single property comparison
-      if (orderByClause.length === 1) {
-        const direction = orderByClause[0]!.direction
-        return direction === `desc` ? descComparator(a, b) : ascComparator(a, b)
       }
+      return arrayA.length - arrayB.length
+    }
 
-      return ascComparator(a, b)
+    // Single property comparison
+    if (orderByClause.length === 1) {
+      const clause = orderByClause[0]!
+      const compareFn = makeComparator(clause.compareOptions)
+      return compareFn(a, b)
     }
-  }
 
-  const comparator = makeComparator()
+    return defaultComparator(a, b)
+  }
 
   // Use fractional indexing and return the tuple [value, index]
   return pipeline.pipe(

packages/db/src/query/ir.ts

@@ -2,6 +2,7 @@
 This is the intermediate representation of the query.
 */
 
+import type { CompareOptions } from "./builder/types"
 import type { CollectionImpl } from "../collection"
 import type { NamespacedRow } from "../types"
 
@@ -48,7 +49,7 @@ export type OrderBy = Array<OrderByClause>
 
 export type OrderByClause = {
   expression: BasicExpression
-  direction: OrderByDirection
+  compareOptions: CompareOptions
 }
 
 export type OrderByDirection = `asc` | `desc`

packages/db/src/utils/comparison.ts

@@ -1,3 +1,5 @@
+import type { CompareOptions } from "../query/builder/types"
+
 // WeakMap to store stable IDs for objects
 const objectIds = new WeakMap<object, number>()
 let nextObjectId = 1
@@ -19,21 +21,26 @@ function getObjectId(obj: object): number {
  * Handles null/undefined, strings, arrays, dates, objects, and primitives
  * Always sorts null/undefined values first
  */
-export const ascComparator = (a: any, b: any): number => {
+export const ascComparator = (a: any, b: any, opts: CompareOptions): number => {
+  const { nulls } = opts
+
   // Handle null/undefined
   if (a == null && b == null) return 0
-  if (a == null) return -1
-  if (b == null) return 1
+  if (a == null) return nulls === `first` ? -1 : 1
+  if (b == null) return nulls === `first` ? 1 : -1
 
   // if a and b are both strings, compare them based on locale
   if (typeof a === `string` && typeof b === `string`) {
-    return a.localeCompare(b)
+    if (opts.stringSort === `locale`) {
+      return a.localeCompare(b, opts.locale, opts.localeOptions)
+    }
+    // For lexical sort we rely on direct comparison for primitive values
   }
 
   // if a and b are both arrays, compare them element by element
   if (Array.isArray(a) && Array.isArray(b)) {
     for (let i = 0; i < Math.min(a.length, b.length); i++) {
-      const result = ascComparator(a[i], b[i])
+      const result = ascComparator(a[i], b[i], opts)
       if (result !== 0) {
         return result
       }
@@ -74,6 +81,32 @@ export const ascComparator = (a: any, b: any): number => {
  * Descending comparator function for ordering values
  * Handles null/undefined as largest values (opposite of ascending)
  */
-export const descComparator = (a: unknown, b: unknown): number => {
-  return ascComparator(b, a)
+export const descComparator = (
+  a: unknown,
+  b: unknown,
+  opts: CompareOptions
+): number => {
+  return ascComparator(b, a, {
+    ...opts,
+    nulls: opts.nulls === `first` ? `last` : `first`,
+  })
+}
+
+export function makeComparator(
+  opts: CompareOptions
+): (a: any, b: any) => number {
+  return (a, b) => {
+    if (opts.direction === `asc`) {
+      return ascComparator(a, b, opts)
+    } else {
+      return descComparator(a, b, opts)
+    }
+  }
 }
+
+/** Default comparator orders values in ascending order with nulls first and locale string comparison. */
+export const defaultComparator = makeComparator({
+  direction: `asc`,
+  nulls: `first`,
+  stringSort: `locale`,
+})