Sum and average (#6952)

Refactoring aggregations in Firestore to support sum and average. Sum and average are not public.

---------

Co-authored-by: MarkDuckworth <MarkDuckworth@users.noreply.github.com>

.changeset/popular-items-hide.md

@@ -0,0 +1,5 @@
+---
+"@firebase/firestore": patch
+---
+
+Refactoring the aggregation implementation to support future aggregate functions.

common/api-review/firestore-lite.api.md

@@ -19,11 +19,11 @@ export type AddPrefixToKeys<Prefix extends string, T extends Record<string, unkn
 
 // @public
 export class AggregateField<T> {
-    type: string;
+    readonly type = "AggregateField";
 }
 
 // @public
-export type AggregateFieldType = AggregateField<number>;
+export type AggregateFieldType = AggregateField<number | null>;
 
 // @public
 export class AggregateQuerySnapshot<T extends AggregateSpec> {

common/api-review/firestore.api.md

@@ -19,11 +19,11 @@ export type AddPrefixToKeys<Prefix extends string, T extends Record<string, unkn
 
 // @public
 export class AggregateField<T> {
-    type: string;
+    readonly type = "AggregateField";
 }
 
 // @public
-export type AggregateFieldType = AggregateField<number>;
+export type AggregateFieldType = AggregateField<number | null>;
 
 // @public
 export class AggregateQuerySnapshot<T extends AggregateSpec> {

packages/firestore/lite/index.ts

@@ -29,15 +29,21 @@ registerFirestore();
 
 export {
   aggregateQuerySnapshotEqual,
-  getCount
+  getCount,
+  getAggregate,
+  count,
+  sum,
+  average,
+  aggregateFieldEqual
 } from '../src/lite-api/aggregate';
 
 export {
   AggregateField,
   AggregateFieldType,
   AggregateSpec,
   AggregateSpecData,
-  AggregateQuerySnapshot
+  AggregateQuerySnapshot,
+  AggregateType
 } from '../src/lite-api/aggregate_types';
 
 export { FirestoreSettings as Settings } from '../src/lite-api/settings';

packages/firestore/src/api.ts

@@ -17,15 +17,21 @@
 
 export {
   aggregateQuerySnapshotEqual,
-  getCountFromServer
+  getCountFromServer,
+  getAggregateFromServer,
+  count,
+  sum,
+  average,
+  aggregateFieldEqual
 } from './api/aggregate';
 
 export {
   AggregateField,
   AggregateFieldType,
   AggregateSpec,
   AggregateSpecData,
-  AggregateQuerySnapshot
+  AggregateQuerySnapshot,
+  AggregateType
 } from './lite-api/aggregate_types';
 
 export { FieldPath, documentId } from './api/field_path';

packages/firestore/src/api/aggregate.ts

@@ -15,18 +15,26 @@
  * limitations under the License.
  */
 
-import { Query } from '../api';
-import { firestoreClientRunCountQuery } from '../core/firestore_client';
-import {
-  AggregateField,
-  AggregateQuerySnapshot
-} from '../lite-api/aggregate_types';
+import { AggregateField, AggregateSpec, Query } from '../api';
+import { AggregateImpl } from '../core/aggregate';
+import { firestoreClientRunAggregateQuery } from '../core/firestore_client';
+import { count } from '../lite-api/aggregate';
+import { AggregateQuerySnapshot } from '../lite-api/aggregate_types';
+import { AggregateAlias } from '../model/aggregate_alias';
+import { ObjectValue } from '../model/object_value';
 import { cast } from '../util/input_validation';
+import { mapToArray } from '../util/obj';
 
 import { ensureFirestoreConfigured, Firestore } from './database';
 import { ExpUserDataWriter } from './reference_impl';
 
-export { aggregateQuerySnapshotEqual } from '../lite-api/aggregate';
+export {
+  aggregateQuerySnapshotEqual,
+  count,
+  sum,
+  average,
+  aggregateFieldEqual
+} from '../lite-api/aggregate';
 
 /**
  * Calculates the number of documents in the result set of the given query,
@@ -52,8 +60,89 @@ export { aggregateQuerySnapshotEqual } from '../lite-api/aggregate';
 export function getCountFromServer(
   query: Query<unknown>
 ): Promise<AggregateQuerySnapshot<{ count: AggregateField<number> }>> {
+  const countQuerySpec: { count: AggregateField<number> } = {
+    count: count()
+  };
+
+  return getAggregateFromServer(query, countQuerySpec);
+}
+
+/**
+ * Calculates the specified aggregations over the documents in the result
+ * set of the given query, without actually downloading the documents.
+ *
+ * Using this function to perform aggregations is efficient because only the
+ * final aggregation values, not the documents' data, is downloaded. This
+ * function can even perform aggregations of the documents if the result set
+ * would be prohibitively large to download entirely (e.g. thousands of documents).
+ *
+ * The result received from the server is presented, unaltered, without
+ * considering any local state. That is, documents in the local cache are not
+ * taken into consideration, neither are local modifications not yet
+ * synchronized with the server. Previously-downloaded results, if any, are not
+ * used: every request using this source necessarily involves a round trip to
+ * the server.
+ *
+ * @param query The query whose result set to aggregate over.
+ * @param aggregateSpec An `AggregateSpec` object that specifies the aggregates
+ * to perform over the result set. The AggregateSpec specifies aliases for each
+ * aggregate, which can be used to retrieve the aggregate result.
+ * @example
+ * ```typescript
+ * const aggregateSnapshot = await getAggregateFromServer(query, {
+ *   countOfDocs: count(),
+ *   totalHours: sum('hours'),
+ *   averageScore: average('score')
+ * });
+ *
+ * const countOfDocs: number = aggregateSnapshot.data().countOfDocs;
+ * const totalHours: number = aggregateSnapshot.data().totalHours;
+ * const averageScore: number | null = aggregateSnapshot.data().averageScore;
+ * ```
+ * @internal TODO (sum/avg) remove when public
+ */
+export function getAggregateFromServer<T extends AggregateSpec>(
+  query: Query<unknown>,
+  aggregateSpec: T
+): Promise<AggregateQuerySnapshot<T>> {
   const firestore = cast(query.firestore, Firestore);
   const client = ensureFirestoreConfigured(firestore);
+
+  const internalAggregates = mapToArray(aggregateSpec, (aggregate, alias) => {
+    return new AggregateImpl(
+      new AggregateAlias(alias),
+      aggregate._aggregateType,
+      aggregate._internalFieldPath
+    );
+  });
+
+  // Run the aggregation and convert the results
+  return firestoreClientRunAggregateQuery(
+    client,
+    query._query,
+    internalAggregates
+  ).then(aggregateResult =>
+    convertToAggregateQuerySnapshot(firestore, query, aggregateResult)
+  );
+}
+
+/**
+ * Converts the core aggregration result to an `AggregateQuerySnapshot`
+ * that can be returned to the consumer.
+ * @param query
+ * @param aggregateResult Core aggregation result
+ * @internal
+ */
+function convertToAggregateQuerySnapshot<T extends AggregateSpec>(
+  firestore: Firestore,
+  query: Query<unknown>,
+  aggregateResult: ObjectValue
+): AggregateQuerySnapshot<T> {
   const userDataWriter = new ExpUserDataWriter(firestore);
-  return firestoreClientRunCountQuery(client, query, userDataWriter);
+  const querySnapshot = new AggregateQuerySnapshot<T>(
+    query,
+    userDataWriter,
+    aggregateResult
+  );
+  return querySnapshot;
 }

packages/firestore/src/core/aggregate.ts

@@ -0,0 +1,45 @@
+/**
+ * @license
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { AggregateAlias } from '../model/aggregate_alias';
+import { FieldPath } from '../model/path';
+
+/**
+ * Union type representing the aggregate type to be performed.
+ * @internal
+ */
+export type AggregateType = 'count' | 'avg' | 'sum';
+
+/**
+ * Represents an Aggregate to be performed over a query result set.
+ */
+export interface Aggregate {
+  readonly fieldPath?: FieldPath;
+  readonly alias: AggregateAlias;
+  readonly aggregateType: AggregateType;
+}
+
+/**
+ * Concrete implementation of the Aggregate type.
+ */
+export class AggregateImpl implements Aggregate {
+  constructor(
+    readonly alias: AggregateAlias,
+    readonly aggregateType: AggregateType,
+    readonly fieldPath?: FieldPath
+  ) {}
+}