Serving OR Queries from the client-side index.

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

@@ -314,7 +314,9 @@ public IndexOffset getMinOffset(String collectionGroup) {
   @Override
   public IndexType getIndexType(Target target) {
     IndexType result = IndexType.FULL;
-    for (Target subTarget : getSubTargets(target)) {
+    List<Target> subTargets = getSubTargets(target);
+
+    for (Target subTarget : subTargets) {
       FieldIndex index = getFieldIndex(subTarget);
       if (index == null) {
         result = IndexType.NONE;
@@ -325,6 +327,17 @@ public IndexType getIndexType(Target target) {
         result = IndexType.PARTIAL;
       }
     }
+
+    // OR queries have more than one sub-target (one sub-target per DNF term). We currently consider
+    // OR queries that have a `limit` to have a partial index. For such queries we perform sorting
+    // and apply the limit in memory as a post-processing step.
+    // TODO(orquery): If we have a FULL index *and* we have the index that can be used for sorting
+    //  all DNF branches on the same value, we can improve performance by performing a JOIN in SQL.
+    //  See b/235224019 for more information.
+    if (target.hasLimit() && subTargets.size() > 1 && result == IndexType.FULL) {
+      return IndexType.PARTIAL;
+    }
+
     return result;
   }
 
@@ -509,9 +522,23 @@ public List<DocumentKey> getDocumentsMatchingTarget(Target target) {
       bindings.addAll(Arrays.asList(subQueryAndBindings).subList(1, subQueryAndBindings.length));
     }
 
-    String queryString =
-        "SELECT DISTINCT document_key FROM (" + TextUtils.join(" UNION ", subQueries) + ")";
-
+    // We are constructing:
+    // SELECT DISTINCT document_key FROM (
+    //   (SELECT ...) UNION (SELECT ...) UNION (SELECT ...)
+    //   ORDER BY ...
+    // )
+    // LIMIT ...
+    //
+    // Note: SQLite does not allow performing ORDER BY on each union clause. The ORDER BY must come
+    // after the last union clause. Also note that LIMIT must be applied *after* the DISTINCT
+    // operator has been performed. When dealing with multiple sub-targets, it's possible that the
+    // same document_key appears multiple times.
+    String unionSubTargets =
+        TextUtils.join(" UNION ", subQueries)
+            + "ORDER BY directional_value, document_key "
+            + (target.getKeyOrder().equals(Direction.ASCENDING) ? "asc " : "desc ");
+
+    String queryString = "SELECT DISTINCT document_key FROM (" + unionSubTargets + ")";
     if (target.hasLimit()) {
       queryString = queryString + " LIMIT " + target.getLimit();
     }
@@ -557,11 +584,8 @@ private Object[] generateQueryAndBindings(
     statement.append("AND directional_value ").append(lowerBoundOp).append(" ? ");
     statement.append("AND directional_value ").append(upperBoundOp).append(" ? ");
 
-    // Create the UNION statement by repeating the above generated statement. We can then add
-    // ordering and a limit clause.
+    // Create the UNION statement by repeating the above generated statement.
     StringBuilder sql = repeatSequence(statement, statementCount, " UNION ");
-    sql.append("ORDER BY directional_value, document_key ");
-    sql.append(target.getKeyOrder().equals(Direction.ASCENDING) ? "asc " : "desc ");
 
     if (notIn != null) {
       // Wrap the statement in a NOT-IN call.

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

@@ -482,5 +482,45 @@ public void canPerformOrQueriesUsingFullCollectionScan() throws Exception {
     DocumentSet result5 =
         expectFullCollectionScan(() -> runQuery(query5, MISSING_LAST_LIMBO_FREE_SNAPSHOT));
     assertEquals(docSet(query5.comparator(), doc3), result5);
+
+    // Test with limits (implicit order by ASC): (a==1) || (b > 0) LIMIT 2
+    Query query6 =
+        query("coll").filter(orFilters(filter("a", "==", 1), filter("b", ">", 0))).limitToFirst(2);
+    DocumentSet result6 =
+        expectFullCollectionScan(() -> runQuery(query6, MISSING_LAST_LIMBO_FREE_SNAPSHOT));
+    assertEquals(docSet(query6.comparator(), doc1, doc2), result6);
+
+    // Test with limits (implicit order by DESC): (a==1) || (b > 0) LIMIT_TO_LAST 2
+    Query query7 =
+        query("coll").filter(orFilters(filter("a", "==", 1), filter("b", ">", 0))).limitToLast(2);
+    DocumentSet result7 =
+        expectFullCollectionScan(() -> runQuery(query7, MISSING_LAST_LIMBO_FREE_SNAPSHOT));
+    assertEquals(docSet(query7.comparator(), doc3, doc4), result7);
+
+    // Test with limits (explicit order by ASC): (a==2) || (b == 1) ORDER BY a LIMIT 1
+    Query query8 =
+        query("coll")
+            .filter(orFilters(filter("a", "==", 2), filter("b", "==", 1)))
+            .limitToFirst(1)
+            .orderBy(orderBy("a", "asc"));
+    DocumentSet result8 =
+        expectFullCollectionScan(() -> runQuery(query8, MISSING_LAST_LIMBO_FREE_SNAPSHOT));
+    assertEquals(docSet(query8.comparator(), doc5), result8);
+
+    // Test with limits (explicit order by DESC): (a==2) || (b == 1) ORDER BY a LIMIT_TO_LAST 1
+    Query query9 =
+        query("coll")
+            .filter(orFilters(filter("a", "==", 2), filter("b", "==", 1)))
+            .limitToLast(1)
+            .orderBy(orderBy("a", "asc"));
+    DocumentSet result9 =
+        expectFullCollectionScan(() -> runQuery(query9, MISSING_LAST_LIMBO_FREE_SNAPSHOT));
+    assertEquals(docSet(query9.comparator(), doc2), result9);
+
+    // Test with limits without orderBy (the __name__ ordering is the tie breaker).
+    Query query10 =
+        query("coll").filter(orFilters(filter("a", "==", 2), filter("b", "==", 1))).limitToFirst(1);
+    DocumentSet result10 = expectFullCollectionScan(() -> runQuery(query10, SnapshotVersion.NONE));
+    assertEquals(docSet(query10.comparator(), doc2), result10);
   }
 }

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

@@ -25,6 +25,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
 import static com.google.firebase.firestore.testutil.TestUtil.key;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
+import static com.google.firebase.firestore.testutil.TestUtil.orFilters;
 import static com.google.firebase.firestore.testutil.TestUtil.orderBy;
 import static com.google.firebase.firestore.testutil.TestUtil.path;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
@@ -1007,90 +1008,148 @@ public void testPartialIndexAndFullIndex() throws Exception {
     indexManager.addFieldIndex(fieldIndex("coll", "c", Kind.ASCENDING, "d", Kind.ASCENDING));
 
     Query query1 = query("coll").filter(filter("a", "==", 1));
-    validateIsFullIndex(query1);
+    validateIndexType(query1, IndexManager.IndexType.FULL);
 
     Query query2 = query("coll").filter(filter("b", "==", 1));
-    validateIsFullIndex(query2);
+    validateIndexType(query2, IndexManager.IndexType.FULL);
 
     Query query3 = query("coll").filter(filter("a", "==", 1)).orderBy(orderBy("a"));
-    validateIsFullIndex(query3);
+    validateIndexType(query3, IndexManager.IndexType.FULL);
 
     Query query4 = query("coll").filter(filter("b", "==", 1)).orderBy(orderBy("b"));
-    validateIsFullIndex(query4);
+    validateIndexType(query4, IndexManager.IndexType.FULL);
 
     Query query5 = query("coll").filter(filter("a", "==", 1)).filter(filter("b", "==", 1));
-    validateIsPartialIndex(query5);
+    validateIndexType(query5, IndexManager.IndexType.PARTIAL);
 
     Query query6 = query("coll").filter(filter("a", "==", 1)).orderBy(orderBy("b"));
-    validateIsPartialIndex(query6);
+    validateIndexType(query6, IndexManager.IndexType.PARTIAL);
 
     Query query7 = query("coll").filter(filter("b", "==", 1)).orderBy(orderBy("a"));
-    validateIsPartialIndex(query7);
+    validateIndexType(query7, IndexManager.IndexType.PARTIAL);
 
     Query query8 = query("coll").filter(filter("c", "==", 1)).filter(filter("d", "==", 1));
-    validateIsFullIndex(query8);
+    validateIndexType(query8, IndexManager.IndexType.FULL);
 
     Query query9 =
         query("coll")
             .filter(filter("c", "==", 1))
             .filter(filter("d", "==", 1))
             .orderBy(orderBy("c"));
-    validateIsFullIndex(query9);
+    validateIndexType(query9, IndexManager.IndexType.FULL);
 
     Query query10 =
         query("coll")
             .filter(filter("c", "==", 1))
             .filter(filter("d", "==", 1))
             .orderBy(orderBy("d"));
-    validateIsFullIndex(query10);
+    validateIndexType(query10, IndexManager.IndexType.FULL);
 
     Query query11 =
         query("coll")
             .filter(filter("c", "==", 1))
             .filter(filter("d", "==", 1))
             .orderBy(orderBy("c"))
             .orderBy(orderBy("d"));
-    validateIsFullIndex(query11);
+    validateIndexType(query11, IndexManager.IndexType.FULL);
 
     Query query12 =
         query("coll")
             .filter(filter("c", "==", 1))
             .filter(filter("d", "==", 1))
             .orderBy(orderBy("d"))
             .orderBy(orderBy("c"));
-    validateIsFullIndex(query12);
+    validateIndexType(query12, IndexManager.IndexType.FULL);
 
     Query query13 =
         query("coll")
             .filter(filter("c", "==", 1))
             .filter(filter("d", "==", 1))
             .orderBy(orderBy("e"));
-    validateIsPartialIndex(query13);
+    validateIndexType(query13, IndexManager.IndexType.PARTIAL);
 
     Query query14 = query("coll").filter(filter("c", "==", 1)).filter(filter("d", "<=", 1));
-    validateIsFullIndex(query14);
+    validateIndexType(query14, IndexManager.IndexType.FULL);
 
     Query query15 =
         query("coll")
             .filter(filter("c", "==", 1))
             .filter(filter("d", ">", 1))
             .orderBy(orderBy("d"));
-    validateIsFullIndex(query15);
+    validateIndexType(query15, IndexManager.IndexType.FULL);
   }
 
-  private void validateIsPartialIndex(Query query) {
-    validateIndex(query, false);
-  }
+  @Test
+  public void testIndexTypeForOrQueries() throws Exception {
+    indexManager.addFieldIndex(fieldIndex("coll", "a", Kind.ASCENDING));
+    indexManager.addFieldIndex(fieldIndex("coll", "b", Kind.ASCENDING));
+    indexManager.addFieldIndex(fieldIndex("coll", "b", Kind.ASCENDING, "a", Kind.ASCENDING));
+
+    // OR query without orderBy without limit which has missing sub-target indexes.
+    Query query1 = query("coll").filter(orFilters(filter("a", "==", 1), filter("c", "==", 1)));
+    validateIndexType(query1, IndexManager.IndexType.NONE);
+
+    // OR query with explicit orderBy without limit which has missing sub-target indexes.
+    Query query2 =
+        query("coll")
+            .filter(orFilters(filter("a", "==", 1), filter("c", "==", 1)))
+            .orderBy(orderBy("c"));
+    validateIndexType(query2, IndexManager.IndexType.NONE);
+
+    // OR query with implicit orderBy without limit which has missing sub-target indexes.
+    Query query3 = query("coll").filter(orFilters(filter("a", "==", 1), filter("c", ">", 1)));
+    validateIndexType(query3, IndexManager.IndexType.NONE);
+
+    // OR query with explicit orderBy with limit which has missing sub-target indexes.
+    Query query4 =
+        query("coll")
+            .filter(orFilters(filter("a", "==", 1), filter("c", "==", 1)))
+            .orderBy(orderBy("c"))
+            .limitToFirst(2);
+    validateIndexType(query4, IndexManager.IndexType.NONE);
 
-  private void validateIsFullIndex(Query query) {
-    validateIndex(query, true);
+    // OR query with implicit orderBy with limit which has missing sub-target indexes.
+    Query query5 =
+        query("coll").filter(orFilters(filter("a", "==", 1), filter("c", ">", 1))).limitToFirst(2);
+    validateIndexType(query5, IndexManager.IndexType.NONE);
+
+    // OR query without orderBy without limit which has all sub-target indexes.
+    Query query6 = query("coll").filter(orFilters(filter("a", "==", 1), filter("b", "==", 1)));
+    validateIndexType(query6, IndexManager.IndexType.FULL);
+
+    // OR query with explicit orderBy without limit which has all sub-target indexes.
+    Query query7 =
+        query("coll")
+            .filter(orFilters(filter("a", "==", 1), filter("b", "==", 1)))
+            .orderBy(orderBy("a"));
+    validateIndexType(query7, IndexManager.IndexType.FULL);
+
+    // OR query with implicit orderBy without limit which has all sub-target indexes.
+    Query query8 = query("coll").filter(orFilters(filter("a", ">", 1), filter("b", "==", 1)));
+    validateIndexType(query8, IndexManager.IndexType.FULL);
+
+    // OR query without orderBy with limit which has all sub-target indexes.
+    Query query9 =
+        query("coll").filter(orFilters(filter("a", "==", 1), filter("b", "==", 1))).limitToFirst(2);
+    validateIndexType(query9, IndexManager.IndexType.PARTIAL);
+
+    // OR query with explicit orderBy with limit which has all sub-target indexes.
+    Query query10 =
+        query("coll")
+            .filter(orFilters(filter("a", "==", 1), filter("b", "==", 1)))
+            .orderBy(orderBy("a"))
+            .limitToFirst(2);
+    validateIndexType(query10, IndexManager.IndexType.PARTIAL);
+
+    // OR query with implicit orderBy with limit which has all sub-target indexes.
+    Query query11 =
+        query("coll").filter(orFilters(filter("a", ">", 1), filter("b", "==", 1))).limitToFirst(2);
+    validateIndexType(query11, IndexManager.IndexType.PARTIAL);
   }
 
-  private void validateIndex(Query query, boolean validateFullIndex) {
+  private void validateIndexType(Query query, IndexManager.IndexType expected) {
     IndexManager.IndexType indexType = indexManager.getIndexType(query.toTarget());
-    assertEquals(
-        indexType,
-        validateFullIndex ? IndexManager.IndexType.FULL : IndexManager.IndexType.PARTIAL);
+    assertEquals(indexType, expected);
   }
 
   private void verifySequenceNumber(