Skip to content

Commit

Permalink
feat(cubesql): Add separate ungrouped_scan flag to wrapper replacer c…
Browse files Browse the repository at this point in the history 
…ontext

It is used to track whether this query actually wraps ungrouped scan in from as opposed to old flag, which is used as push to cube enabler

Changes in cost are necessary because now that it is tracked, we are getting proper values in `wrapped_select_ungrouped_scan` cost component, and it turns any ungrouped scan under wrapper more expensive than plan with same amount of wrappers (usually it's just 1 anyway), but with more nodes outside wrapper
  • Loading branch information
@mcheshkov
mcheshkov committed Jan 27, 2025
1 parent 20bfbb8 commit 24a04dc
Show file tree
Hide file tree
Showing 23 changed files with 266 additions and 56 deletions.
2 changes: 1 addition & 1 deletion rust/cubesql/cubesql/src/compile/rewrite/cost.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -255,8 +255,8 @@ pub struct CubePlanCost {
non_pushed_down_limit_sort: i64,
joins: usize,
wrapper_nodes: i64,
wrapped_select_ungrouped_scan: usize,
ast_size_outside_wrapper: usize,
wrapped_select_ungrouped_scan: usize,
filters: i64,
structure_points: i64,
filter_members: i64,
Expand Down
10 changes: 9 additions & 1 deletion rust/cubesql/cubesql/src/compile/rewrite/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -469,6 +469,13 @@ crate::plan_to_language! {
// Known qualifiers of grouped subqueries
// Used to allow to rewrite columns from them even with push to Cube enabled
grouped_subqueries: Vec<String>,
// When `member` is logical plan this means it is actually ungrouped, even when push_to_cube is disabled.
// When `member` is expression it just acts as a pull-through from pushdown.
// It will be filled by every wrapper replacer producer rule, essentially same way as
// ungrouped_scan flag in wrapped_select is filled:
// fixed false for aggregation, copy inner value for projection.
// This flag should make roundtrip from top to bottom and back.
ungrouped_scan: bool,
},
WrapperPushdownReplacer {
member: Arc<LogicalPlan>,
Expand Down Expand Up @@ -1974,9 +1981,10 @@ fn wrapper_replacer_context(
in_projection: impl Display,
cube_members: impl Display,
grouped_subqueries: impl Display,
ungrouped_scan: impl Display,
) -> String {
format!(
"(WrapperReplacerContext {alias_to_cube} {push_to_cube} {in_projection} {cube_members} {grouped_subqueries})",
"(WrapperReplacerContext {alias_to_cube} {push_to_cube} {in_projection} {cube_members} {grouped_subqueries} {ungrouped_scan})",
)
}

Expand Down
44 changes: 44 additions & 0 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs
View file

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate_function.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
)],
"?distinct",
Expand All @@ -51,6 +52,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
self.transform_agg_fun_expr("?fun", "?distinct", "?alias_to_cube"),
Expand Down
3 changes: 3 additions & 0 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/binary_expr.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
"?op",
Expand All @@ -44,6 +45,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
),
Expand All @@ -55,6 +57,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
self.transform_binary_expr("?op", "?alias_to_cube"),
Expand Down
4 changes: 4 additions & 0 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/case.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
wrapper_pullup_replacer(
Expand All @@ -43,6 +44,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
wrapper_pullup_replacer(
Expand All @@ -53,6 +55,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
),
Expand All @@ -64,6 +67,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
self.transform_case_expr("?alias_to_cube"),
Expand Down
6 changes: 6 additions & 0 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/column.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
wrapper_pullup_replacer(
Expand All @@ -35,6 +36,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
),
Expand All @@ -50,6 +52,7 @@ impl WrapperRules {
"WrapperReplacerContextInProjection:true",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
wrapper_pullup_replacer(
Expand All @@ -60,6 +63,7 @@ impl WrapperRules {
"WrapperReplacerContextInProjection:true",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
self.pushdown_simple_measure("?name", "?cube_members"),
Expand All @@ -75,6 +79,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
wrapper_pullup_replacer(
Expand All @@ -85,6 +90,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
self.pushdown_dimension(
Expand Down
19 changes: 17 additions & 2 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/cube_scan_wrapper.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ use crate::{
transforming_rewrite, wrapper_pullup_replacer, wrapper_replacer_context,
CubeScanAliasToCube, CubeScanLimit, CubeScanOffset, CubeScanUngrouped, LogicalPlanLanguage,
WrapperReplacerContextAliasToCube, WrapperReplacerContextGroupedSubqueries,
WrapperReplacerContextPushToCube,
WrapperReplacerContextPushToCube, WrapperReplacerContextUngroupedScan,
},
var, var_iter,
copy_flag, var, var_iter,
};
use egg::Subst;

Expand Down Expand Up @@ -49,6 +49,7 @@ impl WrapperRules {
"WrapperReplacerContextInProjection:false",
"?members",
"?grouped_subqueries_out",
"?ungrouped_scan_out",
),
),
"CubeScanWrapperFinalized:false",
Expand All @@ -62,6 +63,7 @@ impl WrapperRules {
"?alias_to_cube_out",
"?push_to_cube_out",
"?grouped_subqueries_out",
"?ungrouped_scan_out",
),
),
rewrite(
Expand All @@ -85,6 +87,7 @@ impl WrapperRules {
alias_to_cube_var_out: &'static str,
push_to_cube_out_var: &'static str,
grouped_subqueries_out_var: &'static str,
ungrouped_scan_out_var: &'static str,
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
let members_var = var!(members_var);
let alias_to_cube_var = var!(alias_to_cube_var);
Expand All @@ -94,6 +97,7 @@ impl WrapperRules {
let alias_to_cube_var_out = var!(alias_to_cube_var_out);
let push_to_cube_out_var = var!(push_to_cube_out_var);
let grouped_subqueries_out_var = var!(grouped_subqueries_out_var);
let ungrouped_scan_out_var = var!(ungrouped_scan_out_var);
move |egraph, subst| {
let mut has_no_limit_or_offset = true;
for limit in var_iter!(egraph[subst[limit_var]], CubeScanLimit).cloned() {
Expand All @@ -103,6 +107,17 @@ impl WrapperRules {
has_no_limit_or_offset &= offset.is_none();
}

if !copy_flag!(
egraph,
subst,
ungrouped_cube_var,
CubeScanUngrouped,
ungrouped_scan_out_var,
WrapperReplacerContextUngroupedScan
) {
return false;
}

if let Some(_) = egraph[subst[members_var]].data.member_name_to_expr {
for alias_to_cube in
var_iter!(egraph[subst[alias_to_cube_var]], CubeScanAliasToCube).cloned()
Expand Down
3 changes: 3 additions & 0 deletions rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/extract.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
wrapper_pullup_replacer(
Expand All @@ -35,6 +36,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
],
Expand All @@ -50,6 +52,7 @@ impl WrapperRules {
"?in_projection",
"?cube_members",
"?grouped_subqueries",
"?ungrouped_scan",
),
),
self.transform_date_part_expr("?alias_to_cube"),
Expand Down
Loading

0 comments on commit 24a04dc

Please sign in to comment.