@@ -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