Convert query to pipeline

Author: tom-andersen

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/QueryToPipelineTest.java

@@ -38,6 +38,8 @@
 import com.google.firebase.firestore.FirebaseFirestoreSettings;
 import com.google.firebase.firestore.ListenerRegistration;
 import com.google.firebase.firestore.MetadataChanges;
+import com.google.firebase.firestore.PipelineResult;
+import com.google.firebase.firestore.PipelineSnapshot;
 import com.google.firebase.firestore.Query;
 import com.google.firebase.firestore.QuerySnapshot;
 import com.google.firebase.firestore.Source;
@@ -465,6 +467,15 @@ public static List<Map<String, Object>> querySnapshotToValues(QuerySnapshot quer
     return res;
   }
 
+  public static List<Map<String, Object>> pipelineSnapshotToValues(
+      PipelineSnapshot pipelineSnapshot) {
+    List<Map<String, Object>> res = new ArrayList<>();
+    for (PipelineResult result : pipelineSnapshot) {
+      res.add(result.getData());
+    }
+    return res;
+  }
+
   public static List<String> querySnapshotToIds(QuerySnapshot querySnapshot) {
     List<String> res = new ArrayList<>();
     for (DocumentSnapshot doc : querySnapshot) {
@@ -473,6 +484,15 @@ public static List<String> querySnapshotToIds(QuerySnapshot querySnapshot) {
     return res;
   }
 
+  public static List<String> pipelineSnapshotToIds(PipelineSnapshot pipelineResults) {
+    List<String> res = new ArrayList<>();
+    for (PipelineResult result : pipelineResults) {
+      DocumentReference ref = result.getRef();
+      res.add(ref == null ? null : ref.getId());
+    }
+    return res;
+  }
+
   public static void disableNetwork(FirebaseFirestore firestore) {
     if (firestoreStatus.get(firestore)) {
       waitFor(firestore.disableNetwork());
@@ -537,4 +557,23 @@ public static void checkOnlineAndOfflineResultsMatch(Query query, String... expe
       assertEquals(expected, querySnapshotToIds(docsFromCache));
     }
   }
+
+  /**
+   * Checks that running the query while online (against the backend/emulator) results in the same
+   * documents as running the query while offline. If `expectedDocs` is provided, it also checks
+   * that both online and offline query result is equal to the expected documents.
+   *
+   * @param query The query to check
+   * @param expectedDocs Ordered list of document keys that are expected to match the query
+   */
+  public static void checkQueryAndPipelineResultsMatch(Query query, String... expectedDocs) {
+    QuerySnapshot docsFromQuery = waitFor(query.get(Source.SERVER));
+    PipelineSnapshot docsFromPipeline = waitFor(query.pipeline().execute());
+
+    assertEquals(querySnapshotToIds(docsFromQuery), pipelineSnapshotToIds(docsFromPipeline));
+    List<String> expected = asList(expectedDocs);
+    if (!expected.isEmpty()) {
+      assertEquals(expected, querySnapshotToIds(docsFromQuery));
+    }
+  }
 }

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/testutil/IntegrationTestUtil.java

@@ -130,6 +130,7 @@ public Task<DocumentReference> add(@NonNull Object data) {
   }
 
   @NonNull
+  @Override
   public Pipeline pipeline() {
     return new Pipeline(firestore, firestore.getUserDataReader(), new CollectionSource(getPath()));
   }

firebase-firestore/src/main/java/com/google/firebase/firestore/CollectionReference.java

@@ -20,6 +20,7 @@ import com.google.common.collect.FluentIterable
 import com.google.common.collect.ImmutableList
 import com.google.firebase.firestore.model.DocumentKey
 import com.google.firebase.firestore.model.SnapshotVersion
+import com.google.firebase.firestore.model.Values
 import com.google.firebase.firestore.pipeline.AddFieldsStage
 import com.google.firebase.firestore.pipeline.AggregateStage
 import com.google.firebase.firestore.pipeline.AggregateWithAlias
@@ -123,9 +124,9 @@ internal constructor(
 
   fun where(condition: BooleanExpr): Pipeline = append(WhereStage(condition))
 
-  fun offset(offset: Long): Pipeline = append(OffsetStage(offset))
+  fun offset(offset: Int): Pipeline = append(OffsetStage(offset))
 
-  fun limit(limit: Long): Pipeline = append(LimitStage(limit))
+  fun limit(limit: Int): Pipeline = append(LimitStage(limit))
 
   fun distinct(vararg groups: Selectable): Pipeline = append(DistinctStage(groups))
 
@@ -175,12 +176,15 @@ internal constructor(
     append(UnnestStage(selectable))
 
   private inner class ObserverSnapshotTask : PipelineResultObserver {
+    private val userDataWriter =
+      UserDataWriter(firestore, DocumentSnapshot.ServerTimestampBehavior.DEFAULT)
     private val taskCompletionSource = TaskCompletionSource<PipelineSnapshot>()
     private val results: ImmutableList.Builder<PipelineResult> = ImmutableList.builder()
     override fun onDocument(key: DocumentKey?, data: Map<String, Value>, version: SnapshotVersion) {
       results.add(
         PipelineResult(
           firestore,
+          userDataWriter,
           if (key == null) null else DocumentReference(key, firestore),
           data,
           version
@@ -249,20 +253,52 @@ class PipelineSource internal constructor(private val firestore: FirebaseFiresto
 }
 
 class PipelineSnapshot
-internal constructor(private val executionTime: SnapshotVersion, val results: List<PipelineResult>)
+internal constructor(
+  private val executionTime: SnapshotVersion,
+  val results: List<PipelineResult>
+) : Iterable<PipelineResult> {
+  override fun iterator() = results.iterator()
+}
 
 class PipelineResult
 internal constructor(
   private val firestore: FirebaseFirestore,
+  private val userDataWriter: UserDataWriter,
   val ref: DocumentReference?,
   private val fields: Map<String, Value>,
   private val version: SnapshotVersion,
 ) {
 
-  fun getData(): Map<String, Any?> = userDataWriter().convertObject(fields)
+  /**
+   * Returns the ID of the document represented by this result. Returns null if this result is not
+   * corresponding to a Firestore document.
+   */
+  fun getId(): String? = ref?.id
+
+  fun getData(): Map<String, Any?> = userDataWriter.convertObject(fields)
+
+  private fun extractNestedValue(fieldPath: FieldPath): Value? {
+    val segments = fieldPath.internalPath.iterator()
+    if (!segments.hasNext()) {
+      return Values.encodeValue(fields)
+    }
+    val firstSegment = segments.next()
+    if (!fields.containsKey(firstSegment)) {
+      return null
+    }
+    var value: Value? = fields[firstSegment]
+    for (segment in segments) {
+      if (value == null || !value.hasMapValue()) {
+        return null
+      }
+      value = value.mapValue.getFieldsOrDefault(segment, null)
+    }
+    return value
+  }
+
+  fun get(field: String): Any? = get(FieldPath.fromDotSeparatedPath(field))
 
-  private fun userDataWriter(): UserDataWriter =
-    UserDataWriter(firestore, DocumentSnapshot.ServerTimestampBehavior.DEFAULT)
+  fun get(fieldPath: FieldPath): Any? = userDataWriter.convertValue(extractNestedValue(fieldPath))
 
   override fun toString() = "PipelineResult{ref=$ref, version=$version}, data=${getData()}"
 }

firebase-firestore/src/main/java/com/google/firebase/firestore/Pipeline.kt

@@ -1241,6 +1241,11 @@ public AggregateQuery aggregate(
     return new AggregateQuery(this, fields);
   }
 
+  @NonNull
+  public Pipeline pipeline() {
+    return query.toPipeline(firestore, firestore.getUserDataReader());
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

firebase-firestore/src/main/java/com/google/firebase/firestore/Query.java

@@ -230,7 +230,7 @@ private ObjectValue convertAndParseDocumentData(Object input, ParseContext conte
 
     Object converted = CustomClassMapper.convertToPlainJavaTypes(input);
     Value parsedValue = parseData(converted, context);
-    if (parsedValue.getValueTypeCase() != Value.ValueTypeCase.MAP_VALUE) {
+    if (!parsedValue.hasMapValue()) {
       throw new IllegalArgumentException(badDocReason + "of type: " + Util.typeName(input));
     }
     return new ObjectValue(parsedValue);

firebase-firestore/src/main/java/com/google/firebase/firestore/UserDataReader.java

@@ -78,7 +78,7 @@ public Object convertValue(Value value) {
       case TYPE_ORDER_BOOLEAN:
         return value.getBooleanValue();
       case TYPE_ORDER_NUMBER:
-        return value.getValueTypeCase().equals(Value.ValueTypeCase.INTEGER_VALUE)
+        return value.hasIntegerValue()
             ? (Object) value.getIntegerValue() // Cast to Object to prevent type coercion to double
             : (Object) value.getDoubleValue();
       case TYPE_ORDER_STRING:

firebase-firestore/src/main/java/com/google/firebase/firestore/UserDataWriter.java

@@ -17,6 +17,7 @@
 import android.text.TextUtils;
 import androidx.annotation.Nullable;
 import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.pipeline.BooleanExpr;
 import com.google.firebase.firestore.util.Function;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -167,6 +168,19 @@ public String getCanonicalId() {
     return builder.toString();
   }
 
+  @Override
+  BooleanExpr toPipelineExpr() {
+    BooleanExpr[] booleanExprs = filters.stream().map(Filter::toPipelineExpr).toArray(BooleanExpr[]::new);
+    switch (operator) {
+      case AND:
+        return new BooleanExpr("and", booleanExprs);
+      case OR:
+        return new BooleanExpr("or", booleanExprs);
+    }
+    // Handle OPERATOR_UNSPECIFIED and UNRECOGNIZED cases as needed
+    throw new IllegalArgumentException("Unsupported operator: " + operator);
+  }
+
   @Override
   public String toString() {
     return getCanonicalId();

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

@@ -14,11 +14,15 @@
 
 package com.google.firebase.firestore.core;
 
+import static com.google.firebase.firestore.pipeline.Function.and;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
+import static java.lang.Double.isNaN;
 
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.FieldPath;
 import com.google.firebase.firestore.model.Values;
+import com.google.firebase.firestore.pipeline.BooleanExpr;
+import com.google.firebase.firestore.pipeline.Field;
 import com.google.firebase.firestore.util.Assert;
 import com.google.firestore.v1.Value;
 import java.util.Arrays;
@@ -172,6 +176,49 @@ public List<Filter> getFilters() {
     return Collections.singletonList(this);
   }
 
+  @Override
+  BooleanExpr toPipelineExpr() {
+    Field x = new Field(field);
+    BooleanExpr exists = x.exists();
+    switch (operator) {
+      case LESS_THAN:
+        return and(exists, x.lt(value));
+      case LESS_THAN_OR_EQUAL:
+        return and(exists, x.lte(value));
+      case EQUAL:
+        if (value.hasNullValue()) {
+          return and(exists, x.isNull());
+        } else if (value.hasDoubleValue() && isNaN(value.getDoubleValue())) {
+          return and(exists, x.isNan());
+        } else {
+          return and(exists, x.eq(value));
+        }
+      case NOT_EQUAL:
+        if (value.hasNullValue()) {
+          return and(exists, x.isNotNull());
+        } else if (value.hasDoubleValue() && isNaN(value.getDoubleValue())) {
+          return and(exists, x.isNotNan());
+        } else {
+          return and(exists, x.neq(value));
+        }
+      case GREATER_THAN:
+        return and(exists, x.gt(value));
+      case GREATER_THAN_OR_EQUAL:
+        return and(exists, x.gte(value));
+      case ARRAY_CONTAINS:
+        return and(exists, x.arrayContains(value));
+      case ARRAY_CONTAINS_ANY:
+        return and(exists, x.arrayContainsAny(value.getArrayValue().getValuesList()));
+      case IN:
+        return and(exists, x.eqAny(value.getArrayValue().getValuesList()));
+      case NOT_IN:
+        return and(exists, x.notEqAny(value.getArrayValue().getValuesList()));
+      default:
+        // Handle OPERATOR_UNSPECIFIED and UNRECOGNIZED cases as needed
+        throw new IllegalArgumentException("Unsupported operator: " + operator);
+    }
+  }
+
   @Override
   public String toString() {
     return getCanonicalId();

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

@@ -15,6 +15,7 @@
 package com.google.firebase.firestore.core;
 
 import com.google.firebase.firestore.model.Document;
+import com.google.firebase.firestore.pipeline.BooleanExpr;
 import java.util.List;
 
 public abstract class Filter {
@@ -29,4 +30,6 @@ public abstract class Filter {
 
   /** Returns a list of all filters that are contained within this filter */
   public abstract List<Filter> getFilters();
+
+  abstract BooleanExpr toPipelineExpr();
 }