Skip to content
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

feat(incremental): copy multiple tables in parallel (#1237) #1413

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(incremental): copy multiple tables in parallel (#1237) #1413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
6 changes: 6 additions & 0 deletions .changes/unreleased/Features-20241126-000241.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
kind: Features
body: copy tables and partitions in parallel
time: 2024-11-26T00:02:41.54479+01:00
custom:
Author: AxelThevenot
Issue: "1237"
43 changes: 29 additions & 14 deletions dbt/adapters/bigquery/connections.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -402,17 +402,14 @@ def standard_to_legacy(table):
_, iterator = self.raw_execute(sql, use_legacy_sql=True)
return self.get_table_from_response(iterator)

def copy_bq_table(self, source, destination, write_disposition) -> None:
def copy_bq_table(self, source, destination, write_disposition, partition_ids=None) -> None:
conn = self.get_thread_connection()
client: Client = conn.handle

# -------------------------------------------------------------------------------
# BigQuery allows to use copy API using two different formats:
# 1. client.copy_table(source_table_id, destination_table_id)
# where source_table_id = "your-project.source_dataset.source_table"
# 2. client.copy_table(source_table_ids, destination_table_id)
# where source_table_ids = ["your-project.your_dataset.your_table_name", ...]
# Let's use uniform function call and always pass list there
# BigQuery allows to use copy API on the same table in parallel
# so each source (and if partition of each source if given) is copied
# into the destination table in parallel.
# -------------------------------------------------------------------------------
if type(source) is not list:
source = [source]
Expand All @@ -436,14 +433,32 @@ def copy_bq_table(self, source, destination, write_disposition) -> None:
", ".join(source_ref.path for source_ref in source_ref_array),
destination_ref.path,
)

with self.exception_handler(msg):
copy_job = client.copy_table(
source_ref_array,
destination_ref,
job_config=CopyJobConfig(write_disposition=write_disposition),
retry=self._retry.create_reopen_with_deadline(conn),
)
copy_job.result(timeout=self._retry.create_job_execution_timeout(fallback=300))

copy_jobs = []

# Runs all the copy jobs in parallel
for source_ref in source_ref_array:

for partition_id in partition_ids or [None]:
source_ref_partition = (
f"{source_ref}${partition_id}" if partition_id else source_ref
)
destination_ref_partition = (
f"{destination_ref}${partition_id}" if partition_id else destination_ref
)
copy_job = client.copy_table(
source_ref_partition,
destination_ref_partition,
job_config=CopyJobConfig(write_disposition=write_disposition),
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will there be more than one element in source_ref_partition?

If we ever be in the scenario where we have source_ref_array greater than one and write_disposition set to WRITE_TRUNCATE, we'll be overwriting the same data.

retry=self._retry.create_reopen_with_deadline(conn),
)
copy_jobs.append(copy_job)
Comment on lines +441 to +457
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would being explicit here clarify the logic?

if partition_ids:
    copy_jobs = [
        client.copy_table(
            f"{source_ref}${partition_id}",
            f"{destination_ref}${partition_id}",
            job_config=CopyJobConfig(write_disposition=write_disposition),
            retry=self._retry.create_reopen_with_deadline(conn),
        )
        for partition_id in partition_ids
        for source_ref in source_ref_array
    ]
else:
    copy_jobs = [client.copy_table(
        source_ref_array,
        destination_ref,
        job_config=CopyJobConfig(write_disposition=write_disposition),
        retry=self._retry.create_reopen_with_deadline(conn),
    )]


# Waits for the jobs to finish
for copy_job in copy_jobs:
copy_job.result(timeout=self._retry.create_job_execution_timeout(fallback=300))

def write_dataframe_to_table(
self,
Expand Down
4 changes: 2 additions & 2 deletions dbt/adapters/bigquery/impl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -409,7 +409,7 @@ def _agate_to_schema(
return bq_schema

@available.parse(lambda *a, **k: "")
def copy_table(self, source, destination, materialization):
def copy_table(self, source, destination, materialization, partition_ids=None):
if materialization == "incremental":
write_disposition = WRITE_APPEND
elif materialization == "table":
Expand All @@ -421,7 +421,7 @@ def copy_table(self, source, destination, materialization):
f"{materialization}"
)

self.connections.copy_bq_table(source, destination, write_disposition)
self.connections.copy_bq_table(source, destination, write_disposition, partition_ids)

return "COPY TABLE with materialization: {}".format(materialization)

Expand Down
18 changes: 11 additions & 7 deletions dbt/include/bigquery/macros/materializations/incremental_strategy/insert_overwrite.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,23 +18,27 @@

{% macro bq_copy_partitions(tmp_relation, target_relation, partitions, partition_by) %}

{% set partition_ids = [] %}

{% for partition in partitions %}
{% if partition_by.data_type == 'int64' %}
{% set partition = partition | as_text %}
{% elif partition_by.granularity == 'hour' %}
{% set partition = partition.strftime("%Y%m%d%H") %}
{% set partition = partition.strftime('%Y%m%d%H') %}
{% elif partition_by.granularity == 'day' %}
{% set partition = partition.strftime("%Y%m%d") %}
{% set partition = partition.strftime('%Y%m%d') %}
{% elif partition_by.granularity == 'month' %}
{% set partition = partition.strftime("%Y%m") %}
{% set partition = partition.strftime('%Y%m') %}
{% elif partition_by.granularity == 'year' %}
{% set partition = partition.strftime("%Y") %}
{% set partition = partition.strftime('%Y') %}
{% endif %}
{% set tmp_relation_partitioned = api.Relation.create(database=tmp_relation.database, schema=tmp_relation.schema, identifier=tmp_relation.table ~ '$' ~ partition, type=tmp_relation.type) %}
{% set target_relation_partitioned = api.Relation.create(database=target_relation.database, schema=target_relation.schema, identifier=target_relation.table ~ '$' ~ partition, type=target_relation.type) %}
{% do adapter.copy_table(tmp_relation_partitioned, target_relation_partitioned, "table") %}

{% do partition_ids.append(partition) %}

{% endfor %}

{% do adapter.copy_table(tmp_relation, target_relation, 'table', partition_ids) %}

{% endmacro %}

{% macro bq_insert_overwrite_sql(
Expand Down