sql_cli/query_plan/
qualify_to_where_transformer.rs

1//! QUALIFY to WHERE clause transformer
2//!
3//! This transformer converts QUALIFY clauses (Snowflake-style window function filtering)
4//! into WHERE clauses that work with lifted window functions.
5//!
6//! # Problem
7//!
8//! Users want to filter on window function results without writing subqueries:
9//! ```sql
10//! SELECT region, sales_amount,
11//!        ROW_NUMBER() OVER (PARTITION BY region ORDER BY sales_amount DESC) AS rn
12//! FROM sales
13//! QUALIFY rn <= 3  -- Cleaner than WHERE in a subquery!
14//! ```
15//!
16//! # Solution
17//!
18//! The transformer works in conjunction with ExpressionLifter:
19//! 1. ExpressionLifter runs first, lifting window functions to CTE
20//! 2. QualifyToWhereTransformer runs second, converting QUALIFY → WHERE
21//!
22//! # Algorithm
23//!
24//! 1. Check if statement has a QUALIFY clause
25//! 2. If yes, move it to WHERE clause
26//! 3. If WHERE already exists, combine with AND
27//! 4. Set qualify field to None
28//!
29//! # Example Transformation
30//!
31//! After ExpressionLifter:
32//! ```sql
33//! WITH _lifted AS (
34//!     SELECT region, sales_amount,
35//!            ROW_NUMBER() OVER (...) AS rn
36//!     FROM sales
37//! )
38//! SELECT * FROM _lifted
39//! QUALIFY rn <= 3
40//! ```
41//!
42//! After QualifyToWhereTransformer:
43//! ```sql
44//! WITH _lifted AS (
45//!     SELECT region, sales_amount,
46//!            ROW_NUMBER() OVER (...) AS rn
47//!     FROM sales
48//! )
49//! SELECT * FROM _lifted
50//! WHERE rn <= 3
51//! ```
52
53use crate::query_plan::pipeline::ASTTransformer;
54use crate::sql::parser::ast::{Condition, LogicalOp, SelectStatement, WhereClause};
55use anyhow::Result;
56use tracing::debug;
57
58/// Transformer that converts QUALIFY clauses to WHERE clauses
59pub struct QualifyToWhereTransformer;
60
61impl QualifyToWhereTransformer {
62    pub fn new() -> Self {
63        Self
64    }
65}
66
67impl Default for QualifyToWhereTransformer {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73impl ASTTransformer for QualifyToWhereTransformer {
74    fn name(&self) -> &str {
75        "QualifyToWhereTransformer"
76    }
77
78    fn description(&self) -> &str {
79        "Converts QUALIFY clauses to WHERE clauses after window function lifting"
80    }
81
82    fn transform(&mut self, mut stmt: SelectStatement) -> Result<SelectStatement> {
83        // Only process if there's a QUALIFY clause
84        if stmt.qualify.is_none() {
85            return Ok(stmt);
86        }
87
88        let qualify_expr = stmt.qualify.take().unwrap();
89
90        debug!("Converting QUALIFY clause to WHERE");
91
92        // If there's already a WHERE clause, combine them with AND
93        if let Some(mut where_clause) = stmt.where_clause.take() {
94            debug!("Combining QUALIFY with existing WHERE using AND");
95
96            // Add AND connector to the last condition in existing WHERE
97            if let Some(last_condition) = where_clause.conditions.last_mut() {
98                last_condition.connector = Some(LogicalOp::And);
99            }
100
101            // Add the QUALIFY expression as a new condition
102            where_clause.conditions.push(Condition {
103                expr: qualify_expr,
104                connector: None,
105            });
106
107            stmt.where_clause = Some(where_clause);
108        } else {
109            // No existing WHERE clause, create new one
110            debug!("Creating new WHERE clause from QUALIFY");
111
112            stmt.where_clause = Some(WhereClause {
113                conditions: vec![Condition {
114                    expr: qualify_expr,
115                    connector: None,
116                }],
117            });
118        }
119
120        Ok(stmt)
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127    use crate::sql::parser::ast::{ColumnRef, QuoteStyle, SqlExpression};
128
129    #[test]
130    fn test_qualify_to_where_simple() {
131        let mut stmt = SelectStatement::default();
132
133        // Add QUALIFY: rn <= 3
134        stmt.qualify = Some(SqlExpression::BinaryOp {
135            left: Box::new(SqlExpression::Column(ColumnRef::unquoted("rn".to_string()))),
136            op: "<=".to_string(),
137            right: Box::new(SqlExpression::NumberLiteral("3".to_string())),
138        });
139
140        let mut transformer = QualifyToWhereTransformer::new();
141        let result = transformer.transform(stmt).unwrap();
142
143        // QUALIFY should be None
144        assert!(result.qualify.is_none());
145
146        // WHERE should exist
147        assert!(result.where_clause.is_some());
148
149        let where_clause = result.where_clause.unwrap();
150        assert_eq!(where_clause.conditions.len(), 1);
151    }
152
153    #[test]
154    fn test_qualify_combined_with_existing_where() {
155        let mut stmt = SelectStatement::default();
156
157        // Add existing WHERE: region = 'North'
158        stmt.where_clause = Some(WhereClause {
159            conditions: vec![Condition {
160                expr: SqlExpression::BinaryOp {
161                    left: Box::new(SqlExpression::Column(ColumnRef::unquoted(
162                        "region".to_string(),
163                    ))),
164                    op: "=".to_string(),
165                    right: Box::new(SqlExpression::StringLiteral("North".to_string())),
166                },
167                connector: None,
168            }],
169        });
170
171        // Add QUALIFY: rn <= 3
172        stmt.qualify = Some(SqlExpression::BinaryOp {
173            left: Box::new(SqlExpression::Column(ColumnRef::unquoted("rn".to_string()))),
174            op: "<=".to_string(),
175            right: Box::new(SqlExpression::NumberLiteral("3".to_string())),
176        });
177
178        let mut transformer = QualifyToWhereTransformer::new();
179        let result = transformer.transform(stmt).unwrap();
180
181        // QUALIFY should be None
182        assert!(result.qualify.is_none());
183
184        // WHERE should have 2 conditions connected by AND
185        let where_clause = result.where_clause.unwrap();
186        assert_eq!(where_clause.conditions.len(), 2);
187        assert!(matches!(
188            where_clause.conditions[0].connector,
189            Some(LogicalOp::And)
190        ));
191    }
192}
sql_cli/query_plan/qualify_to_where_transformer.rs

sql_cli/query_plan/
qualify_to_where_transformer.rs
