Optimize Athena partition deletions for insert_overwrite strategy (#1558)

juhoautio-rovio · colin-k-rogers · Copilot · web-flow · commit afd64a3caf3b · 2026-04-25T18:59:48.000Z
Co-authored-by: Colin Rogers &lt;111200756+colin-rogers-dbt@users.noreply.github.com&gt;
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/dbt-athena/.changes/unreleased/Under the Hood-20260223-225659.yaml b/dbt-athena/.changes/unreleased/Under the Hood-20260223-225659.yaml
@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Optimize Athena partition deletions for insert_overwrite strategy
+time: 2026-02-23T22:56:59.91236+02:00
+custom:
+    Author: juhoautio-rovio
+    Issue: "1125"
diff --git a/dbt-athena/src/dbt/adapters/athena/impl.py b/dbt-athena/src/dbt/adapters/athena/impl.py
@@ -9,7 +9,20 @@
 from functools import lru_cache
 from textwrap import dedent
 from threading import Lock
-from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    FrozenSet,
+    Generator,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 from urllib.parse import urlparse
 from uuid import uuid4
 
@@ -54,6 +67,7 @@
 from dbt.adapters.athena.s3 import S3DataNaming
 from dbt.adapters.athena.utils import (
     AthenaCatalogType,
+    chunk_iterable,
     clean_sql_comment,
     ellipsis_comment,
     get_catalog_id,
@@ -132,6 +146,9 @@ class AthenaConfig(AdapterConfig):
 class AthenaAdapter(SQLAdapter):
     BATCH_CREATE_PARTITION_API_LIMIT = 100
     BATCH_DELETE_PARTITION_API_LIMIT = 25
+    BATCH_DELETE_S3_OBJECTS_API_LIMIT = 1000
+    PARTITION_PROCESSING_CHUNK_SIZE = 1000
+    GET_PARTITIONS_API_EXPRESSION_MAX_LENGTH = 2048
     INTEGER_MAX_VALUE_32_BIT_SIGNED = 0x7FFFFFFF
 
     ConnectionManager = AthenaConnectionManager
@@ -401,7 +418,9 @@ def get_glue_table_location(self, relation: AthenaRelation) -> Optional[str]:
         return None
 
     @available
-    def clean_up_partitions(self, relation: AthenaRelation, where_condition: str) -> None:
+    def clean_up_partitions(
+        self, relation: AthenaRelation, where_condition: Union[str, List[str]]
+    ) -> None:
         conn = self.connections.get_thread_connection()
         creds = conn.credentials
         client = conn.handle
@@ -415,24 +434,79 @@ def clean_up_partitions(self, relation: AthenaRelation, where_condition: str) ->
                 region_name=client.region_name,
                 config=get_boto3_config(num_retries=creds.effective_num_retries),
             )
-        paginator = glue_client.get_paginator("get_partitions")
-        partition_params = {
-            "CatalogId": catalog_id,
-            "DatabaseName": relation.schema,
-            "TableName": relation.identifier,
-            "Expression": where_condition,
-            "ExcludeColumnSchema": True,
-        }
-        partition_pg = paginator.paginate(**partition_params)
-        partitions = partition_pg.build_full_result().get("Partitions")
-        for partition in partitions:
-            self.delete_from_s3(partition["StorageDescriptor"]["Location"])
-            glue_client.delete_partition(
-                CatalogId=catalog_id,
-                DatabaseName=relation.schema,
-                TableName=relation.identifier,
-                PartitionValues=partition["Values"],
-            )
+
+        where_conditions = (
+            [where_condition] if isinstance(where_condition, str) else where_condition
+        )
+
+        def join_or_conditions(conditions: List[str]) -> str:
+            return " or ".join(conditions)
+
+        def get_partition_expressions() -> Generator[List[str], None, None]:
+            current_chunk: List[str] = []
+            for condition in where_conditions:
+                condition_with_brackets = f"({condition})"
+                if (
+                    len(condition_with_brackets)
+                    > AthenaAdapter.GET_PARTITIONS_API_EXPRESSION_MAX_LENGTH
+                ):
+                    raise DbtRuntimeError(
+                        f"Partition condition exceeds the Glue API expression limit of "
+                        f"{AthenaAdapter.GET_PARTITIONS_API_EXPRESSION_MAX_LENGTH} characters: "
+                        f"'{condition_with_brackets[:100]}...'"
+                    )
+                if current_chunk and (
+                    len(join_or_conditions(current_chunk + [condition_with_brackets]))
+                    > AthenaAdapter.GET_PARTITIONS_API_EXPRESSION_MAX_LENGTH
+                ):
+                    yield current_chunk
+                    current_chunk = []
+                current_chunk.append(condition_with_brackets)
+            if current_chunk:
+                yield current_chunk
+
+        def iter_partitions() -> Generator[Dict[str, Any], None, None]:
+            paginator = glue_client.get_paginator("get_partitions")
+            for expression_chunk in get_partition_expressions():
+                expression = join_or_conditions(expression_chunk)
+                partition_params = {
+                    "CatalogId": catalog_id,
+                    "DatabaseName": relation.schema,
+                    "TableName": relation.identifier,
+                    "Expression": expression,
+                    "ExcludeColumnSchema": True,
+                }
+                for page in paginator.paginate(**partition_params):
+                    yield from page.get("Partitions", [])
+
+        def delete_partition_chunk(partitions):
+            self.bulk_delete_from_s3([p["StorageDescriptor"]["Location"] for p in partitions])
+
+            for glue_batch in get_chunks(
+                partitions, AthenaAdapter.BATCH_DELETE_PARTITION_API_LIMIT
+            ):
+                response = glue_client.batch_delete_partition(
+                    CatalogId=catalog_id,
+                    DatabaseName=relation.schema,
+                    TableName=relation.identifier,
+                    PartitionsToDelete=[{"Values": p["Values"]} for p in glue_batch],
+                )
+                if errors := response.get("Errors"):
+                    for err in errors:
+                        LOGGER.error(
+                            f"Failed to delete Glue partition: Values='{err['PartitionValues']}', "
+                            f"Code='{err['ErrorDetail']['ErrorCode']}', "
+                            f"Message='{err['ErrorDetail']['ErrorMessage']}'"
+                        )
+                    raise DbtRuntimeError(
+                        f"Failed to delete {len(errors)} partition(s) from Glue table "
+                        f"'{relation.schema}.{relation.identifier}'"
+                    )
+
+        for partition_params_chunk in chunk_iterable(
+            iter_partitions(), AthenaAdapter.PARTITION_PROCESSING_CHUNK_SIZE
+        ):
+            delete_partition_chunk(partition_params_chunk)
 
     @available
     def clean_up_table(self, relation: AthenaRelation) -> None:
@@ -538,6 +612,68 @@ def delete_from_s3(self, s3_path: str) -> None:
         else:
             LOGGER.debug("S3 path does not exist")
 
+    def bulk_delete_from_s3(self, s3_paths: List[str]) -> None:
+        if not s3_paths:
+            LOGGER.debug("No S3 paths provided for deletion")
+            return
+
+        conn = self.connections.get_thread_connection()
+        creds = conn.credentials
+        client = conn.handle
+        s3_resource = client.session.resource(
+            "s3",
+            region_name=client.region_name,
+            config=get_boto3_config(num_retries=creds.effective_num_retries),
+        )
+
+        # Group paths by bucket to support partitions spread across multiple buckets
+        paths_by_bucket: Dict[str, List[str]] = {}
+        for s3_path in s3_paths:
+            bucket, _ = self._parse_s3_path(s3_path)
+            paths_by_bucket.setdefault(bucket, []).append(s3_path)
+
+        def filter_objects_by_prefixes(
+            s3_bucket: Any, bucket_paths: List[str]
+        ) -> Generator[Any, None, None]:
+            for s3_path in bucket_paths:
+                LOGGER.debug(f"Listing files for deletion: {s3_path}")
+                _, prefix = self._parse_s3_path(s3_path)
+                yield from s3_bucket.objects.filter(Prefix=prefix)
+
+        def chunk_object_keys(objects_iter) -> Generator[List[Dict[str, str]], None, None]:
+            chunk = []
+            for obj in objects_iter:
+                chunk.append({"Key": obj.key})
+                if len(chunk) >= AthenaAdapter.BATCH_DELETE_S3_OBJECTS_API_LIMIT:
+                    yield chunk
+                    chunk = []
+            if chunk:
+                yield chunk
+
+        for bucket_name, bucket_paths in paths_by_bucket.items():
+            s3_bucket = s3_resource.Bucket(bucket_name)
+            for object_keys in chunk_object_keys(
+                filter_objects_by_prefixes(s3_bucket, bucket_paths)
+            ):
+                if object_keys:
+                    LOGGER.debug(f"Calling delete_objects for {len(object_keys)} objects")
+                    response = s3_bucket.delete_objects(Delete={"Objects": object_keys})
+                    deleted_count = len(response.get("Deleted", []))
+                    error_count = len(response.get("Errors", []))
+                    LOGGER.debug(
+                        f"delete_objects result: {deleted_count} deleted, {error_count} errors"
+                    )
+                    if errors := response.get("Errors"):
+                        for err in errors:
+                            LOGGER.error(
+                                f"Failed to delete S3 object: Key='{err['Key']}', "
+                                f"Code='{err['Code']}', Message='{err['Message']}', "
+                                f"Bucket='{bucket_name}'"
+                            )
+                        raise DbtRuntimeError(
+                            f"Failed to delete {len(errors)} object(s) from S3 bucket '{bucket_name}'"
+                        )
+
     @staticmethod
     def _parse_s3_path(s3_path: str) -> Tuple[str, str]:
         """
diff --git a/dbt-athena/src/dbt/adapters/athena/utils.py b/dbt-athena/src/dbt/adapters/athena/utils.py
@@ -1,7 +1,7 @@
 import json
 import re
 from enum import Enum
-from typing import Any, Generator, List, Optional, TypeVar
+from typing import Any, Generator, Iterable, List, Optional, TypeVar
 
 from mypy_boto3_athena.type_defs import DataCatalogTypeDef
 
@@ -62,6 +62,18 @@ def get_chunks(lst: List[T], n: int) -> Generator[List[T], None, None]:
         yield lst[i : i + n]
 
 
+def chunk_iterable(iterable: Iterable[T], n: int) -> Generator[List[T], None, None]:
+    """Yield successive n-sized chunks from any iterable, including generators."""
+    chunk = []
+    for item in iterable:
+        chunk.append(item)
+        if len(chunk) >= n:
+            yield chunk
+            chunk = []
+    if chunk:
+        yield chunk
+
+
 def ellipsis_comment(s: str, max_len: int = 255) -> str:
     """Ellipsis string if it exceeds max length"""
     return f"{s[:(max_len - 3)]}..." if len(s) > max_len else s
diff --git a/dbt-athena/src/dbt/include/athena/macros/materializations/models/incremental/helpers.sql b/dbt-athena/src/dbt/include/athena/macros/materializations/models/incremental/helpers.sql
@@ -111,11 +111,9 @@
       {%- do single_partition.append(partitioned_by[loop.index0] + '=' + value) -%}
     {%- endfor -%}
     {%- set single_partition_expression = single_partition | join(' and ') -%}
-    {%- do partitions.append('(' + single_partition_expression + ')') -%}
-  {%- endfor -%}
-  {%- for i in range(partitions | length) %}
-    {%- do adapter.clean_up_partitions(target_relation, partitions[i]) -%}
+    {%- do partitions.append(single_partition_expression) -%}
   {%- endfor -%}
+  {%- do adapter.clean_up_partitions(target_relation, partitions) -%}
 {%- endmacro %}
 
 {% macro remove_partitions_from_columns(columns_with_partitions, partition_keys) %}
diff --git a/dbt-athena/tests/functional/adapter/test_unique_tmp_table_suffix.py b/dbt-athena/tests/functional/adapter/test_unique_tmp_table_suffix.py
@@ -6,6 +6,7 @@
     extract_running_create_statements,
 )
 
+from dbt.adapters.athena.impl import AthenaAdapter
 from dbt.contracts.results import RunStatus
 from dbt.tests.util import run_dbt
 
@@ -19,7 +20,17 @@
 }}
 select
     random() as rnd,
-    cast(from_iso8601_date('{{ var('logical_date') }}') as date) as date_column
+    cast(date_column as date) as date_column
+from (
+    values (
+        sequence(
+            from_iso8601_date('{{ var('start_date') }}'),
+            from_iso8601_date('{{ var('end_date') }}'),
+            interval '1' day
+        )
+    )
+) as t1(date_array)
+cross join unnest(date_array) as t2(date_column)
 """
 
 
@@ -28,7 +39,7 @@ class TestUniqueTmpTableSuffix:
     def models(self):
         return {"unique_tmp_table_suffix.sql": models__unique_tmp_table_suffix_sql}
 
-    def test__unique_tmp_table_suffix(self, project, capsys):
+    def test__unique_tmp_table_suffix(self, project, monkeypatch, capsys):
         relation_name = "unique_tmp_table_suffix"
         model_run_result_row_count_query = (
             f"select count(*) as records from {project.test_schema}.{relation_name}"
@@ -44,7 +55,7 @@ def test__unique_tmp_table_suffix(self, project, capsys):
                 "--select",
                 relation_name,
                 "--vars",
-                '{"logical_date": "2024-01-01"}',
+                '{"start_date": "2024-01-01", "end_date": "2024-01-01"}',
                 "--log-level",
                 "debug",
                 "--log-format",
@@ -70,19 +81,15 @@ def test__unique_tmp_table_suffix(self, project, capsys):
             re.search(expected_unique_table_name_re, first_model_run_result_table_name)
         )
 
-        records_count_first_run = project.run_sql(model_run_result_row_count_query, fetch="all")[
-            0
-        ][0]
-
-        assert records_count_first_run == 1
+        assert project.run_sql(model_run_result_row_count_query, fetch="all")[0][0] == 1
 
         incremental_model_run = run_dbt(
             [
                 "run",
                 "--select",
                 relation_name,
                 "--vars",
-                '{"logical_date": "2024-01-02"}',
+                '{"start_date": "2024-01-02", "end_date": "2024-01-02"}',
                 "--log-level",
                 "debug",
                 "--log-format",
@@ -110,13 +117,18 @@ def test__unique_tmp_table_suffix(self, project, capsys):
 
         assert first_model_run_result_table_name != incremental_model_run_result_table_name
 
+        # Write 4 partitions with a monkeypatched expression limit to force chunking.
+        # Each "(date_column='2024-01-0X')" is ~27 chars, so 2 fit per chunk (58 < 60),
+        # requiring 2 Glue GetPartitions API calls for 4 partitions.
+        monkeypatch.setattr(AthenaAdapter, "GET_PARTITIONS_API_EXPRESSION_MAX_LENGTH", 60)
+
         incremental_model_run_2 = run_dbt(
             [
                 "run",
                 "--select",
                 relation_name,
                 "--vars",
-                '{"logical_date": "2024-01-03"}',
+                '{"start_date": "2024-01-01", "end_date": "2024-01-04"}',
                 "--log-level",
                 "debug",
                 "--log-format",
@@ -136,5 +148,5 @@ def test__unique_tmp_table_suffix(self, project, capsys):
         )[0]
 
         assert incremental_model_run_result_table_name != incremental_model_run_result_table_name_2
-
         assert first_model_run_result_table_name != incremental_model_run_result_table_name_2
+        assert project.run_sql(model_run_result_row_count_query, fetch="all")[0][0] == 4
diff --git a/dbt-athena/tests/unit/test_adapter.py b/dbt-athena/tests/unit/test_adapter.py
