recce-nightly 0.62.0.20250417__py3-none-any.whl → 1.30.0.20251221__py3-none-any.whl

recce/tasks/profile.py

@@ -1,13 +1,141 @@
-import textwrap
 from typing import List
 
 from pydantic import BaseModel
 
-from .core import Task, TaskResultDiffer, CheckValidator
-from .dataframe import DataFrame
 from ..core import default_context
 from ..exceptions import RecceException
 from ..models import Check
+from .core import CheckValidator, Task, TaskResultDiffer
+from .dataframe import DataFrame
+
+PROFILE_COLUMN_JINJA_TEMPLATE = r"""
+{# Conditions -------------------------------------------- #}
+{%- set is_struct = column_type.startswith('struct') -%}
+{%- set is_numeric =
+    column_type.startswith('int') or
+    column_type.startswith('float') or
+    'numeric' in column_type or
+    'number' in column_type or
+    'double' in column_type or
+    'bigint' in column_type
+-%}
+{%- set is_date_or_time =
+    column_type.startswith('date') or
+    column_type.startswith('timestamp')
+-%}
+{%- set is_logical = column_type.startswith('bool') -%}
+
+{%- if db_type == 'sqlserver' -%}
+    {%- set is_numeric = column_type in [
+        "bigint", "numeric", "smallint", "decimal", "int",
+        "tinyint", "money", "float", "real"
+    ]-%}
+{%- elif db_type == 'athena' -%}
+    {%- set is_numeric =
+        "int" in column_type or
+        "float" in column_type or
+        "decimal" in column_type or
+        "double" in column_type
+    -%}
+{%- endif -%}
+
+{# General Agg ------------------------------------------- #}
+{%- set agg_row_count = 'cast(count(*) as ' ~ dbt.type_bigint() ~ ')' -%}
+{%- set agg_not_null_proportion =
+        'sum(case when ' ~ adapter.quote(column_name) ~ ' is null '
+        ~ 'then 0 '
+        ~ 'else 1 end) / '
+        ~ 'cast(count(*) as ' ~ dbt.type_numeric() ~ ')'
+-%}
+{%- set agg_distinct_proportion =
+        'count(distinct ' ~ adapter.quote(column_name) ~') / '
+        ~ 'cast(count(*) as ' ~ dbt.type_numeric() ~ ')'
+-%}
+{%- set agg_distinct_count = 'count(distinct ' ~ adapter.quote(column_name) ~ ')' -%}
+{%- set agg_is_unique =      'count(distinct ' ~ adapter.quote(column_name) ~ ') = count(*)' -%}
+{%- set agg_min =            'cast(null as ' ~ dbt.type_string() ~ ')' -%}
+{%- set agg_max =            'cast(null as ' ~ dbt.type_string() ~ ')' -%}
+{%- set agg_avg =            'cast(null as ' ~ dbt.type_numeric() ~ ')' -%}
+{%- set agg_median =         'cast(null as ' ~ dbt.type_numeric() ~ ')' -%}
+
+
+{%- if is_struct -%}
+    {%- set agg_distinct_proportion = 'cast(null as ' ~ dbt.type_numeric() ~ ')' -%}
+    {%- set agg_distinct_count = 'cast(null as ' ~ dbt.type_numeric() ~ ')' -%}
+    {%- set agg_is_unique = 'null' -%}
+{%- endif -%}
+
+
+{%- if (is_numeric or is_date_or_time) and (not is_struct) -%}
+    {%- set agg_min =
+        'cast(min(' ~ adapter.quote(column_name) ~ ') as ' ~ dbt.type_string() ~ ')'
+    -%}
+    {%- set agg_max =
+        'cast(max(' ~ adapter.quote(column_name) ~ ') as ' ~ dbt.type_string() ~ ')'
+    -%}
+{%- endif -%}
+
+
+{%- if is_numeric and not is_struct -%}
+    {%- set agg_avg = 'avg(' ~ adapter.quote(column_name) ~ ')' -%}
+
+    {%- if db_type == 'bigquery' -%}
+        {%- set agg_median = 'approx_quantiles(' ~ adapter.quote(column_name) ~ ', 100)[offset(50)]' -%}
+    {%- elif db_type == 'postgres' -%}
+        {%- set agg_median = 'percentile_cont(0.5) within group (order by ' ~ adapter.quote(column_name) ~ ')' -%}
+    {%- elif db_type == 'redshift' -%}
+        {%- set agg_median =
+            '(select percentile_cont(0.5) within group (order by '
+            ~ adapter.quote(column_name) ~ ') from ' ~ relation ~ ')' -%}
+    {%- elif db_type == 'athena' -%}
+        {%- set agg_median = 'approx_percentile( ' ~ adapter.quote(column_name) ~ ', 0.5)' -%}
+    {%- elif db_type == 'sqlserver' -%}
+        {%- set agg_median = 'percentile_cont(' ~ adapter.quote(column_name) ~ ', 0.5) over ()' -%}
+    {%- else -%}
+        {%- set agg_median = 'median(' ~ adapter.quote(column_name) ~ ')' -%}
+    {%- endif -%}
+{%- elif is_logical -%}
+    {%- set agg_avg = 'avg(case when ' ~ adapter.quote(column_name) ~ ' then 1 else 0 end)' -%}
+{%- endif -%}
+
+
+{# Overwrite Agg ----------------------------------------- #}
+
+{# DRC-663: Support bigquery array type }
+{%- set is_array = column_type.startswith('array') -%}
+{%- if db_type == 'bigquery' and is_array -%}
+    {%- set agg_distinct_proportion = 'cast(null as ' ~ dbt.type_numeric() ~ ')' -%}
+    {%- set agg_distinct_count = 'cast(null as ' ~ dbt.type_numeric() ~ ')' -%}
+    {%- set agg_is_unique = 'null' -%}
+    {%- set agg_min =
+        'cast(min(array_length(' ~ adapter.quote(column_name) ~ ')) as ' ~ dbt.type_string() ~ ')'
+    -%}
+    {%- set agg_max =
+        'cast(max(array_length(' ~ adapter.quote(column_name) ~ ')) as ' ~ dbt.type_string() ~ ')'
+    -%}
+    {%- set agg_avg = 'avg(array_length(' ~ adapter.quote(column_name) ~ '))' -%}
+    {%- set agg_median =
+        'approx_quantiles(array_length(' ~ adapter.quote(column_name) ~ '), 100)[offset(50)]'
+    -%}
+{%- endif -%}
+
+
+{# Main Query -------------------------------------------- #}
+
+select
+    '{{ column_name }}' as column_name,
+    nullif('{{ column_type }}', '') as data_type,
+    {{ agg_row_count }} as row_count,
+    {{ agg_not_null_proportion }} as not_null_proportion,
+    {{ agg_distinct_proportion }} as distinct_proportion,
+    {{ agg_distinct_count }} as distinct_count,
+    {{ agg_is_unique }} as is_unique,
+    {{ agg_min }} as min,
+    {{ agg_max }} as max,
+    {{ agg_avg }} as avg,
+    {{ agg_median }} as median
+from {{ relation }}
+"""
 
 
 class ProfileParams(BaseModel):
