Where tests

Author: tom-andersen

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

@@ -25,6 +25,7 @@ import com.google.firebase.firestore.pipeline.AddFieldsStage
 import com.google.firebase.firestore.pipeline.AggregateFunction
 import com.google.firebase.firestore.pipeline.AggregateStage
 import com.google.firebase.firestore.pipeline.AggregateWithAlias
+import com.google.firebase.firestore.pipeline.BaseStage
 import com.google.firebase.firestore.pipeline.BooleanExpr
 import com.google.firebase.firestore.pipeline.CollectionGroupSource
 import com.google.firebase.firestore.pipeline.CollectionSource
@@ -38,19 +39,18 @@ import com.google.firebase.firestore.pipeline.Field
 import com.google.firebase.firestore.pipeline.FindNearestStage
 import com.google.firebase.firestore.pipeline.FunctionExpr
 import com.google.firebase.firestore.pipeline.InternalOptions
-import com.google.firebase.firestore.pipeline.Stage
 import com.google.firebase.firestore.pipeline.LimitStage
 import com.google.firebase.firestore.pipeline.OffsetStage
 import com.google.firebase.firestore.pipeline.Ordering
 import com.google.firebase.firestore.pipeline.PipelineOptions
+import com.google.firebase.firestore.pipeline.RawStage
 import com.google.firebase.firestore.pipeline.RealtimePipelineOptions
 import com.google.firebase.firestore.pipeline.RemoveFieldsStage
 import com.google.firebase.firestore.pipeline.ReplaceStage
 import com.google.firebase.firestore.pipeline.SampleStage
 import com.google.firebase.firestore.pipeline.SelectStage
 import com.google.firebase.firestore.pipeline.Selectable
 import com.google.firebase.firestore.pipeline.SortStage
-import com.google.firebase.firestore.pipeline.BaseStage
 import com.google.firebase.firestore.pipeline.UnionStage
 import com.google.firebase.firestore.pipeline.UnnestStage
 import com.google.firebase.firestore.pipeline.WhereStage
