Skip to content

Commit addde0d

Browse files
paveltiunovpaveltiunov
authored
fix(cubesql): Prioritize ungrouped aggregate scans over ungrouped projection scans so most of the members can be pushed down without wrapping (#7865)
1 parent 82217b2 commit addde0d

File tree

11 files changed

+105
-8
lines changed

11 files changed

+105
-8
lines changed

‎rust/cubesql/cubesql/egraph-debug-template/src/index.js‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ const toRegularNode = (n) => ({
2828
const toEdge = (n) => ({
2929
...n,
3030
id: `${n.source}->${n.target}`,
31-
style: n.source.match(new RegExp(`^${n.target}-`)) ? { stroke: '#f00', 'stroke-width': 10 } : undefined
31+
style: n.source.indexOf(`${n.target}-`) === 0 ? { stroke: '#f00', 'stroke-width': 10 } : undefined
3232
});
3333
const initialNodes = data.combos.map(toGroupNode).concat(data.nodes.map(toRegularNode));
3434
const initialEdges = data.edges.map(toEdge);

‎rust/cubesql/cubesql/src/compile/mod.rs‎

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4705,12 +4705,18 @@ from
47054705
)
47064706
.await;
47074707

4708+
let physical_plan = query_plan.as_physical_plan().await.unwrap();
4709+
println!(
4710+
"Physical plan: {}",
4711+
displayable(physical_plan.as_ref()).indent()
4712+
);
4713+
47084714
let logical_plan = query_plan.as_logical_plan();
47094715
assert_eq!(
47104716
logical_plan.find_cube_scan().request,
47114717
V1LoadRequestQuery {
47124718
measures: Some(vec![]),
4713-
dimensions: Some(vec![]),
4719+
dimensions: Some(vec!["KibanaSampleDataEcommerce.order_date".to_string()]),
47144720
segments: Some(vec![]),
47154721
time_dimensions: None,
47164722
order: None,
@@ -18849,7 +18855,13 @@ limit
1884918855
order: None,
1885018856
limit: None,
1885118857
offset: None,
18852-
filters: None,
18858+
filters: Some(vec![V1LoadRequestQueryFilterItem {
18859+
member: Some("KibanaSampleDataEcommerce.order_date".to_string(),),
18860+
operator: Some("beforeDate".to_string(),),
18861+
values: Some(vec!["2014-05-31T00:00:00.000Z".to_string()]),
18862+
or: None,
18863+
and: None,
18864+
}]),
1885318865
ungrouped: Some(true),
1885418866
}
1885518867
)
@@ -19961,6 +19973,44 @@ limit
1996119973
);
1996219974
}
1996319975

