add test and comment

apache · Aug 27, 2024 · 303c62b · 303c62b
1 parent 56bfec6
commit 303c62b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
@@ -3044,7 +3044,13 @@ public Write<T> withPropagateSuccessfulStorageApiWrites(
     /**
      * If called, then all successful writes will be propagated to {@link WriteResult} and
      * accessible via the {@link WriteResult#getSuccessfulStorageApiInserts} method. The predicate
-     * allows filtering out columns from appearing in the resulting PCollection.
+     * allows filtering out columns from appearing in the resulting PCollection. The argument to the
+     * predicate is the name of the field to potentially be included in the output. Nested fields
+     * will be presented using . notation - e.g. a.b.c.
+     *
+     * <p>The predicate will be invoked repeatedly for every field in every message, so it is
+     * recommended that it be as lightweight as possible. e.g. looking up fields in a hash table
+     * instead of searching a list of field names.
      */
     public Write<T> withPropagateSuccessfulStorageApiWrites(Predicate<String> columnsToPropagate) {
       return toBuilder()

diff --git a/...cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/...cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java
@@ -1603,14 +1603,23 @@ private void storageWriteWithSuccessHandling(boolean columnSubset) throws Except
     List<TableRow> elements =
         IntStream.range(0, 30)
             .mapToObj(Integer::toString)
-            .map(i -> new TableRow().set("number", i).set("string", i))
+            .map(
+                i ->
+                    new TableRow()
+                        .set("number", i)
+                        .set("string", i)
+                        .set("nested", new TableRow().set("number", i)))
             .collect(Collectors.toList());
 
     List<TableRow> expectedSuccessElements = elements;
     if (columnSubset) {
       expectedSuccessElements =
           elements.stream()
-              .map(tr -> new TableRow().set("number", tr.get("number")))
+              .map(
+                  tr ->
+                      new TableRow()
+                          .set("number", tr.get("number"))
+                          .set("nested", new TableRow().set("number", tr.get("number"))))
               .collect(Collectors.toList());
     }
 
@@ -1619,7 +1628,13 @@ private void storageWriteWithSuccessHandling(boolean columnSubset) throws Except
             .setFields(
                 ImmutableList.of(
                     new TableFieldSchema().setName("number").setType("INTEGER"),
-                    new TableFieldSchema().setName("string").setType("STRING")));
+                    new TableFieldSchema().setName("string").setType("STRING"),
+                    new TableFieldSchema()
+                        .setName("nested")
+                        .setType("RECORD")
+                        .setFields(
+                            ImmutableList.of(
+                                new TableFieldSchema().setName("number").setType("INTEGER")))));
 
     TestStream<TableRow> testStream =
         TestStream.create(TableRowJsonCoder.of())
@@ -1645,7 +1660,8 @@ private void storageWriteWithSuccessHandling(boolean columnSubset) throws Except
     if (columnSubset) {
       write =
           write.withPropagateSuccessfulStorageApiWrites(
-              (Serializable & Predicate<String>) s -> s.equals("number"));
+              (Serializable & Predicate<String>)
+                  s -> s.equals("number") || s.equals("nested") || s.equals("nested.number"));
     }
     if (useStreaming) {
       if (useStorageApiApproximate) {