Skip to content

Commit 5bd01f8

Browse files
committed
feat(athena): add build_with_subquery config to skip tmp table staging
Signed-off-by: Daisuke Taniwaki <daisuketaniwaki@gmail.com>
1 parent 77a0fe3 commit 5bd01f8

8 files changed

Lines changed: 556 additions & 47 deletions

File tree

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Features
2+
body: build_with_subquery config option to skip tmp table staging for incremental merge and append strategies
3+
time: 2026-04-06T22:54:58.739286+09:00
4+
custom:
5+
Author: dtaniwaki
6+
Issue: "1829"

dbt-athena/src/dbt/adapters/athena/impl.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,9 @@ class AthenaConfig(AdapterConfig):
104104
force_batch: Skip creating the table as ctas and run the operation directly in batch insert mode.
105105
unique_tmp_table_suffix: Enforce the use of a unique id as tmp table suffix instead of __dbt_tmp.
106106
temp_schema: Define in which schema to create temporary tables used in incremental runs.
107+
build_with_subquery: Use a subquery directly instead of staging data into __dbt_tmp.
108+
Creates an empty tmp table (WITH NO DATA) for schema comparison, then applies via subquery.
109+
Supported for iceberg merge and append strategies. Incompatible with force_batch.
107110
"""
108111

109112
work_group: Optional[str] = None
@@ -127,6 +130,7 @@ class AthenaConfig(AdapterConfig):
127130
force_batch: bool = False
128131
unique_tmp_table_suffix: bool = False
129132
temp_schema: Optional[str] = None
133+
build_with_subquery: bool = False
130134

131135

132136
class AthenaAdapter(SQLAdapter):

dbt-athena/src/dbt/include/athena/macros/materializations/models/incremental/helpers.sql

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
target_relation,
5454
existing_relation,
5555
force_batch,
56+
source_sql=none,
5657
statement_name="main"
5758
)
5859
%}
@@ -65,18 +66,32 @@
6566
{% if force_batch %}
6667
{% do batch_incremental_insert(tmp_relation, target_relation, dest_cols_csv) %}
6768
{% else %}
68-
{%- set insert_full -%}
69-
insert into {{ target_relation }} ({{ dest_cols_csv }})
70-
(
71-
select {{ dest_cols_csv }}
72-
from {{ tmp_relation }}
73-
);
74-
{%- endset -%}
69+
{%- if source_sql is not none -%}
70+
{%- set insert_full -%}
71+
insert into {{ target_relation }} ({{ dest_cols_csv }})
72+
(
73+
select {{ dest_cols_csv }}
74+
from ({{ source_sql }}) _dbt_sbq
75+
);
76+
{%- endset -%}
77+
{%- else -%}
78+
{%- set insert_full -%}
79+
insert into {{ target_relation }} ({{ dest_cols_csv }})
80+
(
81+
select {{ dest_cols_csv }}
82+
from {{ tmp_relation }}
83+
);
84+
{%- endset -%}
85+
{%- endif -%}
7586

7687
{%- set query_result = adapter.run_query_with_partitions_limit_catching(insert_full) -%}
7788
{%- do log('QUERY RESULT: ' ~ query_result) -%}
7889
{%- if query_result == 'TOO_MANY_OPEN_PARTITIONS' -%}
90+
{%- if source_sql is not none -%}
91+
{% do exceptions.raise_compiler_error('Runtime error: TOO_MANY_OPEN_PARTITIONS encountered with build_with_subquery=True. Disable build_with_subquery to enable automatic batching.') %}
92+
{%- else -%}
7993
{% do batch_incremental_insert(tmp_relation, target_relation, dest_cols_csv) %}
94+
{%- endif -%}
8095
{%- endif -%}
8196
{%- endif -%}
8297

dbt-athena/src/dbt/include/athena/macros/materializations/models/incremental/incremental.sql

Lines changed: 89 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
{% set lf_grants = config.get('lf_grants') %}
1111
{% set partitioned_by = config.get('partitioned_by') %}
1212
{% set force_batch = config.get('force_batch', False) | as_bool -%}
13+
{% set build_with_subquery = config.get('build_with_subquery', False) | as_bool -%}
1314
{% set unique_tmp_table_suffix = config.get('unique_tmp_table_suffix', False) | as_bool -%}
1415
{% set temp_schema = config.get('temp_schema') %}
1516
{% set target_relation = this.incorporate(type='table') %}
@@ -46,6 +47,18 @@
4647
{% endif %}
4748
{% endif %}
4849

50+
{% if build_with_subquery %}
51+
{% if model_language == 'python' %}
52+
{% do exceptions.raise_compiler_error('build_with_subquery is not supported with Python models.') %}
53+
{% endif %}
54+
{% if force_batch %}
55+
{% do exceptions.raise_compiler_error('build_with_subquery is incompatible with force_batch. Batching requires data in the tmp table.') %}
56+
{% endif %}
57+
{% if strategy == 'insert_overwrite' %}
58+
{% do exceptions.raise_compiler_error('build_with_subquery is not supported with insert_overwrite strategy.') %}
59+
{% endif %}
60+
{% endif %}
61+
4962
{{ run_hooks(pre_hooks, inside_transaction=False) }}
5063

5164
-- `BEGIN` happens here:
@@ -146,22 +159,36 @@
146159

147160
-- Append Strategy --
148161
{% elif strategy == 'append' %}
149-
{% if old_tmp_relation is not none %}
150-
{% do drop_relation(old_tmp_relation) %}
162+
{% if build_with_subquery %}
163+
{% if old_tmp_relation is not none %}
164+
{% do drop_relation(old_tmp_relation) %}
165+
{% endif %}
166+
{%- set empty_sql = 'SELECT * FROM (' ~ compiled_code ~ ') _dbt_sbq WITH NO DATA' -%}
167+
{% do run_query(create_table_as(True, tmp_relation, empty_sql)) %}
168+
169+
{% set build_sql = incremental_insert(
170+
on_schema_change, tmp_relation, target_relation, existing_relation, false, source_sql=compiled_code
171+
)
172+
%}
173+
{% do to_drop.append(tmp_relation) %}
174+
{% else %}
175+
{% if old_tmp_relation is not none %}
176+
{% do drop_relation(old_tmp_relation) %}
177+
{% endif %}
178+
{% set query_result = safe_create_table_as(True, tmp_relation, compiled_code, model_language, force_batch) -%}
179+
{%- if model_language == 'python' -%}
180+
{% call statement('create_table', language=model_language) %}
181+
{{ query_result }}
182+
{% endcall %}
183+
{%- endif -%}
184+
{% set build_sql = incremental_insert(
185+
on_schema_change, tmp_relation, target_relation, existing_relation, force_batch
186+
)
187+
%}
188+
{% do to_drop.append(tmp_relation) %}
151189
{% endif %}
152-
{% set query_result = safe_create_table_as(True, tmp_relation, compiled_code, model_language, force_batch) -%}
153-
{%- if model_language == 'python' -%}
154-
{% call statement('create_table', language=model_language) %}
155-
{{ query_result }}
156-
{% endcall %}
157-
{%- endif -%}
158-
{% set build_sql = incremental_insert(
159-
on_schema_change, tmp_relation, target_relation, existing_relation, force_batch
160-
)
161-
%}
162-
{% do to_drop.append(tmp_relation) %}
163190

164-
-- Iceberge Merge Stategy --
191+
-- Iceberg Merge Strategy --
165192
{% elif strategy == 'merge' and table_type == 'iceberg' %}
166193
{% set unique_key = config.get('unique_key') %}
167194
{% set incremental_predicates = config.get('incremental_predicates') %}
@@ -186,29 +213,55 @@
186213
{% do exceptions.raise_compiler_error(inc_predicates_not_list) %}
187214
{% endif %}
188215
{% endif %}
189-
{% if old_tmp_relation is not none %}
190-
{% do drop_relation(old_tmp_relation) %}
216+
217+
{% if build_with_subquery %}
218+
-- Create empty tmp table for schema comparison (no data scan)
219+
{% if old_tmp_relation is not none %}
220+
{% do drop_relation(old_tmp_relation) %}
221+
{% endif %}
222+
{%- set empty_sql = 'SELECT * FROM (' ~ compiled_code ~ ') _dbt_sbq WITH NO DATA' -%}
223+
{% do run_query(create_table_as(True, tmp_relation, empty_sql)) %}
224+
225+
{% set build_sql = iceberg_merge(
226+
on_schema_change=on_schema_change,
227+
tmp_relation=tmp_relation,
228+
target_relation=target_relation,
229+
unique_key=unique_key,
230+
incremental_predicates=incremental_predicates,
231+
existing_relation=existing_relation,
232+
delete_condition=delete_condition,
233+
update_condition=update_condition,
234+
insert_condition=insert_condition,
235+
force_batch=false,
236+
source_sql=compiled_code,
237+
)
238+
%}
239+
{% do to_drop.append(tmp_relation) %}
240+
{% else %}
241+
{% if old_tmp_relation is not none %}
242+
{% do drop_relation(old_tmp_relation) %}
243+
{% endif %}
244+
{% set query_result = safe_create_table_as(True, tmp_relation, compiled_code, model_language, force_batch) -%}
245+
{%- if model_language == 'python' -%}
246+
{% call statement('create_table', language=model_language) %}
247+
{{ query_result }}
248+
{% endcall %}
249+
{%- endif -%}
250+
{% set build_sql = iceberg_merge(
251+
on_schema_change=on_schema_change,
252+
tmp_relation=tmp_relation,
253+
target_relation=target_relation,
254+
unique_key=unique_key,
255+
incremental_predicates=incremental_predicates,
256+
existing_relation=existing_relation,
257+
delete_condition=delete_condition,
258+
update_condition=update_condition,
259+
insert_condition=insert_condition,
260+
force_batch=force_batch,
261+
)
262+
%}
263+
{% do to_drop.append(tmp_relation) %}
191264
{% endif %}
192-
{% set query_result = safe_create_table_as(True, tmp_relation, compiled_code, model_language, force_batch) -%}
193-
{%- if model_language == 'python' -%}
194-
{% call statement('create_table', language=model_language) %}
195-
{{ query_result }}
196-
{% endcall %}
197-
{%- endif -%}
198-
{% set build_sql = iceberg_merge(
199-
on_schema_change=on_schema_change,
200-
tmp_relation=tmp_relation,
201-
target_relation=target_relation,
202-
unique_key=unique_key,
203-
incremental_predicates=incremental_predicates,
204-
existing_relation=existing_relation,
205-
delete_condition=delete_condition,
206-
update_condition=update_condition,
207-
insert_condition=insert_condition,
208-
force_batch=force_batch,
209-
)
210-
%}
211-
{% do to_drop.append(tmp_relation) %}
212265
{% endif %}
213266

214267
{% call statement("main", language=model_language) %}

dbt-athena/src/dbt/include/athena/macros/materializations/models/incremental/merge.sql

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@
7777
update_condition,
7878
insert_condition,
7979
force_batch,
80+
source_sql=none,
8081
statement_name="main"
8182
)
8283
%}
@@ -143,9 +144,15 @@
143144
{%- if force_batch -%}
144145
{% do batch_iceberg_merge(tmp_relation, target_relation, merge_part, dest_cols_csv) %}
145146
{%- else -%}
146-
{%- set src_part -%}
147-
merge into {{ target_relation }} as target using {{ tmp_relation }} as src
148-
{%- endset -%}
147+
{%- if source_sql is not none -%}
148+
{%- set src_part -%}
149+
merge into {{ target_relation }} as target using ({{ source_sql }}) as src
150+
{%- endset -%}
151+
{%- else -%}
152+
{%- set src_part -%}
153+
merge into {{ target_relation }} as target using {{ tmp_relation }} as src
154+
{%- endset -%}
155+
{%- endif -%}
149156
{%- set merge_full -%}
150157
{{ src_part }}
151158
{{ merge_part }}
@@ -154,7 +161,11 @@
154161
{%- set query_result = adapter.run_query_with_partitions_limit_catching(merge_full) -%}
155162
{%- do log('QUERY RESULT: ' ~ query_result) -%}
156163
{%- if query_result == 'TOO_MANY_OPEN_PARTITIONS' -%}
157-
{% do batch_iceberg_merge(tmp_relation, target_relation, merge_part, dest_cols_csv) %}
164+
{%- if source_sql is not none -%}
165+
{% do exceptions.raise_compiler_error('Runtime error: TOO_MANY_OPEN_PARTITIONS encountered with build_with_subquery=True. Disable build_with_subquery to enable automatic batching.') %}
166+
{%- else -%}
167+
{% do batch_iceberg_merge(tmp_relation, target_relation, merge_part, dest_cols_csv) %}
168+
{%- endif -%}
158169
{%- endif -%}
159170
{%- endif -%}
160171

0 commit comments

Comments
 (0)