19976+
#[tokio::test]
19977+
async fn test_case_wrapper_distinct() {
19978+
if !Rewriter::sql_push_down_enabled() {
19979+
return;
19980+
}
19981+
init_logger();
19982+
19983+
let query_plan = convert_select_to_query_plan(
19984+
r#"SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END, COUNT(DISTINCT countDistinct) mp
19985+
FROM KibanaSampleDataEcommerce a
19986+
WHERE
19987+
(
19988+
(
19989+
( a.order_date ) >= '2024-01-01'
19990+
AND ( a.order_date ) < '2024-02-01'
19991+
)
19992+
)
19993+
GROUP BY 1"#
19994+
.to_string(),
19995+
DatabaseProtocol::PostgreSQL,
19996+
)
19997+
.await;
19998+
19999+
let logical_plan = query_plan.as_logical_plan();
20000+
assert!(logical_plan
20001+
.find_cube_scan_wrapper()
20002+
.wrapped_sql
20003+
.unwrap()
20004+
.sql
20005+
.contains("CASE WHEN"));
20006+
20007+
let physical_plan = query_plan.as_physical_plan().await.unwrap();
20008+
println!(
20009+
"Physical plan: {}",
20010+
displayable(physical_plan.as_ref()).indent()
20011+
);
20012+
}
20013+
1996420014
#[tokio::test]
1996520015
async fn test_case_wrapper_alias_with_order() {
1996620016
if !Rewriter::sql_push_down_enabled() {

‎rust/cubesql/cubesql/src/compile/rewrite/cost.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use crate::{
33
rewrite::{
44
rules::utils::granularity_str_to_int_order, CubeScanUngrouped, CubeScanWrapped,
55
DimensionName, LogicalPlanLanguage, MemberErrorPriority, ScalarUDFExprFun,
6-
TimeDimensionGranularity,
6+
TimeDimensionGranularity, WrappedSelectUngroupedScan,
77
},
88
MetaContext,
99
},
@@ -44,6 +44,7 @@ pub struct CubePlanCost {
4444
non_pushed_down_window: i64,
4545
ungrouped_aggregates: usize,
4646
wrapper_nodes: i64,
47+
wrapped_select_ungrouped_scan: usize,
4748
ast_size_outside_wrapper: usize,
4849
filters: i64,
4950
structure_points: i64,
@@ -134,6 +135,8 @@ impl CubePlanCost {
134135
+ other.ast_size_outside_wrapper,
135136
ungrouped_aggregates: self.ungrouped_aggregates + other.ungrouped_aggregates,
136137
wrapper_nodes: self.wrapper_nodes + other.wrapper_nodes,
138+
wrapped_select_ungrouped_scan: self.wrapped_select_ungrouped_scan
139+
+ other.wrapped_select_ungrouped_scan,
137140
cube_scan_nodes: self.cube_scan_nodes + other.cube_scan_nodes,
138141
time_dimensions_used_as_dimensions: self.time_dimensions_used_as_dimensions
139142
+ other.time_dimensions_used_as_dimensions,
@@ -197,6 +200,7 @@ impl CubePlanCost {
197200
CubePlanState::Wrapper => 0,
198201
} + self.ungrouped_aggregates,
199202
wrapper_nodes: self.wrapper_nodes,
203+
wrapped_select_ungrouped_scan: self.wrapped_select_ungrouped_scan,
200204
cube_scan_nodes: self.cube_scan_nodes,
201205
ast_size_without_alias: self.ast_size_without_alias,
202206
ast_size: self.ast_size,
@@ -364,6 +368,11 @@ impl CostFunction<LogicalPlanLanguage> for BestCubePlan {
364368
_ => 0,
365369
};
366370

371+
let wrapped_select_ungrouped_scan = match enode {
372+
LogicalPlanLanguage::WrappedSelectUngroupedScan(WrappedSelectUngroupedScan(true)) => 1,
373+
_ => 0,
374+
};
375+
367376
let initial_cost = CubePlanCostAndState {
368377
cost: CubePlanCost {
369378
replacers: this_replacers,
@@ -380,6 +389,7 @@ impl CostFunction<LogicalPlanLanguage> for BestCubePlan {
380389
structure_points,
381390
ungrouped_aggregates: 0,
382391
wrapper_nodes,
392+
wrapped_select_ungrouped_scan,
383393
empty_wrappers: 0,
384394
ast_size_outside_wrapper: 0,
385395
ast_size_inside_wrapper,

‎rust/cubesql/cubesql/src/compile/rewrite/mod.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,7 @@ crate::plan_to_language! {
268268
order_expr: Vec<Expr>,
269269
alias: Option<String>,
270270
ungrouped: bool,
271+
ungrouped_scan: bool,
271272
},
272273
WrappedSelectJoin {
273274
input: Arc<LogicalPlan>,
@@ -788,9 +789,10 @@ fn wrapped_select(
788789
order_expr: impl Display,
789790
alias: impl Display,
790791
ungrouped: impl Display,
792+
ungrouped_scan: impl Display,
791793
) -> String {
792794
format!(
793-
"(WrappedSelect {} {} {} {} {} {} {} {} {} {} {} {} {} {})",
795+
"(WrappedSelect {} {} {} {} {} {} {} {} {} {} {} {} {} {} {})",
794796
select_type,
795797
projection_expr,
796798
group_expr,
@@ -805,6 +807,7 @@ fn wrapped_select(
805807
order_expr,
806808
alias,
807809
ungrouped,
810+
ungrouped_scan
808811
)
809812
}
810813

‎rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/aggregate.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ impl WrapperRules {
9898
),
9999
"WrappedSelectAlias:None",
100100
"?select_ungrouped",
101+
"WrappedSelectUngroupedScan:false",
101102
),
102103
"CubeScanWrapperFinalized:false",
103104
),

‎rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/filter.rs‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use crate::{
77
wrapped_select_joins_empty_tail, wrapped_select_order_expr_empty_tail,
88
wrapped_select_projection_expr_empty_tail, wrapped_select_window_expr_empty_tail,
99
wrapper_pullup_replacer, wrapper_pushdown_replacer, LogicalPlanLanguage,
10-
WrappedSelectUngrouped, WrapperPullupReplacerUngrouped,
10+
WrappedSelectUngrouped, WrappedSelectUngroupedScan, WrapperPullupReplacerUngrouped,
1111
},
1212
var, var_iter,
1313
};
@@ -191,10 +191,11 @@ impl WrapperRules {
191191
),
192192
"WrappedSelectAlias:None",
193193
"?select_ungrouped",
194+
"?select_ungrouped_scan",
194195
),
195196
"CubeScanWrapperFinalized:false",
196197
),
197-
self.transform_filter("?ungrouped", "?select_ungrouped"),
198+
self.transform_filter("?ungrouped", "?select_ungrouped", "?select_ungrouped_scan"),
198199
)]);
199200

200201
Self::list_pushdown_pullup_rules(
@@ -209,9 +210,11 @@ impl WrapperRules {
209210
&self,
210211
ungrouped_var: &'static str,
211212
select_ungrouped_var: &'static str,
213+
select_ungrouped_scan_var: &'static str,
212214
) -> impl Fn(&mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>, &mut Subst) -> bool {
213215
let ungrouped_var = var!(ungrouped_var);
214216
let select_ungrouped_var = var!(select_ungrouped_var);
217+
let select_ungrouped_scan_var = var!(select_ungrouped_scan_var);
215218
move |egraph, subst| {
216219
for ungrouped in
217220
var_iter!(egraph[subst[ungrouped_var]], WrapperPullupReplacerUngrouped).cloned()
@@ -222,6 +225,13 @@ impl WrapperRules {
222225
WrappedSelectUngrouped(ungrouped),
223226
)),
224227
);
228+
229+
subst.insert(
230+
select_ungrouped_scan_var,
231+
egraph.add(LogicalPlanLanguage::WrappedSelectUngroupedScan(
232+
WrappedSelectUngroupedScan(ungrouped),
233+
)),
234+
);
225235
return true;
226236
}
227237
false

‎rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/limit.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ impl WrapperRules {
3232
"?order_expr",
3333
"?select_alias",
3434
"?select_ungrouped",
35+
"?select_ungrouped_scan",
3536
),
3637
"?alias_to_cube",
3738
"?ungrouped",
@@ -58,6 +59,7 @@ impl WrapperRules {
5859
"?order_expr",
5960
"?select_alias",
6061
"?select_ungrouped",
62+
"?select_ungrouped_scan",
6163
),
6264
"?alias_to_cube",
6365
"?ungrouped",

‎rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/order.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ impl WrapperRules {
2828
wrapped_select_order_expr_empty_tail(),
2929
"?select_alias",
3030
"?select_ungrouped",
31+
"?select_ungrouped_scan",
3132
),
3233
"?alias_to_cube",
3334
"?ungrouped",
@@ -95,6 +96,7 @@ impl WrapperRules {
9596
),
9697
"?select_alias",
9798
"?select_ungrouped",
99+
"?select_ungrouped_scan",
98100
),
99101
"CubeScanWrapperFinalized:false",
100102
),

