Do not apply SQL LIMIT if a partial index is used. (#3583)

* Do not apply SQL LIMIT if a partial index is used.

* Alternate implementation.

* Address feedback.

* Address feedback.

* Fix bug in `servedByIndex`.

* Consider array-contains segments when counting matching segments.

* Rename TargetIndexType to IndexType.

Author: ehsannas

firebase-firestore/src/main/java/com/google/firebase/firestore/core/Target.java

@@ -14,6 +14,8 @@
 
 package com.google.firebase.firestore.core;
 
+import static com.google.firebase.firestore.core.FieldFilter.Operator.ARRAY_CONTAINS;
+import static com.google.firebase.firestore.core.FieldFilter.Operator.ARRAY_CONTAINS_ANY;
 import static com.google.firebase.firestore.model.Values.max;
 import static com.google.firebase.firestore.model.Values.min;
 
@@ -29,8 +31,10 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Set;
 
 /**
  * A Target represents the WatchTarget representation of a Query, which is used by the LocalStore
@@ -369,6 +373,37 @@ public Direction getKeyOrder() {
     return this.orderBys.get(this.orderBys.size() - 1).getDirection();
   }
 
+  /** Returns the number of segments of a perfect index for this target. */
+  public int getSegmentCount() {
+    Set<FieldPath> fields = new HashSet<>();
+    boolean hasArraySegment = false;
+    for (Filter filter : filters) {
+      for (FieldFilter subFilter : filter.getFlattenedFilters()) {
+        // __name__ is not an explicit segment of any index, so we don't need to count it.
+        if (subFilter.getField().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 (subFilter.getOperator().equals(ARRAY_CONTAINS)
+            || subFilter.getOperator().equals(ARRAY_CONTAINS_ANY)) {
+          hasArraySegment = true;
+        } else {
+          fields.add(subFilter.getField());
+        }
+      }
+    }
+    for (OrderBy orderBy : orderBys) {
+      // __name__ is not an explicit segment of any index, so we don't need to count it.
+      if (!orderBy.getField().isKeyField()) {
+        fields.add(orderBy.getField());
+      }
+    }
+    return fields.size() + (hasArraySegment ? 1 : 0);
+  }
+
   /** Returns a canonical string representing this target. */
   public String getCanonicalId() {
     if (memoizedCanonicalId != null) {

firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexManager.java

@@ -32,6 +32,21 @@
  * Collection Group queries.
  */
 public interface IndexManager {
+  /** Represents the index state as it relates to a particular target. */
+  public 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
+  }
 
   /** Initializes the IndexManager. */
   void start();
@@ -83,12 +98,8 @@ public interface IndexManager {
   /** Returns the minimum offset for the given collection group. */
   IndexOffset getMinOffset(String collectionGroup);
 
-  /**
-   * Returns an index that can be used to serve the provided target. Returns {@code null} if no
-   * index is configured.
-   */
-  @Nullable
-  FieldIndex getFieldIndex(Target target);
+  /** Returns the type of index (if any) that can be used to serve the given target */
+  IndexType getIndexType(Target target);
 
   /**
    * Returns the documents that match the given target based on the provided index or {@code null}

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryIndexManager.java

@@ -60,13 +60,6 @@ public void deleteFieldIndex(FieldIndex index) {
     // Field indices are not supported with memory persistence.
   }
 
-  @Nullable
-  @Override
-  public FieldIndex getFieldIndex(Target target) {
-    // Field indices are not supported with memory persistence.
-    return null;
-  }
-
   @Override
   @Nullable
   public List<DocumentKey> getDocumentsMatchingTarget(Target target) {
@@ -108,6 +101,11 @@ public IndexOffset getMinOffset(String collectionGroup) {
     return IndexOffset.NONE;
   }
 
+  @Override
+  public IndexType getIndexType(Target target) {
+    return IndexType.NONE;
+  }
+
   @Override
   public void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents) {
     // Field indices are not supported with memory persistence.

firebase-firestore/src/main/java/com/google/firebase/firestore/local/QueryEngine.java

@@ -21,6 +21,7 @@
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.core.Target;
+import com.google.firebase.firestore.local.IndexManager.IndexType;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex;
@@ -110,11 +111,27 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return null;
     }
 
-    List<DocumentKey> keys = indexManager.getDocumentsMatchingTarget(target);
-    if (keys == null) {
+    IndexType indexType = indexManager.getIndexType(target);
+
+    if (indexType.equals(IndexType.NONE)) {
+      // The target cannot be served from any index.
       return null;
     }
 
+    if (indexType.equals(IndexType.PARTIAL)) {
+      // We cannot apply a limit for targets that are served using a partial index.
+      // If a partial index will be used to serve the target, the query may return a superset of
+      // documents that match the target (e.g. if the index doesn't include all the target's
+      // filters), or may return the correct set of documents in the wrong order (e.g. if the index
+      // doesn't include a segment for one of the orderBys). Therefore a limit should not be applied
+      // in such cases.
+      query = query.limitToFirst(Target.NO_LIMIT);
+      target = query.toTarget();
+    }
+
+    List<DocumentKey> keys = indexManager.getDocumentsMatchingTarget(target);
+    hardAssert(keys != null, "index manager must return results for partial and full indexes.");
+
     ImmutableSortedMap<DocumentKey, Document> indexedDocuments =
         localDocumentsView.getDocuments(keys);
     IndexOffset offset = indexManager.getMinOffset(target);

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteIndexManager.java

@@ -23,6 +23,7 @@
 import static java.lang.Math.max;
 
 import android.text.TextUtils;
+import android.util.Pair;
 import androidx.annotation.Nullable;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
@@ -311,11 +312,30 @@ public IndexOffset getMinOffset(String collectionGroup) {
     return getMinOffset(fieldIndexes);
   }
 
+  @Override
+  public IndexType getIndexType(Target target) {
+    for (Target subTarget : getSubTargets(target)) {
+      Pair<FieldIndex, Integer> indexAndSegmentCount = getFieldIndexAndSegmentCount(subTarget);
+      if (indexAndSegmentCount == null) {
+        return IndexType.NONE;
+      }
+
+      Integer indexSegmentCount = indexAndSegmentCount.second;
+      if (indexSegmentCount < subTarget.getSegmentCount()) {
+        return IndexType.PARTIAL;
+      }
+    }
+    return IndexType.FULL;
+  }
+
   @Override
   public IndexOffset getMinOffset(Target target) {
     List<FieldIndex> fieldIndexes = new ArrayList<>();
     for (Target subTarget : getSubTargets(target)) {
-      fieldIndexes.add(getFieldIndex(subTarget));
+      Pair<FieldIndex, Integer> indexAndSegmentCount = getFieldIndexAndSegmentCount(subTarget);
+      if (indexAndSegmentCount != null) {
+        fieldIndexes.add(indexAndSegmentCount.first);
+      }
     }
     return getMinOffset(fieldIndexes);
   }
@@ -448,11 +468,12 @@ public List<DocumentKey> getDocumentsMatchingTarget(Target target) {
     List<Object> bindings = new ArrayList<>();
 
     for (Target subTarget : getSubTargets(target)) {
-      FieldIndex fieldIndex = getFieldIndex(subTarget);
-      if (fieldIndex == null) {
+      Pair<FieldIndex, Integer> indexAndSegmentCount = getFieldIndexAndSegmentCount(subTarget);
+      if (indexAndSegmentCount == null) {
         return null;
       }
 
+      FieldIndex fieldIndex = indexAndSegmentCount.first;
       @Nullable List<Value> arrayValues = subTarget.getArrayValues(fieldIndex);
       @Nullable Collection<Value> notInValues = subTarget.getNotInValues(fieldIndex);
       @Nullable Bound lowerBound = subTarget.getLowerBound(fieldIndex);
@@ -491,7 +512,8 @@ public List<DocumentKey> getDocumentsMatchingTarget(Target target) {
 
     String queryString =
         "SELECT DISTINCT document_key FROM (" + TextUtils.join(" UNION ", subQueries) + ")";
-    if (target.getLimit() != -1) {
+
+    if (target.hasLimit()) {
       queryString = queryString + " LIMIT " + target.getLimit();
     }
 
@@ -547,9 +569,6 @@ private Object[] generateQueryAndBindings(
     StringBuilder sql = repeatSequence(statement, statementCount, " UNION ");
     sql.append("ORDER BY directional_value, document_key ");
     sql.append(target.getKeyOrder().equals(Direction.ASCENDING) ? "asc " : "desc ");
-    if (target.getLimit() != -1) {
-      sql.append("LIMIT ").append(target.getLimit()).append(" ");
-    }
 
     if (notIn != null) {
       // Wrap the statement in a NOT-IN call.
@@ -607,9 +626,13 @@ private Object[] fillBounds(
     return bindArgs;
   }
 
+  /**
+   * Returns an index that can be used to serve the provided target, as well as the number of index
+   * segments that matched the target filters and orderBys. Returns {@code null} if no index is
+   * configured.
+   */
   @Nullable
-  @Override
-  public FieldIndex getFieldIndex(Target target) {
+  public Pair<FieldIndex, Integer> getFieldIndexAndSegmentCount(Target target) {
     hardAssert(started, "IndexManager not started");
 
     TargetIndexMatcher targetIndexMatcher = new TargetIndexMatcher(target);
@@ -623,21 +646,20 @@ public FieldIndex getFieldIndex(Target target) {
       return null;
     }
 
-    List<FieldIndex> matchingIndexes = new ArrayList<>();
+    // Return the index with the most number of segments.
+    int maxNumMatchingSegments = -1;
+    FieldIndex matchingIndex = null;
     for (FieldIndex fieldIndex : collectionIndexes) {
-      boolean matches = targetIndexMatcher.servedByIndex(fieldIndex);
-      if (matches) {
-        matchingIndexes.add(fieldIndex);
+      int numMatchingSegments = targetIndexMatcher.servedByIndex(fieldIndex);
+      if (numMatchingSegments > maxNumMatchingSegments) {
+        maxNumMatchingSegments = numMatchingSegments;
+        matchingIndex = fieldIndex;
       }
     }
-
-    if (matchingIndexes.isEmpty()) {
+    if (maxNumMatchingSegments == -1) {
       return null;
     }
-
-    // Return the index with the most number of segments
-    return Collections.max(
-        matchingIndexes, (l, r) -> Integer.compare(l.getSegments().size(), r.getSegments().size()));
+    return new Pair<>(matchingIndex, maxNumMatchingSegments);
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/model/TargetIndexMatcher.java

@@ -111,7 +111,8 @@ public TargetIndexMatcher(Target target) {
   }
 
   /**
-   * Returns whether the index can be used to serve the TargetIndexMatcher's target.
+   * Returns the number of index segments that match the TargetIndexMatcher's target. Returns -1 if
+   * the index cannot be used to serve the target.
    *
    * <p>An index is considered capable of serving the target when:
    *
@@ -131,21 +132,26 @@ public TargetIndexMatcher(Target target) {
    *
    * @throws AssertionError if the index is for a different collection
    */
-  public boolean servedByIndex(FieldIndex index) {
+  public int servedByIndex(FieldIndex index) {
     hardAssert(index.getCollectionGroup().equals(collectionId), "Collection IDs do not match");
+    int numMatchingSegments = 0;
 
     // If there is an array element, find a matching filter.
     FieldIndex.Segment arraySegment = index.getArraySegment();
-    if (arraySegment != null && !hasMatchingEqualityFilter(arraySegment)) {
-      return false;
+    if (arraySegment != null) {
+      if (hasMatchingEqualityFilter(arraySegment)) {
+        ++numMatchingSegments;
+      } else {
+        return -1;
+      }
     }
 
     Iterator<OrderBy> orderBys = this.orderBys.iterator();
     List<FieldIndex.Segment> segments = index.getDirectionalSegments();
     int segmentIndex = 0;
 
     // Process all equalities first. Equalities can appear out of order.
-    for (; segmentIndex < segments.size(); ++segmentIndex) {
+    for (; segmentIndex < segments.size(); ++segmentIndex, ++numMatchingSegments) {
       // We attempt to greedily match all segments to equality filters. If a filter matches an
       // index segment, we can mark the segment as used. Since it is not possible to use the same
       // field path in both an equality and inequality/oderBy clause, we do not have to consider the
@@ -162,28 +168,29 @@ public boolean servedByIndex(FieldIndex index) {
     // filters and we do not need to map any segments to the target's inequality and orderBy
     // clauses.
     if (segmentIndex == segments.size()) {
-      return true;
+      return numMatchingSegments;
     }
 
     // If there is an inequality filter, the next segment must match both the filter and the first
     // orderBy clause.
     if (inequalityFilter != null) {
       FieldIndex.Segment segment = segments.get(segmentIndex);
       if (!matchesFilter(inequalityFilter, segment) || !matchesOrderBy(orderBys.next(), segment)) {
-        return false;
+        return -1;
       }
       ++segmentIndex;
+      ++numMatchingSegments;
     }
 
     // All remaining segments need to represent the prefix of the target's orderBy.
-    for (; segmentIndex < segments.size(); ++segmentIndex) {
+    for (; segmentIndex < segments.size(); ++segmentIndex, ++numMatchingSegments) {
       FieldIndex.Segment segment = segments.get(segmentIndex);
       if (!orderBys.hasNext() || !matchesOrderBy(orderBys.next(), segment)) {
-        return false;
+        return -1;
       }
     }
 
-    return true;
+    return numMatchingSegments;
   }
 
   private boolean hasMatchingEqualityFilter(FieldIndex.Segment segment) {

firebase-firestore/src/test/java/com/google/firebase/firestore/local/SQLiteIndexManagerTest.java

@@ -524,7 +524,7 @@ public void testArrayContainsDoesNotMatchNonArray() {
   public void testNoMatchingFilter() {
     setUpSingleValueFilter();
     Query query = query("coll").filter(filter("unknown", "==", true));
-    assertNull(indexManager.getFieldIndex(query.toTarget()));
+    assertEquals(indexManager.getIndexType(query.toTarget()), IndexManager.IndexType.NONE);
     assertNull(indexManager.getDocumentsMatchingTarget(query.toTarget()));
   }
 
@@ -1002,7 +1002,7 @@ private void addDoc(String key, Map<String, Object> data) {
 
   private void verifyResults(Query query, String... documents) {
     Target target = query.toTarget();
-    Iterable<DocumentKey> results = indexManager.getDocumentsMatchingTarget(target);
+    List<DocumentKey> results = indexManager.getDocumentsMatchingTarget(target);
     assertNotNull("Target cannot be served from index.", results);
     List<DocumentKey> keys =
         Arrays.stream(documents).map(TestUtil::key).collect(Collectors.toList());