From bdd6a7dde1d7d41cab7eb46a70c50c4d95827f2d Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 17:57:03 +0200 Subject: [PATCH 1/3] [Java] BigQueryIO: set table labels from dynamic destinations Add TableDestination.withLabels() and TableDestinationCoderV4. Streaming inserts and the Storage Write API set labels when they create the table. FILE_LOADS patches them after the load or copy job. withTableLabels() switches to(SerializableFunction) writes to the new coder. Fixes #26386. --- CHANGES.md | 1 + .../beam/sdk/io/gcp/bigquery/BatchLoads.java | 10 +- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 54 ++++- .../sdk/io/gcp/bigquery/BigQueryServices.java | 11 + .../io/gcp/bigquery/BigQueryServicesImpl.java | 22 ++ .../io/gcp/bigquery/CreateTableHelpers.java | 5 + .../bigquery/DynamicDestinationsHelpers.java | 32 ++- .../sdk/io/gcp/bigquery/PrepareWrite.java | 29 +++ .../sdk/io/gcp/bigquery/TableDestination.java | 73 ++++++- .../gcp/bigquery/TableDestinationCoderV4.java | 64 ++++++ .../beam/sdk/io/gcp/bigquery/WriteRename.java | 10 + .../beam/sdk/io/gcp/bigquery/WriteTables.java | 11 + .../io/gcp/testing/FakeDatasetService.java | 22 ++ .../io/gcp/bigquery/BigQueryIOWriteTest.java | 197 ++++++++++++++++++ .../bigquery/TableDestinationCoderV4Test.java | 71 +++++++ 15 files changed, 596 insertions(+), 16 deletions(-) create mode 100644 sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4.java create mode 100644 sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4Test.java diff --git a/CHANGES.md b/CHANGES.md index aa9f5aadb040..5f2bd496085a 100644 --- a/CHANGES.md +++ b/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 diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java index d12af24e407d..e46343e98938 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java @@ -188,6 +188,7 @@ class BatchLoads RowWriterFactory rowWriterFactory, @Nullable String kmsKey, boolean clusteringEnabled, + boolean labelsEnabled, boolean useAvroLogicalTypes, @Nullable String tempDataset, BadRecordRouter badRecordRouter, @@ -216,8 +217,13 @@ class BatchLoads 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; } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java index 0b8263ecb1e4..1c545dc53f28 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java @@ -2794,6 +2794,7 @@ public static Write 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 getJsonClustering(); + abstract boolean getTableLabelsEnabled(); + abstract CreateDisposition getCreateDisposition(); abstract WriteDisposition getWriteDisposition(); @@ -3137,6 +3140,8 @@ abstract Builder setAvroSchemaFactory( abstract Builder setJsonClustering(ValueProvider clustering); + abstract Builder setTableLabelsEnabled(boolean tableLabelsEnabled); + abstract Builder setCreateDisposition(CreateDisposition createDisposition); abstract Builder setWriteDisposition(WriteDisposition writeDisposition); @@ -3345,7 +3350,8 @@ public Write to(ValueProvider tableSpec) { * *

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 to( SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> tableFunction) { @@ -3358,7 +3364,9 @@ public Write to( * *

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 to(DynamicDestinations dynamicDestinations) { checkArgument(dynamicDestinations != null, "dynamicDestinations can not be null"); @@ -3550,6 +3558,23 @@ public Write 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)}. + * + *

Destinations in {@link WriteResult#getSuccessfulTableLoads()} keep their labels, except + * for tables that {@link Method#FILE_LOADS} wrote through temporary tables. + * + *

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 withTableLabels() { + return toBuilder().setTableLabelsEnabled(true).build(); + } + /** Specifies whether the table should be created if it does not exist. */ public Write 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 dynamicDestinations = getDynamicDestinations(); if (dynamicDestinations == null) { ValueProvider 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 WriteResult expandTyped( PCollection input, DynamicDestinations dynamicDestinations) { boolean optimizeWrites = getOptimizeWrites(); @@ -4431,7 +4473,10 @@ private 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 WriteResult continueExpandTyped( rowWriterFactory, getKmsKey(), getJsonClustering() != null, + getTableLabelsEnabled(), getUseAvroLogicalTypes(), getWriteTempDataset(), getBadRecordRouter(), diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java index 0288caeb2c49..ed6b43e3ad52 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServices.java +++ b/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 @@ 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 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; diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java index 81ba6ce3ca72..63f596fc9b58 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java +++ b/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 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) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java index 52ac77532b1d..75c6617c5bd7 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/CreateTableHelpers.java +++ b/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 labels = tableDestination.getLabels(); + if (labels != null) { + table.setLabels(labels); + } + if (kmsKey != null) { table.setEncryptionConfiguration(new EncryptionConfiguration().setKmsKeyName(kmsKey)); } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java index 4cd5d25202a3..9fae46f7609e 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/DynamicDestinationsHelpers.java +++ b/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 extends DynamicDestinations, TableDestination> tableFunction; private final boolean clusteringEnabled; + private final boolean labelsEnabled; TableFunctionDestinations( SerializableFunction<@Nullable ValueInSingleWindow, 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 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 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 getDestinationCoder() { + // keep the inner coder when it already encodes labels, otherwise V3 would drop them + @Nullable Coder 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()); } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/PrepareWrite.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/PrepareWrite.java index a5d4c0e43215..4dc8a1def25f 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/PrepareWrite.java +++ b/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 extends PTransform, PCollection>> { private DynamicDestinations dynamicDestinations; private SerializableFunction formatFunction; + private boolean destinationCoderDropsLabels; public PrepareWrite( DynamicDestinations dynamicDestinations, SerializableFunction 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 dynamicDestinations, + SerializableFunction 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, diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestination.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestination.java index 6b81c2322a51..33ad9365d1f2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestination.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestination.java @@ -21,8 +21,11 @@ import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TimePartitioning; import java.io.Serializable; +import java.util.Map; import java.util.Objects; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSortedMap; import org.checkerframework.checker.nullness.qual.Nullable; /** Encapsulates a BigQuery table destination. */ @@ -32,6 +35,10 @@ public class TableDestination implements Serializable { private final @Nullable String tableDescription; private final @Nullable String jsonTimePartitioning; private final @Nullable String jsonClustering; + // sorted so that TableDestinationCoderV4 writes the same bytes for equal label sets. runners + // group load jobs by the encoded destination, so a hash-ordered map could split one table + // into several groups + private final @Nullable ImmutableSortedMap labels; public TableDestination(String tableSpec, @Nullable String tableDescription) { this(tableSpec, tableDescription, (String) null, (String) null); @@ -106,15 +113,67 @@ public TableDestination( @Nullable String tableDescription, @Nullable String jsonTimePartitioning, @Nullable String jsonClustering) { + this(tableSpec, tableDescription, jsonTimePartitioning, jsonClustering, null); + } + + private TableDestination( + String tableSpec, + @Nullable String tableDescription, + @Nullable String jsonTimePartitioning, + @Nullable String jsonClustering, + @Nullable Map labels) { this.tableSpec = tableSpec; this.tableDescription = tableDescription; this.jsonTimePartitioning = jsonTimePartitioning; this.jsonClustering = jsonClustering; + if (labels == null) { + this.labels = null; + } else { + // ImmutableSortedMap rejects nulls with a bare NullPointerException, so name the bad entry + for (Map.Entry label : labels.entrySet()) { + Preconditions.checkArgument( + label.getKey() != null && label.getValue() != null, + "table labels may not contain null keys or values, but got %s for table %s", + labels, + tableSpec); + } + this.labels = ImmutableSortedMap.copyOf(labels); + } } public TableDestination withTableReference(TableReference tableReference) { return new TableDestination( - tableReference, tableDescription, jsonTimePartitioning, jsonClustering); + BigQueryHelpers.toTableSpec(tableReference), + tableDescription, + jsonTimePartitioning, + jsonClustering, + labels); + } + + /** + * Returns a copy of this destination with the given BigQuery table labels. + * + *

With {@link BigQueryIO.Write.Method#STREAMING_INSERTS} and {@link + * BigQueryIO.Write.Method#STORAGE_WRITE_API} the labels are set only when BigQueryIO creates the + * table. With {@link BigQueryIO.Write.Method#FILE_LOADS} they are patched onto the table after + * every load, the same way the table description is, so they are also added to tables that + * already exist. Labels already on the table that are not in this map are kept. + * + *

Return the same labels for every element written to one table. Destinations that differ only + * in their labels are grouped separately, so the table would get one load job per label set and, + * with {@link BigQueryIO.Write.WriteDisposition#WRITE_TRUNCATE}, each job would replace the data + * of the others. + * + *

Labels are only encoded by {@link TableDestinationCoderV4}. When using {@link + * BigQueryIO.Write#to(org.apache.beam.sdk.transforms.SerializableFunction)}, also set {@link + * BigQueryIO.Write#withTableLabels()}. When using {@link + * BigQueryIO.Write#to(DynamicDestinations)} with {@link TableDestination} as the destination + * type, return {@link TableDestinationCoderV4} from {@link + * DynamicDestinations#getDestinationCoder()}. + */ + public TableDestination withLabels(@Nullable Map labels) { + return new TableDestination( + tableSpec, tableDescription, jsonTimePartitioning, jsonClustering, labels); } /** Return the tablespec in [project:].dataset.tableid format. */ @@ -174,12 +233,19 @@ public TableReference getTableReference() { return tableDescription; } + public @Nullable Map getLabels() { + return labels; + } + @Override public String toString() { String toString = "tableSpec: " + tableSpec; if (tableDescription != null) { toString += " tableDescription: " + tableDescription; } + if (labels != null) { + toString += " labels: " + labels; + } return toString; } @@ -191,11 +257,12 @@ public boolean equals(@Nullable Object o) { TableDestination other = (TableDestination) o; return Objects.equals(this.tableSpec, other.tableSpec) && Objects.equals(this.tableDescription, other.tableDescription) - && Objects.equals(this.jsonTimePartitioning, other.jsonTimePartitioning); + && Objects.equals(this.jsonTimePartitioning, other.jsonTimePartitioning) + && Objects.equals(this.labels, other.labels); } @Override public int hashCode() { - return Objects.hash(tableSpec, tableDescription, jsonTimePartitioning); + return Objects.hash(tableSpec, tableDescription, jsonTimePartitioning, labels); } } diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4.java new file mode 100644 index 000000000000..522f82f463d0 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.io.gcp.bigquery; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.util.Map; +import org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.MapCoder; +import org.apache.beam.sdk.coders.NullableCoder; +import org.apache.beam.sdk.coders.StringUtf8Coder; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A {@link Coder} for {@link TableDestination} that includes time partitioning, clustering and + * table label information. Users must opt in to this version of the coder by setting {@link + * BigQueryIO.Write#withTableLabels()}, otherwise {@link TableDestinationCoderV2} or {@link + * TableDestinationCoderV3} will be used and labels will be discarded. + */ +public class TableDestinationCoderV4 extends AtomicCoder { + private static final TableDestinationCoderV4 INSTANCE = new TableDestinationCoderV4(); + private static final Coder<@Nullable Map> labelsCoder = + NullableCoder.of(MapCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())); + + public static TableDestinationCoderV4 of() { + return INSTANCE; + } + + @Override + public void encode(TableDestination value, OutputStream outStream) throws IOException { + TableDestinationCoderV3.of().encode(value, outStream); + labelsCoder.encode(value.getLabels(), outStream); + } + + @Override + public TableDestination decode(InputStream inStream) throws IOException { + TableDestination destination = TableDestinationCoderV3.of().decode(inStream); + Map labels = labelsCoder.decode(inStream); + return destination.withLabels(labels); + } + + // MapCoder on its own is not deterministic because it writes entries in the map's iteration + // order. TableDestination stores labels in a sorted map, so equal label sets always encode to + // the same bytes + @Override + public void verifyDeterministic() throws NonDeterministicException {} +} diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java index 061e66024e29..baa7c7abcbe1 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteRename.java @@ -236,6 +236,16 @@ public void finishBundle(FinishBundleContext c) throws Exception { .setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), pendingJob.tableDestination.getTableDescription()); } + // copy jobs have no setting for destination table labels, so they are patched + // onto the final table once the copy is done + Map labels = pendingJob.tableDestination.getLabels(); + if (labels != null) { + TableReference ref = pendingJob.tableDestination.getTableReference(); + datasetService.patchTableLabels( + ref.clone() + .setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), + labels); + } c.output( KV.of(pendingJob.tableDestination, pendingJob.tempTables), pendingJob.window.maxTimestamp(), diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java index a95bfc5f7c87..1e8d0243c916 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/WriteTables.java @@ -364,6 +364,17 @@ public void finishBundle(FinishBundleContext c) throws Exception { BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), pendingJob.tableDestination.getTableDescription()); } + // temp tables are deleted after the copy into the final table, which + // WriteRename labels instead + Map labels = pendingJob.tableDestination.getLabels(); + if (!tempTable && labels != null) { + TableReference ref = pendingJob.tableReference; + datasetService.patchTableLabels( + ref.clone() + .setTableId( + BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), + labels); + } Result result = new AutoValue_WriteTables_Result( BigQueryHelpers.toJsonString(pendingJob.tableReference), diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java index 67dc802f1181..91f6cefb54b2 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeDatasetService.java @@ -742,6 +742,28 @@ public Table patchTableDescription( } } + @Override + public Table patchTableLabels(TableReference tableReference, Map labels) + throws IOException, InterruptedException { + validateWholeTableReference(tableReference); + synchronized (FakeDatasetService.class) { + TableContainer tableContainer = + getTableContainer( + tableReference.getProjectId(), + tableReference.getDatasetId(), + tableReference.getTableId()); + Table table = tableContainer.getTable(); + // the real tables.patch merges the given labels into the existing ones + Map merged = new HashMap<>(); + if (table.getLabels() != null) { + merged.putAll(table.getLabels()); + } + merged.putAll(labels); + table.setLabels(merged); + return table; + } + } + @Override public Table patchTableSchema(TableReference tableReference, TableSchema newSchema) throws IOException, InterruptedException { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java index 62e57ec848cb..340b858229ef 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOWriteTest.java @@ -679,6 +679,203 @@ public void testClusteringTableFunction() throws Exception { assertEquals(clustering, table.getClustering()); } + @Test + public void testTableLabelsTableFunction() throws Exception { + assumeTrue(!useStorageApi); + TableRow row1 = new TableRow().set("date", "2018-01-01").set("number", "1"); + TableRow row2 = new TableRow().set("date", "2018-01-02").set("number", "2"); + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("date").setType("DATE"), + new TableFieldSchema().setName("number").setType("INTEGER"))); + + WriteResult result = + p.apply(Create.of(row1, row2)) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination( + "project-id:dataset-id.table-" + vsw.getValue().get("number"), + null, + null, + new Clustering().setFields(ImmutableList.of("date"))) + .withLabels( + ImmutableMap.of( + "number", (String) vsw.getValue().get("number")))) + .withTestServices(fakeBqServices) + .withMethod(useStreaming ? Method.STREAMING_INSERTS : Method.FILE_LOADS) + .withSchema(schema) + .withClustering() + .withTableLabels() + .withoutValidation()); + if (!useStreaming) { + PAssert.that(result.getSuccessfulTableLoads()) + .containsInAnyOrder( + new TableDestination( + "project-id:dataset-id.table-1", + null, + null, + new Clustering().setFields(ImmutableList.of("date"))) + .withLabels(ImmutableMap.of("number", "1")), + new TableDestination( + "project-id:dataset-id.table-2", + null, + null, + new Clustering().setFields(ImmutableList.of("date"))) + .withLabels(ImmutableMap.of("number", "2"))); + } + p.run(); + + for (String number : ImmutableList.of("1", "2")) { + Table table = + fakeDatasetService.getTable( + BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-" + number)); + assertEquals(ImmutableMap.of("number", number), table.getLabels()); + assertEquals(new Clustering().setFields(ImmutableList.of("date")), table.getClustering()); + } + } + + @Test + public void testTableLabelsWithTempTablesKeepsExistingLabels() throws Exception { + assumeTrue(!useStorageApi); + assumeTrue(!useStreaming); + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of(new TableFieldSchema().setName("number").setType("INTEGER"))); + fakeDatasetService.createTable( + new Table() + .setTableReference(BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-id")) + .setSchema(schema) + .setLabels(ImmutableMap.of("team", "data"))); + List elements = Lists.newArrayList(); + for (int i = 0; i < 30; ++i) { + elements.add(new TableRow().set("number", i)); + } + + p.apply(Create.of(elements)) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination("project-id:dataset-id.table-id", null) + .withLabels(ImmutableMap.of("env", "prod"))) + .withSchema(schema) + .withTestServices(fakeBqServices) + .withMethod(Method.FILE_LOADS) + .withWriteDisposition(Write.WriteDisposition.WRITE_APPEND) + .withTableLabels() + // one file per partition forces temp tables and a copy job into the final table + .withMaxBytesPerPartition(1) + .withMaxFilesPerPartition(1) + .withoutValidation()); + p.run(); + + Table table = + fakeDatasetService.getTable( + BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-id")); + assertEquals(ImmutableMap.of("team", "data", "env", "prod"), table.getLabels()); + assertThat( + fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"), + containsInAnyOrder(Iterables.toArray(elements, TableRow.class))); + } + + @Test + public void testTableLabelsTableFunctionRequiresWithTableLabels() throws Exception { + assumeTrue(!useStorageApi); + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of(new TableFieldSchema().setName("number").setType("INTEGER"))); + + p.apply(Create.of(ImmutableList.of(new TableRow().set("number", "1")))) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination("project-id:dataset-id.table-id", null) + .withLabels(ImmutableMap.of("env", "prod"))) + .withTestServices(fakeBqServices) + .withMethod(useStreaming ? Method.STREAMING_INSERTS : Method.FILE_LOADS) + .withSchema(schema) + .withoutValidation()); + + thrown.expectMessage("withTableLabels()"); + p.run(); + } + + @Test + public void testTableLabelsDynamicDestinationsRequireCoderV4() throws Exception { + assumeTrue(!useStorageApi); + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of(new TableFieldSchema().setName("number").setType("INTEGER"))); + // TableSchema is not Serializable, and the anonymous DynamicDestinations below is serialized + // with the pipeline, so it captures the schema as JSON + String jsonSchema = BigQueryHelpers.toJsonString(schema); + + p.apply(Create.of(ImmutableList.of(new TableRow().set("number", "1")))) + .apply( + BigQueryIO.writeTableRows() + .to( + new DynamicDestinations() { + @Override + public TableDestination getDestination( + ValueInSingleWindow element) { + return new TableDestination("project-id:dataset-id.table-id", null) + .withLabels(ImmutableMap.of("env", "prod")); + } + + @Override + public TableDestination getTable(TableDestination destination) { + return destination; + } + + @Override + public TableSchema getSchema(TableDestination destination) { + return BigQueryHelpers.fromJsonString(jsonSchema, TableSchema.class); + } + + @Override + public Coder getDestinationCoder() { + return TableDestinationCoderV3.of(); + } + }) + .withTestServices(fakeBqServices) + .withMethod(useStreaming ? Method.STREAMING_INSERTS : Method.FILE_LOADS) + .withoutValidation()); + + thrown.expectMessage("TableDestinationCoderV4"); + p.run(); + } + + @Test + public void testWithTableLabelsRejectsSingleTable() throws Exception { + p.enableAbandonedNodeEnforcement(false); + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("withTableLabels() only applies to to(SerializableFunction)"); + p.apply(Create.of(ImmutableList.of(new TableRow().set("number", "1")))) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.table-id") + .withTestServices(fakeBqServices) + .withTableLabels() + .withoutValidation()); + } + + @Test + public void testTableLabelsRejectNullValues() { + Map labels = new HashMap<>(); + labels.put("env", null); + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("table labels may not contain null keys or values"); + new TableDestination("project-id:dataset-id.table-id", null).withLabels(labels); + } + public void runStreamingFileLoads(String tableRef, boolean useTempTables, boolean useTempDataset) throws Exception { assumeTrue(!useStorageApi); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4Test.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4Test.java new file mode 100644 index 000000000000..586955034e19 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4Test.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.io.gcp.bigquery; + +import com.google.api.services.bigquery.model.Clustering; +import com.google.api.services.bigquery.model.TimePartitioning; +import java.util.LinkedHashMap; +import java.util.Map; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.testing.CoderProperties; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Test case for {@link TableDestinationCoderV4}. */ +@RunWith(JUnit4.class) +public class TableDestinationCoderV4Test { + + private static final Coder TEST_CODER = TableDestinationCoderV4.of(); + + @Test + public void testDecodeEncodeEqualWithLabels() throws Exception { + TableDestination value = + new TableDestination( + "project:dataset.table", + "description", + new TimePartitioning().setType("DAY").setField("date"), + new Clustering().setFields(ImmutableList.of("date"))) + .withLabels(ImmutableMap.of("env", "prod", "team", "data")); + + CoderProperties.coderDecodeEncodeEqual(TEST_CODER, value); + } + + @Test + public void testDecodeEncodeEqualWithoutLabels() throws Exception { + CoderProperties.coderDecodeEncodeEqual( + TEST_CODER, new TableDestination("project:dataset.table", null)); + } + + @Test + public void testLabelInsertionOrderDoesNotChangeEncoding() throws Exception { + Map envFirst = new LinkedHashMap<>(); + envFirst.put("env", "prod"); + envFirst.put("team", "data"); + Map teamFirst = new LinkedHashMap<>(); + teamFirst.put("team", "data"); + teamFirst.put("env", "prod"); + + CoderProperties.coderDeterministic( + TEST_CODER, + new TableDestination("project:dataset.table", null).withLabels(envFirst), + new TableDestination("project:dataset.table", null).withLabels(teamFirst)); + } +} From 4a4b8e3e3de74184c3d3d141e9d80171bebfa66e Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 20:23:46 +0200 Subject: [PATCH 2/3] Add table_labels_enabled to BigQueryIO write translation --- .../sdk/io/gcp/bigquery/BigQueryIOTranslation.java | 10 ++++++++++ .../sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java | 1 + 2 files changed, 11 insertions(+) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java index 9e81b675a755..cefdbe0fd163 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java +++ b/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 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) { diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java index 4391a66e27ee..426f8c513c25 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslationTest.java @@ -132,6 +132,7 @@ public class BigQueryIOTranslationTest { WRITE_TRANSFORM_SCHEMA_MAPPING.put("getAutoSharding", "auto_sharding"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getPropagateSuccessful", "propagate_successful"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getAutoSchemaUpdate", "auto_schema_update"); + WRITE_TRANSFORM_SCHEMA_MAPPING.put("getTableLabelsEnabled", "table_labels_enabled"); WRITE_TRANSFORM_SCHEMA_MAPPING.put( "getAutoSchemaUpdateStrictTimeout", "auto_schema_update_strict_timeout_ms"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getWriteProtosClass", "write_protos_class"); From a6ef0308565a414077bdce2a072dc67c0503885e Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Tue, 27 Oct 2026 21:13:18 +0200 Subject: [PATCH 3/3] Retrigger CI