diff --git a/src/query_plan/column_dependency_lifter.rs b/src/query_plan/column_dependency_lifter.rs deleted file mode 100644 index 55f6cb35..00000000 --- a/src/query_plan/column_dependency_lifter.rs +++ /dev/null @@ -1,331 +0,0 @@ -use crate::sql::parser::ast::{ - ColumnRef, SelectStatement, SelectItem, SqlExpression, OrderByColumn, CTE, CTEType -}; -use std::collections::{HashMap, HashSet}; - -/// Column Dependency Lifter - Rewrites queries to handle column alias dependencies -/// -/// This preprocessor automatically generates CTEs when a computed column is referenced -/// in the same SELECT statement (e.g., in window functions, GROUP BY, or ORDER BY). -/// -/// Example transformation: -/// ```sql -/// -- Input (illegal - can't reference 'root' alias in PARTITION BY): -/// SELECT -/// CASE WHEN CONTAINS(id, '|') THEN SUBSTRING_AFTER(id, '|', 1) ELSE id END AS root, -/// ROW_NUMBER() OVER (PARTITION BY root) as rn -/// FROM table -/// -/// -- Output (legal - CTE provides 'root' for use): -/// WITH __lifted_cols AS ( -/// SELECT -/// *, -/// CASE WHEN CONTAINS(id, '|') THEN SUBSTRING_AFTER(id, '|', 1) ELSE id END AS root -/// FROM table -/// ) -/// SELECT -/// root, -/// ROW_NUMBER() OVER (PARTITION BY root) as rn -/// FROM __lifted_cols -/// ``` -pub struct ColumnDependencyLifter { - cte_counter: usize, -} - -impl ColumnDependencyLifter { - pub fn new() -> Self { - Self { cte_counter: 0 } - } - - /// Check if a statement needs column dependency lifting and rewrite if necessary - pub fn lift_dependencies(&mut self, statement: &mut SelectStatement) -> bool { - // Find all column aliases defined in SELECT - let column_aliases = self.extract_column_aliases(statement); - - // Find all references to these aliases in window functions, GROUP BY, etc. - let dependencies = self.find_alias_dependencies(statement, &column_aliases); - - if dependencies.is_empty() { - return false; - } - - // Generate a CTE to compute the dependent columns - let lifted_cte = self.generate_lifted_cte(statement, &dependencies); - - // Rewrite the main query to use the CTE - self.rewrite_query_with_cte(statement, lifted_cte, &dependencies); - - true - } - - /// Extract all column aliases from SELECT items - fn extract_column_aliases(&self, statement: &SelectStatement) -> HashMap { - let mut aliases = HashMap::new(); - - for item in &statement.select_items { - match item { - SelectItem::Expression { expr, alias } => { - aliases.insert(alias.clone(), expr.clone()); - } - SelectItem::Column(name) => { - // Simple columns might be aliased implicitly - if let SqlExpression::Column(col) = SqlExpression::Column(name.clone()) { - if col != name { - aliases.insert(name.clone(), SqlExpression::Column(col)); - } - } - } - _ => {} - } - } - - aliases - } - - /// Find dependencies on column aliases in window functions, GROUP BY, etc. - fn find_alias_dependencies( - &self, - statement: &SelectStatement, - aliases: &HashMap - ) -> HashSet { - let mut dependencies = HashSet::new(); - - // Check window functions in SELECT items - for item in &statement.select_items { - if let SelectItem::Expression { expr, .. } = item { - self.find_deps_in_expression(expr, aliases, &mut dependencies); - } - } - - // Check GROUP BY - if let Some(group_by) = &statement.group_by { - for col in group_by { - if aliases.contains_key(col) { - dependencies.insert(col.clone()); - } - } - } - - // Check ORDER BY - if let Some(order_by) = &statement.order_by { - for order_col in order_by { - if let SqlExpression::Column(col) = &order_col.column { - if aliases.contains_key(col) { - dependencies.insert(col.clone()); - } - } - } - } - - // Check HAVING clause - if let Some(having) = &statement.having { - self.find_deps_in_expression(having, aliases, &mut dependencies); - } - - dependencies - } - - /// Recursively find dependencies in an expression - fn find_deps_in_expression( - &self, - expr: &SqlExpression, - aliases: &HashMap, - dependencies: &mut HashSet - ) { - match expr { - SqlExpression::Window { func: _, args, partition_by, order_by, .. } => { - // Check args - for arg in args { - self.find_deps_in_expression(arg, aliases, dependencies); - } - - // Check PARTITION BY - if let Some(partition) = partition_by { - for col in partition { - if aliases.contains_key(col) { - dependencies.insert(col.clone()); - } - } - } - - // Check ORDER BY - if let Some(order) = order_by { - for order_col in order { - if let SqlExpression::Column(col) = &order_col.column { - if aliases.contains_key(col) { - dependencies.insert(col.clone()); - } - } - } - } - } - SqlExpression::Column(col) => { - // Check if this column is actually an alias - if aliases.contains_key(col) { - dependencies.insert(col.clone()); - } - } - SqlExpression::FunctionCall { args, .. } => { - for arg in args { - self.find_deps_in_expression(arg, aliases, dependencies); - } - } - SqlExpression::BinaryOp { left, right, .. } => { - self.find_deps_in_expression(left, aliases, dependencies); - self.find_deps_in_expression(right, aliases, dependencies); - } - SqlExpression::Case { when_clauses, else_clause } => { - for (cond, result) in when_clauses { - self.find_deps_in_expression(cond, aliases, dependencies); - self.find_deps_in_expression(result, aliases, dependencies); - } - if let Some(else_expr) = else_clause { - self.find_deps_in_expression(else_expr, aliases, dependencies); - } - } - _ => {} - } - } - - /// Generate a CTE that computes the dependent columns - fn generate_lifted_cte( - &mut self, - statement: &SelectStatement, - dependencies: &HashSet - ) -> CTE { - self.cte_counter += 1; - let cte_name = format!("__lifted_{}", self.cte_counter); - - // Build SELECT items for the CTE - let mut cte_select_items = vec![SelectItem::Star]; // Include all original columns - - // Add computed columns that are dependencies - for item in &statement.select_items { - if let SelectItem::Expression { expr, alias } = item { - if dependencies.contains(alias) { - cte_select_items.push(SelectItem::Expression { - expr: expr.clone(), - alias: alias.clone(), - }); - } - } - } - - // Build the CTE query - let cte_query = SelectStatement { - distinct: false, - columns: vec!["*".to_string()], // Legacy field - select_items: cte_select_items, - from_table: statement.from_table.clone(), - from_subquery: statement.from_subquery.clone(), - from_function: statement.from_function.clone(), - from_alias: statement.from_alias.clone(), - joins: statement.joins.clone(), - where_clause: statement.where_clause.clone(), - order_by: None, // Don't need ordering in the CTE - group_by: None, // GROUP BY goes in outer query - having: None, // HAVING goes in outer query - limit: None, - offset: None, - ctes: vec![], // No nested CTEs in the generated one - }; - - CTE { - name: cte_name, - cte_type: CTEType::Standard(cte_query), - } - } - - /// Rewrite the main query to use the generated CTE - fn rewrite_query_with_cte( - &self, - statement: &mut SelectStatement, - cte: CTE, - dependencies: &HashSet - ) { - let cte_name = cte.name.clone(); - - // Remove the computed expressions from SELECT items if they're dependencies - let mut new_select_items = Vec::new(); - for item in &statement.select_items { - match item { - SelectItem::Expression { expr: _, alias } if dependencies.contains(alias) => { - // Replace with simple column reference - new_select_items.push(SelectItem::Column(alias.clone())); - } - _ => { - new_select_items.push(item.clone()); - } - } - } - - // Update the statement - statement.select_items = new_select_items; - statement.from_table = Some(cte_name); - statement.from_subquery = None; - statement.from_function = None; - // Keep the original alias if there was one - - // Add the CTE to the statement - statement.ctes.insert(0, cte); - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_simple_dependency_lifting() { - let mut lifter = ColumnDependencyLifter::new(); - - // Create a query that uses an alias in PARTITION BY - let mut statement = SelectStatement { - select_items: vec![ - SelectItem::Expression { - expr: SqlExpression::Case { - when_clauses: vec![( - SqlExpression::FunctionCall { - name: "CONTAINS".to_string(), - args: vec![ - SqlExpression::Column(ColumnRef::unquoted("id".to_string())), - SqlExpression::String("|".to_string()), - ], - }, - SqlExpression::FunctionCall { - name: "SUBSTRING_AFTER".to_string(), - args: vec![ - SqlExpression::Column(ColumnRef::unquoted("id".to_string())), - SqlExpression::String("|".to_string()), - SqlExpression::Integer(1), - ], - }, - )], - else_clause: Some(Box::new(SqlExpression::Column(ColumnRef::unquoted("id".to_string())))), - }, - alias: "root".to_string(), - }, - SelectItem::Expression { - expr: SqlExpression::Window { - func: "ROW_NUMBER".to_string(), - args: vec![], - partition_by: Some(vec!["root".to_string()]), - order_by: None, - frame: None, - }, - alias: "rn".to_string(), - }, - ], - from_table: Some("test_table".to_string()), - ..Default::default() - }; - - // Apply lifting - let lifted = lifter.lift_dependencies(&mut statement); - - assert!(lifted); - assert_eq!(statement.ctes.len(), 1); - assert!(statement.ctes[0].name.starts_with("__lifted_")); - assert_eq!(statement.from_table, Some(statement.ctes[0].name.clone())); - } -} \ No newline at end of file diff --git a/src/query_plan/cte_hoister.rs b/src/query_plan/cte_hoister.rs index 2a043e00..4ca60df6 100644 --- a/src/query_plan/cte_hoister.rs +++ b/src/query_plan/cte_hoister.rs @@ -58,8 +58,8 @@ impl CTEHoister { /// Recursively hoist CTEs from a SELECT statement fn hoist_from_statement(&mut self, mut statement: SelectStatement) -> SelectStatement { // Hoist from subquery in FROM clause - if let Some(subquery) = statement.from_subquery.take() { - let rewritten_sub = self.hoist_from_statement(*subquery); + statement.map_from_subquery(|subquery| { + let rewritten_sub = self.hoist_from_statement(subquery); // If the subquery has CTEs, hoist them for cte in rewritten_sub.ctes.clone() { @@ -67,11 +67,11 @@ impl CTEHoister { } // Return the subquery without its CTEs (they're hoisted) - statement.from_subquery = Some(Box::new(SelectStatement { + SelectStatement { ctes: Vec::new(), ..rewritten_sub - })); - } + } + }); // Hoist from CTEs in this statement let local_ctes = statement.ctes.drain(..).collect::>(); @@ -415,6 +415,59 @@ impl CTEHoister { mod tests { use super::*; + /// Regression: hoisting a derived table must rewrite the `from_source` + /// copy, not only the legacy `from_subquery`. + /// + /// The parser fills both with clones of the same subquery, and the executor + /// reads `from_source`. Rewriting one alone left a stale pre-hoist copy in + /// the field that actually gets executed. + /// + /// This currently produces the right answer either way — the stale copy is + /// a self-contained statement that still carries its own CTEs — so the + /// desync is latent, not a live wrong-results bug. The assertion pins the + /// two representations together before the correlated-subquery work starts + /// depending on `from_source` being authoritative. + #[test] + fn test_derived_table_hoisting_updates_from_source() { + use crate::sql::parser::ast::TableSource; + use crate::sql::recursive_parser::Parser; + + let mut parser = Parser::new( + "SELECT symbol FROM (WITH x AS (SELECT symbol FROM trades) SELECT symbol FROM x) sub", + ); + let stmt = parser.parse().expect("query should parse"); + + let hoisted = CTEHoister::hoist_ctes(stmt); + + // The inner CTE was lifted to the top level. + assert_eq!(hoisted.ctes.len(), 1, "inner CTE should be hoisted"); + assert_eq!(hoisted.ctes[0].name, "x"); + + // Both representations must show the CTE-stripped subquery. Before the + // fix, from_source still held the original with `ctes.len() == 1`. + #[allow(deprecated)] + let legacy = hoisted + .from_subquery + .as_ref() + .expect("from_subquery should be present"); + assert!(legacy.ctes.is_empty(), "legacy copy should be stripped"); + + match hoisted.from_source { + Some(TableSource::DerivedTable { + ref query, + ref alias, + }) => { + assert!( + query.ctes.is_empty(), + "from_source holds a stale pre-hoist subquery with {} CTE(s)", + query.ctes.len() + ); + assert_eq!(alias, "sub", "derived-table alias must survive the rewrite"); + } + ref other => panic!("expected a DerivedTable from_source, got {other:?}"), + } + } + #[test] fn test_simple_cte_hoisting() { // Test that a simple nested CTE gets hoisted diff --git a/src/query_plan/expression_lifter.rs b/src/query_plan/expression_lifter.rs index dd390f13..79007918 100644 --- a/src/query_plan/expression_lifter.rs +++ b/src/query_plan/expression_lifter.rs @@ -184,7 +184,7 @@ impl ExpressionLifter { lifted_ctes.push(cte); // Update the main query to reference the CTE - stmt.from_table = Some(lift_expr.suggested_name); + stmt.set_from_table(lift_expr.suggested_name); // Replace the complex WHERE expression with a simple column reference use crate::sql::parser::ast::Condition; @@ -354,16 +354,9 @@ impl ExpressionLifter { } stmt.select_items = new_select_items; - // Set from_source to reference the CTE (preferred) - stmt.from_source = Some(crate::sql::parser::ast::TableSource::Table( - cte_name.clone(), - )); - // Also set deprecated field for backward compatibility - #[allow(deprecated)] - { - stmt.from_table = Some(cte_name.clone()); - stmt.from_subquery = None; - } + // Point the FROM clause at the CTE (keeps from_source and the legacy + // fields in sync) + stmt.set_from_table(cte_name.clone()); stmt.where_clause = None; // Already in the CTE CTE { diff --git a/src/query_plan/in_operator_lifter.rs b/src/query_plan/in_operator_lifter.rs index 12e6451a..6b4939af 100644 --- a/src/query_plan/in_operator_lifter.rs +++ b/src/query_plan/in_operator_lifter.rs @@ -215,7 +215,7 @@ impl InOperatorLifter { stmt.ctes.push(cte); // Update the FROM clause to use the CTE - stmt.from_table = Some(cte_name); + stmt.set_from_table(cte_name); true } @@ -290,4 +290,49 @@ mod tests { } )); } + + /// Regression: the lifter must repoint `from_source` at the lifted CTE, not + /// just the legacy `from_table`. + /// + /// The executor reads `from_source` first and only falls back to + /// `from_table`, so rewriting one without the other left `from_source` + /// pointing at the original base table — which has no lifted column. + /// + /// Parsed here rather than hand-built: the parser populates both fields, + /// which is what makes the desync reachable. A fixture with + /// `from_source: None` cannot catch this. + #[test] + fn test_rewrite_query_repoints_from_source_to_cte() { + use crate::sql::recursive_parser::Parser; + + let mut parser = Parser::new("SELECT symbol FROM trades WHERE LOWER(symbol) IN ('aapl')"); + let mut stmt = parser.parse().expect("query should parse"); + + // Precondition: the parser sets both representations. + #[allow(deprecated)] + { + assert_eq!(stmt.from_table.as_deref(), Some("trades")); + } + assert!(matches!( + stmt.from_source, + Some(TableSource::Table(ref t)) if t == "trades" + )); + + assert!( + InOperatorLifter::new().rewrite_query(&mut stmt), + "LOWER(col) IN (...) should trigger lifting" + ); + + #[allow(deprecated)] + let rewritten = stmt.from_table.clone().expect("from_table should be set"); + assert_eq!(rewritten, "trades_lifted"); + + match stmt.from_source { + Some(TableSource::Table(ref t)) => assert_eq!( + t, &rewritten, + "from_source must follow from_table to the lifted CTE" + ), + ref other => panic!("expected from_source to name the lifted CTE, got {other:?}"), + } + } } diff --git a/src/query_plan/into_clause_remover.rs b/src/query_plan/into_clause_remover.rs index 44dc90a2..e0c5c49b 100644 --- a/src/query_plan/into_clause_remover.rs +++ b/src/query_plan/into_clause_remover.rs @@ -41,9 +41,7 @@ impl IntoClauseRemover { statement.into_table = None; // Remove from subquery in FROM clause - if let Some(subquery) = statement.from_subquery.take() { - statement.from_subquery = Some(Box::new(Self::remove_from_statement(*subquery))); - } + statement.map_from_subquery(Self::remove_from_statement); // Remove from JOIN subqueries statement.joins = statement diff --git a/src/sql/parser/ast.rs b/src/sql/parser/ast.rs index 23f670e4..12b75e95 100644 --- a/src/sql/parser/ast.rs +++ b/src/sql/parser/ast.rs @@ -515,6 +515,62 @@ impl Default for SelectStatement { } } +impl SelectStatement { + /// Point the FROM clause at a named table, updating `from_source` and the + /// legacy fields together. + /// + /// Both representations must move in lockstep: the executor + /// (`query_engine.rs`) reads `from_source` first and only falls back to + /// `from_table`, so a rewrite that touches just the legacy field is + /// silently discarded whenever `from_source` is populated. + pub fn set_from_table(&mut self, table: String) { + self.from_source = Some(TableSource::Table(table.clone())); + #[allow(deprecated)] + { + self.from_table = Some(table); + self.from_subquery = None; + self.from_function = None; + } + } + + /// Replace the derived table (subquery) in the FROM clause, updating + /// `from_source` and the legacy field together. + /// + /// The parser populates `from_subquery` and `from_source::DerivedTable` + /// with clones of the same subquery, so rewriting only the former leaves a + /// stale copy in `from_source` — which is the one the executor reads. The + /// existing derived-table alias is preserved. + pub fn set_from_subquery(&mut self, query: Box) { + #[allow(deprecated)] + let alias = match &self.from_source { + Some(TableSource::DerivedTable { alias, .. }) => alias.clone(), + _ => self.from_alias.clone().unwrap_or_default(), + }; + self.from_source = Some(TableSource::DerivedTable { + query: query.clone(), + alias, + }); + #[allow(deprecated)] + { + self.from_subquery = Some(query); + } + } + + /// Rewrite the derived table (subquery) in the FROM clause in place, if + /// there is one, keeping `from_source` and the legacy field in sync. + /// + /// Prefer this over taking `from_subquery` and reassigning it: the + /// take-then-restore pattern leaves the `from_source` copy stale in + /// between, and the alias is easy to drop on the way through. + pub fn map_from_subquery(&mut self, f: impl FnOnce(SelectStatement) -> SelectStatement) { + #[allow(deprecated)] + let taken = self.from_subquery.take(); + if let Some(subquery) = taken { + self.set_from_subquery(Box::new(f(*subquery))); + } + } +} + /// INTO clause for creating temporary tables #[derive(Debug, Clone, PartialEq)] pub struct IntoTable { diff --git a/tests/integration/test_examples.py b/tests/integration/test_examples.py index 94734ddf..67062751 100755 --- a/tests/integration/test_examples.py +++ b/tests/integration/test_examples.py @@ -324,7 +324,9 @@ def capture_expectation(cli_path: str, example_name: str, expectations_dir: Path print("To test: ./tests/integration/test_examples.py", example_name) def main(): - cli_path = './target/release/sql-cli' + # Windows builds the binary with a .exe suffix (same resolution as + # tests/comparison/engines.py); without this the suite aborts on win32. + cli_path = './target/release/sql-cli.exe' if sys.platform == 'win32' else './target/release/sql-cli' examples_dir = Path('examples') expectations_dir = Path('examples/expectations')