‎rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/projection.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@ use crate::{
66
wrapped_select_having_expr_empty_tail, wrapped_select_joins_empty_tail,
77
wrapped_select_order_expr_empty_tail, wrapped_select_window_expr_empty_tail,
88
wrapper_pullup_replacer, wrapper_pushdown_replacer, LogicalPlanLanguage, ProjectionAlias,
9-
WrappedSelectAlias, WrappedSelectUngrouped, WrapperPullupReplacerUngrouped,
9+
WrappedSelectAlias, WrappedSelectUngrouped, WrappedSelectUngroupedScan,
10+
WrapperPullupReplacerUngrouped,
1011
},
1112
var, var_iter,
1213
};
@@ -92,6 +93,7 @@ impl WrapperRules {
9293
),
9394
"?select_alias",
9495
"?select_ungrouped",
96+
"?select_ungrouped_scan",
9597
),
9698
"CubeScanWrapperFinalized:false",
9799
),
@@ -100,6 +102,7 @@ impl WrapperRules {
100102
"?ungrouped",
101103
"?select_alias",
102104
"?select_ungrouped",
105+
"?select_ungrouped_scan",
103106
),
104107
)]);
105108

@@ -117,11 +120,13 @@ impl WrapperRules {
117120
ungrouped_var: &'static str,
118121
select_alias_var: &'static str,
119122
select_ungrouped_var: &'static str,
123+
select_ungrouped_scan_var: &'static str,
120124
) -> impl Fn(&mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>, &mut Subst) -> bool {
121125
let projection_alias_var = var!(projection_alias_var);
122126
let ungrouped_var = var!(ungrouped_var);
123127
let select_alias_var = var!(select_alias_var);
124128
let select_ungrouped_var = var!(select_ungrouped_var);
129+
let select_ungrouped_scan_var = var!(select_ungrouped_scan_var);
125130
move |egraph, subst| {
126131
for projection_alias in
127132
var_iter!(egraph[subst[projection_alias_var]], ProjectionAlias).cloned()
@@ -135,6 +140,12 @@ impl WrapperRules {
135140
WrappedSelectUngrouped(ungrouped),
136141
)),
137142
);
143+
subst.insert(
144+
select_ungrouped_scan_var,
145+
egraph.add(LogicalPlanLanguage::WrappedSelectUngroupedScan(
146+
WrappedSelectUngroupedScan(ungrouped),
147+
)),
148+
);
138149
subst.insert(
139150
select_alias_var,
140151
egraph.add(LogicalPlanLanguage::WrappedSelectAlias(WrappedSelectAlias(

‎rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/window.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ impl WrapperRules {
2727
"?order_expr",
2828
"?select_alias",
2929
"?select_ungrouped",
30+
"?select_ungrouped_scan",
3031
),
3132
"?alias_to_cube",
3233
"?ungrouped",
@@ -95,6 +96,7 @@ impl WrapperRules {
9596
),
9697
"?select_alias",
9798
"?select_ungrouped",
99+
"?select_ungrouped_scan",
98100
),
99101
"CubeScanWrapperFinalized:false",
100102
),

0 commit comments

Comments
 (0)