Improve the performance of the at_least_one test (#776)

* Update at_least_one.sql * Update CHANGELOG.md * Update at_least_one.sql * Update at_least_one.sql Fix ambiguous col issue * Update at_least_one.sql ensure case match * Change requests * Fix integration test * Update changelog * Restore original test * Remove ambiguous column names --------- Co-authored-by: Doug Beatty <[email protected]>
dbt-labs · Jun 2, 2023 · 5b36260 · 5b36260
1 parent e9d23b5
commit 5b36260
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 8 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,7 +10,7 @@
 
 # Unreleased
 ## New features
-* ZZZ by @YYY in https://github.com/dbt-labs/dbt-utils/pull/XXX
+* Improve the performance of the `at_least_one` test by pruning early. This is especially helpful when running against external tables. By @joshuahuntley in https://github.com/dbt-labs/dbt-utils/pull/775
 ## Fixes
 * Fix legacy links in README by @dbeatty10 in https://github.com/dbt-labs/dbt-utils/pull/796
 ## Quality of life
@@ -48,8 +48,8 @@
 
 # dbt utils v1.0
 
-## Migration Guide 
-The full migration guide is at https://docs.getdbt.com/guides/migration/versions/upgrading-to-dbt-utils-v1.0 
+## Migration Guide
+The full migration guide is at https://docs.getdbt.com/guides/migration/versions/upgrading-to-dbt-utils-v1.0
 
 ## New features
 - New macro `get_single_value` ([#696](https://github.com/dbt-labs/dbt-utils/pull/696))
@@ -63,7 +63,7 @@ The full migration guide is at https://docs.getdbt.com/guides/migration/versions
 ## Fixes
 - `union()` now includes/excludes columns case-insensitively
 - The `expression_is_true test` doesn’t output * unless storing failures, a cost improvement for BigQuery ([#683](https://github.com/dbt-labs/dbt-utils/issues/683), [#686](https://github.com/dbt-labs/dbt-utils/pull/686))
-- Updated the `slugify` macro to prepend "_" to column names beginning with a number since most databases do not allow names to begin with numbers. 
+- Updated the `slugify` macro to prepend "_" to column names beginning with a number since most databases do not allow names to begin with numbers.
 
 ## Under the hood
 - Remove deprecated table argument from `unpivot` ([#671](https://github.com/dbt-labs/dbt-utils/pull/671))

diff --git a/integration_tests/models/generic_tests/schema.yml b/integration_tests/models/generic_tests/schema.yml
@@ -89,14 +89,14 @@ seeds:
           upper_bound_column: valid_to
           partition_by: subscription_id
           zero_length_range_allowed: true
- 
+
   - name: data_unique_combination_of_columns
     tests:
       - dbt_utils.unique_combination_of_columns:
           combination_of_columns:
             - month
             - product
-  
+
   - name: data_cardinality_equality_a
     columns:
       - name: same_name
@@ -191,7 +191,7 @@ models:
             - first_name
             - last_name
             - email
-  
+
   - name: test_fewer_rows_than
     tests:
       - dbt_utils.fewer_rows_than:

diff --git a/macros/generic_tests/at_least_one.sql b/macros/generic_tests/at_least_one.sql
@@ -4,20 +4,38 @@
 
 {% macro default__test_at_least_one(model, column_name, group_by_columns) %}
 
+{% set pruned_cols = [column_name] %}
+
 {% if group_by_columns|length() > 0 %}
+
   {% set select_gb_cols = group_by_columns|join(' ,') + ', ' %}
   {% set groupby_gb_cols = 'group by ' + group_by_columns|join(',') %}
+  {% set pruned_cols = group_by_columns %}
+
+  {% if column_name not in pruned_cols %}
+    {% do pruned_cols.append(column_name) %}
+  {% endif %}
+
 {% endif %}
 
+{% set select_pruned_cols = pruned_cols|join(' ,') %}
+
 select *
 from (
+    with pruned_rows as (
+      select
+        {{ select_pruned_cols }}
+      from {{ model }}
+      where {{ column_name }} is not null
+      limit 1
+    )
     select
         {# In TSQL, subquery aggregate columns need aliases #}
         {# thus: a filler col name, 'filler_column' #}
       {{select_gb_cols}}
       count({{ column_name }}) as filler_column
 
-    from {{ model }}
+    from pruned_rows
 
     {{groupby_gb_cols}}