[Java] BigQueryIO: set table labels from dynamic destinations

CHANGES.md

@@ -68,6 +68,7 @@
 * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)).
 * Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)).
 * (Java/Python/Go) DebeziumIO: `readAsJson()` can now include `op` and the remaining Debezium `source` fields (e.g. MySQL `file`, `pos`, `row`, `ts_ms`; PostgreSQL `lsn`, `txId`, `ts_ms`) in `metadata` via `withFullMetadata(true)` / `full_metadata=True` / `debeziumio.FullMetadata()`. Off by default ([#40701](https://github.com/apache/beam/issues/40701)).
+* (Java) BigQueryIO: dynamic destinations can now set BigQuery table labels with `TableDestination.withLabels(...)`. When writing with `to(SerializableFunction)`, also set `withTableLabels()`, which switches to the new `TableDestinationCoderV4` ([#26386](https://github.com/apache/beam/issues/26386)).
 
 ## New Features / Improvements
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java

@@ -188,6 +188,7 @@ class BatchLoads<DestinationT, ElementT>
       RowWriterFactory<ElementT, DestinationT> rowWriterFactory,
       @Nullable String kmsKey,
       boolean clusteringEnabled,
+      boolean labelsEnabled,
       boolean useAvroLogicalTypes,
       @Nullable String tempDataset,
       BadRecordRouter badRecordRouter,
@@ -216,8 +217,13 @@ class BatchLoads<DestinationT, ElementT>
     this.rowWriterFactory = rowWriterFactory;
     schemaUpdateOptions = Collections.emptySet();
     this.tempDataset = tempDataset;
-    this.tableDestinationCoder =
-        clusteringEnabled ? TableDestinationCoderV3.of() : TableDestinationCoderV2.of();
+    if (labelsEnabled) {
+      this.tableDestinationCoder = TableDestinationCoderV4.of();
+    } else if (clusteringEnabled) {
+      this.tableDestinationCoder = TableDestinationCoderV3.of();
+    } else {
+      this.tableDestinationCoder = TableDestinationCoderV2.of();
+    }
     this.badRecordRouter = badRecordRouter;
     this.badRecordErrorHandler = badRecordErrorHandler;
   }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java

@@ -2794,6 +2794,7 @@ public static <T> Write<T> write() {
         .setIgnoreUnknownValues(false)
         .setIgnoreInsertIds(false)
         .setUseAvroLogicalTypes(false)
+        .setTableLabelsEnabled(false)
         .setMaxFilesPerPartition(BatchLoads.DEFAULT_MAX_FILES_PER_PARTITION)
         .setMaxBytesPerPartition(BatchLoads.DEFAULT_MAX_BYTES_PER_PARTITION)
         .setOptimizeWrites(false)
@@ -3016,6 +3017,8 @@ public enum Method {
 
     abstract @Nullable ValueProvider<String> getJsonClustering();
 
+    abstract boolean getTableLabelsEnabled();
+
     abstract CreateDisposition getCreateDisposition();
 
     abstract WriteDisposition getWriteDisposition();
@@ -3137,6 +3140,8 @@ abstract Builder<T> setAvroSchemaFactory(
 
       abstract Builder<T> setJsonClustering(ValueProvider<String> clustering);
 
+      abstract Builder<T> setTableLabelsEnabled(boolean tableLabelsEnabled);
+
       abstract Builder<T> setCreateDisposition(CreateDisposition createDisposition);
 
       abstract Builder<T> setWriteDisposition(WriteDisposition writeDisposition);
@@ -3345,7 +3350,8 @@ public Write<T> to(ValueProvider<String> tableSpec) {
      *
      * <p>If the function produces destinations configured with clustering fields, ensure that
      * {@link #withClustering()} is also set so that the clustering configurations get properly
-     * encoded and decoded.
+     * encoded and decoded. Likewise, if it produces destinations with labels set via {@link
+     * TableDestination#withLabels(Map)}, ensure that {@link #withTableLabels()} is also set.
      */
     public Write<T> to(
         SerializableFunction<@Nullable ValueInSingleWindow<T>, TableDestination> tableFunction) {
@@ -3358,7 +3364,9 @@ public Write<T> to(
      *
      * <p>If any of the returned destinations are configured with clustering fields, ensure that the
      * passed {@link DynamicDestinations} object returns {@link TableDestinationCoderV3} when {@link
-     * DynamicDestinations#getDestinationCoder()} is called.
+     * DynamicDestinations#getDestinationCoder()} is called. If the destination type is {@link
+     * TableDestination} and the returned destinations have labels set via {@link
+     * TableDestination#withLabels(Map)}, return {@link TableDestinationCoderV4} instead.
      */
     public Write<T> to(DynamicDestinations<T, ?> dynamicDestinations) {
       checkArgument(dynamicDestinations != null, "dynamicDestinations can not be null");
@@ -3550,6 +3558,23 @@ public Write<T> withClustering() {
       return withClustering(new Clustering());
     }
 
+    /**
+     * Allows writing table labels when {@link #to(SerializableFunction)} is used. The returned
+     * {@link TableDestination} objects should set the labels per table with {@link
+     * TableDestination#withLabels(Map)}.
+     *
+     * <p>Destinations in {@link WriteResult#getSuccessfulTableLoads()} keep their labels, except
+     * for tables that {@link Method#FILE_LOADS} wrote through temporary tables.
+     *
+     * <p>Setting this option enables use of {@link TableDestinationCoderV4} which encodes time
+     * partitioning, clustering and label information. As with {@link #withClustering()}, pipelines
+     * using an older coder must be drained before setting this option, since {@link
+     * TableDestinationCoderV4} will not be able to read state written with a previous version.
+     */
+    public Write<T> withTableLabels() {
+      return toBuilder().setTableLabelsEnabled(true).build();
+    }
+
     /** Specifies whether the table should be created if it does not exist. */
     public Write<T> withCreateDisposition(CreateDisposition createDisposition) {
       checkArgument(createDisposition != null, "createDisposition can not be null");
@@ -4258,6 +4283,14 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) {
                 + " There is no need to call BigQueryIO.Write.withTimePartitioning.");
       }
 
+      if (getTableLabelsEnabled()) {
+        checkArgument(
+            getTableFunction() != null,
+            "withTableLabels() only applies to to(SerializableFunction). With"
+                + " to(DynamicDestinations), return TableDestinationCoderV4 from"
+                + " DynamicDestinations.getDestinationCoder() instead.");
+      }
+
       DynamicDestinations<T, ?> dynamicDestinations = getDynamicDestinations();
       if (dynamicDestinations == null) {
         ValueProvider<String> jsonTableRef = getJsonTableRef();
@@ -4274,7 +4307,9 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) {
           // set.
           tableDestinations =
               new TableFunctionDestinations<>(
-                  checkStateNotNull(tableFunction), jsonClustering != null);
+                  checkStateNotNull(tableFunction),
+                  jsonClustering != null,
+                  getTableLabelsEnabled());
         }
 
         // Wrap with a DynamicDestinations class that will provide a schema. There might be no
@@ -4307,6 +4342,13 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) {
       return expandTyped(input, dynamicDestinations);
     }
 
+    /** Whether {@code coder} is one of the {@link TableDestination} coders without labels. */
+    private static boolean coderDropsTableLabels(Coder<?> coder) {
+      return coder instanceof TableDestinationCoder
+          || coder instanceof TableDestinationCoderV2
+          || coder instanceof TableDestinationCoderV3;
+    }
+
     private <DestinationT> WriteResult expandTyped(
         PCollection<T> input, DynamicDestinations<T, DestinationT> dynamicDestinations) {
       boolean optimizeWrites = getOptimizeWrites();
@@ -4431,7 +4473,10 @@ private <DestinationT> WriteResult expandTyped(
           input
               .apply(
                   "PrepareWrite",
-                  new PrepareWrite<>(dynamicDestinations, SerializableFunctions.identity()))
+                  new PrepareWrite<>(
+                      dynamicDestinations,
+                      SerializableFunctions.identity(),
+                      coderDropsTableLabels(destinationCoder)))
               .setCoder(KvCoder.of(destinationCoder, input.getCoder()));
 
       return continueExpandTyped(
@@ -4558,6 +4603,7 @@ private <DestinationT> WriteResult continueExpandTyped(
                 rowWriterFactory,
                 getKmsKey(),
                 getJsonClustering() != null,
+                getTableLabelsEnabled(),
                 getUseAvroLogicalTypes(),
                 getWriteTempDataset(),
                 getBadRecordRouter(),

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java

@@ -483,6 +483,7 @@ static class BigQueryIOWriteTranslator implements TransformPayloadTranslator<Wri
             .addNullableByteArrayField("bad_record_error_handler")
             .addNullableByteArrayField("bad_record_router")
             .addNullableInt64Field("auto_schema_update_strict_timeout_ms")
+            .addNullableBooleanField("table_labels_enabled")
             .build();
 
     public static final String BIGQUERY_WRITE_TRANSFORM_URN =
@@ -643,6 +644,7 @@ public Row toConfigRow(Write<?> transform) {
       fieldValues.put("use_beam_schema", transform.getUseBeamSchema());
       fieldValues.put("auto_sharding", transform.getAutoSharding());
       fieldValues.put("auto_schema_update", transform.getAutoSchemaUpdate());
+      fieldValues.put("table_labels_enabled", transform.getTableLabelsEnabled());
       org.joda.time.Duration autoSchemaUpdateTimeout = transform.getAutoSchemaUpdateStrictTimeout();
       if (autoSchemaUpdateTimeout != null) {
         fieldValues.put(
@@ -938,6 +940,14 @@ public Write<?> fromConfigRow(Row configRow, PipelineOptions options) {
         if (autoSchemaUpdate != null) {
           builder = builder.setAutoSchemaUpdate(autoSchemaUpdate);
         }
+        Boolean tableLabelsEnabled = configRow.getBoolean("table_labels_enabled");
+        // rows written by SDKs without table labels have no value here, and the builder has no
+        // default for this required property, so building the transform would fail without it
+        if (tableLabelsEnabled != null) {
+          builder = builder.setTableLabelsEnabled(tableLabelsEnabled);
+        } else {
+          builder = builder.setTableLabelsEnabled(false);
+        }
         Long autoSchemaUpdateStrictTimeoutMillis =
             configRow.getInt64("auto_schema_update_strict_timeout_ms");
         if (autoSchemaUpdateStrictTimeoutMillis != null) {

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java

@@ -48,6 +48,7 @@
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.List;
+import java.util.Map;
 import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.values.FailsafeValueInSingleWindow;
 import org.apache.beam.sdk.values.ValueInSingleWindow;
@@ -203,6 +204,16 @@ <T> long insertAll(
     Table patchTableDescription(TableReference tableReference, @Nullable String tableDescription)
         throws IOException, InterruptedException;
 
+    /**
+     * Patch BigQuery {@link Table} labels. Labels in the map are added or overwritten, labels
+     * already on the table that are not in the map are kept.
+     */
+    default Table patchTableLabels(TableReference tableReference, Map<String, String> labels)
+        throws IOException, InterruptedException {
+      throw new UnsupportedOperationException(
+          getClass().getName() + " does not support patching table labels");
+    }
+
     Table patchTableSchema(
         TableReference tableReference, com.google.api.services.bigquery.model.TableSchema newSchema)
         throws IOException, InterruptedException;

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java

@@ -1565,6 +1565,28 @@ public Table patchTableDescription(
           ALWAYS_RETRY);
     }
 
+    @Override
+    public Table patchTableLabels(TableReference tableReference, Map<String, String> labels)
+        throws IOException, InterruptedException {
+      Table table = new Table().setLabels(labels);
+
+      return executeWithRetries(
+          client
+              .tables()
+              .patch(
+                  tableReference.getProjectId(),
+                  tableReference.getDatasetId(),
+                  tableReference.getTableId(),
+                  table)
+              .setPrettyPrint(false),
+          String.format(
+              "Unable to patch table labels: %s, aborting after %d retries.",
+              tableReference, MAX_RPC_RETRIES),
+          Sleeper.DEFAULT,
+          createDefaultBackoff(),
+          ALWAYS_RETRY);
+    }
+
     @Override
     public Table patchTableSchema(
         TableReference tableReference, com.google.api.services.bigquery.model.TableSchema newSchema)

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java

@@ -204,6 +204,11 @@ private static void tryCreateTable(
           table.setClustering(clustering);
         }
 
+        Map<String, String> labels = tableDestination.getLabels();
+        if (labels != null) {
+          table.setLabels(labels);
+        }
+
         if (kmsKey != null) {
           table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey));
         }

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java

@@ -121,12 +121,15 @@ static class TableFunctionDestinations<T> extends DynamicDestinations<T, TableDe
     private final SerializableFunction<@Nullable ValueInSingleWindow<T>, TableDestination>
         tableFunction;
     private final boolean clusteringEnabled;
+    private final boolean labelsEnabled;
 
     TableFunctionDestinations(
         SerializableFunction<@Nullable ValueInSingleWindow<T>, TableDestination> tableFunction,
-        boolean clusteringEnabled) {
+        boolean clusteringEnabled,
+        boolean labelsEnabled) {
       this.tableFunction = tableFunction;
       this.clusteringEnabled = clusteringEnabled;
+      this.labelsEnabled = labelsEnabled;
     }
 
     @Override
@@ -152,7 +155,9 @@ public TableDestination getTable(TableDestination destination) {
 
     @Override
     public Coder<TableDestination> getDestinationCoder() {
-      if (clusteringEnabled) {
+      if (labelsEnabled) {
+        return TableDestinationCoderV4.of();
+      } else if (clusteringEnabled) {
         return TableDestinationCoderV3.of();
       } else {
         return TableDestinationCoderV2.of();
@@ -358,11 +363,20 @@ public TableDestination getDestination(@Nullable ValueInSingleWindow<T> element)
           evaluatedClustering != null ? evaluatedClustering : destination.getJsonClustering();
 
       return new TableDestination(
-          destination.getTableSpec(), destination.getTableDescription(), partitioning, clustering);
+              destination.getTableSpec(),
+              destination.getTableDescription(),
+              partitioning,
+              clustering)
+          .withLabels(destination.getLabels());
     }
 
     @Override
     public Coder<TableDestination> getDestinationCoder() {
+      // keep the inner coder when it already encodes labels, otherwise V3 would drop them
+      @Nullable Coder<TableDestination> innerCoder = super.getDestinationCoder();
+      if (innerCoder instanceof TableDestinationCoderV4) {
+        return TableDestinationCoderV4.of();
+      }
       if (jsonClustering != null) {
         return TableDestinationCoderV3.of();
       } else {
@@ -499,11 +513,15 @@ public TableDestination getTable(DestinationT destination) {
       if (existingTable == null) {
         return wrappedDestination;
       } else {
+        // partitioning and clustering must match the existing table so that temp tables can be
+        // copied into it. labels play no part in that, so keep the requested ones so that they
+        // are still patched onto the table after the copy
         return new TableDestination(
-            wrappedDestination.getTableSpec(),
-            existingTable.getDescription(),
-            existingTable.getTimePartitioning(),
-            existingTable.getClustering());
+                wrappedDestination.getTableSpec(),
+                existingTable.getDescription(),
+                existingTable.getTimePartitioning(),
+                existingTable.getClustering())
+            .withLabels(wrappedDestination.getLabels());
       }
     }
 

sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/PrepareWrite.java

@@ -43,12 +43,26 @@ public class PrepareWrite<InputT, DestinationT extends @NonNull Object, OutputT>
     extends PTransform<PCollection<InputT>, PCollection<KV<DestinationT, OutputT>>> {
   private DynamicDestinations<InputT, DestinationT> dynamicDestinations;
   private SerializableFunction<InputT, OutputT> formatFunction;
+  private boolean destinationCoderDropsLabels;
 
   public PrepareWrite(
       DynamicDestinations<InputT, DestinationT> dynamicDestinations,
       SerializableFunction<InputT, OutputT> formatFunction) {
+    this(dynamicDestinations, formatFunction, false);
+  }
+
+  /**
+   * @param destinationCoderDropsLabels whether the destination coder has no field for {@link
+   *     TableDestination#getLabels()}. When set, a destination that carries labels fails the
+   *     element instead of losing its labels when the destination is encoded.
+   */
+  public PrepareWrite(
+      DynamicDestinations<InputT, DestinationT> dynamicDestinations,
+      SerializableFunction<InputT, OutputT> formatFunction,
+      boolean destinationCoderDropsLabels) {
     this.dynamicDestinations = dynamicDestinations;
     this.formatFunction = formatFunction;
+    this.destinationCoderDropsLabels = destinationCoderDropsLabels;
   }
 
   @Override
@@ -75,6 +89,21 @@ public void processElement(
                             + "but %s returned null on element %s",
                         dynamicDestinations,
                         element);
+                    // this is the last point before the destination is encoded. after that the
+                    // labels are gone and every later step sees a destination without them
+                    if (destinationCoderDropsLabels
+                        && tableDestination instanceof TableDestination) {
+                      checkArgument(
+                          ((TableDestination) tableDestination).getLabels() == null,
+                          "%s returned a TableDestination with labels, but its destination coder"
+                              + " cannot encode labels, so they would be dropped. With"
+                              + " to(SerializableFunction), set withTableLabels() on"
+                              + " BigQueryIO.write(). With to(DynamicDestinations), return"
+                              + " TableDestinationCoderV4 from getDestinationCoder()."
+                              + " Destination: %s",
+                          dynamicDestinations,
+                          tableDestination);
+                    }
                     OutputT outputValue = formatFunction.apply(element);
                     Preconditions.checkArgumentNotNull(
                         outputValue,