@@ -33,14 +161,14 @@ class ProfileDiffTask(Task):
 
     def execute(self):
         import agate
+
         from recce.adapter.dbt_adapter import DbtAdapter, merge_tables
+
         dbt_adapter: DbtAdapter = default_context().adapter
 
         model: str = self.params.model
         selected_columns: List[str] = self.params.columns
 
-        self._verify_dbt_profiler(dbt_adapter)
-
         with dbt_adapter.connection_named("query"):
             self.connection = dbt_adapter.get_thread_connection()
 
@@ -58,7 +186,7 @@ class ProfileDiffTask(Task):
             tables: List[agate.Table] = []
 
             for column in base_columns:
-                self.update_progress(message=f'[Base] Profile column: {column.name}', percentage=completed / total)
+                self.update_progress(message=f"[Base] Profile column: {column.name}", percentage=completed / total)
                 relation = dbt_adapter.create_relation(model, base=True)
                 response, table = self._profile_column(dbt_adapter, relation, column)
                 tables.append(table)
@@ -68,7 +196,7 @@ class ProfileDiffTask(Task):
 
             tables: List[agate.Table] = []
             for column in curr_columns:
-                self.update_progress(message=f'[Current] Profile column: {column.column}', percentage=completed / total)
+                self.update_progress(message=f"[Current] Profile column: {column.column}", percentage=completed / total)
                 relation = dbt_adapter.create_relation(model, base=False)
                 response, table = self._profile_column(dbt_adapter, relation, column)
                 tables.append(table)
