From 2121346ce4658f58aa6c477c35d3917f298f0e44 Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Sun, 25 Oct 2026 17:09:31 +0200 Subject: [PATCH 1/2] [Java] BigQueryIO: support integer-range partitioning in Write Add Write.withRangePartitioning(...) for single tables, and withRangePartitioning() for table functions and DynamicDestinations that set a RangePartitioning on each TableDestination. Range partitioning travels in a new TableDestinationCoderV4, used only when the option is set, so existing pipelines keep their coder. Turning it on requires draining and restarting a running pipeline. FILE_LOADS temp tables now copy the range partitioning of an existing destination table, so the copy jobs into it succeed. Fixes #27069 --- CHANGES.md | 1 + ...ltCoderCloudObjectTranslatorRegistrar.java | 2 + .../beam/sdk/io/gcp/bigquery/BatchLoads.java | 4 +- .../sdk/io/gcp/bigquery/BigQueryHelpers.java | 8 + .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 132 ++++++- .../gcp/bigquery/BigQueryIOTranslation.java | 14 + .../io/gcp/bigquery/CreateTableHelpers.java | 6 + .../bigquery/DynamicDestinationsHelpers.java | 91 +++-- .../sdk/io/gcp/bigquery/PrepareWrite.java | 35 ++ .../sdk/io/gcp/bigquery/TableDestination.java | 68 +++- .../gcp/bigquery/TableDestinationCoderV4.java | 64 ++++ .../gcp/bigquery/UpdateSchemaDestination.java | 7 + .../beam/sdk/io/gcp/bigquery/WriteTables.java | 6 + .../sdk/io/gcp/testing/FakeJobService.java | 10 + .../bigquery/BigQueryIOTranslationTest.java | 33 ++ .../io/gcp/bigquery/BigQueryIOWriteTest.java | 329 +++++++++++++++++- .../bigquery/BigQueryRangePartitioningIT.java | 229 ++++++++++++ .../bigquery/TableDestinationCoderV4Test.java | 107 ++++++ 18 files changed, 1099 insertions(+), 47 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/BigQueryRangePartitioningIT.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 8fa59e7d5ea0..fb2221f357ef 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: added `Write.withRangePartitioning(...)` to create integer-range partitioned tables. For dynamic destinations, set a `RangePartitioning` in each `TableDestination` and call `withRangePartitioning()`. The option uses a new `TableDestinationCoderV4`, so to turn it on, drain the running pipeline and start a new one; in-place update won't work. `FILE_LOADS` temp tables now also copy the range partitioning of an existing destination table ([#27069](https://github.com/apache/beam/issues/27069)). ## New Features / Improvements diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/DefaultCoderCloudObjectTranslatorRegistrar.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/DefaultCoderCloudObjectTranslatorRegistrar.java index 7137eabf1f04..d5035f4d757b 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/DefaultCoderCloudObjectTranslatorRegistrar.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/DefaultCoderCloudObjectTranslatorRegistrar.java @@ -45,6 +45,7 @@ import org.apache.beam.sdk.io.gcp.bigquery.RowMutation; import org.apache.beam.sdk.io.gcp.bigquery.TableDestinationCoderV2; import org.apache.beam.sdk.io.gcp.bigquery.TableDestinationCoderV3; +import org.apache.beam.sdk.io.gcp.bigquery.TableDestinationCoderV4; import org.apache.beam.sdk.io.gcp.bigquery.TableRowJsonCoder; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; @@ -109,6 +110,7 @@ public class DefaultCoderCloudObjectTranslatorRegistrar StringUtf8Coder.class, TableDestinationCoderV2.class, TableDestinationCoderV3.class, + TableDestinationCoderV4.class, TableRowJsonCoder.class, TextualIntegerCoder.class, VarIntCoder.class, 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..8cd133556318 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 rangePartitioningEnabled, boolean useAvroLogicalTypes, @Nullable String tempDataset, BadRecordRouter badRecordRouter, @@ -217,7 +218,8 @@ class BatchLoads schemaUpdateOptions = Collections.emptySet(); this.tempDataset = tempDataset; this.tableDestinationCoder = - clusteringEnabled ? TableDestinationCoderV3.of() : TableDestinationCoderV2.of(); + DynamicDestinationsHelpers.tableDestinationCoder( + clusteringEnabled, rangePartitioningEnabled); 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/BigQueryHelpers.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java index 3e3b216eae84..de290b404613 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryHelpers.java @@ -32,6 +32,7 @@ import com.google.api.services.bigquery.model.JobConfigurationLoad; import com.google.api.services.bigquery.model.JobReference; import com.google.api.services.bigquery.model.JobStatus; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; @@ -828,6 +829,13 @@ public String apply(TimePartitioning partitioning) { } } + static class RangePartitioningToJson implements SerializableFunction { + @Override + public String apply(RangePartitioning partitioning) { + return toJsonString(partitioning); + } + } + static Clustering clusteringFromJsonFields(String jsonStringClustering) { JsonElement jsonClustering = JsonParser.parseString(jsonStringClustering); 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 07fa1a320e92..491aa28261d3 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 @@ -33,6 +33,7 @@ import com.google.api.services.bigquery.model.JobConfigurationQuery; import com.google.api.services.bigquery.model.JobReference; import com.google.api.services.bigquery.model.JobStatistics; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableCell; import com.google.api.services.bigquery.model.TableConstraints; @@ -94,6 +95,7 @@ import org.apache.beam.sdk.io.fs.ResolveOptions; import org.apache.beam.sdk.io.fs.ResourceId; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.JsonTableRefToTableRef; +import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.RangePartitioningToJson; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.TableRefToJson; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.TableSchemaToJsonSchema; import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.TableSpecToTableRef; @@ -431,6 +433,25 @@ * .withTimePartitioning(new TimePartitioning().setField("time"))); * } * + *

Writing to integer-range partitioned BigQuery tables

+ * + *

