Serve OR queries from the index.

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

@@ -26,7 +26,6 @@
 import com.google.firebase.firestore.util.Logger;
 import java.util.Collection;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -133,7 +132,8 @@ private int writeEntriesForCollectionGroup(
       String collectionGroup, int documentsRemainingUnderCap) {
     // Use the earliest offset of all field indexes to query the local cache.
     Collection<FieldIndex> fieldIndexes = indexManager.getFieldIndexes(collectionGroup);
-    IndexOffset existingOffset = getExistingOffset(fieldIndexes);
+    hardAssert(!fieldIndexes.isEmpty(), "Updating collection without indexes");
+    IndexOffset existingOffset = indexManager.getLeastRecentIndexOffset(fieldIndexes);
 
     LocalDocumentsResult nextBatch =
         localDocumentsView.getNextDocuments(
@@ -161,24 +161,6 @@ private IndexOffset getNewOffset(IndexOffset existingOffset, LocalDocumentsResul
         Math.max(lookupResult.getBatchId(), existingOffset.getLargestBatchId()));
   }
 
-  /** Returns the lowest offset for the provided index group. */
-  private IndexOffset getExistingOffset(Collection<FieldIndex> fieldIndexes) {
-    hardAssert(!fieldIndexes.isEmpty(), "Updating collection without indexes");
-
-    Iterator<FieldIndex> it = fieldIndexes.iterator();
-    IndexOffset minOffset = it.next().getIndexState().getOffset();
-    int minBatchId = minOffset.getLargestBatchId();
-    while (it.hasNext()) {
-      IndexOffset newOffset = it.next().getIndexState().getOffset();
-      if (newOffset.compareTo(minOffset) < 0) {
-        minOffset = newOffset;
-      }
-      minBatchId = Math.max(newOffset.getLargestBatchId(), minBatchId);
-    }
-
-    return IndexOffset.create(minOffset.getReadTime(), minOffset.getDocumentKey(), minBatchId);
-  }
-
   @VisibleForTesting
   void setMaxDocumentsToProcess(int newMax) {
     maxDocumentsToProcess = newMax;

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

@@ -74,6 +74,18 @@ public interface IndexManager {
   /** Returns all configured field indexes. */
   Collection<FieldIndex> getFieldIndexes();
 
+  /** Returns whether we can serve the given target from the index. */
+  boolean canServeFromIndex(Target target);
+
+  /**
+   * Iterates over all field indexes that are used to serve the given target, and returns the least
+   * recent offset of them all. Asserts that the target can be served from index.
+   */
+  FieldIndex.IndexOffset getLeastRecentIndexOffset(Target target);
+
+  /** Returns the lowest offset for the provided index group. */
+  FieldIndex.IndexOffset getLeastRecentIndexOffset(Collection<FieldIndex> fieldIndexes);
+
   /**
    * Returns an index that can be used to serve the provided target. Returns {@code null} if no
    * index is configured.
@@ -82,7 +94,7 @@ public interface IndexManager {
   FieldIndex getFieldIndex(Target target);
 
   /** Returns the documents that match the given target based on the provided index. */
-  Set<DocumentKey> getDocumentsMatchingTarget(FieldIndex fieldIndex, Target target);
+  Set<DocumentKey> getDocumentsMatchingTarget(Target target);
 
   /** Returns the next collection group to update. Returns {@code null} if no group exists. */
   @Nullable

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

@@ -69,7 +69,7 @@ public FieldIndex getFieldIndex(Target target) {
 
   @Override
   @Nullable
-  public Set<DocumentKey> getDocumentsMatchingTarget(FieldIndex fieldIndex, Target target) {
+  public Set<DocumentKey> getDocumentsMatchingTarget(Target target) {
     // Field indices are not supported with memory persistence.
     return Collections.emptySet();
   }
@@ -98,6 +98,21 @@ public Collection<FieldIndex> getFieldIndexes() {
     return Collections.emptyList();
   }
 
+  @Override
+  public boolean canServeFromIndex(Target target) {
+    return false;
+  }
+
+  @Override
+  public FieldIndex.IndexOffset getLeastRecentIndexOffset(Target target) {
+    return FieldIndex.IndexOffset.NONE;
+  }
+
+  @Override
+  public FieldIndex.IndexOffset getLeastRecentIndexOffset(Collection<FieldIndex> fieldIndexes) {
+    return FieldIndex.IndexOffset.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

@@ -105,22 +105,21 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
    */
   private @Nullable ImmutableSortedMap<DocumentKey, Document> performQueryUsingIndex(
       Query query, Target target) {
-    // TODO(orquery): Update this condition when we are able to serve or queries from the index.
-    if (query.matchesAllDocuments() || query.containsCompositeFilters()) {
+    if (query.matchesAllDocuments()) {
       // Don't use index queries that can be executed by scanning the collection.
       return null;
     }
 
-    FieldIndex fieldIndex = indexManager.getFieldIndex(query.toTarget());
-    if (fieldIndex == null) {
+    if (!indexManager.canServeFromIndex(target)) {
       return null;
     }
 
-    Set<DocumentKey> keys = indexManager.getDocumentsMatchingTarget(fieldIndex, target);
+    Set<DocumentKey> keys = indexManager.getDocumentsMatchingTarget(target);
     ImmutableSortedMap<DocumentKey, Document> indexedDocuments =
         localDocumentsView.getDocuments(keys);
+
     return appendRemainingResults(
-        values(indexedDocuments), query, fieldIndex.getIndexState().getOffset());
+        values(indexedDocuments), query, indexManager.getLeastRecentIndexOffset(target));
   }
 
   /**
@@ -131,8 +130,7 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       Query query,
       ImmutableSortedSet<DocumentKey> remoteKeys,
       SnapshotVersion lastLimboFreeSnapshotVersion) {
-    // TODO(orquery): Update this condition when we are able to serve or queries from the index.
-    if (query.matchesAllDocuments() || query.containsCompositeFilters()) {
+    if (query.matchesAllDocuments()) {
       // Don't use index queries that can be executed by scanning the collection.
       return null;
     }
@@ -240,7 +238,11 @@ private ImmutableSortedMap<DocumentKey, Document> appendRemainingResults(
     ImmutableSortedMap<DocumentKey, Document> remainingResults =
         localDocumentsView.getDocumentsMatchingQuery(query, offset);
     for (Document entry : indexedResults) {
-      remainingResults = remainingResults.insert(entry.getKey(), entry);
+      // For OR queries, it is possible that a document that's been indexed also shows up in
+      // "remaining results" since we use the least recent IndexOffset of all DNF terms.
+      if (!remainingResults.containsKey(entry.getKey())) {
+        remainingResults = remainingResults.insert(entry.getKey(), entry);
+      }
     }
     return remainingResults;
   }

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

@@ -21,11 +21,14 @@
 import static com.google.firebase.firestore.util.Util.repeatSequence;
 import static java.lang.Math.max;
 
+import android.text.TextUtils;
 import androidx.annotation.Nullable;
+import com.google.common.collect.ObjectArrays;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.Bound;
+import com.google.firebase.firestore.core.CompositeFilter;
 import com.google.firebase.firestore.core.FieldFilter;
 import com.google.firebase.firestore.core.Filter;
 import com.google.firebase.firestore.core.Target;
@@ -41,10 +44,13 @@
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.model.TargetIndexMatcher;
 import com.google.firebase.firestore.util.Logger;
+import com.google.firebase.firestore.util.LogicUtils;
 import com.google.firestore.admin.v1.Index;
+import com.google.firestore.v1.StructuredQuery;
 import com.google.firestore.v1.Value;
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
@@ -66,6 +72,12 @@ final class SQLiteIndexManager implements IndexManager {
   private final LocalSerializer serializer;
   private final String uid;
 
+  /**
+   * Maps from a target to its equivalent list of sub-targets. Each sub-target contains only one
+   * term from the target's disjunctive normal form (DNF).
+   */
+  private final Map<Target, List<Target>> targetToDnfSubTargets = new HashMap<>();
+
   /**
    * An in-memory copy of the index entries we've already written since the SDK launched. Used to
    * avoid re-writing the same entry repeatedly.
@@ -272,6 +284,80 @@ public Collection<FieldIndex> getFieldIndexes() {
     return allIndices;
   }
 
+  @Override
+  public boolean canServeFromIndex(Target target) {
+    for (Target subTarget : getSubTargets(target)) {
+      // If any of the sub-queries cannot be served from the index, the target as a whole cannot be
+      // served from the index.
+      if (getFieldIndex(subTarget) == null) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  @Override
+  public FieldIndex.IndexOffset getLeastRecentIndexOffset(Collection<FieldIndex> fieldIndexes) {
+    hardAssert(
+        !fieldIndexes.isEmpty(),
+        "Found empty index group when looking for least recent index offset.");
+
+    Iterator<FieldIndex> it = fieldIndexes.iterator();
+    FieldIndex.IndexOffset minOffset = it.next().getIndexState().getOffset();
+    int minBatchId = minOffset.getLargestBatchId();
+    while (it.hasNext()) {
+      FieldIndex.IndexOffset newOffset = it.next().getIndexState().getOffset();
+      if (newOffset.compareTo(minOffset) < 0) {
+        minOffset = newOffset;
+      }
+      minBatchId = Math.max(newOffset.getLargestBatchId(), minBatchId);
+    }
+
+    return FieldIndex.IndexOffset.create(
+        minOffset.getReadTime(), minOffset.getDocumentKey(), minBatchId);
+  }
+
+  @Override
+  public FieldIndex.IndexOffset getLeastRecentIndexOffset(Target target) {
+    hardAssert(
+        canServeFromIndex(target),
+        "Cannot find least recent index offset if target cannot be served from index.");
+    List<FieldIndex> fieldIndexes = new ArrayList<>();
+    for (Target subTarget : getSubTargets(target)) {
+      fieldIndexes.add(getFieldIndex(subTarget));
+    }
+    return getLeastRecentIndexOffset(fieldIndexes);
+  }
+
+  private List<Target> getSubTargets(Target target) {
+    if (targetToDnfSubTargets.containsKey(target)) {
+      return targetToDnfSubTargets.get(target);
+    }
+    List<Target> subTargets = new ArrayList<>();
+    if (target.getFilters().isEmpty()) {
+      subTargets.add(target);
+    } else {
+      // There is an implicit AND operation between all the filters stored in the target.
+      List<Filter> dnf =
+          LogicUtils.DnfTransform(
+              new CompositeFilter(
+                  target.getFilters(), StructuredQuery.CompositeFilter.Operator.AND));
+      for (Filter term : dnf) {
+        subTargets.add(
+            new Target(
+                target.getPath(),
+                target.getCollectionGroup(),
+                term.getFilters(),
+                target.getOrderBy(),
+                target.getLimit(),
+                target.getStartAt(),
+                target.getEndAt()));
+      }
+    }
+    targetToDnfSubTargets.put(target, subTargets);
+    return subTargets;
+  }
+
   /**
    * Stores the index in the memoized indexes table and updates {@link #nextIndexToUpdate}, {@link
    * #memoizedMaxIndexId} and {@link #memoizedMaxSequenceNumber}.
@@ -364,41 +450,60 @@ private SortedSet<IndexEntry> getExistingIndexEntries(
   }
 
   @Override
-  public Set<DocumentKey> getDocumentsMatchingTarget(FieldIndex fieldIndex, Target target) {
+  public Set<DocumentKey> getDocumentsMatchingTarget(Target target) {
     hardAssert(started, "IndexManager not started");
 
-    @Nullable List<Value> arrayValues = target.getArrayValues(fieldIndex);
-    @Nullable List<Value> notInValues = target.getNotInValues(fieldIndex);
-    @Nullable Bound lowerBound = target.getLowerBound(fieldIndex);
-    @Nullable Bound upperBound = target.getUpperBound(fieldIndex);
-
-    if (Logger.isDebugEnabled()) {
-      Logger.debug(
-          TAG,
-          "Using index '%s' to execute '%s' (Arrays: %s, Lower bound: %s, Upper bound: %s)",
-          fieldIndex,
-          target,
-          arrayValues,
-          lowerBound,
-          upperBound);
-    }
-
-    Object[] lowerBoundEncoded = encodeBound(fieldIndex, target, lowerBound);
-    String lowerBoundOp = lowerBound != null && lowerBound.isInclusive() ? ">=" : ">";
-    Object[] upperBoundEncoded = encodeBound(fieldIndex, target, upperBound);
-    String upperBoundOp = upperBound != null && upperBound.isInclusive() ? "<=" : "<";
-    Object[] notInEncoded = encodeValues(fieldIndex, target, notInValues);
-
-    SQLitePersistence.Query query =
-        generateQuery(
-            target,
-            fieldIndex.getIndexId(),
+    List<String> subQueries = new ArrayList<>();
+    List<Object> bindings = new ArrayList<>();
+
+    for (Target subTarget : getSubTargets(target)) {
+      FieldIndex fieldIndex = getFieldIndex(subTarget);
+      @Nullable List<Value> arrayValues = subTarget.getArrayValues(fieldIndex);
+      @Nullable List<Value> notInValues = subTarget.getNotInValues(fieldIndex);
+      @Nullable Bound lowerBound = subTarget.getLowerBound(fieldIndex);
+      @Nullable Bound upperBound = subTarget.getUpperBound(fieldIndex);
+
+      if (Logger.isDebugEnabled()) {
+        Logger.debug(
+            TAG,
+            "Using index '%s' to execute '%s' (Arrays: %s, Lower bound: %s, Upper bound: %s)",
+            fieldIndex,
+            subTarget,
             arrayValues,
-            lowerBoundEncoded,
-            lowerBoundOp,
-            upperBoundEncoded,
-            upperBoundOp,
-            notInEncoded);
+            lowerBound,
+            upperBound);
+      }
+
+      Object[] lowerBoundEncoded = encodeBound(fieldIndex, subTarget, lowerBound);
+      String lowerBoundOp = lowerBound != null && lowerBound.isInclusive() ? ">=" : ">";
+      Object[] upperBoundEncoded = encodeBound(fieldIndex, subTarget, upperBound);
+      String upperBoundOp = upperBound != null && upperBound.isInclusive() ? "<=" : "<";
+      Object[] notInEncoded = encodeValues(fieldIndex, subTarget, notInValues);
+
+      Object[] subQueryAndBindings =
+          generateQueryAndBindings(
+              subTarget,
+              fieldIndex.getIndexId(),
+              arrayValues,
+              lowerBoundEncoded,
+              lowerBoundOp,
+              upperBoundEncoded,
+              upperBoundOp,
+              notInEncoded);
+      subQueries.add(String.valueOf(subQueryAndBindings[0]));
+      bindings.addAll(Arrays.asList(subQueryAndBindings).subList(1, subQueryAndBindings.length));
+    }
+
+    String queryString =
+        subQueries.size() == 1
+            ?
+            // If there's only one subQuery, just execute the one subQuery.
+            subQueries.get(0)
+            :
+            // Construct "subQuery1 UNION subQuery2 UNION ... LIMIT N"
+            TextUtils.join(" UNION ", subQueries) + " LIMIT " + target.getLimit();
+
+    SQLitePersistence.Query query = db.query(queryString).binding(bindings.toArray());
 
     Set<DocumentKey> result = new HashSet<>();
     query.forEach(
@@ -408,8 +513,11 @@ public Set<DocumentKey> getDocumentsMatchingTarget(FieldIndex fieldIndex, Target
     return result;
   }
 
-  /** Returns a SQL query on 'index_entries' that unions all bounds. */
-  private SQLitePersistence.Query generateQuery(
+  /**
+   * Constructs a SQL query on 'index_entries' that unions all bounds. Returns an array with SQL
+   * query string as the first element, followed by binding arguments.
+   */
+  private Object[] generateQueryAndBindings(
       Target target,
       int indexId,
       @Nullable List<Value> arrayValues,
@@ -461,7 +569,7 @@ private SQLitePersistence.Query generateQuery(
     // Fill in the bind ("question marks") variables.
     Object[] bindArgs =
         fillBounds(statementCount, indexId, arrayValues, lowerBounds, upperBounds, notIn);
-    return db.query(sql.toString()).binding(bindArgs);
+    return ObjectArrays.concat(sql.toString(), bindArgs);
   }
 
   /** Returns the bind arguments for all {@code statementCount} statements. */