@@ -159,10 +159,10 @@ private constructor(
    * This method provides a way to call stages that are supported by the Firestore backend but that
    * are not implemented in the SDK version being used.
    *
-   * @param stage An [Stage] object that specifies stage name and parameters.
+   * @param rawStage An [RawStage] object that specifies stage name and parameters.
    * @return A new [Pipeline] object with this stage appended to the stage list.
    */
-  fun addStage(stage: Stage): Pipeline = append(stage)
+  fun addStage(rawStage: RawStage): Pipeline = append(rawStage)
 
   /**
    * Adds new fields to outputs from previous stages.

firebase-firestore/src/main/java/com/google/firebase/firestore/core/pipeline.kt

@@ -92,7 +92,7 @@ abstract class Expr internal constructor() {
               }
               .toTypedArray()
           )
-        is List<*> -> ListOfExprs(value.map(toExpr).toTypedArray())
+        is List<*> -> array(value)
         else -> null
       }
     }
@@ -1018,7 +1018,7 @@ abstract class Expr internal constructor() {
      */
     @JvmStatic
     fun eqAny(expression: Expr, values: List<Any>): BooleanExpr =
-      eqAny(expression, ListOfExprs(toArrayOfExprOrConstant(values)))
+      eqAny(expression, array(values))
 
     /**
      * Creates an expression that checks if an [expression], when evaluated, is equal to any of the
@@ -1043,7 +1043,7 @@ abstract class Expr internal constructor() {
      */
     @JvmStatic
     fun eqAny(fieldName: String, values: List<Any>): BooleanExpr =
-      eqAny(fieldName, ListOfExprs(toArrayOfExprOrConstant(values)))
+      eqAny(fieldName, array(values))
 
     /**
      * Creates an expression that checks if a field's value is equal to any of the elements of
@@ -1068,7 +1068,7 @@ abstract class Expr internal constructor() {
      */
     @JvmStatic
     fun notEqAny(expression: Expr, values: List<Any>): BooleanExpr =
-      notEqAny(expression, ListOfExprs(toArrayOfExprOrConstant(values)))
+      notEqAny(expression, array(values))
 
     /**
      * Creates an expression that checks if an [expression], when evaluated, is not equal to all the
@@ -1093,7 +1093,7 @@ abstract class Expr internal constructor() {
      */
     @JvmStatic
     fun notEqAny(fieldName: String, values: List<Any>): BooleanExpr =
-      notEqAny(fieldName, ListOfExprs(toArrayOfExprOrConstant(values)))
+      notEqAny(fieldName, array(values))
 
     /**
      * Creates an expression that checks if a field's value is not equal to all of the elements of
@@ -2776,7 +2776,7 @@ abstract class Expr internal constructor() {
      */
     @JvmStatic
     fun arrayContainsAll(array: Expr, values: List<Any>) =
-      arrayContainsAll(array, ListOfExprs(toArrayOfExprOrConstant(values)))
+      arrayContainsAll(array, array(values))
 
     /**
      * Creates an expression that checks if [array] contains all elements of [arrayExpression].
@@ -2802,7 +2802,7 @@ abstract class Expr internal constructor() {
         "array_contains_all",
         evaluateArrayContainsAll,
         arrayFieldName,
-        ListOfExprs(toArrayOfExprOrConstant(values))
+        array(values)
       )
 
     /**
@@ -2829,7 +2829,7 @@ abstract class Expr internal constructor() {
         "array_contains_any",
         evaluateArrayContainsAny,
         array,
-        ListOfExprs(toArrayOfExprOrConstant(values))
+        array(values)
       )
 
     /**
@@ -2856,7 +2856,7 @@ abstract class Expr internal constructor() {
         "array_contains_any",
         evaluateArrayContainsAny,
         arrayFieldName,
-        ListOfExprs(toArrayOfExprOrConstant(values))
+        array(values)
       )
 
     /**
@@ -4197,17 +4197,6 @@ class Field internal constructor(private val fieldPath: ModelFieldPath) : Select
   }
 }
 
-internal class ListOfExprs(private val expressions: Array<out Expr>) : Expr() {
-  override fun toProto(userDataReader: UserDataReader): Value =
-    encodeValue(expressions.map { it.toProto(userDataReader) })
-
-  override fun evaluateContext(
-    context: EvaluationContext
-  ): (input: MutableDocument) -> EvaluateResult {
-    TODO("Not yet implemented")
-  }
-}
-
 /**
  * This class defines the base class for Firestore [Pipeline] functions, which can be evaluated
  * within pipeline execution.
@@ -4378,7 +4367,7 @@ internal constructor(name: String, function: EvaluateFunction, params: Array<out
  *
  * You create [Ordering] instances using the [ascending] and [descending] helper methods.
  */
-class Ordering private constructor(val expr: Expr, private val dir: Direction) {
+class Ordering private constructor(val expr: Expr, internal val dir: Direction) {
   companion object {
 
     /**
@@ -4416,12 +4405,9 @@ class Ordering private constructor(val expr: Expr, private val dir: Direction) {
     fun descending(fieldName: String): Ordering = Ordering(field(fieldName), Direction.DESCENDING)
   }
 
-  private class Direction private constructor(val proto: Value) {
-    private constructor(protoString: String) : this(encodeValue(protoString))
-    companion object {
-      val ASCENDING = Direction("ascending")
-      val DESCENDING = Direction("descending")
-    }
+  internal enum class Direction(val proto: Value) {
+    ASCENDING(encodeValue("ascending")),
+    DESCENDING(encodeValue("descending"))
   }
 
   internal fun toProto(userDataReader: UserDataReader): Value =

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/expressions.kt

@@ -29,9 +29,14 @@ import com.google.firestore.v1.Pipeline
 import com.google.firestore.v1.Value
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.filter
-
-abstract class BaseStage<T : BaseStage<T>>
-internal constructor(protected val name: String, internal val options: InternalOptions) {
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.toList
+
+sealed class BaseStage<T : BaseStage<T>>(
+  protected val name: String,
+  internal val options: InternalOptions
+) {
   internal fun toProtoStage(userDataReader: UserDataReader): Pipeline.Stage {
     val builder = Pipeline.Stage.newBuilder()
     builder.setName(name)
@@ -106,32 +111,32 @@ internal constructor(protected val name: String, internal val options: InternalO
  * This class provides a way to call stages that are supported by the Firestore backend but that are
  * not implemented in the SDK version being used.
  */
-class Stage
+class RawStage
 private constructor(
   name: String,
   private val arguments: List<GenericArg>,
   options: InternalOptions = InternalOptions.EMPTY
-) : BaseStage<Stage>(name, options) {
+) : BaseStage<RawStage>(name, options) {
   companion object {
     /**
      * Specify name of stage
      *
      * @param name The unique name of the stage to add.
-     * @return [Stage] with specified parameters.
+     * @return [RawStage] with specified parameters.
      */
-    @JvmStatic fun ofName(name: String) = Stage(name, emptyList(), InternalOptions.EMPTY)
+    @JvmStatic fun ofName(name: String) = RawStage(name, emptyList(), InternalOptions.EMPTY)
   }
 
-  override fun self(options: InternalOptions) = Stage(name, arguments, options)
+  override fun self(options: InternalOptions) = RawStage(name, arguments, options)
 
   /**
    * Specify arguments to stage.
    *
    * @param arguments A list of ordered parameters to configure the stage's behavior.
-   * @return [Stage] with specified parameters.
+   * @return [RawStage] with specified parameters.
    */
-  fun withArguments(vararg arguments: Any): Stage =
-    Stage(name, arguments.map(GenericArg::from), options)
+  fun withArguments(vararg arguments: Any): RawStage =
+    RawStage(name, arguments.map(GenericArg::from), options)
 
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
     arguments.asSequence().map { it.toProto(userDataReader) }
@@ -546,6 +551,43 @@ internal constructor(
   override fun self(options: InternalOptions) = SortStage(orders, options)
   override fun args(userDataReader: UserDataReader): Sequence<Value> =
     orders.asSequence().map { it.toProto(userDataReader) }
+
+  override fun evaluate(
+    context: EvaluationContext,
+    inputs: Flow<MutableDocument>
+  ): Flow<MutableDocument> {
+    val evaluates: Array<EvaluateDocument> =
+      orders.map { it.expr.evaluateContext(context) }.toTypedArray()
+    val directions: Array<Ordering.Direction> = orders.map { it.dir }.toTypedArray()
+    return flow {
+      inputs
+        // For each document, lazily evaluate order expression values.
+        .map { doc ->
+          val orderValues =
+            evaluates
+              .map { lazy(LazyThreadSafetyMode.PUBLICATION) { it(doc).value ?: Values.MIN_VALUE } }
+              .toTypedArray<Lazy<Value>>()
+          Pair(doc, orderValues)
+        }
+        .toList()
+        .sortedWith(
+          Comparator { px, py ->
+            val x = px.second
+            val y = py.second
+            directions.forEachIndexed<Ordering.Direction> { i, dir ->
+              val r =
+                when (dir) {
+                  Ordering.Direction.ASCENDING -> Values.compare(x[i].value, y[i].value)
+                  Ordering.Direction.DESCENDING -> Values.compare(y[i].value, x[i].value)
+                }
+              if (r != 0) return@Comparator r
+            }
+            0
+          }
+        )
+        .forEach { p -> emit(p.first) }
+    }
+  }
 }
 
 internal class DistinctStage

firebase-firestore/src/main/java/com/google/firebase/firestore/pipeline/stage.kt

@@ -1,4 +1,18 @@
-package com.google.firebase.firestore.core
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.firebase.firestore.pipeline
 
 import com.google.common.truth.Truth.assertThat
 import com.google.firebase.firestore.RealtimePipelineSource
@@ -10,7 +24,10 @@ import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
 import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
 
+@RunWith(RobolectricTestRunner::class)
 internal class PipelineTests {
 
   @Test