Qualify segment filters in CTEs (#55)

* Qualify segment filters in CTEs

* Format segment filter helper

* Avoid qualifying subquery columns

Author: nicosuave

sidemantic/sql/generator.py

@@ -439,6 +439,33 @@ def _resolve_segments(self, segments: list[str]) -> list[str]:
         Raises:
             ValueError: If segment not found
         """
+
+        def qualify_unaliased_columns(filter_sql: str, model_alias: str) -> str:
+            """Qualify unaliased columns in segment filters with model alias."""
+            try:
+                parsed = sqlglot.parse_one(filter_sql, dialect=self.dialect)
+            except Exception:
+                return filter_sql
+
+            def visit(node: exp.Expression) -> None:
+                if isinstance(node, exp.Subquery):
+                    return
+
+                if isinstance(node, exp.Column) and not node.table:
+                    node.set("table", model_alias)
+
+                for arg in node.args.values():
+                    if isinstance(arg, exp.Expression):
+                        visit(arg)
+                    elif isinstance(arg, list):
+                        for item in arg:
+                            if isinstance(item, exp.Expression):
+                                visit(item)
+
+            visit(parsed)
+
+            return parsed.sql(dialect=self.dialect)
+
         filters = []
         for seg_ref in segments:
             # Parse model.segment format
@@ -457,6 +484,7 @@ def _resolve_segments(self, segments: list[str]) -> list[str]:
             # Get SQL expression with model alias replaced
             # Use model_cte as the alias (consistent with CTE naming)
             filter_sql = segment.get_sql(f"{model_name}_cte")
+            filter_sql = qualify_unaliased_columns(filter_sql, f"{model_name}_cte")
             filters.append(filter_sql)
 
         return filters

tests/optimizations/test_predicate_pushdown.py

@@ -3,7 +3,8 @@
 import sqlglot
 from sqlglot import exp
 
-from sidemantic import Dimension, Metric, Model
+from sidemantic import Dimension, Metric, Model, Segment
+from sidemantic.sql.generator import SQLGenerator
 
 
 def test_single_model_filter_pushdown(layer):
@@ -283,6 +284,42 @@ def test_segment_filters_pushed_down(layer):
     assert "completed" in where_sql
 
 
+def test_segment_filter_skips_subquery_columns(layer):
+    """Test that segment filter qualification does not touch subquery columns."""
+    model = Model(
+        name="orders",
+        table="orders_table",
+        primary_key="id",
+        metrics=[
+            Metric(name="count", agg="count"),
+        ],
+        segments=[
+            Segment(name="in_other", sql="id in (select id from other_table where flag = 'y')"),
+        ],
+    )
+
+    layer.add_model(model)
+
+    generator = SQLGenerator(layer.graph)
+    filters = generator._resolve_segments(["orders.in_other"])
+    assert len(filters) == 1
+
+    filter_sql = filters[0]
+    parsed = sqlglot.parse_one(filter_sql)
+
+    assert any(col.table == "orders_cte" for col in parsed.find_all(exp.Column))
+
+    subquery = None
+    for subquery_def in parsed.find_all(exp.Subquery):
+        subquery = subquery_def
+        break
+
+    assert subquery is not None
+
+    for col in subquery.find_all(exp.Column):
+        assert not col.table
+
+
 def test_metric_level_filters_not_pushed(layer):
     """Test that metric-level filters are applied via CASE WHEN, not pushed to CTE.
 

tests/queries/test_basic.py

@@ -2,8 +2,10 @@
 
 import duckdb
 import pytest
+import sqlglot
+from sqlglot import exp
 
-from sidemantic import Dimension, Metric, Model, Relationship, SemanticLayer
+from sidemantic import Dimension, Metric, Model, Relationship, Segment, SemanticLayer
 from tests.utils import df_rows
 
 
@@ -361,6 +363,53 @@ def test_count_distinct_without_sql_uses_primary_key(layer):
     assert "COUNT(DISTINCT" in sql
 
 
+def test_count_distinct_with_segment_filter_without_model_placeholder(layer):
+    """Test count_distinct with segment filters that omit {model} placeholders."""
+    layer = SemanticLayer()
+
+    location = Model(
+        name="location",
+        table="dim_location",
+        primary_key="sk_location_id",
+        dimensions=[
+            Dimension(name="city", type="categorical"),
+        ],
+        metrics=[
+            Metric(name="count", agg="count_distinct"),  # No sql field
+        ],
+        segments=[
+            Segment(name="lockers_3000", sql="zipcode = '3000'"),
+        ],
+    )
+
+    layer.add_model(location)
+
+    sql = layer.compile(
+        metrics=["location.count"],
+        dimensions=["location.city"],
+        segments=["location.lockers_3000"],
+    )
+
+    # Should still use primary key for count_distinct
+    assert "sk_location_id AS count_raw" in sql
+    assert "count AS count_raw" not in sql
+
+    parsed = sqlglot.parse_one(sql)
+    cte = None
+    for cte_def in parsed.find_all(exp.CTE):
+        if cte_def.alias == "location_cte":
+            cte = cte_def
+            break
+
+    assert cte is not None
+
+    where_clause = cte.this.find(exp.Where)
+    assert where_clause is not None
+    where_sql = where_clause.sql()
+    assert "zipcode" in where_sql
+    assert "3000" in where_sql
+
+
 def test_count_distinct_with_explicit_sql(layer):
     """Test that count_distinct with explicit sql uses that column."""
     layer = SemanticLayer()