Add tests for column pruning CTEs.

dbt-labs · Nov 4, 2024 · 37d8c5b · 37d8c5b
1 parent e326ad3
commit 37d8c5b
Show file tree

Hide file tree

Showing 6 changed files with 542 additions and 0 deletions.
diff --git a/...s_metricflow/snapshots/test_cte_column_pruner.py/str/test_multi_child_pruning__result.txt b/...s_metricflow/snapshots/test_cte_column_pruner.py/str/test_multi_child_pruning__result.txt
@@ -0,0 +1,50 @@
+optimizer:
+  SqlColumnPrunerOptimizer
+
+sql_before_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+      , test_table_alias.col_1 AS cte_source_0__col_1
+      , test_table_alias.col_1 AS cte_source_0__col_2
+    FROM test_schema.test_table test_table_alias
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+    , right_source_alias.right_source__col_1 AS top_level__col_1
+  FROM cte_source_0 cte_source_0_alias
+  INNER JOIN (
+    -- Joined sub-query
+    SELECT
+      cte_source_0_alias_in_right_source.cte_source_0__col_0 AS right_source__col_0
+      , cte_source_0_alias_in_right_source.cte_source_0__col_1 AS right_source__col_1
+    FROM cte_source_0 cte_source_0_alias_in_right_source
+  ) right_source_alias
+  ON
+    cte_source_0_alias.cte_source_0__col_1 = right_source_alias.right_source__col_1
+
+sql_after_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+      , test_table_alias.col_1 AS cte_source_0__col_1
+    FROM test_schema.test_table test_table_alias
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+    , right_source_alias.right_source__col_1 AS top_level__col_1
+  FROM cte_source_0 cte_source_0_alias
+  INNER JOIN (
+    -- Joined sub-query
+    SELECT
+      cte_source_0_alias_in_right_source.cte_source_0__col_1 AS right_source__col_1
+    FROM cte_source_0 cte_source_0_alias_in_right_source
+  ) right_source_alias
+  ON
+    cte_source_0_alias.cte_source_0__col_1 = right_source_alias.right_source__col_1
diff --git a/tests_metricflow/snapshots/test_cte_column_pruner.py/str/test_nested_pruning__result.txt b/tests_metricflow/snapshots/test_cte_column_pruner.py/str/test_nested_pruning__result.txt
@@ -0,0 +1,44 @@
+optimizer:
+  SqlColumnPrunerOptimizer
+
+sql_before_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+      , test_table_alias.col_1 AS cte_source_0__col_1
+    FROM test_schema.test_table test_table_alias
+  )
+
+  , cte_source_1 AS (
+    -- CTE source 1
+    SELECT
+      cte_source_0_alias.cte_source_0__col_0 AS cte_source_1__col_0
+      , cte_source_0_alias.cte_source_0__col_0 AS cte_source_1__col_1
+    FROM cte_source_0 cte_source_0_alias
+  )
+
+  SELECT
+    cte_source_1_alias.cte_source_1__col_0 AS top_level__col_0
+  FROM cte_source_1 cte_source_1_alias
+
+sql_after_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+    FROM test_schema.test_table test_table_alias
+  )
+
+  , cte_source_1 AS (
+    -- CTE source 1
+    SELECT
+      cte_source_0_alias.cte_source_0__col_0 AS cte_source_1__col_0
+    FROM cte_source_0 cte_source_0_alias
+  )
+
+  SELECT
+    cte_source_1_alias.cte_source_1__col_0 AS top_level__col_0
+  FROM cte_source_1 cte_source_1_alias
diff --git a/tests_metricflow/snapshots/test_cte_column_pruner.py/str/test_no_pruning__result.txt b/tests_metricflow/snapshots/test_cte_column_pruner.py/str/test_no_pruning__result.txt
@@ -0,0 +1,28 @@
+optimizer:
+  SqlColumnPrunerOptimizer
+
+sql_before_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+    FROM test_schema.test_table test_table_alias
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+  FROM cte_source_0 cte_source_0_alias
+
+sql_after_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+    FROM test_schema.test_table test_table_alias
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+  FROM cte_source_0 cte_source_0_alias
diff --git a/tests_metricflow/snapshots/test_cte_column_pruner.py/str/test_simple_pruning__result.txt b/tests_metricflow/snapshots/test_cte_column_pruner.py/str/test_simple_pruning__result.txt
@@ -0,0 +1,29 @@
+optimizer:
+  SqlColumnPrunerOptimizer
+
+sql_before_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+      , test_table_alias.col_0 AS cte_source_0__col_1
+    FROM test_schema.test_table test_table_alias
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+  FROM cte_source_0 cte_source_0_alias
+
+sql_after_optimizing:
+  -- Top-level SELECT
+  WITH cte_source_0 AS (
+    -- CTE source 0
+    SELECT
+      test_table_alias.col_0 AS cte_source_0__col_0
+    FROM test_schema.test_table test_table_alias
+  )
+
+  SELECT
+    cte_source_0_alias.cte_source_0__col_0 AS top_level__col_0
+  FROM cte_source_0 cte_source_0_alias
diff --git a/tests_metricflow/sql/optimizer/check_optimizer.py b/tests_metricflow/sql/optimizer/check_optimizer.py
@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+import logging
+
+from _pytest.fixtures import FixtureRequest
+from metricflow_semantics.mf_logging.formatting import indent
+from metricflow_semantics.mf_logging.lazy_formattable import LazyFormat
+from metricflow_semantics.test_helpers.config_helpers import MetricFlowTestConfiguration
+from metricflow_semantics.test_helpers.snapshot_helpers import assert_str_snapshot_equal
+
+from metricflow.sql.optimizer.sql_query_plan_optimizer import SqlQueryPlanOptimizer
+from metricflow.sql.render.sql_plan_renderer import SqlQueryPlanRenderer
+from metricflow.sql.sql_plan import SqlQueryPlan, SqlSelectStatementNode
+
+logger = logging.getLogger(__name__)
+
+
+def assert_optimizer_result_snapshot_equal(
+    request: FixtureRequest,
+    mf_test_configuration: MetricFlowTestConfiguration,
+    optimizer: SqlQueryPlanOptimizer,
+    sql_plan_renderer: SqlQueryPlanRenderer,
+    select_statement: SqlSelectStatementNode,
+) -> None:
+    """Helper to assert that the SQL snapshot of the optimizer result is the same as the stored one."""
+    sql_before_optimizing = sql_plan_renderer.render_sql_query_plan(SqlQueryPlan(select_statement)).sql
+    logger.debug(
+        LazyFormat(
+            "Optimizing SELECT statement",
+            select_statement=select_statement.structure_text(),
+            sql_before_optimizing=sql_before_optimizing,
+        )
+    )
+
+    column_pruned_select_node = optimizer.optimize(select_statement)
+    sql_after_optimizing = sql_plan_renderer.render_sql_query_plan(SqlQueryPlan(column_pruned_select_node)).sql
+    logger.debug(
+        LazyFormat(
+            "Optimized SQL",
+            sql_before_optimizing=sql_before_optimizing,
+            sql_after_optimizing=sql_after_optimizing,
+        )
+    )
+    snapshot_str = "\n".join(
+        [
+            "optimizer:",
+            indent(optimizer.__class__.__name__),
+            "",
+            "sql_before_optimizing:",
+            indent(sql_before_optimizing),
+            "",
+            "sql_after_optimizing:",
+            indent(sql_after_optimizing),
+        ]
+    )
+    assert_str_snapshot_equal(
+        request=request,
+        mf_test_configuration=mf_test_configuration,
+        snapshot_id="result",
+        snapshot_str=snapshot_str,
+    )