Skip to content

Commit 513a9d8

Browse files
mcheshkovmcheshkov
committed
feat(cubesql): Allow to use CubeScan(ungrouped=true) in wrappers without push to Cube
* New wrapped_no_push_to_cube cost component, necessary when two wrappers does not differ at all, only by push_to_cube flag
1 parent 77c3634 commit 513a9d8

File tree

4 files changed

+108
-16
lines changed

4 files changed

+108
-16
lines changed

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

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2323,7 +2323,7 @@ from
23232323
assert_eq!(
23242324
logical_plan.find_cube_scan().request,
23252325
V1LoadRequestQuery {
2326-
measures: Some(vec![]),
2326+
measures: Some(vec!["KibanaSampleDataEcommerce.sumPrice".to_string(),]),
23272327
dimensions: Some(vec![]),
23282328
segments: Some(vec![]),
23292329
order: Some(vec![]),
@@ -2750,7 +2750,7 @@ limit
27502750

27512751
let query_plan = convert_select_to_query_plan_with_meta(
27522752
r#"
2753-
SELECT MIN(a.sixteen_charchar), MAX(a.sixteen_charchar_foo), MAX(a.sixteen_charchar_bar) FROM (SELECT * FROM SixteenChar) a
2753+
SELECT MIN(a.sixteen_charchar), MAX(a.sixteen_charchar_foo), MAX(a.sixteen_charchar_bar), MAX(a.sixteen_charchar_baz) FROM (SELECT * FROM SixteenChar) a
27542754
"#
27552755
.to_string(),
27562756
get_sixteen_char_member_cube(),
@@ -2766,7 +2766,12 @@ limit
27662766
assert_eq!(
27672767
query_plan.as_logical_plan().find_cube_scan().request,
27682768
V1LoadRequestQuery {
2769-
measures: Some(vec![]),
2769+
measures: Some(vec![
2770+
"SixteenChar.sixteen_charchar".to_string(),
2771+
"SixteenChar.sixteen_charchar_foo".to_string(),
2772+
"SixteenChar.sixteen_charchar_bar".to_string(),
2773+
"SixteenChar.sixteen_charchar_baz".to_string(),
2774+
]),
27702775
dimensions: Some(vec![]),
27712776
segments: Some(vec![]),
27722777
order: Some(vec![]),
@@ -2781,15 +2786,15 @@ limit
27812786
.wrapped_sql
27822787
.unwrap()
27832788
.sql
2784-
.contains("sixteen_charchar_1"));
2789+
.contains("\"max_a_sixteen_ch_1\""));
27852790

27862791
assert!(query_plan
27872792
.as_logical_plan()
27882793
.find_cube_scan_wrapper()
27892794
.wrapped_sql
27902795
.unwrap()
27912796
.sql
2792-
.contains("sixteen_charchar_2"));
2797+
.contains("\"max_a_sixteen_ch_2\""));
27932798
}
27942799

27952800
#[tokio::test]
@@ -9569,11 +9574,7 @@ ORDER BY "source"."str0" ASC
95699574
assert_eq!(
95709575
query_plan.as_logical_plan().find_cube_scan().request,
95719576
V1LoadRequestQuery {
9572-
measures: Some(vec![
9573-
"Logs.agentCount".to_string(),
9574-
"Logs.agentCountApprox".to_string(),
9575-
"KibanaSampleDataEcommerce.count".to_string()
9576-
]),
9577+
measures: Some(vec!["Logs.agentCountApprox".to_string(),]),
95779578
dimensions: Some(vec![
95789579
"KibanaSampleDataEcommerce.taxful_total_price".to_string(),
95799580
]),
@@ -13732,11 +13733,7 @@ ORDER BY "source"."str0" ASC
1373213733
.wrapped_sql
1373313734
.unwrap()
1373413735
.sql;
13735-
if Rewriter::top_down_extractor_enabled() {
13736-
assert!(sql.contains("LIMIT 1000"));
13737-
} else {
13738-
assert!(sql.contains("\"limit\":1000"));
13739-
}
13736+
assert!(sql.contains("LIMIT 1000"));
1374013737
assert!(sql.contains("% 7"));
1374113738

1374213739
let physical_plan = query_plan.as_physical_plan().await.unwrap();

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

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use crate::{
44
compile::rewrite::{
55
rules::utils::granularity_str_to_int_order, CubeScanUngrouped, CubeScanWrapped,
66
DimensionName, LogicalPlanLanguage, MemberErrorPriority, ScalarUDFExprFun,
7-
TimeDimensionGranularity, WrappedSelectUngroupedScan,
7+
TimeDimensionGranularity, WrappedSelectPushToCube, WrappedSelectUngroupedScan,
88
},
99
transport::{MetaContext, V1CubeMetaDimensionExt},
1010
};
@@ -186,6 +186,11 @@ impl BestCubePlan {
186186
_ => 0,
187187
};
188188

189+
let wrapped_no_push_to_cube = match enode {
190+
LogicalPlanLanguage::WrappedSelectPushToCube(WrappedSelectPushToCube(false)) => 1,
191+
_ => 0,
192+
};
193+
189194
let unwrapped_subqueries = match enode {
190195
LogicalPlanLanguage::Subquery(_) => 1,
191196
_ => 0,
@@ -218,6 +223,7 @@ impl BestCubePlan {
218223
ast_size: 1,
219224
ungrouped_nodes,
220225
unwrapped_subqueries,
226+
wrapped_no_push_to_cube,
221227
}
222228
}
223229
}
@@ -267,6 +273,10 @@ pub struct CubePlanCost {
267273
ast_size: usize,
268274
ast_size_inside_wrapper: usize,
269275
ungrouped_nodes: usize,
276+
// This is necessary for queries like `SELECT AVG(avgMeasure) FROM cube;`
277+
// Query like that can be represented by WrappedSelect with both push_to_cube=true and push_to_cube=false
278+
// And this component is the only thing to differ those
279+
wrapped_no_push_to_cube: usize,
270280
}
271281

272282
#[derive(Debug, Clone, Eq, Hash, PartialEq)]
@@ -387,6 +397,7 @@ impl CubePlanCost {
387397
ast_size_inside_wrapper: self.ast_size_inside_wrapper + other.ast_size_inside_wrapper,
388398
ungrouped_nodes: self.ungrouped_nodes + other.ungrouped_nodes,
389399
unwrapped_subqueries: self.unwrapped_subqueries + other.unwrapped_subqueries,
400+
wrapped_no_push_to_cube: self.wrapped_no_push_to_cube + other.wrapped_no_push_to_cube,
390401
}
391402
}
392403

@@ -465,6 +476,7 @@ impl CubePlanCost {
465476
ast_size: self.ast_size,
466477
ast_size_inside_wrapper: self.ast_size_inside_wrapper,
467478
ungrouped_nodes: self.ungrouped_nodes,
479+
wrapped_no_push_to_cube: self.wrapped_no_push_to_cube,
468480
}
469481
}
470482
}

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

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ use egg::Subst;
1414
impl WrapperRules {
1515
pub fn cube_scan_wrapper_rules(&self, rules: &mut Vec<CubeRewrite>) {
1616
rules.extend(vec![
17+
// This rule would copy ungrouped from CubeScan to push_to_cube in WrapperPullupReplacer
18+
// So, if this CubeScan is ungrouped, then next wrapper would try to push to cube
1719
transforming_rewrite(
1820
"wrapper-cube-scan-wrap",
1921
cube_scan(
@@ -59,6 +61,51 @@ impl WrapperRules {
5961
"?ungrouped_scan_out",
6062
),
6163
),
64+
// This rule would add WrapperPullupReplacer without push_to_cube for ungrouped CubeScan
65+
// So when this CubeScan is ungrouped, next wrapped could use it without pushing to Cube
66+
// It's important for stuff like Aggregate(MAX(avgMeasure), CubeScan(ungrouped=true))
67+
transforming_rewrite(
68+
"wrapper-ungrouped-cube-scan-as-grouped-wrap",
69+
cube_scan(
70+
"?alias_to_cube",
71+
"?members",
72+
"?filters",
73+
"?order",
74+
"?limit",
75+
"?offset",
76+
"CubeScanSplit:false",
77+
"?can_pushdown_join",
78+
"CubeScanWrapped:false",
79+
"CubeScanUngrouped:true",
80+
),
81+
cube_scan_wrapper(
82+
wrapper_pullup_replacer(
83+
cube_scan(
84+
"?alias_to_cube",
85+
"?members",
86+
"?filters",
87+
"?order",
88+
"?limit",
89+
"?offset",
90+
"CubeScanSplit:false",
91+
"?can_pushdown_join",
92+
"CubeScanWrapped:true",
93+
"CubeScanUngrouped:true",
94+
),
95+
"?alias_to_cube_out",
96+
"WrapperPullupReplacerPushToCube:false",
97+
"WrapperPullupReplacerUngroupedScan:true",
98+
"WrapperPullupReplacerInProjection:false",
99+
"?members",
100+
),
101+
"CubeScanWrapperFinalized:false",
102+
),
103+
self.transform_wrap_ungropued_cube_scan_as_grouped(
104+
"?members",
105+
"?alias_to_cube",
106+
"?alias_to_cube_out",
107+
),
108+
),
62109
rewrite(
63110
"wrapper-finalize-pull-up-replacer",
64111
cube_scan_wrapper(
@@ -126,4 +173,32 @@ impl WrapperRules {
126173
false
127174
}
128175
}
176+
177+
fn transform_wrap_ungropued_cube_scan_as_grouped(
178+
&self,
179+
members_var: &'static str,
180+
alias_to_cube_var: &'static str,
181+
alias_to_cube_var_out: &'static str,
182+
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
183+
let members_var = var!(members_var);
184+
let alias_to_cube_var = var!(alias_to_cube_var);
185+
let alias_to_cube_var_out = var!(alias_to_cube_var_out);
186+
move |egraph, subst| {
187+
if let Some(_) = egraph[subst[members_var]].data.member_name_to_expr {
188+
for alias_to_cube in
189+
var_iter!(egraph[subst[alias_to_cube_var]], CubeScanAliasToCube).cloned()
190+
{
191+
subst.insert(
192+
alias_to_cube_var_out,
193+
egraph.add(LogicalPlanLanguage::WrapperPullupReplacerAliasToCube(
194+
WrapperPullupReplacerAliasToCube(alias_to_cube),
195+
)),
196+
);
197+
return true;
198+
}
199+
}
200+
201+
false
202+
}
203+
}
129204
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -476,6 +476,14 @@ pub fn get_sixteen_char_member_cube() -> Vec<CubeMeta> {
476476
agg_type: Some("count".to_string()),
477477
meta: None,
478478
},
479+
CubeMetaMeasure {
480+
name: "SixteenChar.sixteen_charchar_baz".to_string(),
481+
title: None,
482+
description: None,
483+
r#type: "number".to_string(),
484+
agg_type: Some("sum".to_string()),
485+
meta: None,
486+
},
479487
],
480488
segments: vec![],
481489
joins: None,

0 commit comments

Comments
 (0)