Fix: Aggregate query order-by normalization (#7507)

MarkDuckworth · web-flow · commit b395277f3de5 · 2023-08-04T08:53:28.000-06:00
Fix aggregate query order-by normalization to support future aggregate operations.
diff --git a/.changeset/quiet-gorillas-smoke.md b/.changeset/quiet-gorillas-smoke.md
@@ -0,0 +1,5 @@
+---
+"@firebase/firestore": patch
+---
+
+Refactored aggregate query order-by normalization to support future aggregate operations.
diff --git a/packages/firestore/src/core/query.ts b/packages/firestore/src/core/query.ts
@@ -43,7 +43,7 @@ export const enum LimitType {
 /**
  * The Query interface defines all external properties of a query.
  *
- * QueryImpl implements this interface to provide memoization for `queryOrderBy`
+ * QueryImpl implements this interface to provide memoization for `queryNormalizedOrderBy`
  * and `queryToTarget`.
  */
 export interface Query {
@@ -65,11 +65,18 @@ export interface Query {
  * Visible for testing.
  */
 export class QueryImpl implements Query {
-  memoizedOrderBy: OrderBy[] | null = null;
+  memoizedNormalizedOrderBy: OrderBy[] | null = null;
 
-  // The corresponding `Target` of this `Query` instance.
+  // The corresponding `Target` of this `Query` instance, for use with
+  // non-aggregate queries.
   memoizedTarget: Target | null = null;
 
+  // The corresponding `Target` of this `Query` instance, for use with
+  // aggregate queries. Unlike targets for non-aggregate queries,
+  // aggregate query targets do not contain normalized order-bys, they only
+  // contain explicit order-bys.
+  memoizedAggregateTarget: Target | null = null;
+
   /**
    * Initializes a Query with a path and optional additional query constraints.
    * Path must currently be empty if this is a collection group query.
@@ -86,13 +93,13 @@ export class QueryImpl implements Query {
   ) {
     if (this.startAt) {
       debugAssert(
-        this.startAt.position.length <= queryOrderBy(this).length,
+        this.startAt.position.length <= queryNormalizedOrderBy(this).length,
         'Bound is longer than orderBy'
       );
     }
     if (this.endAt) {
       debugAssert(
-        this.endAt.position.length <= queryOrderBy(this).length,
+        this.endAt.position.length <= queryNormalizedOrderBy(this).length,
         'Bound is longer than orderBy'
       );
     }
@@ -211,14 +218,16 @@ export function isCollectionGroupQuery(query: Query): boolean {
 }
 
 /**
- * Returns the implicit order by constraint that is used to execute the Query,
- * which can be different from the order by constraints the user provided (e.g.
- * the SDK and backend always orders by `__name__`).
+ * Returns the normalized order-by constraint that is used to execute the Query,
+ * which can be different from the order-by constraints the user provided (e.g.
+ * the SDK and backend always orders by `__name__`). The normalized order-by
+ * includes implicit order-bys in addition to the explicit user provided
+ * order-bys.
  */
-export function queryOrderBy(query: Query): OrderBy[] {
+export function queryNormalizedOrderBy(query: Query): OrderBy[] {
   const queryImpl = debugCast(query, QueryImpl);
-  if (queryImpl.memoizedOrderBy === null) {
-    queryImpl.memoizedOrderBy = [];
+  if (queryImpl.memoizedNormalizedOrderBy === null) {
+    queryImpl.memoizedNormalizedOrderBy = [];
 
     const inequalityField = getInequalityFilterField(queryImpl);
     const firstOrderByField = getFirstOrderByField(queryImpl);
@@ -227,9 +236,9 @@ export function queryOrderBy(query: Query): OrderBy[] {
       // inequality filter field for it to be a valid query.
       // Note that the default inequality field and key ordering is ascending.
       if (!inequalityField.isKeyField()) {
-        queryImpl.memoizedOrderBy.push(new OrderBy(inequalityField));
+        queryImpl.memoizedNormalizedOrderBy.push(new OrderBy(inequalityField));
       }
-      queryImpl.memoizedOrderBy.push(
+      queryImpl.memoizedNormalizedOrderBy.push(
         new OrderBy(FieldPath.keyField(), Direction.ASCENDING)
       );
     } else {
@@ -241,76 +250,103 @@ export function queryOrderBy(query: Query): OrderBy[] {
       );
       let foundKeyOrdering = false;
       for (const orderBy of queryImpl.explicitOrderBy) {
-        queryImpl.memoizedOrderBy.push(orderBy);
+        queryImpl.memoizedNormalizedOrderBy.push(orderBy);
         if (orderBy.field.isKeyField()) {
           foundKeyOrdering = true;
         }
       }
       if (!foundKeyOrdering) {
         // The order of the implicit key ordering always matches the last
-        // explicit order by
+        // explicit order-by
         const lastDirection =
           queryImpl.explicitOrderBy.length > 0
             ? queryImpl.explicitOrderBy[queryImpl.explicitOrderBy.length - 1]
                 .dir
             : Direction.ASCENDING;
-        queryImpl.memoizedOrderBy.push(
+        queryImpl.memoizedNormalizedOrderBy.push(
           new OrderBy(FieldPath.keyField(), lastDirection)
         );
       }
     }
   }
-  return queryImpl.memoizedOrderBy;
+  return queryImpl.memoizedNormalizedOrderBy;
 }
 
 /**
- * Converts this `Query` instance to it's corresponding `Target` representation.
+ * Converts this `Query` instance to its corresponding `Target` representation.
  */
 export function queryToTarget(query: Query): Target {
   const queryImpl = debugCast(query, QueryImpl);
   if (!queryImpl.memoizedTarget) {
-    if (queryImpl.limitType === LimitType.First) {
-      queryImpl.memoizedTarget = newTarget(
-        queryImpl.path,
-        queryImpl.collectionGroup,
-        queryOrderBy(queryImpl),
-        queryImpl.filters,
-        queryImpl.limit,
-        queryImpl.startAt,
-        queryImpl.endAt
-      );
-    } else {
-      // Flip the orderBy directions since we want the last results
-      const orderBys = [] as OrderBy[];
-      for (const orderBy of queryOrderBy(queryImpl)) {
-        const dir =
-          orderBy.dir === Direction.DESCENDING
-            ? Direction.ASCENDING
-            : Direction.DESCENDING;
-        orderBys.push(new OrderBy(orderBy.field, dir));
-      }
+    queryImpl.memoizedTarget = _queryToTarget(
+      queryImpl,
+      queryNormalizedOrderBy(query)
+    );
+  }
 
-      // We need to swap the cursors to match the now-flipped query ordering.
-      const startAt = queryImpl.endAt
-        ? new Bound(queryImpl.endAt.position, queryImpl.endAt.inclusive)
-        : null;
-      const endAt = queryImpl.startAt
-        ? new Bound(queryImpl.startAt.position, queryImpl.startAt.inclusive)
-        : null;
-
-      // Now return as a LimitType.First query.
-      queryImpl.memoizedTarget = newTarget(
-        queryImpl.path,
-        queryImpl.collectionGroup,
-        orderBys,
-        queryImpl.filters,
-        queryImpl.limit,
-        startAt,
-        endAt
-      );
-    }
+  return queryImpl.memoizedTarget;
+}
+
+/**
+ * Converts this `Query` instance to its corresponding `Target` representation,
+ * for use within an aggregate query. Unlike targets for non-aggregate queries,
+ * aggregate query targets do not contain normalized order-bys, they only
+ * contain explicit order-bys.
+ */
+export function queryToAggregateTarget(query: Query): Target {
+  const queryImpl = debugCast(query, QueryImpl);
+
+  if (!queryImpl.memoizedAggregateTarget) {
+    // Do not include implicit order-bys for aggregate queries.
+    queryImpl.memoizedAggregateTarget = _queryToTarget(
+      queryImpl,
+      query.explicitOrderBy
+    );
+  }
+
+  return queryImpl.memoizedAggregateTarget;
+}
+
+function _queryToTarget(queryImpl: QueryImpl, orderBys: OrderBy[]): Target {
+  if (queryImpl.limitType === LimitType.First) {
+    return newTarget(
+      queryImpl.path,
+      queryImpl.collectionGroup,
+      orderBys,
+      queryImpl.filters,
+      queryImpl.limit,
+      queryImpl.startAt,
+      queryImpl.endAt
+    );
+  } else {
+    // Flip the orderBy directions since we want the last results
+    orderBys = orderBys.map(orderBy => {
+      const dir =
+        orderBy.dir === Direction.DESCENDING
+          ? Direction.ASCENDING
+          : Direction.DESCENDING;
+      return new OrderBy(orderBy.field, dir);
+    });
+
+    // We need to swap the cursors to match the now-flipped query ordering.
+    const startAt = queryImpl.endAt
+      ? new Bound(queryImpl.endAt.position, queryImpl.endAt.inclusive)
+      : null;
+    const endAt = queryImpl.startAt
+      ? new Bound(queryImpl.startAt.position, queryImpl.startAt.inclusive)
+      : null;
+
+    // Now return as a LimitType.First query.
+    return newTarget(
+      queryImpl.path,
+      queryImpl.collectionGroup,
+      orderBys,
+      queryImpl.filters,
+      queryImpl.limit,
+      startAt,
+      endAt
+    );
   }
-  return queryImpl.memoizedTarget!;
 }
 
 export function queryWithAddedFilter(query: Query, filter: Filter): Query {
@@ -461,14 +497,14 @@ function queryMatchesPathAndCollectionGroup(
  * in the results.
  */
 function queryMatchesOrderBy(query: Query, doc: Document): boolean {
-  // We must use `queryOrderBy()` to get the list of all orderBys (both implicit and explicit).
+  // We must use `queryNormalizedOrderBy()` to get the list of all orderBys (both implicit and explicit).
   // Note that for OR queries, orderBy applies to all disjunction terms and implicit orderBys must
   // be taken into account. For example, the query "a > 1 || b==1" has an implicit "orderBy a" due
   // to the inequality, and is evaluated as "a > 1 orderBy a || b==1 orderBy a".
   // A document with content of {b:1} matches the filters, but does not match the orderBy because
   // it's missing the field 'a'.
-  for (const orderBy of queryOrderBy(query)) {
-    // order by key always matches
+  for (const orderBy of queryNormalizedOrderBy(query)) {
+    // order-by key always matches
     if (!orderBy.field.isKeyField() && doc.data.field(orderBy.field) === null) {
       return false;
     }
@@ -489,13 +525,13 @@ function queryMatchesFilters(query: Query, doc: Document): boolean {
 function queryMatchesBounds(query: Query, doc: Document): boolean {
   if (
     query.startAt &&
-    !boundSortsBeforeDocument(query.startAt, queryOrderBy(query), doc)
+    !boundSortsBeforeDocument(query.startAt, queryNormalizedOrderBy(query), doc)
   ) {
     return false;
   }
   if (
     query.endAt &&
-    !boundSortsAfterDocument(query.endAt, queryOrderBy(query), doc)
+    !boundSortsAfterDocument(query.endAt, queryNormalizedOrderBy(query), doc)
   ) {
     return false;
   }
@@ -526,7 +562,7 @@ export function newQueryComparator(
 ): (d1: Document, d2: Document) => number {
   return (d1: Document, d2: Document): number => {
     let comparedOnKeyField = false;
-    for (const orderBy of queryOrderBy(query)) {
+    for (const orderBy of queryNormalizedOrderBy(query)) {
       const comp = compareDocs(orderBy, d1, d2);
       if (comp !== 0) {
         return comp;
diff --git a/packages/firestore/src/lite-api/query.ts b/packages/firestore/src/lite-api/query.ts
@@ -33,7 +33,7 @@ import {
   isCollectionGroupQuery,
   LimitType,
   Query as InternalQuery,
-  queryOrderBy,
+  queryNormalizedOrderBy,
   queryWithAddedFilter,
   queryWithAddedOrderBy,
   queryWithEndAt,
@@ -907,7 +907,7 @@ export function newQueryBoundFromDocument(
   // the provided document. Without the key (by using the explicit sort
   // orders), multiple documents could match the position, yielding duplicate
   // results.
-  for (const orderBy of queryOrderBy(query)) {
+  for (const orderBy of queryNormalizedOrderBy(query)) {
     if (orderBy.field.isKeyField()) {
       components.push(refValue(databaseId, doc.key));
     } else {
diff --git a/packages/firestore/src/remote/datastore.ts b/packages/firestore/src/remote/datastore.ts
@@ -18,7 +18,7 @@
 import { CredentialsProvider } from '../api/credentials';
 import { User } from '../auth/user';
 import { Aggregate } from '../core/aggregate';
-import { Query, queryToTarget } from '../core/query';
+import { queryToAggregateTarget, Query, queryToTarget } from '../core/query';
 import { Document } from '../model/document';
 import { DocumentKey } from '../model/document_key';
 import { Mutation } from '../model/mutation';
@@ -248,7 +248,7 @@ export async function invokeRunAggregationQueryRpc(
   const datastoreImpl = debugCast(datastore, DatastoreImpl);
   const { request, aliasMap } = toRunAggregationQueryRequest(
     datastoreImpl.serializer,
-    queryToTarget(query),
+    queryToAggregateTarget(query),
     aggregates
   );
 
diff --git a/packages/firestore/test/unit/core/query.test.ts b/packages/firestore/test/unit/core/query.test.ts
@@ -27,12 +27,13 @@ import {
   newQueryForPath,
   Query,
   queryMatches,
-  queryOrderBy,
+  queryNormalizedOrderBy,
   queryWithAddedFilter,
   queryWithEndAt,
   queryWithLimit,
   queryWithStartAt,
   stringifyQuery,
+  queryToAggregateTarget,
   queryToTarget,
   QueryImpl,
   queryEquals,
@@ -852,6 +853,57 @@ describe('Query', () => {
     assertQueryMatches(query5, [doc3], [doc1, doc2, doc4, doc5]);
   });
 
+  it('generates appropriate order-bys for aggregate and non-aggregate targets', () => {
+    const col = newQueryForPath(ResourcePath.fromString('collection'));
+
+    // Build two identical queries
+    const query1 = queryWithAddedFilter(col, filter('foo', '>', 1));
+    const query2 = queryWithAddedFilter(col, filter('foo', '>', 1));
+
+    // Compute an aggregate and non-aggregate target from the queries
+    const aggregateTarget = queryToAggregateTarget(query1);
+    const target = queryToTarget(query2);
+
+    expect(aggregateTarget.orderBy.length).to.equal(0);
+    expect(target.orderBy.length).to.equal(2);
+    expect(target.orderBy[0].dir).to.equal('asc');
+    expect(target.orderBy[0].field.canonicalString()).to.equal('foo');
+    expect(target.orderBy[1].dir).to.equal('asc');
+    expect(target.orderBy[1].field.canonicalString()).to.equal('__name__');
+  });
+
+  it('generated order-bys are not affected by previously memoized targets', () => {
+    const col = newQueryForPath(ResourcePath.fromString('collection'));
+
+    // Build two identical queries
+    const query1 = queryWithAddedFilter(col, filter('foo', '>', 1));
+    const query2 = queryWithAddedFilter(col, filter('foo', '>', 1));
+
+    // query1 - first to aggregate target, then to non-aggregate target
+    const aggregateTarget1 = queryToAggregateTarget(query1);
+    const target1 = queryToTarget(query1);
+
+    // query2 - first to non-aggregate target, then to aggregate target
+    const target2 = queryToTarget(query2);
+    const aggregateTarget2 = queryToAggregateTarget(query2);
+
+    expect(aggregateTarget1.orderBy.length).to.equal(0);
+
+    expect(aggregateTarget2.orderBy.length).to.equal(0);
+
+    expect(target1.orderBy.length).to.equal(2);
+    expect(target1.orderBy[0].dir).to.equal('asc');
+    expect(target1.orderBy[0].field.canonicalString()).to.equal('foo');
+    expect(target1.orderBy[1].dir).to.equal('asc');
+    expect(target1.orderBy[1].field.canonicalString()).to.equal('__name__');
+
+    expect(target2.orderBy.length).to.equal(2);
+    expect(target2.orderBy[0].dir).to.equal('asc');
+    expect(target2.orderBy[0].field.canonicalString()).to.equal('foo');
+    expect(target2.orderBy[1].dir).to.equal('asc');
+    expect(target2.orderBy[1].field.canonicalString()).to.equal('__name__');
+  });
+
   function assertQueryMatches(
     query: Query,
     matching: MutableDocument[],
@@ -866,7 +918,7 @@ describe('Query', () => {
   }
 
   function assertImplicitOrderBy(query: Query, ...orderBys: OrderBy[]): void {
-    expect(queryOrderBy(query)).to.deep.equal(orderBys);
+    expect(queryNormalizedOrderBy(query)).to.deep.equal(orderBys);
   }
 
   function assertCanonicalId(query: Query, expectedCanonicalId: string): void {