To create tables with integer-range + * partitioning, specify {@link BigQueryIO.Write#withRangePartitioning(RangePartitioning)}. For + * example: + * + * 

{@code
+ * PCollection orders = ...;
+ *
+ * orders.apply(BigQueryIO.write()
+ *         .withSchema(schema)
+ *         .withFormatFunction(order -> new TableRow()...)
+ *         .to("my-project:my_dataset.my_table")
+ *         .withRangePartitioning(new RangePartitioning()
+ *             .setField("customer_id")
+ *             .setRange(new RangePartitioning.Range().setStart(0L).setEnd(1000L).setInterval(10L))));
+ * }
+ * *

Writing different values to different tables

* *

A common use case is to dynamically generate BigQuery table names based on the current value. @@ -2988,6 +3009,8 @@ public enum Method { abstract @Nullable ValueProvider getJsonClustering(); + abstract @Nullable ValueProvider getJsonRangePartitioning(); + abstract CreateDisposition getCreateDisposition(); abstract WriteDisposition getWriteDisposition(); @@ -3109,6 +3132,8 @@ abstract Builder setAvroSchemaFactory( abstract Builder setJsonClustering(ValueProvider clustering); + abstract Builder setJsonRangePartitioning(ValueProvider jsonRangePartitioning); + abstract Builder setCreateDisposition(CreateDisposition createDisposition); abstract Builder setWriteDisposition(WriteDisposition writeDisposition); @@ -3330,7 +3355,8 @@ 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 any of the returned destinations are + * configured with range partitioning, return {@link TableDestinationCoderV4} instead. */ public Write to(DynamicDestinations dynamicDestinations) { checkArgument(dynamicDestinations != null, "dynamicDestinations can not be null"); @@ -3522,6 +3548,63 @@ public Write withClustering() { return withClustering(new Clustering()); } + /** + * Sets the integer-range + * partitioning of tables created by this write. BigQuery doesn't allow both time and range + * partitioning on one table, so this can't be combined with {@link #withTimePartitioning}. + * + *

This form is only for writing to a single table. With {@link #to(SerializableFunction)}, + * call {@link #withRangePartitioning()} instead and set the range partitioning in each returned + * {@link TableDestination}. With {@link #to(DynamicDestinations)}, also return {@link + * TableDestinationCoderV4} from {@link DynamicDestinations#getDestinationCoder()}. + * + *

This option switches the destination coder to {@link TableDestinationCoderV4}, which also + * encodes range partitioning. V4 can't read state written by the older coders, so a running + * pipeline can't be updated in place to turn this on: drain it and start a new one. New + * pipelines can set it whether or not their tables are range partitioned. + */ + public Write withRangePartitioning(RangePartitioning partitioning) { + checkArgument(partitioning != null, "partitioning can not be null"); + return withJsonRangePartitioning( + StaticValueProvider.of(BigQueryHelpers.toJsonString(partitioning))); + } + + /** + * Like {@link #withRangePartitioning(RangePartitioning)} but using a deferred {@link + * ValueProvider}. + */ + public Write withRangePartitioning(ValueProvider partitioning) { + checkArgument(partitioning != null, "partitioning can not be null"); + return withJsonRangePartitioning( + NestedValueProvider.of(partitioning, new RangePartitioningToJson())); + } + + /** + * The same as {@link #withRangePartitioning(RangePartitioning)}, but takes a JSON-serialized + * object. For example: {@code {"field": "id", "range": {"start": "0", "end": "100", "interval": + * "10"}}} + */ + public Write withJsonRangePartitioning(ValueProvider partitioning) { + checkArgument(partitioning != null, "partitioning can not be null"); + return toBuilder().setJsonRangePartitioning(partitioning).build(); + } + + /** + * Turns on range partitioning for dynamic destinations. Each {@link TableDestination} returned + * by {@link #to(SerializableFunction)} sets its own range partitioning. With {@link + * #to(DynamicDestinations)}, also return {@link TableDestinationCoderV4} from {@link + * DynamicDestinations#getDestinationCoder()}. For a single table, use {@link + * #withRangePartitioning(RangePartitioning)}. + * + *

Like {@link #withRangePartitioning(RangePartitioning)}, this switches to {@link + * TableDestinationCoderV4}, so a running pipeline has to be drained and restarted to turn it + * on. + */ + public Write withRangePartitioning() { + return withRangePartitioning(new RangePartitioning()); + } + /** 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"); @@ -4219,6 +4302,22 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { !getUseBeamSchema(), "Auto schema update not supported when using Beam schemas."); } + checkArgument( + !(hasJsonConfig(getJsonTimePartitioning()) && hasJsonConfig(getJsonRangePartitioning())), + "A BigQuery table can be either time partitioned or range partitioned. Only one of" + + " withTimePartitioning and withRangePartitioning may be set."); + + ValueProvider jsonRangePartitioning = getJsonRangePartitioning(); + if (jsonRangePartitioning != null + && (getDynamicDestinations() != null || getTableFunction() != null)) { + checkArgument( + jsonRangePartitioning.isAccessible() + && !ConstantTimePartitioningClusteringDestinations.isJsonConfigPresent( + jsonRangePartitioning), + "When writing to dynamic destinations, set range partitioning in the returned" + + " TableDestination objects and call withRangePartitioning() with no arguments."); + } + if (getJsonTimePartitioning() != null) { checkArgument( getDynamicDestinations() == null, @@ -4240,13 +4339,18 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { if (jsonTableRef != null) { tableDestinations = DynamicDestinationsHelpers.ConstantTableDestinations.fromJsonTableRef( - jsonTableRef, getTableDescription(), jsonClustering != null); + jsonTableRef, + getTableDescription(), + jsonClustering != null, + jsonRangePartitioning != null); } else { // Checked above: exactly one of jsonTableRef, tableFunction or dynamicDestinations is // set. tableDestinations = new TableFunctionDestinations<>( - checkStateNotNull(tableFunction), jsonClustering != null); + checkStateNotNull(tableFunction), + jsonClustering != null, + jsonRangePartitioning != null); } // Wrap with a DynamicDestinations class that will provide a schema. There might be no @@ -4259,12 +4363,15 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { tableDestinations = new SchemaFromViewDestinations<>(tableDestinations, schemaFromView); } - // Wrap with a DynamicDestinations class that will provide the proper TimePartitioning. + // Wrap with a DynamicDestinations class that will provide the proper TimePartitioning, + // Clustering and RangePartitioning. ValueProvider jsonTimePartitioning = getJsonTimePartitioning(); - if (jsonTimePartitioning != null || jsonClustering != null) { + if (jsonTimePartitioning != null + || jsonClustering != null + || jsonRangePartitioning != null) { tableDestinations = new ConstantTimePartitioningClusteringDestinations<>( - tableDestinations, jsonTimePartitioning, jsonClustering); + tableDestinations, jsonTimePartitioning, jsonClustering, jsonRangePartitioning); } List primaryKey = getPrimaryKey(); if (primaryKey != null) { @@ -4403,7 +4510,10 @@ private WriteResult expandTyped( input .apply( "PrepareWrite", - new PrepareWrite<>(dynamicDestinations, SerializableFunctions.identity())) + new PrepareWrite<>( + dynamicDestinations, + SerializableFunctions.identity(), + PrepareWrite.dropsRangePartitioning(destinationCoder))) .setCoder(KvCoder.of(destinationCoder, input.getCoder())); return continueExpandTyped( @@ -4530,6 +4640,7 @@ private WriteResult continueExpandTyped( rowWriterFactory, getKmsKey(), getJsonClustering() != null, + getJsonRangePartitioning() != null, getUseAvroLogicalTypes(), getWriteTempDataset(), getBadRecordRouter(), @@ -4715,6 +4826,13 @@ private boolean hasJsonTypeInSchema(JsonElement schema) { return false; } + /** Returns true if the JSON config is accessible at construction time and not empty. */ + private static boolean hasJsonConfig(@Nullable ValueProvider json) { + return json != null + && json.isAccessible() + && ConstantTimePartitioningClusteringDestinations.isJsonConfigPresent(json); + } + @Override public void populateDisplayData(DisplayData.Builder builder) { super.populateDisplayData(builder); 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..590cf5384d27 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) { if (jsonClustering != null) { fieldValues.put("clustering", jsonClustering.get()); } + ValueProvider jsonRangePartitioning = transform.getJsonRangePartitioning(); + if (jsonRangePartitioning != null) { + fieldValues.put("json_range_partitioning", jsonRangePartitioning.get()); + } if (transform.getCreateDisposition() != null) { fieldValues.put("create_disposition", toByteArray(transform.getCreateDisposition())); } @@ -749,6 +754,15 @@ public Write fromConfigRow(Row configRow, PipelineOptions options) { builder = builder.setJsonClustering(StaticValueProvider.of(jsonClustering)); } } + if (TransformUpgrader.compareVersions(updateCompatibilityBeamVersion, "2.78.0") >= 0) { + // This property was added for Beam 2.78.0 hence not available when + // upgrading the transform from previous Beam versions. + String jsonRangePartitioning = configRow.getString("json_range_partitioning"); + if (jsonRangePartitioning != null) { + builder = + builder.setJsonRangePartitioning(StaticValueProvider.of(jsonRangePartitioning)); + } + } byte[] createDispositionBytes = configRow.getBytes("create_disposition"); if (createDispositionBytes != null) { builder = 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..ca9da9d1e078 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 @@ -27,6 +27,7 @@ import com.google.api.services.bigquery.model.BigLakeConfiguration; import com.google.api.services.bigquery.model.Clustering; import com.google.api.services.bigquery.model.EncryptionConfiguration; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableConstraints; import com.google.api.services.bigquery.model.TableReference; @@ -204,6 +205,11 @@ private static void tryCreateTable( table.setClustering(clustering); } + RangePartitioning rangePartitioning = tableDestination.getRangePartitioning(); + if (rangePartitioning != null) { + table.setRangePartitioning(rangePartitioning); + } + 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..525408d3f9aa 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 @@ -57,36 +57,58 @@ class DynamicDestinationsHelpers { private static final Logger LOG = LoggerFactory.getLogger(DynamicDestinationsHelpers.class); + /** + * Returns the oldest {@link TableDestination} coder that can carry the enabled fields. Picking + * the oldest one keeps existing pipelines update compatible. + */ + static Coder tableDestinationCoder( + boolean clusteringEnabled, boolean rangePartitioningEnabled) { + if (rangePartitioningEnabled) { + return TableDestinationCoderV4.of(); + } else if (clusteringEnabled) { + return TableDestinationCoderV3.of(); + } else { + return TableDestinationCoderV2.of(); + } + } + /** Always returns a constant table destination. */ static class ConstantTableDestinations extends DynamicDestinations { private final ValueProvider tableSpec; private final @Nullable String tableDescription; private final boolean clusteringEnabled; + private final boolean rangePartitioningEnabled; ConstantTableDestinations( ValueProvider tableSpec, @Nullable String tableDescription, - boolean clusteringEnabled) { + boolean clusteringEnabled, + boolean rangePartitioningEnabled) { this.tableSpec = tableSpec; this.tableDescription = tableDescription; this.clusteringEnabled = clusteringEnabled; + this.rangePartitioningEnabled = rangePartitioningEnabled; } static ConstantTableDestinations fromTableSpec( ValueProvider tableSpec, @Nullable String tableDescription, - boolean clusteringEnabled) { - return new ConstantTableDestinations<>(tableSpec, tableDescription, clusteringEnabled); + boolean clusteringEnabled, + boolean rangePartitioningEnabled) { + return new ConstantTableDestinations<>( + tableSpec, tableDescription, clusteringEnabled, rangePartitioningEnabled); } static ConstantTableDestinations fromJsonTableRef( ValueProvider jsonTableRef, @Nullable String tableDescription, - boolean clusteringEnabled) { + boolean clusteringEnabled, + boolean rangePartitioningEnabled) { return new ConstantTableDestinations<>( NestedValueProvider.of(jsonTableRef, new JsonTableRefToTableSpec()), tableDescription, - clusteringEnabled); + clusteringEnabled, + rangePartitioningEnabled); } @Override @@ -108,11 +130,7 @@ public TableDestination getTable(TableDestination destination) { @Override public Coder getDestinationCoder() { - if (clusteringEnabled) { - return TableDestinationCoderV3.of(); - } else { - return TableDestinationCoderV2.of(); - } + return tableDestinationCoder(clusteringEnabled, rangePartitioningEnabled); } } @@ -121,12 +139,15 @@ static class TableFunctionDestinations extends DynamicDestinations, TableDestination> tableFunction; private final boolean clusteringEnabled; + private final boolean rangePartitioningEnabled; TableFunctionDestinations( SerializableFunction<@Nullable ValueInSingleWindow, TableDestination> tableFunction, - boolean clusteringEnabled) { + boolean clusteringEnabled, + boolean rangePartitioningEnabled) { this.tableFunction = tableFunction; this.clusteringEnabled = clusteringEnabled; + this.rangePartitioningEnabled = rangePartitioningEnabled; } @Override @@ -152,11 +173,7 @@ public TableDestination getTable(TableDestination destination) { @Override public Coder getDestinationCoder() { - if (clusteringEnabled) { - return TableDestinationCoderV3.of(); - } else { - return TableDestinationCoderV2.of(); - } + return tableDestinationCoder(clusteringEnabled, rangePartitioningEnabled); } } @@ -295,16 +312,19 @@ static class ConstantTimePartitioningClusteringDestinations private final @Nullable ValueProvider jsonTimePartitioning; private final @Nullable ValueProvider jsonClustering; + private final @Nullable ValueProvider jsonRangePartitioning; // Lazily initialized and cached values. private @Nullable String evaluatedPartitioning = null; private @Nullable String evaluatedClustering = null; + private @Nullable String evaluatedRangePartitioning = null; private final AtomicBoolean initialized = new AtomicBoolean(false); ConstantTimePartitioningClusteringDestinations( DynamicDestinations inner, @Nullable ValueProvider jsonTimePartitioning, - @Nullable ValueProvider jsonClustering) { + @Nullable ValueProvider jsonClustering, + @Nullable ValueProvider jsonRangePartitioning) { super(inner); checkArgument( @@ -313,12 +333,16 @@ static class ConstantTimePartitioningClusteringDestinations && jsonTimePartitioning.get() != null) || (jsonClustering != null && jsonClustering.isAccessible() - && jsonClustering.get() != null), - "at least one of jsonTimePartitioning or jsonClustering must be non-null, accessible " - + "and present"); + && jsonClustering.get() != null) + || (jsonRangePartitioning != null + && jsonRangePartitioning.isAccessible() + && jsonRangePartitioning.get() != null), + "at least one of jsonTimePartitioning, jsonClustering or jsonRangePartitioning must be " + + "non-null, accessible and present"); this.jsonTimePartitioning = jsonTimePartitioning; this.jsonClustering = jsonClustering; + this.jsonRangePartitioning = jsonRangePartitioning; } static boolean isJsonConfigPresent(ValueProvider json) { @@ -340,6 +364,11 @@ private synchronized void evaluateOncePartitioningAndClustering() { this.evaluatedClustering = jsonClustering.get(); } } + if (jsonRangePartitioning != null) { + if (isJsonConfigPresent(jsonRangePartitioning)) { + this.evaluatedRangePartitioning = jsonRangePartitioning.get(); + } + } initialized.set(true); } @@ -356,18 +385,22 @@ public TableDestination getDestination(@Nullable ValueInSingleWindow element) : destination.getJsonTimePartitioning(); String clustering = evaluatedClustering != null ? evaluatedClustering : destination.getJsonClustering(); + String rangePartitioning = + evaluatedRangePartitioning != null + ? evaluatedRangePartitioning + : destination.getJsonRangePartitioning(); return new TableDestination( - destination.getTableSpec(), destination.getTableDescription(), partitioning, clustering); + destination.getTableSpec(), + destination.getTableDescription(), + partitioning, + clustering, + rangePartitioning); } @Override public Coder getDestinationCoder() { - if (jsonClustering != null) { - return TableDestinationCoderV3.of(); - } else { - return TableDestinationCoderV2.of(); - } + return tableDestinationCoder(jsonClustering != null, jsonRangePartitioning != null); } @Override @@ -379,6 +412,9 @@ public String toString() { if (jsonClustering != null) { helper.add("jsonClustering", jsonClustering); } + if (jsonRangePartitioning != null) { + helper.add("jsonRangePartitioning", jsonRangePartitioning); + } return helper.toString(); } } @@ -503,7 +539,8 @@ public TableDestination getTable(DestinationT destination) { wrappedDestination.getTableSpec(), existingTable.getDescription(), existingTable.getTimePartitioning(), - existingTable.getClustering()); + existingTable.getClustering(), + existingTable.getRangePartitioning()); } } 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..80d7f5a003ea 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 @@ -21,6 +21,7 @@ import com.google.api.services.bigquery.model.TableRow; import java.io.IOException; +import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.transforms.ParDo; @@ -43,16 +44,38 @@ public class PrepareWrite extends PTransform, PCollection>> { private DynamicDestinations dynamicDestinations; private SerializableFunction formatFunction; + private boolean destinationCoderDropsRangePartitioning; public PrepareWrite( DynamicDestinations dynamicDestinations, SerializableFunction formatFunction) { + this(dynamicDestinations, formatFunction, false); + } + + /** + * Like {@link #PrepareWrite(DynamicDestinations, SerializableFunction)}, but fails on any {@link + * TableDestination} with range partitioning when {@code destinationCoderDropsRangePartitioning} + * is set, since the destination coder would drop it. + */ + PrepareWrite( + DynamicDestinations dynamicDestinations, + SerializableFunction formatFunction, + boolean destinationCoderDropsRangePartitioning) { this.dynamicDestinations = dynamicDestinations; this.formatFunction = formatFunction; + this.destinationCoderDropsRangePartitioning = destinationCoderDropsRangePartitioning; + } + + /** Returns true if {@code coder} encodes {@link TableDestination} without range partitioning. */ + static boolean dropsRangePartitioning(Coder coder) { + return coder instanceof TableDestinationCoder + || coder instanceof TableDestinationCoderV2 + || coder instanceof TableDestinationCoderV3; } @Override public PCollection> expand(PCollection input) { + final boolean checkRangePartitioning = destinationCoderDropsRangePartitioning; return input.apply( ParDo.of( new DoFn>() { @@ -75,6 +98,18 @@ public void processElement( + "but %s returned null on element %s", dynamicDestinations, element); + if (checkRangePartitioning && tableDestination instanceof TableDestination) { + TableDestination destination = (TableDestination) tableDestination; + checkArgument( + destination.getJsonRangePartitioning() == null, + "%s returned a destination with range partitioning for table %s, but" + + " range partitioning is not enabled, so the destination coder" + + " would drop it. Call withRangePartitioning() on" + + " BigQueryIO.write(). If you provided a custom DynamicDestinations," + + " also return TableDestinationCoderV4 from getDestinationCoder().", + dynamicDestinations, + destination.getTableSpec()); + } 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..f1952ee2b034 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 @@ -18,6 +18,7 @@ package org.apache.beam.sdk.io.gcp.bigquery; import com.google.api.services.bigquery.model.Clustering; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TimePartitioning; import java.io.Serializable; @@ -32,6 +33,7 @@ public class TableDestination implements Serializable { private final @Nullable String tableDescription; private final @Nullable String jsonTimePartitioning; private final @Nullable String jsonClustering; + private final @Nullable String jsonRangePartitioning; public TableDestination(String tableSpec, @Nullable String tableDescription) { this(tableSpec, tableDescription, (String) null, (String) null); @@ -106,15 +108,62 @@ public TableDestination( @Nullable String tableDescription, @Nullable String jsonTimePartitioning, @Nullable String jsonClustering) { + this(tableSpec, tableDescription, jsonTimePartitioning, jsonClustering, (String) null); + } + + /** + * Creates a destination with range partitioning. BigQuery doesn't allow both time and range + * partitioning on one table, so set at most one of {@code timePartitioning} and {@code + * rangePartitioning}. + */ + public TableDestination( + String tableSpec, + @Nullable String tableDescription, + @Nullable TimePartitioning timePartitioning, + @Nullable Clustering clustering, + @Nullable RangePartitioning rangePartitioning) { + this( + tableSpec, + tableDescription, + timePartitioning != null ? BigQueryHelpers.toJsonString(timePartitioning) : null, + clustering != null ? BigQueryHelpers.toJsonString(clustering) : null, + rangePartitioning != null ? BigQueryHelpers.toJsonString(rangePartitioning) : null); + } + + public TableDestination( + TableReference tableReference, + @Nullable String tableDescription, + @Nullable String jsonTimePartitioning, + @Nullable String jsonClustering, + @Nullable String jsonRangePartitioning) { + this( + BigQueryHelpers.toTableSpec(tableReference), + tableDescription, + jsonTimePartitioning, + jsonClustering, + jsonRangePartitioning); + } + + public TableDestination( + String tableSpec, + @Nullable String tableDescription, + @Nullable String jsonTimePartitioning, + @Nullable String jsonClustering, + @Nullable String jsonRangePartitioning) { this.tableSpec = tableSpec; this.tableDescription = tableDescription; this.jsonTimePartitioning = jsonTimePartitioning; this.jsonClustering = jsonClustering; + this.jsonRangePartitioning = jsonRangePartitioning; } public TableDestination withTableReference(TableReference tableReference) { return new TableDestination( - tableReference, tableDescription, jsonTimePartitioning, jsonClustering); + tableReference, + tableDescription, + jsonTimePartitioning, + jsonClustering, + jsonRangePartitioning); } /** Return the tablespec in [project:].dataset.tableid format. */ @@ -170,6 +219,18 @@ public TableReference getTableReference() { } } + public @Nullable String getJsonRangePartitioning() { + return jsonRangePartitioning; + } + + public @Nullable RangePartitioning getRangePartitioning() { + if (jsonRangePartitioning == null) { + return null; + } else { + return BigQueryHelpers.fromJsonString(jsonRangePartitioning, RangePartitioning.class); + } + } + public @Nullable String getTableDescription() { return tableDescription; } @@ -191,11 +252,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.jsonRangePartitioning, other.jsonRangePartitioning); } @Override public int hashCode() { - return Objects.hash(tableSpec, tableDescription, jsonTimePartitioning); + return Objects.hash(tableSpec, tableDescription, jsonTimePartitioning, jsonRangePartitioning); } } 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..268598f0d49e --- /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 org.apache.beam.sdk.coders.AtomicCoder; +import org.apache.beam.sdk.coders.Coder; +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 encodes time partitioning, clustering and range + * partitioning. {@link BigQueryIO.Write} switches to it only when one of its range partitioning + * options is set. Otherwise it uses {@link TableDestinationCoderV3} or {@link + * TableDestinationCoderV2}, which drop range partitioning. + */ +public class TableDestinationCoderV4 extends AtomicCoder { + private static final TableDestinationCoderV4 INSTANCE = new TableDestinationCoderV4(); + private static final Coder<@Nullable String> rangePartitioningCoder = + NullableCoder.of(StringUtf8Coder.of()); + + public static TableDestinationCoderV4 of() { + return INSTANCE; + } + + @Override + public void encode(TableDestination value, OutputStream outStream) throws IOException { + TableDestinationCoderV3.of().encode(value, outStream); + rangePartitioningCoder.encode(value.getJsonRangePartitioning(), outStream); + } + + @Override + public TableDestination decode(InputStream inStream) throws IOException { + TableDestination destination = TableDestinationCoderV3.of().decode(inStream); + String jsonRangePartitioning = rangePartitioningCoder.decode(inStream); + return new TableDestination( + destination.getTableSpec(), + destination.getTableDescription(), + destination.getJsonTimePartitioning(), + destination.getJsonClustering(), + jsonRangePartitioning); + } + + @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/UpdateSchemaDestination.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java index 536d318f0213..5e194d1dfe2b 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/UpdateSchemaDestination.java @@ -22,6 +22,7 @@ import com.google.api.services.bigquery.model.EncryptionConfiguration; import com.google.api.services.bigquery.model.JobConfigurationLoad; import com.google.api.services.bigquery.model.JobReference; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; @@ -169,6 +170,7 @@ public void processElement( tableReference, tableDestination.getTimePartitioning(), tableDestination.getClustering(), + tableDestination.getRangePartitioning(), schema, writeDisposition, createDisposition, @@ -237,6 +239,7 @@ public void finishBundle(FinishBundleContext context) throws Exception { TableReference tableReference, @Nullable TimePartitioning timePartitioning, @Nullable Clustering clustering, + @Nullable RangePartitioning rangePartitioning, @Nullable TableSchema schema, BigQueryIO.Write.WriteDisposition writeDisposition, BigQueryIO.Write.CreateDisposition createDisposition, @@ -291,6 +294,10 @@ public void finishBundle(FinishBundleContext context) throws Exception { loadConfig.setClustering(clustering); } + if (rangePartitioning != null) { + loadConfig.setRangePartitioning(rangePartitioning); + } + if (kmsKey != null) { loadConfig.setDestinationEncryptionConfiguration( new EncryptionConfiguration().setKmsKeyName(kmsKey)); 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..a13cb4595373 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 @@ -23,6 +23,7 @@ import com.google.api.services.bigquery.model.EncryptionConfiguration; import com.google.api.services.bigquery.model.JobConfigurationLoad; import com.google.api.services.bigquery.model.JobReference; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableSchema; import com.google.api.services.bigquery.model.TimePartitioning; @@ -289,6 +290,7 @@ public void processElement( tableReference, tableDestination.getTimePartitioning(), tableDestination.getClustering(), + tableDestination.getRangePartitioning(), tableSchema, partitionFiles, writeDisposition, @@ -471,6 +473,7 @@ private PendingJob startLoad( TableReference ref, @Nullable TimePartitioning timePartitioning, @Nullable Clustering clustering, + @Nullable RangePartitioning rangePartitioning, @Nullable TableSchema schema, List gcsUris, WriteDisposition writeDisposition, @@ -502,6 +505,9 @@ private PendingJob startLoad( if (clustering != null) { loadConfig.setClustering(clustering); } + if (rangePartitioning != null) { + loadConfig.setRangePartitioning(rangePartitioning); + } if (kmsKey != null) { loadConfig.setDestinationEncryptionConfiguration( new EncryptionConfiguration().setKmsKeyName(kmsKey)); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeJobService.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeJobService.java index ac09b11638de..45326b825a24 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeJobService.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/testing/FakeJobService.java @@ -37,6 +37,7 @@ import com.google.api.services.bigquery.model.JobStatistics; import com.google.api.services.bigquery.model.JobStatistics4; import com.google.api.services.bigquery.model.JobStatus; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableReference; import com.google.api.services.bigquery.model.TableRow; @@ -401,6 +402,9 @@ private JobStatus runLoadJob(JobReference jobRef, JobConfigurationLoad load) if (load.getClustering() != null) { existingTable = existingTable.setClustering(load.getClustering()); } + if (load.getRangePartitioning() != null) { + existingTable = existingTable.setRangePartitioning(load.getRangePartitioning()); + } datasetService.createTable(existingTable); } @@ -430,6 +434,7 @@ private JobStatus runCopyJob(JobConfigurationTableCopy copy) } TimePartitioning partitioning = null; Clustering clustering = null; + RangePartitioning rangePartitioning = null; TableSchema schema = null; boolean first = true; List allRows = Lists.newArrayList(); @@ -442,12 +447,16 @@ private JobStatus runCopyJob(JobConfigurationTableCopy copy) if (!Objects.equals(clustering, table.getClustering())) { return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto()); } + if (!Objects.equals(rangePartitioning, table.getRangePartitioning())) { + return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto()); + } if (!Objects.equals(schema, table.getSchema())) { return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto()); } } partitioning = table.getTimePartitioning(); clustering = table.getClustering(); + rangePartitioning = table.getRangePartitioning(); schema = table.getSchema(); first = false; allRows.addAll( @@ -460,6 +469,7 @@ private JobStatus runCopyJob(JobConfigurationTableCopy copy) .setSchema(schema) .setTimePartitioning(partitioning) .setClustering(clustering) + .setRangePartitioning(rangePartitioning) .setEncryptionConfiguration(copy.getDestinationEncryptionConfiguration())); datasetService.insertAll(destination, allRows, null); return new JobStatus().setState("DONE"); 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..7849d70afb47 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 @@ -20,9 +20,11 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import com.google.api.services.bigquery.model.Clustering; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.TableRow; import java.util.Arrays; import java.util.HashMap; @@ -94,6 +96,7 @@ public class BigQueryIOTranslationTest { WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonSchema", "json_schema"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonTimePartitioning", "json_time_partitioning"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonClustering", "clustering"); + WRITE_TRANSFORM_SCHEMA_MAPPING.put("getJsonRangePartitioning", "json_range_partitioning"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getCreateDisposition", "create_disposition"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getWriteDisposition", "write_disposition"); WRITE_TRANSFORM_SCHEMA_MAPPING.put("getSchemaUpdateOptions", "schema_update_options"); @@ -278,6 +281,36 @@ public void testReCreateWriteTransformFromRowTable() { writeTransformFromRow.getJsonClustering().get()); } + @Test + public void testReCreateWriteTransformFromRowRangePartitioning() { + RangePartitioning testRangePartitioning = + new RangePartitioning() + .setField("id") + .setRange(new RangePartitioning.Range().setStart(0L).setEnd(100L).setInterval(10L)); + BigQueryIO.Write writeTransform = + BigQueryIO.write() + .to("dummyproject:dummydataset.dummytable") + .withRangePartitioning(testRangePartitioning); + + BigQueryIOTranslation.BigQueryIOWriteTranslator translator = + new BigQueryIOTranslation.BigQueryIOWriteTranslator(); + Row row = translator.toConfigRow(writeTransform); + + PipelineOptions options = PipelineOptionsFactory.create(); + options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.78.0"); + BigQueryIO.Write writeTransformFromRow = + (BigQueryIO.Write) translator.fromConfigRow(row, options); + assertNotNull(writeTransformFromRow.getJsonRangePartitioning()); + assertEquals( + BigQueryHelpers.toJsonString(testRangePartitioning), + writeTransformFromRow.getJsonRangePartitioning().get()); + + // Range partitioning is not translated when upgrading from a Beam version that predates it. + options.as(StreamingOptions.class).setUpdateCompatibilityVersion("2.77.0"); + writeTransformFromRow = (BigQueryIO.Write) translator.fromConfigRow(row, options); + assertNull(writeTransformFromRow.getJsonRangePartitioning()); + } + @Test public void testWriteTransformRowIncludesAllFields() { // These fields do not represent properties of the transform. 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..8e7f20d13d5e 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 @@ -46,6 +46,7 @@ import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.Job; import com.google.api.services.bigquery.model.JobConfigurationLoad; +import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.Table; import com.google.api.services.bigquery.model.TableCell; import com.google.api.services.bigquery.model.TableDataInsertAllResponse; @@ -596,18 +597,20 @@ void testTimePartitioningAndClustering( } } - void testTimePartitioningAndClusteringWithAllMethods( - Boolean enablePartitioning, Boolean enableClustering) throws Exception { - BigQueryIO.Write.Method method; + private BigQueryIO.Write.Method writeMethodForTestParameters() { if (useStorageApi) { - method = - useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; + return useStorageApiApproximate ? Method.STORAGE_API_AT_LEAST_ONCE : Method.STORAGE_WRITE_API; } else if (useStreaming) { - method = Method.STREAMING_INSERTS; + return Method.STREAMING_INSERTS; } else { - method = Method.FILE_LOADS; + return Method.FILE_LOADS; } - testTimePartitioningAndClustering(method, enablePartitioning, enableClustering); + } + + void testTimePartitioningAndClusteringWithAllMethods( + Boolean enablePartitioning, Boolean enableClustering) throws Exception { + testTimePartitioningAndClustering( + writeMethodForTestParameters(), enablePartitioning, enableClustering); } @Test @@ -679,6 +682,314 @@ public void testClusteringTableFunction() throws Exception { assertEquals(clustering, table.getClustering()); } + private static RangePartitioning testRangePartitioning() { + return new RangePartitioning() + .setField("number") + .setRange(new RangePartitioning.Range().setStart(0L).setEnd(100L).setInterval(10L)); + } + + private static TableSchema rangePartitioningTestSchema() { + return new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("date").setType("DATE"), + new TableFieldSchema().setName("number").setType("INTEGER"))); + } + + void testRangePartitioning(boolean enableClustering) throws Exception { + TableRow row1 = new TableRow().set("date", "2018-01-01").set("number", "1"); + TableRow row2 = new TableRow().set("date", "2018-01-02").set("number", "2"); + + RangePartitioning rangePartitioning = testRangePartitioning(); + Clustering clustering = new Clustering().setFields(ImmutableList.of("date")); + TableSchema schema = rangePartitioningTestSchema(); + + Write writeTransform = + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.table-id") + .withTestServices(fakeBqServices) + .withMethod(writeMethodForTestParameters()) + .withSchema(schema) + .withRangePartitioning(rangePartitioning) + .withoutValidation(); + if (enableClustering) { + writeTransform = writeTransform.withClustering(clustering); + } + + p.apply(Create.of(row1, row2)).apply(writeTransform); + p.run(); + Table table = + fakeDatasetService.getTable( + BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-id")); + + assertEquals(schema, table.getSchema()); + assertEquals(rangePartitioning, table.getRangePartitioning()); + assertNull(table.getTimePartitioning()); + if (enableClustering) { + assertEquals(clustering, table.getClustering()); + } + } + + @Test + public void testRangePartitioningWithoutClustering() throws Exception { + testRangePartitioning(false); + } + + @Test + public void testRangePartitioningWithClustering() throws Exception { + testRangePartitioning(true); + } + + @Test + public void testRangePartitioningTableFunction() throws Exception { + 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 = rangePartitioningTestSchema(); + + p.apply(Create.of(row1, row2)) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> { + String tableSpec = + "project-id:dataset-id.table-" + vsw.getValue().get("number"); + return new TableDestination( + tableSpec, + null, + null, + new Clustering().setFields(ImmutableList.of("date")), + testRangePartitioning()); + }) + .withTestServices(fakeBqServices) + .withMethod(writeMethodForTestParameters()) + .withSchema(schema) + .withClustering() + .withRangePartitioning() + .withoutValidation()); + p.run(); + + for (String tableId : ImmutableList.of("table-1", "table-2")) { + Table table = + fakeDatasetService.getTable( + BigQueryHelpers.parseTableSpec("project-id:dataset-id." + tableId)); + assertEquals(schema, table.getSchema()); + assertEquals(testRangePartitioning(), table.getRangePartitioning()); + assertEquals(new Clustering().setFields(ImmutableList.of("date")), table.getClustering()); + } + } + + @Test + public void testRangePartitioningTableFunctionRequiresOptIn() throws Exception { + p.apply(Create.of(new TableRow().set("date", "2018-01-01").set("number", "1"))) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination( + "project-id:dataset-id.table-id", + null, + null, + null, + testRangePartitioning())) + .withTestServices(fakeBqServices) + .withMethod(writeMethodForTestParameters()) + .withSchema(rangePartitioningTestSchema()) + .withoutValidation()); + + thrown.expect(RuntimeException.class); + thrown.expectMessage("range partitioning is not enabled"); + p.run(); + } + + /** Writes every element to one range partitioned table, with a configurable coder. */ + static class RangePartitionedDestinations + extends DynamicDestinations { + private final Coder destinationCoder; + + RangePartitionedDestinations(Coder destinationCoder) { + this.destinationCoder = destinationCoder; + } + + @Override + public TableDestination getDestination(@Nullable ValueInSingleWindow element) { + return new TableDestination( + "project-id:dataset-id.table-id", null, null, null, testRangePartitioning()); + } + + @Override + public TableDestination getTable(TableDestination destination) { + return destination; + } + + @Override + public TableSchema getSchema(TableDestination destination) { + return rangePartitioningTestSchema(); + } + + @Override + public Coder getDestinationCoder() { + return destinationCoder; + } + } + + @Test + public void testRangePartitioningDynamicDestinations() throws Exception { + p.apply(Create.of(new TableRow().set("date", "2018-01-01").set("number", "1"))) + .apply( + BigQueryIO.writeTableRows() + .to(new RangePartitionedDestinations(TableDestinationCoderV4.of())) + .withTestServices(fakeBqServices) + .withMethod(writeMethodForTestParameters()) + .withRangePartitioning() + .withoutValidation()); + p.run(); + + Table table = + fakeDatasetService.getTable( + BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-id")); + assertEquals(testRangePartitioning(), table.getRangePartitioning()); + } + + @Test + public void testRangePartitioningDynamicDestinationsRequiresV4Coder() throws Exception { + p.apply(Create.of(new TableRow().set("date", "2018-01-01").set("number", "1"))) + .apply( + BigQueryIO.writeTableRows() + .to(new RangePartitionedDestinations(TableDestinationCoderV3.of())) + .withTestServices(fakeBqServices) + .withMethod(writeMethodForTestParameters()) + .withRangePartitioning() + .withoutValidation()); + + thrown.expect(RuntimeException.class); + thrown.expectMessage("also return TableDestinationCoderV4 from getDestinationCoder()"); + p.run(); + } + + @Test + public void testConstantRangePartitioningNotAllowedWithTableFunction() throws Exception { + p.enableAbandonedNodeEnforcement(false); + + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("call withRangePartitioning() with no arguments"); + p.apply(Create.empty(TableRowJsonCoder.of())) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination("project-id:dataset-id.table-id", null)) + .withSchema(rangePartitioningTestSchema()) + .withRangePartitioning(testRangePartitioning()) + .withTestServices(fakeBqServices) + .withoutValidation()); + } + + @Test + public void testTimeAndRangePartitioningNotAllowed() throws Exception { + p.enableAbandonedNodeEnforcement(false); + + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("Only one of withTimePartitioning and withRangePartitioning may be set"); + p.apply(Create.empty(TableRowJsonCoder.of())) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.table-id") + .withSchema(rangePartitioningTestSchema()) + .withTimePartitioning(new TimePartitioning().setType("DAY").setField("date")) + .withRangePartitioning(testRangePartitioning()) + .withTestServices(fakeBqServices) + .withoutValidation()); + } + + @Test + public void testRangePartitioningWithTempTables() throws Exception { + assumeTrue(!useStorageApi); + assumeTrue(!useStreaming); + 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("project-id:dataset-id.table-id") + .withSchema( + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("number").setType("INTEGER")))) + .withRangePartitioning(testRangePartitioning()) + .withTestServices(fakeBqServices) + .withMaxBytesPerPartition(1) + .withMaxFilesPerPartition(1) + .withoutValidation()); + p.run(); + + assertThat( + fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"), + containsInAnyOrder(Iterables.toArray(elements, TableRow.class))); + Table table = + fakeDatasetService.getTable( + BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-id")); + assertEquals(testRangePartitioning(), table.getRangePartitioning()); + for (Job job : fakeJobService.getAllJobs()) { + JobConfigurationLoad load = job.getConfiguration().getLoad(); + if (load != null) { + assertEquals(testRangePartitioning(), load.getRangePartitioning()); + } + } + } + + @Test + public void testTempTablesInheritRangePartitioningOfExistingTable() throws Exception { + assumeTrue(!useStorageApi); + assumeTrue(!useStreaming); + + TableSchema schema = + new TableSchema() + .setFields( + ImmutableList.of( + new TableFieldSchema().setName("str").setType("STRING"), + new TableFieldSchema().setName("number").setType("INTEGER"))); + fakeDatasetService.createTable( + new Table() + .setTableReference(BigQueryHelpers.parseTableSpec("project-id:dataset-id.table-id")) + .setSchema(schema) + .setRangePartitioning(testRangePartitioning())); + + List elements = Lists.newArrayList(); + for (int i = 1; i < 10; i++) { + elements.add(new TableRow().set("str", "a").set("number", i)); + } + + // No withRangePartitioning() here: the temp tables should take their range partitioning from + // the existing destination table. + p.apply(Create.of(elements)) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.table-id") + .withCreateDisposition(CreateDisposition.CREATE_NEVER) + .withTestServices(fakeBqServices) + .withMaxBytesPerPartition(1) + .withMaxFilesPerPartition(1) + .withoutValidation()); + p.run(); + + assertThat( + fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"), + containsInAnyOrder(Iterables.toArray(elements, TableRow.class))); + int tempTableLoads = 0; + for (Job job : fakeJobService.getAllJobs()) { + JobConfigurationLoad load = job.getConfiguration().getLoad(); + if (load != null && !"table-id".equals(load.getDestinationTable().getTableId())) { + assertEquals(testRangePartitioning(), load.getRangePartitioning()); + tempTableLoads++; + } + } + assertTrue("Expected the write to use temp tables", tempTableLoads > 0); + } + public void runStreamingFileLoads(String tableRef, boolean useTempTables, boolean useTempDataset) throws Exception { assumeTrue(!useStorageApi); @@ -3169,7 +3480,7 @@ private void testWritePartition( boolean isSingleton = numTables == 1 && numFilesPerTable == 0; DynamicDestinations dynamicDestinations = new DynamicDestinationsHelpers.ConstantTableDestinations<>( - ValueProvider.StaticValueProvider.of("SINGLETON"), "", false); + ValueProvider.StaticValueProvider.of("SINGLETON"), "", false, false); List> expectedPartitions = Lists.newArrayList(); if (isSingleton) { expectedPartitions.add(ShardedKey.of(new TableDestination("SINGLETON", ""), 1)); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryRangePartitioningIT.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryRangePartitioningIT.java new file mode 100644 index 000000000000..9b2d4e9319c8 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryRangePartitioningIT.java @@ -0,0 +1,229 @@ +/* + * 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 static org.apache.beam.sdk.io.gcp.testing.BigtableUtils.TEMP_DATASET_PREFIX; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; + +import com.google.api.services.bigquery.model.Clustering; +import com.google.api.services.bigquery.model.RangePartitioning; +import com.google.api.services.bigquery.model.Table; +import com.google.api.services.bigquery.model.TableFieldSchema; +import com.google.api.services.bigquery.model.TableReference; +import com.google.api.services.bigquery.model.TableRow; +import com.google.api.services.bigquery.model.TableSchema; +import java.security.SecureRandom; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.extensions.gcp.options.GcpOptions; +import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.io.gcp.testing.BigqueryClient; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.testing.TestPipelineOptions; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.ValueInSingleWindow; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Integration test that writes to integer-range partitioned tables in BigQuery. */ +@RunWith(JUnit4.class) +public class BigQueryRangePartitioningIT { + private static final int NUM_ROWS = 100; + + private static String project; + private static String bigQueryLocation; + private static final BigqueryClient BQ_CLIENT = new BigqueryClient("BigQueryRangePartitioningIT"); + private static final String DATASET_NAME = + TEMP_DATASET_PREFIX + + "range_partitioning_" + + System.currentTimeMillis() + + "_" + + new SecureRandom().nextInt(32); + private static final RangePartitioning RANGE_PARTITIONING = + new RangePartitioning() + .setField("id") + .setRange(new RangePartitioning.Range().setStart(0L).setEnd(100L).setInterval(10L)); + private static final Clustering CLUSTERING = new Clustering().setFields(Arrays.asList("name")); + private static final TableSchema SCHEMA = + new TableSchema() + .setFields( + Arrays.asList( + new TableFieldSchema().setName("id").setType("INTEGER"), + new TableFieldSchema().setName("name").setType("STRING"))); + + private TestPipelineOptions options; + + @BeforeClass + public static void setupTestEnvironment() throws Exception { + project = TestPipeline.testingPipelineOptions().as(GcpOptions.class).getProject(); + bigQueryLocation = + TestPipeline.testingPipelineOptions().as(TestBigQueryOptions.class).getBigQueryLocation(); + BQ_CLIENT.createNewDataset(project, DATASET_NAME, null, bigQueryLocation); + } + + @Before + public void setUp() { + options = TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class); + options.setTempLocation( + FileSystems.matchNewDirectory(options.getTempRoot(), "temp-it").toString()); + } + + @AfterClass + public static void cleanup() { + BQ_CLIENT.deleteDataset(project, DATASET_NAME); + } + + private static List rows() { + List rows = new ArrayList<>(); + for (int i = 0; i < NUM_ROWS; i++) { + rows.add(new TableRow().set("id", i).set("name", "name_" + (i % 7))); + } + return rows; + } + + private Table getTable(String tableName) throws Exception { + return BQ_CLIENT.getTableResource(project, DATASET_NAME, tableName); + } + + private static void assertRangePartitioning(Table table) { + RangePartitioning actual = table.getRangePartitioning(); + assertNotNull("Table is not range partitioned", actual); + assertEquals(RANGE_PARTITIONING.getField(), actual.getField()); + assertEquals(RANGE_PARTITIONING.getRange().getStart(), actual.getRange().getStart()); + assertEquals(RANGE_PARTITIONING.getRange().getEnd(), actual.getRange().getEnd()); + assertEquals(RANGE_PARTITIONING.getRange().getInterval(), actual.getRange().getInterval()); + } + + private void assertRowCount(String tableName) throws Exception { + TableRow response = + Iterables.getOnlyElement( + BQ_CLIENT.queryUnflattened( + String.format("SELECT COUNT(*) FROM `%s.%s.%s`", project, DATASET_NAME, tableName), + project, + true, + true, + bigQueryLocation)); + assertEquals(NUM_ROWS, Integer.parseInt((String) response.get("f0_"))); + } + + @Test + public void testE2EBigQueryRangePartitioning() throws Exception { + String tableName = "range_partitioned_" + System.currentTimeMillis(); + + Pipeline p = Pipeline.create(options); + p.apply(Create.of(rows())) + .apply( + BigQueryIO.writeTableRows() + .to(String.format("%s.%s", DATASET_NAME, tableName)) + .withRangePartitioning(RANGE_PARTITIONING) + .withSchema(SCHEMA) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE)); + p.run().waitUntilFinish(); + + assertRangePartitioning(getTable(tableName)); + assertRowCount(tableName); + } + + @Test + public void testE2EBigQueryRangePartitioningTableFunction() throws Exception { + String tableName = "range_partitioned_table_function_" + System.currentTimeMillis(); + String destination = String.format("%s.%s", DATASET_NAME, tableName); + + Pipeline p = Pipeline.create(options); + p.apply(Create.of(rows())) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination( + destination, null, null, CLUSTERING, RANGE_PARTITIONING)) + .withClustering() + .withRangePartitioning() + .withSchema(SCHEMA) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_TRUNCATE)); + p.run().waitUntilFinish(); + + Table table = getTable(tableName); + assertRangePartitioning(table); + assertEquals(CLUSTERING.getFields(), table.getClustering().getFields()); + assertRowCount(tableName); + } + + @Test + public void testE2EBigQueryRangePartitioningStorageWriteApi() throws Exception { + String tableName = "range_partitioned_storage_api_" + System.currentTimeMillis(); + + Pipeline p = Pipeline.create(options); + p.apply(Create.of(rows())) + .apply( + BigQueryIO.writeTableRows() + .to(String.format("%s.%s", DATASET_NAME, tableName)) + .withMethod(BigQueryIO.Write.Method.STORAGE_WRITE_API) + .withRangePartitioning(RANGE_PARTITIONING) + .withSchema(SCHEMA) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND)); + p.run().waitUntilFinish(); + + assertRangePartitioning(getTable(tableName)); + assertRowCount(tableName); + } + + @Test + public void testE2EBigQueryTempTablesIntoExistingRangePartitionedTable() throws Exception { + String tableName = "range_partitioned_existing_" + System.currentTimeMillis(); + BQ_CLIENT.createNewTable( + project, + DATASET_NAME, + new Table() + .setTableReference( + new TableReference() + .setProjectId(project) + .setDatasetId(DATASET_NAME) + .setTableId(tableName)) + .setSchema(SCHEMA) + .setRangePartitioning(RANGE_PARTITIONING)); + + // The small partition limits make FILE_LOADS go through temp tables and copy jobs. BigQuery + // rejects the copy unless the temp tables have the same range partitioning as the existing + // table, and no withRangePartitioning() is set here. + Pipeline p = Pipeline.create(options); + p.apply(Create.of(rows())) + .apply( + BigQueryIO.writeTableRows() + .to(String.format("%s.%s", DATASET_NAME, tableName)) + .withMaxBytesPerPartition(1) + .withMaxFilesPerPartition(1) + .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER) + .withWriteDisposition(BigQueryIO.Write.WriteDisposition.WRITE_APPEND)); + p.run().waitUntilFinish(); + + assertRangePartitioning(getTable(tableName)); + assertRowCount(tableName); + } +} 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..24f03fe1eb9f --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableDestinationCoderV4Test.java @@ -0,0 +1,107 @@ +/* + * 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 static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import com.google.api.services.bigquery.model.Clustering; +import com.google.api.services.bigquery.model.RangePartitioning; +import com.google.api.services.bigquery.model.TimePartitioning; +import java.util.Arrays; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.testing.CoderProperties; +import org.apache.beam.sdk.util.CoderUtils; +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(); + + private static final RangePartitioning RANGE_PARTITIONING = + new RangePartitioning() + .setField("id") + .setRange(new RangePartitioning.Range().setStart(0L).setEnd(100L).setInterval(10L)); + + @Test + public void testDecodeEncodeEqualWithRangePartitioning() throws Exception { + TableDestination value = + new TableDestination( + "project:dataset.table", + "description", + null, + new Clustering().setFields(Arrays.asList("a", "b")), + RANGE_PARTITIONING); + + CoderProperties.coderDecodeEncodeEqual(TEST_CODER, value); + + TableDestination decoded = + CoderUtils.decodeFromByteArray(TEST_CODER, CoderUtils.encodeToByteArray(TEST_CODER, value)); + assertEquals(RANGE_PARTITIONING, decoded.getRangePartitioning()); + assertEquals(value.getJsonClustering(), decoded.getJsonClustering()); + assertNull(decoded.getJsonTimePartitioning()); + } + + @Test + public void testDecodeEncodeEqualWithoutRangePartitioning() throws Exception { + TableDestination value = + new TableDestination( + "project:dataset.table", + null, + new TimePartitioning().setType("DAY").setField("date"), + new Clustering().setFields(Arrays.asList("a"))); + + CoderProperties.coderDecodeEncodeEqual(TEST_CODER, value); + + TableDestination decoded = + CoderUtils.decodeFromByteArray(TEST_CODER, CoderUtils.encodeToByteArray(TEST_CODER, value)); + assertEquals(value.getJsonTimePartitioning(), decoded.getJsonTimePartitioning()); + assertEquals(value.getJsonClustering(), decoded.getJsonClustering()); + assertNull(decoded.getJsonRangePartitioning()); + } + + @Test + public void testEncodingExtendsV3() throws Exception { + TableDestination value = + new TableDestination( + "project:dataset.table", + "description", + null, + new Clustering().setFields(Arrays.asList("a")), + RANGE_PARTITIONING); + + byte[] v3Bytes = CoderUtils.encodeToByteArray(TableDestinationCoderV3.of(), value); + byte[] v4Bytes = CoderUtils.encodeToByteArray(TEST_CODER, value); + + // V4 writes the V3 encoding followed by the range partitioning. + assertArrayEquals(v3Bytes, Arrays.copyOf(v4Bytes, v3Bytes.length)); + } + + @Test + public void testCoderIsDeterministic() throws Exception { + CoderProperties.coderDeterministic( + TEST_CODER, + new TableDestination("project:dataset.table", null, null, null, RANGE_PARTITIONING), + new TableDestination("project:dataset.table", null, null, null, RANGE_PARTITIONING)); + } +} From e2a5b0b134aef065ecf4f0589fa96237864ee69f Mon Sep 17 00:00:00 2001 From: Maksym Tymoshyk Date: Sun, 25 Oct 2026 20:18:54 +0200 Subject: [PATCH 2/2] Reject range partitioning configs that BigQuery would drop or refuse withRangePartitioning() with no arguments on a single table only turned on the V4 coder, so BigQuery created an unpartitioned table without an error. The write now fails at construction and says to pass a RangePartitioning. A TableDestination with both time and range partitioning used to fail later, inside BigQuery. The constructor now throws IllegalArgumentException. --- .../beam/sdk/io/gcp/bigquery/BigQueryIO.java | 7 ++++ .../sdk/io/gcp/bigquery/TableDestination.java | 10 ++++- .../io/gcp/bigquery/BigQueryIOWriteTest.java | 40 +++++++++++++++++++ 3 files changed, 55 insertions(+), 2 deletions(-) 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 491aa28261d3..30cdc16c0911 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 @@ -4316,6 +4316,13 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) { jsonRangePartitioning), "When writing to dynamic destinations, set range partitioning in the returned" + " TableDestination objects and call withRangePartitioning() with no arguments."); + } else if (jsonRangePartitioning != null && jsonRangePartitioning.isAccessible()) { + checkArgument( + ConstantTimePartitioningClusteringDestinations.isJsonConfigPresent( + jsonRangePartitioning), + "withRangePartitioning() with no arguments only applies to to(SerializableFunction)" + + " and to(DynamicDestinations). To write to a single table, pass a" + + " RangePartitioning."); } if (getJsonTimePartitioning() != null) { 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 f1952ee2b034..16883dddfcb4 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 @@ -17,6 +17,8 @@ */ package org.apache.beam.sdk.io.gcp.bigquery; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + import com.google.api.services.bigquery.model.Clustering; import com.google.api.services.bigquery.model.RangePartitioning; import com.google.api.services.bigquery.model.TableReference; @@ -113,8 +115,8 @@ public TableDestination( /** * Creates a destination with range partitioning. BigQuery doesn't allow both time and range - * partitioning on one table, so set at most one of {@code timePartitioning} and {@code - * rangePartitioning}. + * partitioning on one table, so this throws {@link IllegalArgumentException} if both {@code + * timePartitioning} and {@code rangePartitioning} are set. */ public TableDestination( String tableSpec, @@ -150,6 +152,10 @@ public TableDestination( @Nullable String jsonTimePartitioning, @Nullable String jsonClustering, @Nullable String jsonRangePartitioning) { + checkArgument( + jsonTimePartitioning == null || jsonRangePartitioning == null, + "Table %s has both time and range partitioning, but a BigQuery table can only have one.", + tableSpec); this.tableSpec = tableSpec; this.tableDescription = tableDescription; this.jsonTimePartitioning = jsonTimePartitioning; 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 8e7f20d13d5e..6025eed05d49 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 @@ -801,6 +801,46 @@ public void testRangePartitioningTableFunctionRequiresOptIn() throws Exception { p.run(); } + @Test + public void testTableFunctionWithTimeAndRangePartitioningFails() throws Exception { + p.apply(Create.of(new TableRow().set("date", "2018-01-01").set("number", "1"))) + .apply( + BigQueryIO.writeTableRows() + .to( + (ValueInSingleWindow vsw) -> + new TableDestination( + "project-id:dataset-id.table-id", + null, + new TimePartitioning().setType("DAY").setField("date"), + null, + testRangePartitioning())) + .withRangePartitioning() + .withTestServices(fakeBqServices) + .withMethod(writeMethodForTestParameters()) + .withSchema(rangePartitioningTestSchema()) + .withoutValidation()); + + thrown.expect(RuntimeException.class); + thrown.expectMessage("has both time and range partitioning"); + p.run(); + } + + @Test + public void testNoArgRangePartitioningNotAllowedWithSingleTable() throws Exception { + p.enableAbandonedNodeEnforcement(false); + + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("To write to a single table, pass a RangePartitioning"); + p.apply(Create.empty(TableRowJsonCoder.of())) + .apply( + BigQueryIO.writeTableRows() + .to("project-id:dataset-id.table-id") + .withSchema(rangePartitioningTestSchema()) + .withRangePartitioning() + .withTestServices(fakeBqServices) + .withoutValidation()); + } + /** Writes every element to one range partitioned table, with a configurable coder. */ static class RangePartitionedDestinations extends DynamicDestinations {