fix: return results from getPartitions() in order (#1521)

dev/src/collection-group.ts

@@ -27,6 +27,7 @@ import {validateInteger} from './validate';
 
 import api = protos.google.firestore.v1;
 import {defaultConverter} from './types';
+import {compareArrays} from './order';
 
 /**
  * A `CollectionGroup` refers to all documents that are contained in a
@@ -79,8 +80,7 @@ export class CollectionGroup<T = firestore.DocumentData>
     const tag = requestTag();
     await this.firestore.initializeIfNeeded(tag);
 
-    let lastValues: api.IValue[] | undefined = undefined;
-    let partitionCount = 0;
+    const partitions: Array<api.IValue>[] = [];
 
     if (desiredPartitionCount > 1) {
       // Partition queries require explicit ordering by __name__.
@@ -100,32 +100,36 @@ export class CollectionGroup<T = firestore.DocumentData>
       stream.resume();
 
       for await (const currentCursor of stream) {
-        ++partitionCount;
-        const currentValues = currentCursor.values ?? [];
-        yield new QueryPartition(
-          this._firestore,
-          this._queryOptions.collectionId,
-          this._queryOptions.converter,
-          lastValues,
-          currentValues
-        );
-        lastValues = currentValues;
+        partitions.push(currentCursor.values ?? []);
       }
     }
 
     logger(
       'Firestore.getPartitions',
       tag,
       'Received %d partitions',
-      partitionCount
+      partitions.length
     );
 
+    // Sort the partitions as they may not be ordered if responses are paged.
+    partitions.sort((l, r) => compareArrays(l, r));
+
+    for (let i = 0; i < partitions.length; ++i) {
+      yield new QueryPartition(
+        this._firestore,
+        this._queryOptions.collectionId,
+        this._queryOptions.converter,
+        i > 0 ? partitions[i - 1] : undefined,
+        partitions[i]
+      );
+    }
+
     // Return the extra partition with the empty cursor.
     yield new QueryPartition(
       this._firestore,
       this._queryOptions.collectionId,
       this._queryOptions.converter,
-      lastValues,
+      partitions.pop(),
       undefined
     );
   }

dev/src/order.ts

@@ -179,7 +179,7 @@ function compareGeoPoints(
 /*!
  * @private
  */
-function compareArrays(left: api.IValue[], right: api.IValue[]): number {
+export function compareArrays(left: api.IValue[], right: api.IValue[]): number {
   for (let i = 0; i < left.length && i < right.length; i++) {
     const valueComparison = compare(left[i], right[i]);
     if (valueComparison !== 0) {

dev/test/partition-query.ts

@@ -18,13 +18,13 @@ import {
   QueryPartition,
 } from '@google-cloud/firestore';
 
-import {describe, it, beforeEach, afterEach} from 'mocha';
+import {afterEach, beforeEach, describe, it} from 'mocha';
 import {expect, use} from 'chai';
 import * as chaiAsPromised from 'chai-as-promised';
 import * as extend from 'extend';
 
 import {google} from '../protos/firestore_v1_proto_api';
-import {Firestore} from '../src';
+import {DocumentReference, Firestore} from '../src';
 import {setTimeoutHandler} from '../src/backoff';
 import {
   ApiOverride,
@@ -39,6 +39,8 @@ import api = google.firestore.v1;
 
 const PROJECT_ID = 'test-project';
 const DATABASE_ROOT = `projects/${PROJECT_ID}/databases/(default)`;
+const DOC1 = `${DATABASE_ROOT}/documents/coll/doc1`;
+const DOC2 = `${DATABASE_ROOT}/documents/coll/doc2`;
 
 export function partitionQueryEquals(
   actual: api.IPartitionQueryRequest | undefined,
@@ -102,10 +104,32 @@ describe('Partition Query', () => {
     return partitions;
   }
 
+  function verifyPartition(
+    partition: FirebaseFirestore.QueryPartition,
+    startAt: string | null,
+    endBefore: string | null
+  ) {
+    if (startAt) {
+      expect(
+        partition.startAt?.map(value => (value as DocumentReference).path)
+      ).to.have.members([startAt]);
+    } else {
+      expect(partition.startAt).to.be.undefined;
+    }
+
+    if (endBefore) {
+      expect(
+        partition.endBefore?.map(value => (value as DocumentReference).path)
+      ).to.have.members([endBefore]);
+    } else {
+      expect(partition.endBefore).to.be.undefined;
+    }
+  }
+
   it('requests one less than desired partitions', () => {
     const desiredPartitionsCount = 2;
     const cursorValue = {
-      values: [{referenceValue: 'projects/p1/databases/d1/documents/coll/doc'}],
+      values: [{referenceValue: DOC1}],
     };
 
     const overrides: ApiOverride = {
@@ -156,12 +180,12 @@ describe('Partition Query', () => {
 
     const expectedStartAt: Array<undefined | api.IValue> = [
       undefined,
-      {referenceValue: 'coll/doc1'},
-      {referenceValue: 'coll/doc2'},
+      {referenceValue: DOC1},
+      {referenceValue: DOC2},
     ];
     const expectedEndBefore: Array<undefined | api.IValue> = [
-      {referenceValue: 'coll/doc1'},
-      {referenceValue: 'coll/doc2'},
+      {referenceValue: DOC1},
+      {referenceValue: DOC2},
       undefined,
     ];
 
@@ -174,10 +198,10 @@ describe('Partition Query', () => {
 
         return stream<api.ICursor>(
           {
-            values: [{referenceValue: 'coll/doc1'}],
+            values: [{referenceValue: DOC1}],
           },
           {
-            values: [{referenceValue: 'coll/doc2'}],
+            values: [{referenceValue: DOC2}],
           }
         );
       },
@@ -255,4 +279,32 @@ describe('Partition Query', () => {
       return result[0].toQuery().get();
     });
   });
+
+  it('sorts partitions', () => {
+    const desiredPartitionsCount = 3;
+
+    const overrides: ApiOverride = {
+      partitionQueryStream: () => {
+        return stream<api.ICursor>(
+          {
+            values: [{referenceValue: DOC2}],
+          },
+          {
+            values: [{referenceValue: DOC1}],
+          }
+        );
+      },
+    };
+
+    return createInstance(overrides).then(async firestore => {
+      const query = firestore.collectionGroup('collectionId');
+
+      const partitions = await getPartitions(query, desiredPartitionsCount);
+      expect(partitions.length).to.equal(3);
+
+      verifyPartition(partitions[0], null, 'coll/doc1');
+      verifyPartition(partitions[1], 'coll/doc1', 'coll/doc2');
+      verifyPartition(partitions[2], 'coll/doc2', null);
+    });
+  });
 });