Skip to content

Commit eef26ca

Browse files
authored
Revert "[Python] Optimize BigQuery copy jobs in file loads using multi-source copy" (#39106)
* Revert "[Python] Optimize BigQuery copy jobs in file loads using multi-source copy (#38983)" This reverts commit 8e4ea73. * Trigger postcommit test.
1 parent 5eefc93 commit eef26ca

4 files changed

Lines changed: 136 additions & 251 deletions

File tree

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
{
22
"comment": "Modify this file in a trivial way to cause this test suite to run.",
33
"pr": "37345",
4-
"modification": 50
4+
"modification": 51
55
}

sdks/python/apache_beam/io/gcp/bigquery_file_loads.py

Lines changed: 82 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -491,8 +491,6 @@ class TriggerCopyJobs(beam.DoFn):
491491
"""
492492

493493
TRIGGER_DELETE_TEMP_TABLES = 'TriggerDeleteTempTables'
494-
# https://docs.cloud.google.com/bigquery/quotas#copy_jobs
495-
MAX_SOURCES_PER_COPY_JOB = 1200
496494

497495
def __init__(
498496
self,
@@ -530,90 +528,96 @@ def process(
530528
self, element_list, job_name_prefix=None, unused_schema_mod_jobs=None):
531529
if isinstance(element_list, tuple):
532530
# Allow this for streaming update compatibility while fixing BEAM-24535.
533-
element_list = [element_list]
531+
self.process_one(element_list, job_name_prefix)
532+
else:
533+
for element in element_list:
534+
self.process_one(element, job_name_prefix)
534535

535-
if not element_list:
536-
return
536+
def process_one(self, element, job_name_prefix):
537+
destination, job_reference = element
537538

538-
first_destination = element_list[0][0]
539-
copy_to_reference = bigquery_tools.parse_table_reference(first_destination)
539+
copy_to_reference = bigquery_tools.parse_table_reference(destination)
540540
if copy_to_reference.projectId is None:
541541
copy_to_reference.projectId = vp.RuntimeValueProvider.get_value(
542542
'project', str, '') or self.project
543543

544-
copy_from_references = []
545-
for destination, job_reference in element_list:
546-
copy_from_reference = bigquery_tools.parse_table_reference(destination)
547-
copy_from_reference.tableId = job_reference.jobId
548-
if copy_from_reference.projectId is None:
549-
copy_from_reference.projectId = vp.RuntimeValueProvider.get_value(
550-
'project', str, '') or self.project
551-
copy_from_references.append(copy_from_reference)
544+
copy_from_reference = bigquery_tools.parse_table_reference(destination)
545+
copy_from_reference.tableId = job_reference.jobId
546+
if copy_from_reference.projectId is None:
547+
copy_from_reference.projectId = vp.RuntimeValueProvider.get_value(
548+
'project', str, '') or self.project
552549

553-
full_table_ref = bigquery_tools.get_hashable_destination(copy_to_reference)
550+
_LOGGER.info(
551+
"Triggering copy job from %s to %s",
552+
copy_from_reference,
553+
copy_to_reference)
554554

555-
is_first_time = full_table_ref not in self._observed_tables
556-
if is_first_time:
557-
self._observed_tables.add(full_table_ref)
558-
if self.bq_io_metadata:
559-
Lineage.sinks().add(
560-
'bigquery',
561-
copy_to_reference.projectId,
562-
copy_to_reference.datasetId,
563-
copy_to_reference.tableId)
564-
565-
# Split into chunks of MAX_SOURCES_PER_COPY_JOB
566-
chunks = [
567-
copy_from_references[i:i + self.MAX_SOURCES_PER_COPY_JOB]
568-
for i in range(
569-
0, len(copy_from_references), self.MAX_SOURCES_PER_COPY_JOB)
570-
]
571-
572-
copy_job_name_base = '%s_%s' % (
573-
job_name_prefix,
574-
_bq_uuid(bigquery_tools.get_hashable_destination(copy_to_reference)))
555+
wait_for_job, write_disposition = (
556+
self._determine_write_disposition(copy_to_reference))
557+
558+
if not self.bq_io_metadata:
559+
self.bq_io_metadata = create_bigquery_io_metadata(self._step_name)
575560

576561
project_id = (
577562
copy_to_reference.projectId
578563
if self.load_job_project_id is None else self.load_job_project_id)
564+
copy_job_name = '%s_%s' % (
565+
job_name_prefix,
566+
_bq_uuid(
567+
'%s:%s.%s' % (
568+
copy_from_reference.projectId,
569+
copy_from_reference.datasetId,
570+
copy_from_reference.tableId)))
571+
job_reference = self.bq_wrapper._insert_copy_job(
572+
project_id,
573+
copy_job_name,
574+
copy_from_reference,
575+
copy_to_reference,
576+
create_disposition=self.create_disposition,
577+
write_disposition=write_disposition,
578+
job_labels=self.bq_io_metadata.add_additional_bq_job_labels())
579+
580+
if wait_for_job:
581+
self.bq_wrapper.wait_for_bq_job(job_reference, sleep_duration_sec=10)
582+
self.pending_jobs.append(
583+
GlobalWindows.windowed_value((destination, job_reference)))
579584

580-
for i, chunk in enumerate(chunks):
581-
if i == 0 and is_first_time:
582-
write_disposition = self.write_disposition
583-
# Wait inline only if we have multiple chunks and write disposition is WRITE_TRUNCATE or WRITE_EMPTY.
584-
# This ensures the first chunk initializes the table, and subsequent chunks (WRITE_APPEND) append to it.
585-
wait_for_job = (
586-
self.write_disposition in ('WRITE_TRUNCATE', 'WRITE_EMPTY') and
587-
len(chunks) > 1)
588-
else:
589-
write_disposition = 'WRITE_APPEND'
590-
wait_for_job = False
591-
592-
chunk_job_name = copy_job_name_base
593-
if len(chunks) > 1:
594-
chunk_job_name = f"{copy_job_name_base}_{i}"
595-
596-
_LOGGER.info(
597-
"Triggering copy job %s from %s to %s (write_disposition: %s)",
598-
chunk_job_name, [str(r) for r in chunk],
599-
copy_to_reference,
600-
write_disposition)
601-
602-
job_reference = self.bq_wrapper._insert_copy_job(
603-
project_id,
604-
chunk_job_name,
605-
chunk,
606-
copy_to_reference,
607-
create_disposition=self.create_disposition,
608-
write_disposition=write_disposition,
609-
job_labels=self.bq_io_metadata.add_additional_bq_job_labels()
610-
if self.bq_io_metadata else None)
611-
612-
if wait_for_job:
613-
self.bq_wrapper.wait_for_bq_job(job_reference, sleep_duration_sec=10)
614-
615-
self.pending_jobs.append(
616-
GlobalWindows.windowed_value((first_destination, job_reference)))
585+
def _determine_write_disposition(self, copy_to_reference) -> tuple[bool, str]:
586+
"""
587+
Determines the write disposition for a BigQuery copy job,
588+
based on destination.
589+
590+
When the write_disposition for a job is WRITE_TRUNCATE, multiple copy jobs
591+
to the same destination can interfere with each other, truncate data, and
592+
write to the BigQuery table repeatedly. To prevent this, the first copy job
593+
runs with the user's specified write_disposition, but subsequent jobs must
594+
always use WRITE_APPEND. This ensures that subsequent copy jobs do not
595+
clear out data appended by previous jobs.
596+
597+
Args:
598+
copy_to_reference: The reference to the destination table.
599+
600+
Returns:
601+
A tuple containing a boolean indicating whether to wait for the job to
602+
complete and the write disposition to use for the job.
603+
"""
604+
full_table_ref = '%s:%s.%s' % (
605+
copy_to_reference.projectId,
606+
copy_to_reference.datasetId,
607+
copy_to_reference.tableId)
608+
if full_table_ref not in self._observed_tables:
609+
write_disposition = self.write_disposition
610+
wait_for_job = True
611+
self._observed_tables.add(full_table_ref)
612+
Lineage.sinks().add(
613+
'bigquery',
614+
copy_to_reference.projectId,
615+
copy_to_reference.datasetId,
616+
copy_to_reference.tableId)
617+
else:
618+
wait_for_job = False
619+
write_disposition = 'WRITE_APPEND'
620+
return wait_for_job, write_disposition
617621

618622
def finish_bundle(self):
619623
for windowed_value in self.pending_jobs:
@@ -740,7 +744,7 @@ def process(
740744
else:
741745
try:
742746
schema = bigquery_tools.table_schema_to_dict(
743-
self.bq_wrapper.get_table(
747+
bigquery_tools.BigQueryWrapper().get_table(
744748
project_id=table_reference.projectId,
745749
dataset_id=table_reference.datasetId,
746750
table_id=table_reference.tableId).schema)
@@ -851,8 +855,7 @@ def process(self, element):
851855
if latest_partition.can_accept(file_size):
852856
latest_partition.add(file_path, file_size)
853857
else:
854-
if latest_partition.files:
855-
partitions.append(latest_partition.files)
858+
partitions.append(latest_partition.files)
856859
latest_partition = PartitionFiles.Partition(
857860
self.max_partition_size, self.max_files_per_partition)
858861
latest_partition.add(file_path, file_size)
@@ -1178,13 +1181,12 @@ def _load_data(
11781181
# the truncation happens only once. See
11791182
# https://github.com/apache/beam/issues/24535.
11801183
finished_temp_tables_load_job_ids_list_pc = (
1181-
finished_temp_tables_load_job_ids_pc
1182-
| beam.MapTuple(
1184+
finished_temp_tables_load_job_ids_pc | beam.MapTuple(
11831185
lambda destination, job_reference: (
1184-
bigquery_tools.get_hashable_destination(destination),
1186+
bigquery_tools.parse_table_reference(destination).tableId,
11851187
(destination, job_reference)))
11861188
| beam.GroupByKey()
1187-
| beam.MapTuple(lambda dest, batch: list(batch)))
1189+
| beam.MapTuple(lambda tableId, batch: list(batch)))
11881190
else:
11891191
# Loads can happen in parallel.
11901192
finished_temp_tables_load_job_ids_list_pc = (

0 commit comments

Comments
 (0)