@@ -76,87 +204,23 @@ class ProfileDiffTask(Task):
                 self.check_cancel()
             current = DataFrame.from_agate(merge_tables(tables))
 
-            return ProfileDiffResult(base=base, current=current)
+            if len(base.columns) == 0 and len(current.columns) != 0:
+                base.columns = current.columns
+            elif len(base.columns) != 0 and len(current.columns) == 0:
+                current.columns = base.columns
 
-    def _verify_dbt_profiler(self, dbt_adapter):
-        for macro_name, macro in dbt_adapter.manifest.macros.items():
-            if macro.package_name == 'dbt_profiler':
-                break
-        else:
-            raise RecceException(
-                r"Package 'dbt_profiler' not found. Please refer to the link to install: https://hub.getdbt.com/data-mie/dbt_profiler/")
+            return ProfileDiffResult(base=base, current=current)
 
     def _profile_column(self, dbt_adapter, relation, column):
-        sql_template = textwrap.dedent(r"""
-        select
-        '{{column_name}}' as column_name,
-        nullif('{{column_type}}', '') as data_type,
-        {{ dbt_profiler.measure_row_count(column_name, column_type) }} as row_count,
-        {{ dbt_profiler.measure_not_null_proportion(column_name, column_type) }} as not_null_proportion,
-        {{ dbt_profiler.measure_distinct_proportion(column_name, column_type) }} as distinct_proportion,
-        {{ dbt_profiler.measure_distinct_count(column_name, column_type) }} as distinct_count,
-        {{ dbt_profiler.measure_is_unique(column_name, column_type) }} as is_unique,
-        {{ dbt_profiler.measure_min(column_name, column_type) }} as min,
-        {{ dbt_profiler.measure_max(column_name, column_type) }} as max,
-        {{ dbt_profiler.measure_avg(column_name, column_type) }} as avg,
-        {{ dbt_profiler.measure_median(column_name, column_type) }} as median
-        from
-        {{ relation }}
-        """)
         column_name = column.name
         column_type = column.data_type.lower()
-        db_type = dbt_adapter.adapter.type()
-        if db_type == 'bigquery' and column_type.startswith('array'):
-            # DRC-663: Support bigquery array type
-            sql_template = textwrap.dedent(r"""
-            select
-            '{{column_name}}' as column_name,
-            nullif('{{column_type}}', '') as data_type,
-            {{ dbt_profiler.measure_row_count(column_name, column_type) }} as row_count,
-            {{ dbt_profiler.measure_not_null_proportion(column_name, column_type) }} as not_null_proportion,
-            cast(null as {{ dbt.type_numeric() }}) as distinct_proportion,
-            cast(null as {{ dbt.type_numeric() }}) as distinct_count,
-            null as is_unique,
-            cast(min(ARRAY_LENGTH({{ adapter.quote(column_name) }})) as {{ dbt_profiler.type_string() }}) as min,
-            cast(max(ARRAY_LENGTH({{ adapter.quote(column_name) }})) as {{ dbt_profiler.type_string() }}) as max,
-            avg(ARRAY_LENGTH({{ adapter.quote(column_name) }})) as avg,
-            APPROX_QUANTILES(ARRAY_LENGTH({{ adapter.quote(column_name) }}), 100)[OFFSET(50)] as median,
-            from
-            {{ relation }}
-            """)
-        elif db_type == 'redshift':
-            # DRC-1149: Support redshift median calculation
-            # https://github.com/data-mie/dbt-profiler/pull/89
-            #
-            # Since dbt-profiler 0.8.2, there is the third parameter for measure_median
-            # For sake of compatibility, we use the new way to call the macro only for redshift
-            sql_template = textwrap.dedent(r"""
-            with source_data as (
-              select
-                *
-              from {{ relation }}
-            )
-            select
-            '{{column_name}}' as column_name,
-            nullif('{{column_type}}', '') as data_type,
-            {{ dbt_profiler.measure_row_count(column_name, column_type) }} as row_count,
-            {{ dbt_profiler.measure_not_null_proportion(column_name, column_type) }} as not_null_proportion,
-            {{ dbt_profiler.measure_distinct_proportion(column_name, column_type) }} as distinct_proportion,
-            {{ dbt_profiler.measure_distinct_count(column_name, column_type) }} as distinct_count,
-            {{ dbt_profiler.measure_is_unique(column_name, column_type) }} as is_unique,
-            {{ dbt_profiler.measure_min(column_name, column_type) }} as min,
-            {{ dbt_profiler.measure_max(column_name, column_type) }} as max,
-            {{ dbt_profiler.measure_avg(column_name, column_type) }} as avg,
-            ({{ dbt_profiler.measure_median(column_name, column_type, 'source_data') }}) as median
-            from
-            source_data
-            """)
+        db_type = dbt_adapter.adapter.type().lower()
 
         try:
             sql = dbt_adapter.generate_sql(
-                sql_template,
+                PROFILE_COLUMN_JINJA_TEMPLATE,
                 base=False,  # always false because we use the macro in current manifest
-                context=dict(relation=relation, column_name=column_name, column_type=column_type)
+                context=dict(relation=relation, column_name=column_name, column_type=column_type, db_type=db_type),
             )
         except Exception as e:
             raise RecceException(f"Failed to generate SQL for profiling column: {column_name}") from e
