feat: Added public preview support for full-text search and geo search in Pipelines. (#2346)

Author: MarkDuckworth

java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java

@@ -54,6 +54,7 @@
 import com.google.cloud.firestore.pipeline.stages.RemoveFields;
 import com.google.cloud.firestore.pipeline.stages.ReplaceWith;
 import com.google.cloud.firestore.pipeline.stages.Sample;
+import com.google.cloud.firestore.pipeline.stages.Search;
 import com.google.cloud.firestore.pipeline.stages.Select;
 import com.google.cloud.firestore.pipeline.stages.Sort;
 import com.google.cloud.firestore.pipeline.stages.Stage;
@@ -223,6 +224,21 @@ private Pipeline append(Stage stage) {
     return new Pipeline(this.rpcContext, stages.append(stage));
   }
 
+  /**
+   * Adds a search stage to the Pipeline.
+   *
+   * <p>This must be the first stage of the pipeline.
+   *
+   * <p>A limited set of expressions are supported in the search stage.
+   *
+   * @param searchStage An object that specifies how search is performed.
+   * @return A new {@code Pipeline} object with this stage appended to the stage list.
+   */
+  @BetaApi
+  public Pipeline search(Search searchStage) {
+    return append(searchStage);
+  }
+
   /**
    * Adds new fields to outputs from previous stages.
    *
@@ -1575,9 +1591,9 @@ public void onError(Throwable t) {
   }
 
   @InternalApi
-  private com.google.firestore.v1.Pipeline toProto() {
+  public com.google.firestore.v1.Pipeline toProto() {
     return com.google.firestore.v1.Pipeline.newBuilder()
-        .addAllStages(stages.transform(StageUtils::toStageProto))
+        .addAllStages(stages.transform(Stage::toStageProto))
         .build();
   }
 

java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/expressions/BooleanFunctionExpression.java

@@ -36,6 +36,11 @@ class BooleanFunctionExpression extends BooleanExpression {
     this(new FunctionExpression(name, java.util.Arrays.asList(params)));
   }
 
+  BooleanFunctionExpression(
+      String name, List<? extends Expression> params, java.util.Map<String, Value> options) {
+    this(new FunctionExpression(name, params, options));
+  }
+
   @Override
   Value toProto() {
     return expr.toProto();

java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/expressions/Expression.java

@@ -121,7 +121,7 @@ public static Expression constant(Timestamp value) {
    */
   @BetaApi
   public static BooleanExpression constant(Boolean value) {
-    return equal(new Constant(value), true);
+    return new BooleanConstant(new Constant(value));
   }
 
   /**
@@ -4593,6 +4593,200 @@ public static BooleanExpression isError(Expression expr) {
     return new BooleanFunctionExpression("is_error", expr);
   }
 
+  /**
+   * Evaluates to the distance in meters between the location in the specified field and the query
+   * location.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * @param fieldName Specifies the field in the document which contains the first {@link GeoPoint}
+   *     for distance computation.
+   * @param location Compute distance to this {@link GeoPoint}.
+   * @return A new {@link Expression} representing the geoDistance operation.
+   */
+  @BetaApi
+  public static Expression geoDistance(String fieldName, GeoPoint location) {
+    return geoDistance(field(fieldName), location);
+  }
+
+  /**
+   * Evaluates to the distance in meters between the location in the specified field and the query
+   * location.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * db.pipeline().collection("restaurants").search(
+   *   Search.withQuery("waffles").withSort(geoDistance(field("location"), new GeoPoint(37.0, -122.0)).ascending())
+   * )
+   * }</pre>
+   *
+   * @param field Specifies the field in the document which contains the first {@link GeoPoint} for
+   *     distance computation.
+   * @param location Compute distance to this {@link GeoPoint}.
+   * @return A new {@link Expression} representing the geoDistance operation.
+   */
+  @BetaApi
+  public static Expression geoDistance(Field field, GeoPoint location) {
+    return new FunctionExpression(
+        "geo_distance", java.util.Arrays.asList(field, constant(location)));
+  }
+
+  /**
+   * Perform a full-text search on all indexed search fields in the document.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * db.pipeline().collection("restaurants").search(Search.withQuery(documentMatches("waffles OR pancakes")))
+   * }</pre>
+   *
+   * @param rquery Define the search query using the search domain-specific language (DSL).
+   * @return A new {@link BooleanExpression} representing the documentMatches operation.
+   */
+  @BetaApi
+  public static BooleanExpression documentMatches(String rquery) {
+    return new BooleanFunctionExpression("document_matches", constant(rquery));
+  }
+
+  /**
+   * Perform a full-text search on the specified field.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * db.pipeline().collection("restaurants").search(Search.withQuery(matches("menu", "waffles")))
+   * }</pre>
+   *
+   * @param fieldName Perform search on this field.
+   * @param rquery Define the search query using the search domain-specific language (DSL).
+   */
+  @InternalApi
+  static BooleanExpression matches(String fieldName, String rquery) {
+    return matches(field(fieldName), rquery);
+  }
+
+  /**
+   * Perform a full-text search on the specified field.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * @param field Perform search on this field.
+   * @param rquery Define the search query using the search domain-specific language (DSL).
+   */
+  @InternalApi
+  static BooleanExpression matches(Field field, String rquery) {
+    return new BooleanFunctionExpression("matches", field, constant(rquery));
+  }
+
+  /**
+   * Evaluates to the search score that reflects the topicality of the document to all of the text
+   * predicates (for example: {@code documentMatches}) in the search query.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * db.pipeline().collection("restaurants").search(
+   *   Search.withQuery("waffles").withSort(score().descending())
+   * )
+   * }</pre>
+   *
+   * @return A new {@link Expression} representing the score operation.
+   */
+  @BetaApi
+  public static Expression score() {
+    return new FunctionExpression("score", com.google.common.collect.ImmutableList.of());
+  }
+
+  /**
+   * Evaluates to an HTML-formatted text snippet that highlights terms matching the search query in
+   * {@code <b>bold</b>}.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * db.pipeline().collection("restaurants").search(
+   *   Search.withQuery("waffles").withAddFields(snippet("menu", "waffles").as("snippet"))
+   * )
+   * }</pre>
+   *
+   * @param fieldName Search the specified field for matching terms.
+   * @param rquery Define the search query using the search domain-specific language (DSL).
+   * @return A new {@link Expression} representing the snippet operation.
+   */
+  @BetaApi
+  public static Expression snippet(String fieldName, String rquery) {
+    return new FunctionExpression(
+        "snippet", java.util.Arrays.asList(field(fieldName), constant(rquery)));
+  }
+
+  /**
+   * Evaluates to an HTML-formatted text snippet that highlights terms matching the search query in
+   * {@code <b>bold</b>}.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * db.pipeline().collection("restaurants").search(
+   *   Search.withQuery("waffles").withAddFields(field("menu").snippet("waffles").as("snippet"))
+   * )
+   * }</pre>
+   *
+   * @param rquery Define the search query using the search domain-specific language (DSL).
+   * @return A new {@link Expression} representing the snippet operation.
+   */
+  @BetaApi
+  public final Expression snippet(String rquery) {
+    return new FunctionExpression(
+        "snippet",
+        java.util.Arrays.asList(this, constant(rquery)),
+        java.util.Collections.singletonMap(
+            "query", com.google.cloud.firestore.PipelineUtils.encodeValue(rquery)));
+  }
+
+  @InternalApi
+  static BooleanExpression between(String fieldName, Expression lowerBound, Expression upperBound) {
+    return between(field(fieldName), lowerBound, upperBound);
+  }
+
+  @InternalApi
+  static BooleanExpression between(String fieldName, Object lowerBound, Object upperBound) {
+    return between(fieldName, toExprOrConstant(lowerBound), toExprOrConstant(upperBound));
+  }
+
+  @InternalApi
+  static BooleanExpression between(
+      Expression expression, Expression lowerBound, Expression upperBound) {
+    return new BooleanFunctionExpression("between", expression, lowerBound, upperBound);
+  }
+
+  @InternalApi
+  static BooleanExpression between(Expression expression, Object lowerBound, Object upperBound) {
+    return between(expression, toExprOrConstant(lowerBound), toExprOrConstant(upperBound));
+  }
+
+  @InternalApi
+  public final BooleanExpression between(Expression lowerBound, Expression upperBound) {
+    return Expression.between(this, lowerBound, upperBound);
+  }
+
+  @InternalApi
+  public final BooleanExpression between(Object lowerBound, Object upperBound) {
+    return Expression.between(this, lowerBound, upperBound);
+  }
+
   // Other Utility Functions
   /**
    * Creates an expression that returns the document ID from a path.

java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/expressions/Field.java

@@ -90,6 +90,32 @@ public Value toProto() {
     return Value.newBuilder().setFieldReferenceValue(path.toString()).build();
   }
 
+  /**
+   * Evaluates to the distance in meters between the location specified by this field and the query
+   * location.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * @param location Compute distance to this {@link com.google.cloud.firestore.GeoPoint}.
+   * @return A new {@link Expression} representing the geoDistance operation.
+   */
+  @BetaApi
+  public Expression geoDistance(com.google.cloud.firestore.GeoPoint location) {
+    return Expression.geoDistance(this, location);
+  }
+
+  /**
+   * Perform a full-text search on this field.
+   *
+   * <p>This Expression can only be used within a {@code Search} stage.
+   *
+   * @param rquery Define the search query using the rquery DTS.
+   */
+  @InternalApi
+  BooleanExpression matches(String rquery) {
+    return Expression.matches(this, rquery);
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/expressions/FunctionExpression.java

@@ -18,20 +18,26 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
-import com.google.common.base.Objects;
 import com.google.firestore.v1.Value;
-import java.util.Collections;
 import java.util.List;
 import java.util.stream.Collectors;
 
 @BetaApi
 public class FunctionExpression extends Expression {
   private final String name;
   private final List<Expression> params;
+  private final java.util.Map<String, Value> options;
 
   public FunctionExpression(String name, List<? extends Expression> params) {
+    this(name, params, java.util.Collections.emptyMap());
+  }
+
+  @InternalApi
+  FunctionExpression(
+      String name, List<? extends Expression> params, java.util.Map<String, Value> options) {
     this.name = name;
-    this.params = Collections.unmodifiableList(params);
+    this.params = java.util.Collections.unmodifiableList(params);
+    this.options = java.util.Collections.unmodifiableMap(options);
   }
 
   @InternalApi
@@ -44,7 +50,8 @@ Value toProto() {
                 .addAllArgs(
                     this.params.stream()
                         .map(FunctionUtils::exprToValue)
-                        .collect(Collectors.toList())))
+                        .collect(Collectors.toList()))
+                .putAllOptions(this.options))
         .build();
   }
 
