Skip to content

Add BigQueryIO support for clone table creation #38883

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lalitium wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Add BigQueryIO support for clone table creation #38883

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
48a4664
Add BigQueryIO clone table support
lalitium Jun 10, 2026
258ca80
Fix BigQueryIO clone validation message
lalitium Jun 10, 2026
3b6ffc8
Track BigQuery clone source in write transform translation
lalitium Jun 10, 2026
d7a03f5
updates based on reviews
lalitium Jun 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions ...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BatchLoads.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,7 +429,7 @@ private WriteResult expandTriggered(PCollection<KV<DestinationT, ElementT>> inpu
zeroLoadJobIdPrefixView,
loadJobProjectId,
WriteDisposition.WRITE_APPEND,
CreateDisposition.CREATE_NEVER,
createDisposition,
maxRetryJobs,
kmsKey,
schemaUpdateOptions,
Expand Down Expand Up @@ -535,7 +535,7 @@ public WriteResult expandUntriggered(PCollection<KV<DestinationT, ElementT>> inp
zeroLoadJobIdPrefixView,
loadJobProjectId,
WriteDisposition.WRITE_APPEND,
CreateDisposition.CREATE_NEVER,
createDisposition,
maxRetryJobs,
kmsKey,
schemaUpdateOptions,
Expand Down
60 changes: 57 additions & 3 deletions ...o/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIO.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.JobService;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StorageClient;
import org.apache.beam.sdk.io.gcp.bigquery.BigQuerySourceBase.ExtractResult;
import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.ConstantCloneSourceDestinations;
import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.ConstantSchemaDestinations;
import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.ConstantTimePartitioningClusteringDestinations;
import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinationsHelpers.SchemaFromViewDestinations;
Expand Down Expand Up @@ -2705,6 +2706,8 @@ public enum Method {

abstract @Nullable ValueProvider<String> getJsonSchema();

abstract @Nullable ValueProvider<String> getJsonCloneSourceTableRef();

abstract @Nullable ValueProvider<String> getJsonTimePartitioning();

abstract @Nullable ValueProvider<String> getJsonClustering();
Expand Down Expand Up @@ -2824,6 +2827,8 @@ abstract Builder<T> setAvroSchemaFactory(

abstract Builder<T> setJsonSchema(ValueProvider<String> jsonSchema);

abstract Builder<T> setJsonCloneSourceTableRef(ValueProvider<String> jsonCloneSourceTableRef);

abstract Builder<T> setJsonTimePartitioning(ValueProvider<String> jsonTimePartitioning);

abstract Builder<T> setJsonClustering(ValueProvider<String> clustering);
Expand Down Expand Up @@ -3162,6 +3167,39 @@ public Write<T> withJsonSchema(ValueProvider<String> jsonSchema) {
return toBuilder().setJsonSchema(jsonSchema).build();
}

/**
* Creates missing destination tables as clones of the specified base table.
*
* <p>The clone source is used only if writing to a table that does not already exist and {@link
* CreateDisposition} is set to {@link CreateDisposition#CREATE_IF_NEEDED}. This option is
* mutually exclusive with {@link #withSchema(TableSchema)}, {@link #withJsonSchema(String)},
* and {@link #withSchemaFromView(PCollectionView)}.
*/
public Write<T> withCloneFrom(String tableSpec) {
checkArgument(tableSpec != null, "tableSpec can not be null");
return withCloneFrom(StaticValueProvider.of(tableSpec));
}

/** Same as {@link #withCloneFrom(String)} but using a {@link TableReference}. */
public Write<T> withCloneFrom(TableReference table) {
checkArgument(table != null, "table can not be null");
return withJsonCloneSourceTableRef(
StaticValueProvider.of(BigQueryHelpers.toJsonString(table)));
}

/** Same as {@link #withCloneFrom(String)} but using a deferred {@link ValueProvider}. */
public Write<T> withCloneFrom(ValueProvider<String> tableSpec) {
checkArgument(tableSpec != null, "tableSpec can not be null");
return withJsonCloneSourceTableRef(
NestedValueProvider.of(
NestedValueProvider.of(tableSpec, new TableSpecToTableRef()), new TableRefToJson()));
}

private Write<T> withJsonCloneSourceTableRef(ValueProvider<String> jsonCloneSourceTableRef) {
checkArgument(jsonCloneSourceTableRef != null, "jsonCloneSourceTableRef can not be null");
return toBuilder().setJsonCloneSourceTableRef(jsonCloneSourceTableRef).build();
}

/**
* Allows the schemas for each table to be computed within the pipeline itself.
*
Expand Down Expand Up @@ -3771,6 +3809,15 @@ public WriteResult expand(PCollection<T> input) {
.collect(Collectors.toList())),
"No more than one of jsonSchema, schemaFromView, or dynamicDestinations may be set");

long activeCreationMetadataCount =
java.util.stream.Stream.of(
getJsonSchema(), getSchemaFromView(), getJsonCloneSourceTableRef())
.filter(arg -> arg != null)
.count();
checkArgument(
activeCreationMetadataCount <= 1,
"No more than one of jsonSchema, schemaFromView, or cloneSource may be set");

// Perform some argument checks
BigQueryOptions bqOptions = input.getPipeline().getOptions().as(BigQueryOptions.class);
Write.Method method = resolveMethod(input);
Expand Down Expand Up @@ -3940,6 +3987,11 @@ && getStorageApiTriggeringFrequency(bqOptions) != null) {
new TableConstraints.PrimaryKey().setColumns(getPrimaryKey())));
}
}
if (getJsonCloneSourceTableRef() != null) {
dynamicDestinations =
new ConstantCloneSourceDestinations<>(
(DynamicDestinations<T, Object>) dynamicDestinations, getJsonCloneSourceTableRef());
}
return expandTyped(input, dynamicDestinations);
}

Expand All @@ -3955,6 +4007,7 @@ private <DestinationT> WriteResult expandTyped(
getJsonSchema() != null
|| getDynamicDestinations() != null
|| getSchemaFromView() != null;
boolean hasTableCreationMetadata = hasSchema || getJsonCloneSourceTableRef() != null;

Class<T> writeProtoClass = getWriteProtosClass();
if (getUseBeamSchema()) {
Expand Down Expand Up @@ -3983,7 +4036,7 @@ private <DestinationT> WriteResult expandTyped(
dynamicDestinations,
StaticValueProvider.of(BigQueryHelpers.toJsonString(tableSchema)));
} else if (writeProtoClass != null) {
if (!hasSchema) {
if (!hasTableCreationMetadata) {
try {
@SuppressWarnings({"unchecked", "nullness"})
Descriptors.Descriptor descriptor =
Expand All @@ -4005,8 +4058,9 @@ private <DestinationT> WriteResult expandTyped(
} else {
// Require a schema if creating one or more tables.
checkArgument(
getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED || hasSchema,
"CreateDisposition is CREATE_IF_NEEDED, however no schema was provided.");
getCreateDisposition() != CreateDisposition.CREATE_IF_NEEDED
|| hasTableCreationMetadata,
"CreateDisposition is CREATE_IF_NEEDED, however no schema was provided and no clone source was provided.");
}

Coder<DestinationT> destinationCoder;
Expand Down
10 changes: 10 additions & 0 deletions ...oud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryIOTranslation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,6 +396,7 @@ static class BigQueryIOWriteTranslator implements TransformPayloadTranslator<Wri
.addNullableBooleanField("use_avro_logical_types")
.addNullableByteArrayField("dynamic_destinations")
.addNullableStringField("json_schema")
.addNullableStringField("json_clone_source_table_ref")
.addNullableStringField("json_time_partitioning")
.addNullableStringField("clustering")
.addNullableByteArrayField("create_disposition")
Expand Down Expand Up @@ -496,6 +497,10 @@ public Row toConfigRow(Write<?> transform) {
if (transform.getJsonSchema() != null) {
fieldValues.put("json_schema", transform.getJsonSchema().get());
}
if (transform.getJsonCloneSourceTableRef() != null) {
fieldValues.put(
"json_clone_source_table_ref", transform.getJsonCloneSourceTableRef().get());
}
if (transform.getJsonTimePartitioning() != null) {
fieldValues.put(
"json_time_partitioning", toByteArray(transform.getJsonTimePartitioning().get()));
Expand Down Expand Up @@ -668,6 +673,11 @@ public Write<?> fromConfigRow(Row configRow, PipelineOptions options) {
if (jsonSchema != null) {
builder = builder.setJsonSchema(StaticValueProvider.of(jsonSchema));
}
String jsonCloneSourceTableRef = configRow.getString("json_clone_source_table_ref");
if (jsonCloneSourceTableRef != null) {
builder =
builder.setJsonCloneSourceTableRef(StaticValueProvider.of(jsonCloneSourceTableRef));
}
String jsonTimePartitioning = configRow.getString("json_time_partitioning");
if (jsonTimePartitioning != null) {
builder = builder.setJsonTimePartitioning(StaticValueProvider.of(jsonTimePartitioning));
Expand Down
1 change: 1 addition & 0 deletions ...ud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryResourceNaming.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ public enum JobType {
LOAD,
TEMP_TABLE_LOAD,
COPY,
CLONE,
EXPORT,
QUERY,
SCHEMA_UPDATE,
Expand Down
Loading
Loading