@@ -165,14 +229,15 @@ class ProfileDiffTask(Task):
             return dbt_adapter.execute(sql, fetch=True)
         except Exception as e:
             from recce.adapter.dbt_adapter import dbt_version
-            if dbt_version < 'v1.8':
+
+            if dbt_version < "v1.8":
                 from dbt.exceptions import DbtDatabaseError
             else:
                 from dbt_common.exceptions import DbtDatabaseError
             if isinstance(e, DbtDatabaseError):
-                if str(e).find('100051') >= 0:
+                if str(e).find("100051") >= 0:
                     # Snowflake error '100051 (22012): Division by zero"'
-                    e = RecceException('No profile diff result due to the model is empty.', False)
+                    e = RecceException("No profile diff result due to the model is empty.", False)
             raise e
 
     def cancel(self):
@@ -180,6 +245,7 @@ class ProfileDiffTask(Task):
 
         if self.connection:
             from recce.adapter.dbt_adapter import DbtAdapter
+
             dbt_adapter: DbtAdapter = default_context().adapter
             with dbt_adapter.connection_named("cancel"):
                 dbt_adapter.cancel(self.connection)
@@ -187,7 +253,7 @@ class ProfileDiffTask(Task):
 
 class ProfileDiffResultDiffer(TaskResultDiffer):
     def _check_result_changed_fn(self, result):
-        return self.diff(result['base'], result['current'])
+        return self.diff(result["base"], result["current"])
 
 
 class ProfileCheckValidator(CheckValidator):
@@ -202,14 +268,14 @@ class ProfileCheckValidator(CheckValidator):
 class ProfileTask(ProfileDiffTask):
     def execute(self):
         import agate
+
         from recce.adapter.dbt_adapter import DbtAdapter, merge_tables
+
         dbt_adapter: DbtAdapter = default_context().adapter
 
         model: str = self.params.model
         selected_columns: List[str] = self.params.columns
 
-        self._verify_dbt_profiler(dbt_adapter)
-
         with dbt_adapter.connection_named("query"):
             self.connection = dbt_adapter.get_thread_connection()
             curr_columns = [column for column in dbt_adapter.get_columns(model, base=False)]
@@ -222,7 +288,7 @@ class ProfileTask(ProfileDiffTask):
 
             tables: List[agate.Table] = []
             for column in curr_columns:
-                self.update_progress(message=f'[Current] Profile column: {column.column}', percentage=completed / total)
+                self.update_progress(message=f"[Current] Profile column: {column.column}", percentage=completed / total)
                 relation = dbt_adapter.create_relation(model, base=False)
                 response, table = self._profile_column(dbt_adapter, relation, column)
                 tables.append(table)