diff --git a/rust/cubesql/cubesql/src/compile/rewrite/mod.rs b/rust/cubesql/cubesql/src/compile/rewrite/mod.rs index eef32945cece1..44d8522c8f5f6 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/mod.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/mod.rs @@ -558,6 +558,12 @@ crate::plan_to_language! { // trace_macros!(false); +/// Positions of `WrappedSelect` children, for the rules that reach into the node instead of +/// matching it with a pattern. Keep in sync with the `WrappedSelect` definition above. +pub const WRAPPED_SELECT_SELECT_TYPE: usize = 0; +pub const WRAPPED_SELECT_FROM: usize = 6; +pub const WRAPPED_SELECT_JOINS: usize = 7; + #[macro_export] macro_rules! var_iter { ($eclass:expr, $field_variant:ident) => {{ diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/join.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/join.rs index b1d7acc547906..4e342fd20f312 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/join.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/join.rs @@ -9,14 +9,18 @@ use crate::{ wrapped_select_order_expr_empty_tail, wrapped_select_projection_expr_empty_tail, wrapped_select_subqueries_empty_tail, wrapped_select_window_expr_empty_tail, wrapper_pullup_replacer, wrapper_pushdown_replacer, wrapper_replacer_context, BinaryExprOp, - ColumnExprColumn, CubeEGraph, JoinLeftOn, JoinRightOn, LogicalPlanLanguage, - WrappedSelectJoinJoinType, WrappedSelectPushToCube, WrapperReplacerContextAliasToCube, - WrapperReplacerContextGroupedSubqueries, + ColumnExprColumn, CubeEGraph, CubeScanWrapperFinalized, JoinLeftOn, JoinRightOn, + LogicalPlanLanguage, WrappedSelectJoinJoinType, WrappedSelectPushToCube, + WrappedSelectSelectType, WrappedSelectType, WrapperReplacerContextAliasToCube, + WrapperReplacerContextGroupedSubqueries, WrapperReplacerContextUngroupedScan, + WRAPPED_SELECT_FROM, WRAPPED_SELECT_JOINS, WRAPPED_SELECT_SELECT_TYPE, }, transport::MetaContext, var, var_iter, var_list_iter, }; +use std::collections::HashSet; + use datafusion::{ logical_expr::{Expr, Operator}, logical_plan::Column, @@ -260,6 +264,7 @@ impl WrapperRules { "CubeScanWrapperFinalized:false", ), self.transform_ungrouped_join_grouped( + "?right_input", "?left_cube_members", "?left_on", "?right_on", @@ -479,6 +484,7 @@ impl WrapperRules { "CubeScanWrapperFinalized:false", ), self.transform_grouped_join_grouped( + "?left_input", "?left_on", "?left_push_to_cube", "?right_on", @@ -490,6 +496,209 @@ impl WrapperRules { "?out_push_to_cube", ), ), + // A pivot query builder emits one CTE per property, all joined to the same root + // CTE. The rule above turns the first of those joins into a WrappedSelect; every + // next join is added to that select's join list here, rather than nesting another + // select around it. Nesting would repeat the whole select per join, and the number + // of ways to represent the result grows with every level. + // + // The context of the left select is reused as is, including its + // `grouped_subqueries`: that list is only read by the column rules to pull a column + // qualified by a joined subquery up as a dimension, and only when pushing members + // to Cube. This select does not (`WrapperReplacerContextPushToCube:false` below), + // so the aliases of the subqueries joined here have nothing to do in that list. + transforming_rewrite( + "wrapper-push-down-grouped-join-grouped-chain", + join( + cube_scan_wrapper( + wrapper_pullup_replacer( + // The select built by the rule above, before anything was pushed + // into it: everything but `from` and `joins` is still empty + wrapped_select( + "WrappedSelectSelectType:Projection", + wrapped_select_projection_expr_empty_tail(), + wrapped_select_subqueries_empty_tail(), + wrapped_select_group_expr_empty_tail(), + wrapped_select_aggr_expr_empty_tail(), + wrapped_select_window_expr_empty_tail(), + "?left_from", + "?left_joins", + wrapped_select_filter_expr_empty_tail(), + wrapped_select_having_expr_empty_tail(), + "WrappedSelectLimit:None", + "WrappedSelectOffset:None", + wrapped_select_order_expr_empty_tail(), + "WrappedSelectAlias:None", + "WrappedSelectDistinct:false", + // Only a select that joins subqueries as plain SQL is extended + // here. A push-to-Cube select carries its joins to the Cube + // query as subquery joins, which are rendered under a + // uniqueness assumption this rule can not check + "WrappedSelectPushToCube:false", + "WrappedSelectUngroupedScan:false", + ), + wrapper_replacer_context( + "?left_alias_to_cube", + // A select with joins can not push anything else to Cube, so + // the join condition needs no member resolution and is built + // as plain columns below + "WrapperReplacerContextPushToCube:false", + "?left_in_projection", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + "CubeScanWrapperFinalized:false", + ), + // The joined side is matched as a whole, and its shape is checked in the + // transform: a grouped subquery can be represented in several ways at once, + // and matching them here would produce a copy of this select per + // representation, for every join in a chain + "?right_wrapper", + "?left_on", + "?right_on", + "?in_join_type", + "?join_constraint", + "JoinNullEqualsNull:false", + ), + cube_scan_wrapper( + wrapped_select( + "WrappedSelectSelectType:Projection", + wrapper_pullup_replacer( + wrapped_select_projection_expr_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapper_pullup_replacer( + wrapped_select_subqueries_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapper_pullup_replacer( + wrapped_select_group_expr_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapper_pullup_replacer( + wrapped_select_aggr_expr_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapper_pullup_replacer( + wrapped_select_window_expr_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapper_pullup_replacer( + "?left_from", + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + // The join list is built by the transform, with the new join at the end + // so that joins keep the order they had in the query. It is built + // resolved, in one node: going through push down and pull up per element + // would let the e-graph hold every mix of resolved and unresolved + // elements of the list. + wrapper_pullup_replacer( + "?out_joins", + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapper_pullup_replacer( + wrapped_select_filter_expr_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + wrapped_select_having_expr_empty_tail(), + "WrappedSelectLimit:None", + "WrappedSelectOffset:None", + wrapper_pullup_replacer( + wrapped_select_order_expr_empty_tail(), + wrapper_replacer_context( + "?left_alias_to_cube", + "WrapperReplacerContextPushToCube:false", + "WrapperReplacerContextInProjection:false", + "?left_cube_members", + "?left_grouped_subqueries", + "WrapperReplacerContextUngroupedScan:false", + "?input_data_source", + ), + ), + "WrappedSelectAlias:None", + "WrappedSelectDistinct:false", + "WrappedSelectPushToCube:false", + "WrappedSelectUngroupedScan:false", + ), + "CubeScanWrapperFinalized:false", + ), + self.transform_grouped_join_grouped_chain( + "?left_joins", + "?left_on", + "?right_wrapper", + "?right_on", + "?in_join_type", + "?input_data_source", + "?out_joins", + ), + ), ]); // DataFusion plans complex join conditions as Filter(?join_condition, CrossJoin(...)) @@ -761,6 +970,7 @@ impl WrapperRules { "CubeScanWrapperFinalized:false", ), self.transform_ungrouped_join_grouped_after_check( + "?right_input", "?right_alias_to_cube", "?out_join_type", "?out_grouped_subqueries", @@ -1016,6 +1226,7 @@ impl WrapperRules { fn transform_ungrouped_join_grouped( &self, + right_input_var: &'static str, left_members_var: &'static str, left_on_var: &'static str, right_on_var: &'static str, @@ -1025,6 +1236,7 @@ impl WrapperRules { out_join_type_var: &'static str, out_grouped_subqueries_var: &'static str, ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool { + let right_input_var = var!(right_input_var); let left_members_var = var!(left_members_var); let left_on_var = var!(left_on_var); @@ -1043,6 +1255,10 @@ impl WrapperRules { // It means we don't care about just a "single cube" in LHS, and there's essentially no cubes by this moment in RHS move |egraph, subst| { + if !Self::can_be_subquery_join(egraph, subst[right_input_var]) { + return false; + } + // We are going to generate join with grouped subquery // TODO Do we have to check stuff like `transform_check_subquery_allowed` is checking: // * Both inputs depend on a single data source @@ -1213,15 +1429,21 @@ impl WrapperRules { fn transform_ungrouped_join_grouped_after_check( &self, + right_input_var: &'static str, right_alias_to_cube_var: &'static str, out_join_type_var: &'static str, out_grouped_subqueries_var: &'static str, ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool { + let right_input_var = var!(right_input_var); let right_alias_to_cube_var = var!(right_alias_to_cube_var); let out_join_type_var = var!(out_join_type_var); let out_grouped_subqueries_var = var!(out_grouped_subqueries_var); move |egraph, subst| { + if !Self::can_be_subquery_join(egraph, subst[right_input_var]) { + return false; + } + for right_alias_to_cube in var_iter!( egraph[subst[right_alias_to_cube_var]], WrapperReplacerContextAliasToCube @@ -1262,6 +1484,7 @@ impl WrapperRules { fn transform_grouped_join_grouped( &self, + left_input_var: &'static str, left_on_var: &'static str, left_push_to_cube_var: &'static str, right_on_var: &'static str, @@ -1272,6 +1495,7 @@ impl WrapperRules { out_grouped_subqueries_var: &'static str, out_push_to_cube_var: &'static str, ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool { + let left_input_var = var!(left_input_var); let left_on_var = var!(left_on_var); let left_push_to_cube_var = var!(left_push_to_cube_var); @@ -1288,6 +1512,13 @@ impl WrapperRules { let meta = self.meta_context.clone(); move |egraph, subst| { + // Joins on top of a select that already has joins are handled by + // `wrapper-push-down-grouped-join-grouped-chain`, which keeps them in a single + // select instead of nesting one select per join + if Self::select_has_joins(egraph, subst[left_input_var]) { + return false; + } + // We are going to generate join with grouped subquery // TODO Do we have to check stuff like `transform_check_subquery_allowed` is checking: // * Both inputs depend on a single data source @@ -1376,4 +1607,279 @@ impl WrapperRules { return false; } } + + /// Whether `input` can be joined to a Cube query as a subquery join. + /// + /// A subquery join is rendered by the schema compiler in two ways - counted into the + /// measures over the joined rowset, or computed per distinct primary key - which agree + /// only when the joined subquery is unique on the join keys. A select that is itself a + /// join of subqueries has no such guarantee, and nothing here can check it, so it is + /// refused: the query then plans without pushing the join into the Cube query, or fails + /// with an explicit error, instead of returning numbers that depend on how a measure + /// happens to be classified. + fn can_be_subquery_join(egraph: &CubeEGraph, input: Id) -> bool { + !Self::joins_subqueries(egraph, input, &mut HashSet::new()) + } + + /// Whether `input` produces its rows by joining subqueries, looking through the selects + /// stacked on top of the join, and through replacers when a select below is not pulled up + /// yet. A grouping in between is where the search stops: aggregation collapses the rows a + /// join below it could have duplicated. Note that this does not make the result unique on + /// the join keys - that also needs the join keys to cover the group keys, which nothing + /// here checks (TODO: check key coverage, the pre-existing subquery join path needs it as + /// well). + /// + /// `visited` keeps the walk linear and terminating: e-classes hold many equivalent nodes + /// pointing at shared children, and unions can make the graph cyclic. + fn joins_subqueries(egraph: &CubeEGraph, input: Id, visited: &mut HashSet) -> bool { + if !visited.insert(egraph.find(input)) { + // Already walked, on this path or another one + return false; + } + + for select in var_list_iter!(egraph[input], WrappedSelect) { + let (Some(select_type), Some(from), Some(joins)) = ( + select.get(WRAPPED_SELECT_SELECT_TYPE), + select.get(WRAPPED_SELECT_FROM), + select.get(WRAPPED_SELECT_JOINS), + ) else { + continue; + }; + + if var_iter!(egraph[*select_type], WrappedSelectSelectType) + .any(|select_type| matches!(select_type, WrappedSelectType::Aggregate)) + { + continue; + } + + if var_list_iter!(egraph[*joins], WrappedSelectJoins).any(|joins| !joins.is_empty()) { + return true; + } + + if Self::joins_subqueries(egraph, *from, visited) { + return true; + } + } + + // A select that is still wrapped in a replacer hides the same structure one node deeper + for replacer in var_list_iter!(egraph[input], WrapperPullupReplacer) + .chain(var_list_iter!(egraph[input], WrapperPushdownReplacer)) + { + let Some(member) = replacer.first() else { + continue; + }; + if Self::joins_subqueries(egraph, *member, visited) { + return true; + } + } + + false + } + + /// Whether `input` is a select that already carries joins. + fn select_has_joins(egraph: &CubeEGraph, input: Id) -> bool { + for select in var_list_iter!(egraph[input], WrappedSelect) { + let Some(joins) = select.get(WRAPPED_SELECT_JOINS) else { + continue; + }; + if var_list_iter!(egraph[*joins], WrappedSelectJoins).any(|joins| !joins.is_empty()) { + return true; + } + } + + false + } + + fn transform_grouped_join_grouped_chain( + &self, + left_joins_var: &'static str, + left_on_var: &'static str, + right_wrapper_var: &'static str, + right_on_var: &'static str, + in_join_type_var: &'static str, + input_data_source_var: &'static str, + out_joins_var: &'static str, + ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool { + let left_joins_var = var!(left_joins_var); + let left_on_var = var!(left_on_var); + let right_wrapper_var = var!(right_wrapper_var); + let right_on_var = var!(right_on_var); + let in_join_type_var = var!(in_join_type_var); + let input_data_source_var = var!(input_data_source_var); + let out_joins_var = var!(out_joins_var); + + let meta = self.meta_context.clone(); + + move |egraph, subst| { + // Only extend a select that is already joining something: a select without joins + // is the job of `wrapper-push-down-grouped-join-grouped` + if !var_list_iter!(egraph[subst[left_joins_var]], WrappedSelectJoins) + .any(|joins| !joins.is_empty()) + { + return false; + } + + let Some(right_input) = Self::grouped_join_right_input( + egraph, + subst[right_wrapper_var], + subst[input_data_source_var], + ) else { + return false; + }; + + for left_join_on in var_iter!(egraph[subst[left_on_var]], JoinLeftOn) { + for right_join_on in var_iter!(egraph[subst[right_on_var]], JoinRightOn) { + for in_join_type in + var_list_iter!(egraph[subst[in_join_type_var]], JoinJoinType).cloned() + { + // The select this join is added to is not pushing anything to Cube: + // it is a join of grouped subqueries, where every join type maps to + // SQL directly + if !Self::is_subquery_join_type_supported( + egraph, + subst, + &meta, + input_data_source_var, + &in_join_type.0, + false, + ) { + continue; + } + + let Some(out_join_expr) = Self::build_join_expr( + egraph, + left_join_on.clone(), + right_join_on.clone(), + ) else { + return false; + }; + + let join_type = egraph.add(LogicalPlanLanguage::WrappedSelectJoinJoinType( + WrappedSelectJoinJoinType(in_join_type.0), + )); + let join = egraph.add(LogicalPlanLanguage::WrappedSelectJoin([ + right_input, + out_join_expr, + join_type, + ])); + + let Some(out_joins) = + Self::append_join(egraph, subst[left_joins_var], join) + else { + return false; + }; + + subst.insert(out_joins_var, out_joins); + + return true; + } + } + } + + false + } + } + + /// Add `join` to the end of the `joins` list, keeping the order of the joins already + /// there: a join condition can refer to anything joined before it, but not after. + fn append_join(egraph: &mut CubeEGraph, joins: Id, join: Id) -> Option { + // Unions can make an e-graph list cyclic, and a rule transform must always return + let max_joins = 64; + + let mut list = vec![]; + let mut current = joins; + loop { + if list.len() >= max_joins { + return None; + } + + let mut nodes = var_list_iter!(egraph[current], WrappedSelectJoins).cloned(); + let node = nodes.next()?; + // A single list must have a single representation, otherwise it is not clear + // which one the join is added to + if nodes.next().is_some() { + return None; + } + match node.as_slice() { + [] => break, + [head, tail] => { + list.push(*head); + current = *tail; + } + _ => return None, + } + } + + let mut result = egraph.add(LogicalPlanLanguage::WrappedSelectJoins(vec![])); + result = egraph.add(LogicalPlanLanguage::WrappedSelectJoins(vec![join, result])); + for head in list.into_iter().rev() { + result = egraph.add(LogicalPlanLanguage::WrappedSelectJoins(vec![head, result])); + } + + Some(result) + } + + /// The plan to join, taken out of an unfinalized wrapper over a grouped subquery from the + /// same data source. A grouped subquery can sit in the e-graph in several shapes at once + /// (a Cube query, a select over it, ...), all interchangeable here, so one is picked + /// rather than joined once per shape: the plain Cube query when there is one, so that the + /// generated SQL keeps the fewest levels of nesting. + fn grouped_join_right_input( + egraph: &CubeEGraph, + right_wrapper: Id, + input_data_source: Id, + ) -> Option { + let mut candidates = vec![]; + + for wrapper in var_list_iter!(egraph[right_wrapper], CubeScanWrapper) { + let [input, finalized] = wrapper.as_slice() else { + continue; + }; + if !var_iter!(egraph[*finalized], CubeScanWrapperFinalized).any(|f| !f) { + continue; + } + + for pullup in var_list_iter!(egraph[*input], WrapperPullupReplacer) { + let [member, context] = pullup.as_slice() else { + continue; + }; + + for replacer_context in var_list_iter!(egraph[*context], WrapperReplacerContext) { + // (alias_to_cube, push_to_cube, in_projection, cube_members, + // grouped_subqueries, ungrouped_scan, input_data_source) + let [.., ungrouped_scan, data_source] = replacer_context.as_slice() else { + continue; + }; + // Grouped on both sides, and the same data source, like the rules that + // match those checks in their pattern + if !var_iter!(egraph[*ungrouped_scan], WrapperReplacerContextUngroupedScan) + .any(|ungrouped| !ungrouped) + { + continue; + } + if *data_source != input_data_source { + continue; + } + + candidates.push(*member); + } + } + } + + // Every candidate renders the same rows, so the pick is only about the shape of the + // generated SQL: a plain Cube query keeps it one level shallower than a select over + // one. Beyond that the choice is arbitrary, and sorting only makes it repeatable + // between runs. + candidates.sort(); + candidates.dedup(); + candidates + .iter() + .find(|member| Self::is_cube_scan(egraph, **member)) + .or_else(|| candidates.first()) + .copied() + } + + fn is_cube_scan(egraph: &CubeEGraph, id: Id) -> bool { + var_list_iter!(egraph[id], CubeScan).next().is_some() + } } diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/wrapper_pull_up.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/wrapper_pull_up.rs index 92f9d9c4c9519..8c3a29d830cea 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/wrapper_pull_up.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/wrapper_pull_up.rs @@ -4,9 +4,8 @@ use crate::{ rewriter::{CubeEGraph, CubeRewrite}, rules::{members::MemberRules, wrapper::WrapperRules}, transforming_rewrite, wrapped_select, wrapped_select_having_expr_empty_tail, - wrapped_select_joins_empty_tail, wrapper_pullup_replacer, wrapper_replacer_context, - LogicalPlanLanguage, WrappedSelectAlias, WrappedSelectSelectType, WrappedSelectType, - WrapperReplacerContextAliasToCube, + wrapper_pullup_replacer, wrapper_replacer_context, LogicalPlanLanguage, WrappedSelectAlias, + WrappedSelectSelectType, WrappedSelectType, WrapperReplacerContextAliasToCube, }, var, var_iter, var_list_iter, }; @@ -276,8 +275,7 @@ impl WrapperRules { ), ), wrapper_pullup_replacer( - // TODO handle non-empty joins - wrapped_select_joins_empty_tail(), + "?joins", wrapper_replacer_context( "?alias_to_cube", "?push_to_cube", @@ -351,7 +349,7 @@ impl WrapperRules { "?inner_push_to_cube", "?inner_ungrouped_scan", ), - wrapped_select_joins_empty_tail(), + "?joins", "?filter_expr", wrapped_select_having_expr_empty_tail(), "WrappedSelectLimit:None", @@ -381,6 +379,7 @@ impl WrapperRules { "?projection_expr", "?group_expr", "?aggr_expr", + "?joins", "?inner_select_type", "?inner_projection_expr", "?inner_group_expr", @@ -455,6 +454,7 @@ impl WrapperRules { projection_expr_var: &'static str, _group_expr_var: &'static str, _aggr_expr_var: &'static str, + joins_var: &'static str, inner_select_type_var: &'static str, inner_projection_expr_var: &'static str, _inner_group_expr_var: &'static str, @@ -465,6 +465,7 @@ impl WrapperRules { ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool { let select_type_var = var!(select_type_var); let projection_expr_var = var!(projection_expr_var); + let joins_var = var!(joins_var); let inner_select_type_var = var!(inner_select_type_var); let inner_projection_expr_var = var!(inner_projection_expr_var); let alias_to_cube_var = var!(alias_to_cube_var); @@ -481,6 +482,14 @@ impl WrapperRules { return false; } + // A select that joins something on top of another select is never a no-op + // wrapper around it, no matter how similar the two projections look + if var_list_iter!(egraph[subst[joins_var]], WrappedSelectJoins) + .any(|joins| !joins.is_empty()) + { + return true; + } + for select_type in var_iter!(egraph[subst[select_type_var]], WrappedSelectSelectType).cloned() { diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs index e202347037baa..07e1546f3bd53 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs @@ -2853,3 +2853,452 @@ async fn test_wrapper_only_system_fields() { displayable(physical_plan.as_ref()).indent() ); } + +/// Pivot SQL from a query builder: several conditional aggregations over one +/// fan-out join must push down as a single grouped Cube query, not as an +/// ungrouped scan aggregated in memory. +#[tokio::test] +async fn test_wrapper_conditional_aggregation_over_join() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + SELECT k.customer_gender AS "gender", + MAX(CASE WHEN l.content = 'PropA' THEN l.read END) AS "PropA", + MAX(CASE WHEN l.content = 'PropB' THEN l.read END) AS "PropB" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + WHERE k.customer_gender = 'female' + GROUP BY 1 + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let request = logical_plan.find_cube_scan_wrapped_sql().request; + assert_eq!( + request.ungrouped, None, + "query is grouped in Cube: {:?}", + request + ); + let measures = request.measures.unwrap_or_default(); + assert_eq!( + measures.len(), + 2, + "both conditional aggregations are pushed down as measures: {:?}", + measures + ); + assert!( + measures + .iter() + .all(|measure| measure.contains("MAX(CASE WHEN")), + "measures keep the conditional aggregation: {:?}", + measures + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +} + +/// The same pivot spread over CTEs joined together: the whole query, including the +/// join between the CTEs, must be pushed down into a single Cube query. +#[tokio::test] +async fn test_wrapper_conditional_aggregation_multi_cte() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + WITH t_root AS ( + SELECT k.customer_gender AS "gender", + MAX(CASE WHEN l.content = 'PropA' THEN l.read END) AS "PropA" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + WHERE k.customer_gender = 'female' + GROUP BY 1 + ), + t_prop_b AS ( + SELECT k.customer_gender AS "__j_gender", + MAX(l.content) AS "PropB" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + WHERE l.content = 'PropB' AND k.customer_gender = 'female' + GROUP BY 1 + ) + SELECT t_root."gender", t_root."PropA", t_prop_b."PropB" + FROM t_root + LEFT JOIN t_prop_b ON t_prop_b."__j_gender" = t_root."gender" + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!( + sql.contains("MAX(CASE WHEN"), + "conditional aggregation is pushed down:\n{}", + sql + ); + // The row-preserving side of the LEFT JOIN must stay in `from`, with the other CTE + // joined to it - swapping the sides would drop unmatched rows + let from_position = sql.find(r#") AS "t_root""#).expect(&sql); + let join_position = sql.find(r#"LEFT JOIN"#).expect(&sql); + assert!( + from_position < join_position, + "left CTE stays the from of the LEFT JOIN:\n{}", + sql + ); + assert!( + sql.contains(r#") AS "t_prop_b" ON ("t_root"."gender" = "t_prop_b"."j_gender")"#), + "right CTE is joined as a subquery on the original condition:\n{}", + sql + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + let plan = format!("{}", displayable(physical_plan.as_ref()).indent()); + assert_eq!( + plan.matches("CubeScanExecutionPlan").count(), + 1, + "the join between the CTEs is executed by the data source, not in memory:\n{}", + plan + ); +} + +/// A pivot query builder emits one CTE per property, all joined to the root CTE. +/// Every join in that chain must be pushed down, not just the first one. +#[tokio::test] +async fn test_wrapper_conditional_aggregation_multi_cte_chain() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + WITH t_root AS ( + SELECT k.customer_gender AS "gender", + MAX(CASE WHEN l.content = 'PropA' THEN l.read END) AS "PropA" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + GROUP BY 1 + ), + t_prop_b AS ( + SELECT k.customer_gender AS "__j_b", MAX(l.content) AS "PropB" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + WHERE l.content = 'PropB' + GROUP BY 1 + ), + t_prop_c AS ( + SELECT k.customer_gender AS "__j_c", MAX(l.content) AS "PropC" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + WHERE l.content = 'PropC' + GROUP BY 1 + ) + SELECT t_root."gender", t_root."PropA", t_prop_b."PropB", t_prop_c."PropC" + FROM t_root + LEFT JOIN t_prop_b ON t_prop_b."__j_b" = t_root."gender" + LEFT JOIN t_prop_c ON t_prop_c."__j_c" = t_root."gender" + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + for joined in [ + r#"AS "t_prop_b" ON ("t_root"."gender" = "t_prop_b"."j_b")"#, + r#"AS "t_prop_c" ON ("t_root"."gender" = "t_prop_c"."j_c")"#, + ] { + assert!( + sql.contains(joined), + "both CTE joins are pushed down, missing {}:\n{}", + joined, + sql + ); + } + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + let plan = format!("{}", displayable(physical_plan.as_ref()).indent()); + assert_eq!( + plan.matches("CubeScanExecutionPlan").count(), + 1, + "a chain of CTE joins becomes a single Cube query:\n{}", + plan + ); +} + +/// A join condition can reference a CTE joined earlier in the query, so pushed-down joins +/// must keep the order they had: a subquery can only be referenced after it is joined. +#[tokio::test] +async fn test_wrapper_grouped_join_chain_keeps_join_order() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + WITH t_root AS ( + SELECT k.customer_gender AS "gender", MAX(l.content) AS "a" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + GROUP BY 1 + ), + t_b AS ( + SELECT k.customer_gender AS "jb", MAX(l.content) AS "b" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + GROUP BY 1 + ), + t_c AS ( + SELECT k.customer_gender AS "jc", MIN(l.content) AS "c" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + GROUP BY 1 + ) + SELECT t_root."gender", t_b."b", t_c."c" + FROM t_root + LEFT JOIN t_b ON t_b."jb" = t_root."gender" + LEFT JOIN t_c ON t_c."jc" = t_b."b" + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + let joined_b = sql + .find(r#") AS "t_b" ON ("t_root"."gender" = "t_b"."jb")"#) + .expect(&sql); + let joined_c = sql + .find(r#") AS "t_c" ON ("t_b"."b" = "t_c"."jc")"#) + .expect(&sql); + assert!( + joined_b < joined_c, + "t_b is joined before the condition that references it:\n{}", + sql + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + let plan = format!("{}", displayable(physical_plan.as_ref()).indent()); + assert_eq!( + plan.matches("CubeScanExecutionPlan").count(), + 1, + "the whole chain is one Cube query:\n{}", + plan + ); +} + +/// A pushed-down join of grouped subqueries is not unique on its join keys, so it must not be +/// handed to a Cube query as a subquery join: that rendering counts or deduplicates join fanout +/// depending on how a measure happens to be classified. Refusing to plan the query is the +/// intended outcome - an error is preferable to numbers that depend on the data model. +#[tokio::test] +async fn test_wrapper_grouped_join_is_not_used_as_cube_subquery_join() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query = r#" + WITH m1 AS ( + SELECT customer_gender AS g, sum(sumPrice) AS s + FROM KibanaSampleDataEcommerce GROUP BY 1 + ), + m2 AS ( + SELECT customer_gender AS g2, avg(avgPrice) AS v + FROM KibanaSampleDataEcommerce GROUP BY 1 + ), + joined AS ( + SELECT m2.v AS k, m1.g AS g FROM m1 LEFT JOIN m2 ON m2.g2 = m1.g + ) + SELECT k.customer_gender, MEASURE(k.avgPrice) AS p + FROM KibanaSampleDataEcommerce k + LEFT JOIN joined ON joined.k = k.customer_gender + GROUP BY 1 + "# + .to_string(); + + let meta = crate::compile::test::get_test_tenant_ctx(); + let session = + crate::compile::test::get_test_session(DatabaseProtocol::PostgreSQL, meta.clone()).await; + let query_plan = crate::compile::test::convert_sql_to_cube_query(&query, meta, session).await; + + // Today the query is refused outright. Planning it some other way would be fine too, as + // long as the join of subqueries does not end up as a subquery join of a Cube query - so + // check that, rather than only checking that something failed. + match query_plan { + Err(error) => { + println!("Refused at compile time: {}", error); + } + Ok(query_plan) => { + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + let plan = format!("{}", displayable(physical_plan.as_ref()).indent()); + panic!( + "expected the query to be refused, it planned instead{}:\n{}", + if plan.contains("subqueryJoins") { + " - with the join of subqueries sent to Cube as a subquery join" + } else { + "" + }, + plan + ); + } + } +} + +/// Cube measures can not be computed over a pushed-down join: the aggregation has to stay +/// outside the wrapper, where it fails with the explicit MEASURE error, rather than be folded +/// into a Cube query that would give it aggregation semantics it does not have. +#[tokio::test] +async fn test_wrapper_no_measure_over_grouped_join_chain() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + WITH t_root AS ( + SELECT k.customer_gender AS "g", MAX(l.content) AS "a" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + GROUP BY 1 + ), + t_b AS ( + SELECT k.customer_gender AS "jb", MIN(l.content) AS "b" + FROM KibanaSampleDataEcommerce k + LEFT JOIN Logs l ON k.__cubeJoinField = l.__cubeJoinField + GROUP BY 1 + ) + SELECT t_root."g", MEASURE(t_b."b") AS m + FROM t_root + LEFT JOIN t_b ON t_b."jb" = t_root."g" + GROUP BY 1 + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + assert!( + logical_plan.try_expect_root_cube_scan().is_none(), + "MEASURE() over a pushed-down join is not wrapped:\n{}", + logical_plan.display_indent() + ); +} + +/// A LIMIT between two joins belongs to the join below it. The second join must go on top of +/// the limited select, never into it, or it would join before the rows are picked. +#[tokio::test] +async fn test_wrapper_grouped_join_chain_keeps_limit_between_joins() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + WITH m1 AS ( + SELECT customer_gender AS g, sum(sumPrice) AS s + FROM KibanaSampleDataEcommerce GROUP BY 1 + ), + m2 AS ( + SELECT customer_gender AS g2, avg(avgPrice) AS v + FROM KibanaSampleDataEcommerce GROUP BY 1 + ), + m3 AS ( + SELECT customer_gender AS g3, count(count) AS w + FROM KibanaSampleDataEcommerce GROUP BY 1 + ), + limited AS ( + SELECT m1.g AS g, m2.v AS v FROM m1 LEFT JOIN m2 ON m2.g2 = m1.g LIMIT 5 + ) + SELECT limited.g, limited.v, m3.w + FROM limited + LEFT JOIN m3 ON m3.g3 = limited.g + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + let limit = sql.find("LIMIT 5").expect(&sql); + let second_join = sql + .find(r#"AS "m3" ON ("limited"."g" = "m3"."g3")"#) + .expect(&sql); + assert!( + limit < second_join, + "the second join is applied to the limited rows, not inside the limit:\n{}", + sql + ); +} + +/// Same column names on both sides of a pushed-down join must keep distinct aliases, +/// or the outer projection would read the same column twice. +#[tokio::test] +async fn test_wrapper_grouped_join_wrapped_left_duplicate_names() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan( + r#" + WITH m1 AS ( + SELECT customer_gender AS g, MAX(taxful_total_price) AS v + FROM KibanaSampleDataEcommerce + GROUP BY 1 + ), + m2 AS ( + SELECT customer_gender AS g, MAX(minPrice) AS v + FROM KibanaSampleDataEcommerce + GROUP BY 1 + ) + SELECT COALESCE(m1.g, m2.g) AS g, m1.v AS v1, m2.v AS v2 + FROM m1 + LEFT JOIN m2 ON m1.g = m2.g + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!( + sql.contains(r#"COALESCE("m1"."g", "m1"."g_1")"#), + "join sides keep distinct aliases:\n{}", + sql + ); + assert!( + sql.contains(r#""m1"."v" "v1", "m1"."v_1" "v2""#), + "same-named measures from both sides stay distinct:\n{}", + sql + ); + + let physical_plan = query_plan.as_physical_plan().await.unwrap(); + println!( + "Physical plan: {}", + displayable(physical_plan.as_ref()).indent() + ); +}