add remove fields support

Author: wu-hui

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

@@ -41,12 +41,14 @@
 import com.google.cloud.firestore.pipeline.stages.GenericStage;
 import com.google.cloud.firestore.pipeline.stages.Limit;
 import com.google.cloud.firestore.pipeline.stages.Offset;
+import com.google.cloud.firestore.pipeline.stages.RemoveFields;
 import com.google.cloud.firestore.pipeline.stages.Select;
 import com.google.cloud.firestore.pipeline.stages.Sort;
 import com.google.cloud.firestore.pipeline.stages.Stage;
 import com.google.cloud.firestore.pipeline.stages.StageUtils;
 import com.google.cloud.firestore.pipeline.stages.Where;
 import com.google.common.collect.FluentIterable;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.firestore.v1.Document;
 import com.google.firestore.v1.ExecutePipelineRequest;
@@ -56,6 +58,7 @@
 import io.opencensus.trace.AttributeValue;
 import io.opencensus.trace.Tracing;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -165,6 +168,52 @@ public Pipeline addFields(Selectable... fields) {
     return append(new AddFields(PipelineUtils.selectablesToMap(fields)));
   }
 
+  /**
+   * Remove fields from outputs of previous stages.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * firestore.pipeline().collection("books")
+   *   .removeFields(
+   *     "rating", "cost"
+   *   );
+   * }</pre>
+   *
+   * @param fields The fields to remove.
+   * @return A new Pipeline object with this stage appended to the stage list.
+   */
+  @BetaApi
+  public Pipeline removeFields(String... fields) {
+    return append(
+        new RemoveFields(
+            ImmutableList.<Field>builder()
+                .addAll(Arrays.stream(fields).map(f -> Field.of(f)).iterator())
+                .build()));
+  }
+
+  /**
+   * Remove fields from outputs of previous stages.
+   *
+   * <p>Example:
+   *
+   * <pre>{@code
+   * firestore.pipeline().collection("books")
+   *   .removeFields(
+   *     Field.of("rating"), Field.of("cost")
+   *   );
+   * }</pre>
+   *
+   * @param fields The fields to remove.
+   * @return A new Pipeline object with this stage appended to the stage list.
+   */
+  @BetaApi
+  public Pipeline removeFields(Field... fields) {
+    return append(
+        new RemoveFields(
+            ImmutableList.<Field>builder().addAll(Arrays.stream(fields).iterator()).build()));
+  }
+
   /**
    * Selects or creates a set of fields from the outputs of previous stages.
    *

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/RemoveFields.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2024 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.cloud.firestore.pipeline.stages;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.firestore.pipeline.expressions.Field;
+import com.google.common.collect.ImmutableList;
+
+@InternalApi
+public final class RemoveFields implements Stage {
+
+  private static final String name = "remove_fields";
+  private final ImmutableList<Field> fields;
+
+  @InternalApi
+  public RemoveFields(ImmutableList<Field> fields) {
+    this.fields = fields;
+  }
+
+  @InternalApi
+  public ImmutableList<Field> getFields() {
+    return fields;
+  }
+
+  @Override
+  public String getName() {
+    return name;
+  }
+}

google-cloud-firestore/src/main/java/com/google/cloud/firestore/pipeline/stages/StageUtils.java

@@ -68,6 +68,15 @@ public static com.google.firestore.v1.Pipeline.Stage toStageProto(Stage stage) {
           .setName(addFieldsStage.getName())
           .addArgs(encodeValue(addFieldsStage.getFields()))
           .build();
+    } else if (stage instanceof RemoveFields) {
+      RemoveFields removeFieldsStage = (RemoveFields) stage;
+      return com.google.firestore.v1.Pipeline.Stage.newBuilder()
+          .setName(removeFieldsStage.getName())
+          .addAllArgs(
+              removeFieldsStage.getFields().stream()
+                  .map(f -> encodeValue(f))
+                  .collect(Collectors.toList()))
+          .build();
     } else if (stage instanceof Where) {
       Where whereStage = (Where) stage; // Use wildcard for generic type
       return com.google.firestore.v1.Pipeline.Stage.newBuilder()

google-cloud-firestore/src/test/java/com/google/cloud/firestore/it/ITPipelineTest.java

@@ -364,6 +364,60 @@ public void selectSpecificFields() throws Exception {
                 map("title", "The Handmaid's Tale", "author", "Margaret Atwood")));
   }
 
+  @Test
+  public void addAndRemoveFields() throws Exception {
+    List<PipelineResult> results =
+        firestore
+            .pipeline()
+            .collection(collection.getPath())
+            .addFields(
+                strConcat(Field.of("author"), "_", Field.of("title")).as("author_title"),
+                strConcat(Field.of("title"), "_", Field.of("author")).as("title_author"))
+            .removeFields("title_author", "tags", "awards", "rating", "title")
+            .removeFields(Field.of("published"), Field.of("genre"), Field.of("nestedField"))
+            .sort(Field.of("author_title").ascending())
+            .execute()
+            .get();
+
+    assertThat(data(results))
+        .isEqualTo(
+            Lists.newArrayList(
+                map(
+                    "author_title",
+                    "Douglas Adams_The Hitchhiker's Guide to the Galaxy",
+                    "author",
+                    "Douglas Adams"),
+                map(
+                    "author_title",
+                    "F. Scott Fitzgerald_The Great Gatsby",
+                    "author",
+                    "F. Scott Fitzgerald"),
+                map("author_title", "Frank Herbert_Dune", "author", "Frank Herbert"),
+                map(
+                    "author_title",
+                    "Fyodor Dostoevsky_Crime and Punishment",
+                    "author",
+                    "Fyodor Dostoevsky"),
+                map(
+                    "author_title",
+                    "Gabriel García Márquez_One Hundred Years of Solitude",
+                    "author",
+                    "Gabriel García Márquez"),
+                map("author_title", "George Orwell_1984", "author", "George Orwell"),
+                map("author_title", "Harper Lee_To Kill a Mockingbird", "author", "Harper Lee"),
+                map(
+                    "author_title",
+                    "J.R.R. Tolkien_The Lord of the Rings",
+                    "author",
+                    "J.R.R. Tolkien"),
+                map("author_title", "Jane Austen_Pride and Prejudice", "author", "Jane Austen"),
+                map(
+                    "author_title",
+                    "Margaret Atwood_The Handmaid's Tale",
+                    "author",
+                    "Margaret Atwood")));
+  }
+
   @Test
   public void whereByMultipleConditions() throws Exception {
     List<PipelineResult> results =