Support plumbing for partial index execution (#6143)

schmidt-sebastian · web-flow · commit 15409a7b9581 · 2022-04-14T10:23:39.000-06:00
diff --git a/packages/firestore/src/core/target.ts b/packages/firestore/src/core/target.ts
@@ -41,6 +41,7 @@ import {
 } from '../model/values';
 import { Value as ProtoValue } from '../protos/firestore_proto_api';
 import { debugAssert, debugCast, fail } from '../util/assert';
+import { SortedSet } from '../util/sorted_set';
 import { isNullOrUndefined } from '../util/types';
 
 /**
@@ -484,6 +485,46 @@ function targetGetUpperBoundForField(
   return { value, inclusive };
 }
 
+/** Returns the number of segments of a perfect index for this target. */
+export function targetGetSegmentCount(target: Target): number {
+  let fields = new SortedSet<FieldPath>(FieldPath.comparator);
+  let hasArraySegment = false;
+
+  for (const filter of target.filters) {
+    // TODO(orquery): Use the flattened filters here
+    const fieldFilter = filter as FieldFilter;
+
+    // __name__ is not an explicit segment of any index, so we don't need to
+    // count it.
+    if (fieldFilter.field.isKeyField()) {
+      continue;
+    }
+
+    // ARRAY_CONTAINS or ARRAY_CONTAINS_ANY filters must be counted separately.
+    // For instance, it is possible to have an index for "a ARRAY a ASC". Even
+    // though these are on the same field, they should be counted as two
+    // separate segments in an index.
+    if (
+      fieldFilter.op === Operator.ARRAY_CONTAINS ||
+      fieldFilter.op === Operator.ARRAY_CONTAINS_ANY
+    ) {
+      hasArraySegment = true;
+    } else {
+      fields = fields.add(fieldFilter.field);
+    }
+  }
+
+  for (const orderBy of target.orderBy) {
+    // __name__ is not an explicit segment of any index, so we don't need to
+    // count it.
+    if (!orderBy.field.isKeyField()) {
+      fields = fields.add(orderBy.field);
+    }
+  }
+
+  return fields.size + (hasArraySegment ? 1 : 0);
+}
+
 export abstract class Filter {
   abstract matches(doc: Document): boolean;
 }
diff --git a/packages/firestore/src/local/index_manager.ts b/packages/firestore/src/local/index_manager.ts
@@ -24,6 +24,23 @@ import { ResourcePath } from '../model/path';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
 
+/** Represents the index state as it relates to a particular target. */
+export const enum IndexType {
+  /** Indicates that no index could be found for serving the target. */
+  NONE,
+  /**
+   * Indicates that only a "partial index" could be found for serving the
+   * target. A partial index is one which does not have a segment for every
+   * filter/orderBy in the target.
+   */
+  PARTIAL,
+  /**
+   * Indicates that a "full index" could be found for serving the target. A full
+   * index is one which has a segment for every filter/orderBy in the target.
+   */
+  FULL
+}
+
 /**
  * Represents a set of indexes that are used to execute queries efficiently.
  *
@@ -94,13 +111,13 @@ export interface IndexManager {
   ): PersistencePromise<FieldIndex[]>;
 
   /**
-   * Returns an index that can be used to serve the provided target. Returns
-   * `null` if no index is configured.
+   * Returns the type of index (if any) that can be used to serve the given
+   * target.
    */
-  getFieldIndex(
+  getIndexType(
     transaction: PersistenceTransaction,
     target: Target
-  ): PersistencePromise<FieldIndex | null>;
+  ): PersistencePromise<IndexType>;
 
   /**
    * Returns the documents that match the given target based on the provided
diff --git a/packages/firestore/src/local/indexeddb_index_manager.ts b/packages/firestore/src/local/indexeddb_index_manager.ts
@@ -27,6 +27,7 @@ import {
   targetGetArrayValues,
   targetGetLowerBound,
   targetGetNotInValues,
+  targetGetSegmentCount,
   targetGetUpperBound
 } from '../core/target';
 import { FirestoreIndexValueWriter } from '../index/firestore_index_value_writer';
@@ -59,7 +60,7 @@ import {
   decodeResourcePath,
   encodeResourcePath
 } from './encoded_resource_path';
-import { IndexManager } from './index_manager';
+import { IndexManager, IndexType } from './index_manager';
 import {
   DbCollectionParent,
   DbIndexConfiguration,
@@ -438,7 +439,7 @@ export class IndexedDbIndexManager implements IndexManager {
     );
   }
 
-  getFieldIndex(
+  private getFieldIndex(
     transaction: PersistenceTransaction,
     target: Target
   ): PersistencePromise<FieldIndex | null> {
@@ -449,13 +450,33 @@ export class IndexedDbIndexManager implements IndexManager {
         : target.path.lastSegment();
 
     return this.getFieldIndexes(transaction, collectionGroup).next(indexes => {
-      const matchingIndexes = indexes.filter(i =>
-        targetIndexMatcher.servedByIndex(i)
-      );
+      // Return the index with the most number of segments.
+      let index: FieldIndex | null = null;
+      for (const candidate of indexes) {
+        const matches = targetIndexMatcher.servedByIndex(candidate);
+        if (
+          matches &&
+          (!index || candidate.fields.length > index.fields.length)
+        ) {
+          index = candidate;
+        }
+      }
+      return index;
+    });
+  }
 
-      // Return the index that matches the most number of segments.
-      matchingIndexes.sort((l, r) => r.fields.length - l.fields.length);
-      return matchingIndexes.length > 0 ? matchingIndexes[0] : null;
+  getIndexType(
+    transaction: PersistenceTransaction,
+    target: Target
+  ): PersistencePromise<IndexType> {
+    // TODO(orqueries): We should look at the subtargets here
+    return this.getFieldIndex(transaction, target).next(index => {
+      if (!index) {
+        return IndexType.NONE as IndexType;
+      }
+      return index.fields.length < targetGetSegmentCount(target)
+        ? IndexType.PARTIAL
+        : IndexType.FULL;
     });
   }
 
diff --git a/packages/firestore/src/local/memory_index_manager.ts b/packages/firestore/src/local/memory_index_manager.ts
@@ -23,7 +23,7 @@ import { ResourcePath } from '../model/path';
 import { debugAssert } from '../util/assert';
 import { SortedSet } from '../util/sorted_set';
 
-import { IndexManager } from './index_manager';
+import { IndexManager, IndexType } from './index_manager';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
 
@@ -74,12 +74,12 @@ export class MemoryIndexManager implements IndexManager {
     return PersistencePromise.resolve<DocumentKey[] | null>(null);
   }
 
-  getFieldIndex(
+  getIndexType(
     transaction: PersistenceTransaction,
     target: Target
-  ): PersistencePromise<FieldIndex | null> {
+  ): PersistencePromise<IndexType> {
     // Field indices are not supported with memory persistence.
-    return PersistencePromise.resolve<FieldIndex | null>(null);
+    return PersistencePromise.resolve<IndexType>(IndexType.NONE);
   }
 
   getFieldIndexes(
diff --git a/packages/firestore/test/unit/local/index_manager.test.ts b/packages/firestore/test/unit/local/index_manager.test.ts
@@ -30,6 +30,7 @@ import {
   queryWithStartAt
 } from '../../../src/core/query';
 import { FieldFilter } from '../../../src/core/target';
+import { IndexType } from '../../../src/local/index_manager';
 import { IndexedDbPersistence } from '../../../src/local/indexeddb_persistence';
 import { INDEXING_SCHEMA_VERSION } from '../../../src/local/indexeddb_schema';
 import { Persistence } from '../../../src/local/persistence';
@@ -631,7 +632,9 @@ describe('IndexedDbIndexManager', async () => {
       query('coll'),
       filter('unknown', '==', true)
     );
-    expect(await indexManager.getFieldIndex(queryToTarget(q))).to.be.null;
+    expect(await indexManager.getIndexType(queryToTarget(q))).to.equal(
+      IndexType.NONE
+    );
     expect(await indexManager.getDocumentsMatchingTarget(queryToTarget(q))).to
       .be.null;
   });
@@ -1416,6 +1419,141 @@ describe('IndexedDbIndexManager', async () => {
     );
   });
 
+  it('serves partial and full index', async () => {
+    await indexManager.addFieldIndex(
+      fieldIndex('coll', { fields: [['a', IndexKind.ASCENDING]] })
+    );
+    await indexManager.addFieldIndex(
+      fieldIndex('coll', { fields: [['b', IndexKind.ASCENDING]] })
+    );
+    await indexManager.addFieldIndex(
+      fieldIndex('coll', {
+        fields: [
+          ['c', IndexKind.ASCENDING],
+          ['d', IndexKind.ASCENDING]
+        ]
+      })
+    );
+
+    const query1 = queryWithAddedFilter(query('coll'), filter('a', '==', 1));
+    await validateIsFullIndex(query1);
+
+    const query2 = queryWithAddedFilter(query('coll'), filter('b', '==', 1));
+    await validateIsFullIndex(query2);
+
+    const query3 = queryWithAddedOrderBy(
+      queryWithAddedFilter(query('coll'), filter('a', '==', 1)),
+      orderBy('a')
+    );
+    await validateIsFullIndex(query3);
+
+    const query4 = queryWithAddedOrderBy(
+      queryWithAddedFilter(query('coll'), filter('b', '==', 1)),
+      orderBy('b')
+    );
+    await validateIsFullIndex(query4);
+
+    const query5 = queryWithAddedFilter(
+      queryWithAddedFilter(query('coll'), filter('a', '==', 1)),
+      filter('b', '==', 1)
+    );
+    await validateIsPartialIndex(query5);
+
+    const query6 = queryWithAddedOrderBy(
+      queryWithAddedFilter(query('coll'), filter('a', '==', 1)),
+      orderBy('b')
+    );
+    await validateIsPartialIndex(query6);
+
+    const query7 = queryWithAddedOrderBy(
+      queryWithAddedFilter(query('coll'), filter('b', '==', 1)),
+      orderBy('a')
+    );
+    await validateIsPartialIndex(query7);
+
+    const query8 = queryWithAddedFilter(
+      queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+      filter('d', '==', 1)
+    );
+    await validateIsFullIndex(query8);
+
+    const query9 = queryWithAddedOrderBy(
+      queryWithAddedFilter(
+        queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+        filter('d', '==', 1)
+      ),
+      orderBy('c')
+    );
+    await validateIsFullIndex(query9);
+
+    const query10 = queryWithAddedOrderBy(
+      queryWithAddedFilter(
+        queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+        filter('d', '==', 1)
+      ),
+      orderBy('d')
+    );
+    await validateIsFullIndex(query10);
+
+    const query11 = queryWithAddedOrderBy(
+      queryWithAddedOrderBy(
+        queryWithAddedFilter(
+          queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+          filter('d', '==', 1)
+        ),
+        orderBy('c')
+      ),
+      orderBy('d')
+    );
+    await validateIsFullIndex(query11);
+
+    const query12 = queryWithAddedOrderBy(
+      queryWithAddedOrderBy(
+        queryWithAddedFilter(
+          queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+          filter('d', '==', 1)
+        ),
+        orderBy('d')
+      ),
+      orderBy('c')
+    );
+    await validateIsFullIndex(query12);
+
+    const query13 = queryWithAddedOrderBy(
+      queryWithAddedFilter(
+        queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+        filter('d', '==', 1)
+      ),
+      orderBy('e')
+    );
+    await validateIsPartialIndex(query13);
+
+    const query14 = queryWithAddedFilter(
+      queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+      filter('d', '<=', 1)
+    );
+    await validateIsFullIndex(query14);
+
+    const query15 = queryWithAddedOrderBy(
+      queryWithAddedFilter(
+        queryWithAddedFilter(query('coll'), filter('c', '==', 1)),
+        filter('d', '>', 1)
+      ),
+      orderBy('d')
+    );
+    await validateIsFullIndex(query15);
+  });
+
+  async function validateIsPartialIndex(query: Query): Promise<void> {
+    const indexType = await indexManager.getIndexType(queryToTarget(query));
+    expect(indexType).to.equal(IndexType.PARTIAL);
+  }
+
+  async function validateIsFullIndex(query: Query): Promise<void> {
+    const indexType = await indexManager.getIndexType(queryToTarget(query));
+    expect(indexType).to.equal(IndexType.FULL);
+  }
+
   async function setUpSingleValueFilter(): Promise<void> {
     await indexManager.addFieldIndex(
       fieldIndex('coll', { fields: [['count', IndexKind.ASCENDING]] })
diff --git a/packages/firestore/test/unit/local/test_index_manager.ts b/packages/firestore/test/unit/local/test_index_manager.ts
@@ -16,7 +16,7 @@
  */
 
 import { Target } from '../../../src/core/target';
-import { IndexManager } from '../../../src/local/index_manager';
+import { IndexManager, IndexType } from '../../../src/local/index_manager';
 import { Persistence } from '../../../src/local/persistence';
 import { DocumentMap } from '../../../src/model/collections';
 import { DocumentKey } from '../../../src/model/document_key';
@@ -71,9 +71,9 @@ export class TestIndexManager {
     );
   }
 
-  getFieldIndex(target: Target): Promise<FieldIndex | null> {
-    return this.persistence.runTransaction('getFieldIndex', 'readonly', txn =>
-      this.indexManager.getFieldIndex(txn, target)
+  getIndexType(target: Target): Promise<IndexType> {
+    return this.persistence.runTransaction('getIndexType', 'readonly', txn =>
+      this.indexManager.getIndexType(txn, target)
     );
   }
 