@@ -53,11 +60,13 @@ public boolean equals(Object o) {
     if (this == o) return true;
     if (o == null || getClass() != o.getClass()) return false;
     FunctionExpression that = (FunctionExpression) o;
-    return Objects.equal(name, that.name) && Objects.equal(params, that.params);
+    return java.util.Objects.equals(name, that.name)
+        && java.util.Objects.equals(params, that.params)
+        && java.util.Objects.equals(options, that.options);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hashCode(name, params);
+    return java.util.Objects.hash(name, params, options);
   }
 }

java-firestore/google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/expressions/Selectable.java

@@ -19,4 +19,24 @@
 import com.google.api.core.BetaApi;
 
 @BetaApi
-public interface Selectable {}
+public interface Selectable {
+  /**
+   * Converts an object to a {@link Selectable}.
+   *
+   * @param o The object to convert. Supported types: {@link Selectable}, {@link String}, {@link
+   *     com.google.cloud.firestore.FieldPath}.
+   * @return The converted {@link Selectable}.
+   */
+  @com.google.api.core.InternalApi
+  static Selectable toSelectable(Object o) {
+    if (o instanceof Selectable) {
+      return (Selectable) o;
+    } else if (o instanceof String) {
+      return Expression.field((String) o);
+    } else if (o instanceof com.google.cloud.firestore.FieldPath) {
+      return Expression.field((com.google.cloud.firestore.FieldPath) o);
+    } else {
+      throw new IllegalArgumentException("Unknown Selectable type: " + o.getClass().getName());
+    }
+  }
+}