Expand example and fix.

googleapis · Dec 18, 2024 · 94e5750 · 94e5750
1 parent 6da4709
commit 94e5750
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 16 deletions.
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/Pipeline.java
@@ -505,12 +505,10 @@ public Pipeline distinct(Selectable... selectables) {
    * <pre>{@code
    * // Find books with similar "topicVectors" to the given targetVector
    * firestore.pipeline().collection("books")
-   *     .findNearest("topicVectors", targetVector, FindNearest.DistanceMeasure.cosine(),
-   *        FindNearestOptions
-   *          .builder()
-   *          .limit(10)
-   *          .distanceField("distance")
-   *          .build());
+   *     .findNearest("topicVectors", targetVector, FindNearest.DistanceMeasure.COSINE,
+   *        FindNearestOptions.DEFAULT
+   *          .withLimit(10)
+   *          .withDistanceField("distance"));
    * }</pre>
    *
    * @param fieldName The name of the field containing the vector data. This field should store
@@ -544,12 +542,10 @@ public Pipeline findNearest(
    * // Find books with similar "topicVectors" to the given targetVector
    * firestore.pipeline().collection("books")
    *     .findNearest(
-   *        FindNearest.of(Field.of("topicVectors"), targetVector, FindNearest.DistanceMeasure.cosine()),
-   *        FindNearest
-   *          .builder()
-   *          .limit(10)
-   *          .distanceField("distance")
-   *          .build());
+   *        FindNearest.of(Field.of("topicVectors"), targetVector, FindNearest.DistanceMeasure.COSINE),
+   *        FindNearestOptions.DEFAULT
+   *          .withLimit(10)
+   *          .withDistanceField("distance"));
    * }</pre>
    *
    * @param property The expression that evaluates to a vector value using the stage inputs.

diff --git a/...irestore/src/main/java/com/google/cloud/firestore/pipeline/stages/FindNearestOptions.java b/...irestore/src/main/java/com/google/cloud/firestore/pipeline/stages/FindNearestOptions.java
@@ -40,4 +40,8 @@ public FindNearestOptions withLimit(long limit) {
   public FindNearestOptions withDistanceField(Field distanceField) {
     return with("distance_field", distanceField);
   }
+
+  public FindNearestOptions withDistanceField(String distanceField) {
+    return withDistanceField(Field.of(distanceField));
+  }
 }
diff --git a/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java b/google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java
@@ -54,6 +54,8 @@
 import com.google.cloud.firestore.pipeline.stages.AggregateOptions;
 import com.google.cloud.firestore.pipeline.stages.CollectionHints;
 import com.google.cloud.firestore.pipeline.stages.CollectionOptions;
+import com.google.cloud.firestore.pipeline.stages.FindNearest;
+import com.google.cloud.firestore.pipeline.stages.FindNearestOptions;
 import com.google.cloud.firestore.pipeline.stages.PipelineOptions;
 import com.google.cloud.firestore.pipeline.stages.PipelineOptions.ExecutionMode;
 import com.google.cloud.firestore.pipeline.stages.Sample;
@@ -1173,24 +1175,36 @@ public void testUnnest() throws Exception {
   @Test
   public void testOptions() {
     // This is just example of execute and stage options.
-    // Nothing is executed against firestore.
-
     PipelineOptions opts = PipelineOptions.DEFAULT
         .withIndexRecommendationEnabled()
         .withExecutionMode(ExecutionMode.PROFILE);
 
+    double[] vector = {1.0, 2.0, 3.0};
+
     Pipeline pipeline = firestore.pipeline()
         .collection(
             "/k",
             // Remove Hints overload - can be added later.
-            CollectionOptions.DEFAULT.withHints(CollectionHints.DEFAULT.withForceIndex("abcdef"))
+            CollectionOptions.DEFAULT
+                .withHints(CollectionHints.DEFAULT
+                    .withForceIndex("abcdef")
+                    .with("foo", "bar"))
+                .with("foo", "bar")
         )
+       .findNearest("topicVectors", vector, FindNearest.DistanceMeasure.COSINE,
+           FindNearestOptions.DEFAULT
+              .withLimit(10)
+              .withDistanceField("distance")
+              .with("foo", "bar"))
         .aggregate(
             Aggregate
                 .withAccumulators(avg("rating").as("avg_rating"))
                 .withGroups("genre")
                 .withOptions(AggregateOptions.DEFAULT
-                    .withHints(AggregateHints.DEFAULT.withForceStreamableEnabled()))
+                    .withHints(AggregateHints.DEFAULT
+                        .withForceStreamableEnabled()
+                        .with("foo", "bar"))
+                    .with("foo", "bar"))
         );
 
     pipeline.execute(opts);