Skip to content

Commit a333ee6

Browse files
mcheshkovmcheshkov
committed
feat(cubesql): Add separate ungrouped_scan flags to pushdown and pullup replacers
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
1 parent 48265af commit a333ee6

31 files changed

+744
-61
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -242,8 +242,8 @@ pub struct CubePlanCost {
242242
non_pushed_down_grouping_sets: i64,
243243
non_pushed_down_limit_sort: i64,
244244
wrapper_nodes: i64,
245-
wrapped_select_ungrouped_scan: usize,
246245
ast_size_outside_wrapper: usize,
246+
wrapped_select_ungrouped_scan: usize,
247247
filters: i64,
248248
structure_points: i64,
249249
filter_members: i64,

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

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -465,6 +465,9 @@ crate::plan_to_language! {
465465
// We can use sumMeasure instead of SUM(sumMeasure) ONLY in with push to Cube
466466
// An vice versa, we can't use SUM(sumMeasure) in grouped query to Cube, so it can be allowed ONLY without push to grouped Cube query
467467
push_to_cube: bool,
468+
// This means that `member` is rewritten in context, where data souce (like `from` in WrappedSelect) is actually ungrouped, even when push_to_cube is disabled.
469+
// This flag should be passed from top, by the rule that starts wrapping new logical plan node.
470+
ungrouped_scan: bool,
468471
in_projection: bool,
469472
cube_members: Vec<LogicalPlan>,
470473
},
@@ -480,6 +483,10 @@ crate::plan_to_language! {
480483
// We can use sumMeasure instead of SUM(sumMeasure) ONLY in with push to Cube
481484
// An vice versa, we can't use SUM(sumMeasure) in grouped query to Cube, so it can be allowed ONLY without push to grouped Cube query
482485
push_to_cube: bool,
486+
// When `member` is logical plan this means is is actually ungrouped, even when push_to_cube is disabled.
487+
// When `member` is expression it just acts as a pull-through from pushdown.
488+
// This flag should make roundtrip from top to bottom and back.
489+
ungrouped_scan: bool,
483490
in_projection: bool,
484491
cube_members: Vec<LogicalPlan>,
485492
},
@@ -1945,25 +1952,27 @@ fn wrapper_pushdown_replacer(
19451952
members: impl Display,
19461953
alias_to_cube: impl Display,
19471954
push_to_cube: impl Display,
1955+
ungrouped_scan: impl Display,
19481956
in_projection: impl Display,
19491957
cube_members: impl Display,
19501958
) -> String {
19511959
format!(
1952-
"(WrapperPushdownReplacer {} {} {} {} {})",
1953-
members, alias_to_cube, push_to_cube, in_projection, cube_members
1960+
"(WrapperPushdownReplacer {} {} {} {} {} {})",
1961+
members, alias_to_cube, push_to_cube, ungrouped_scan, in_projection, cube_members
19541962
)
19551963
}
19561964

19571965
fn wrapper_pullup_replacer(
19581966
members: impl Display,
19591967
alias_to_cube: impl Display,
19601968
push_to_cube: impl Display,
1969+
ungrouped_scan: impl Display,
19611970
in_projection: impl Display,
19621971
cube_members: impl Display,
19631972
) -> String {
19641973
format!(
1965-
"(WrapperPullupReplacer {} {} {} {} {})",
1966-
members, alias_to_cube, push_to_cube, in_projection, cube_members
1974+
"(WrapperPullupReplacer {} {} {} {} {} {})",
1975+
members, alias_to_cube, push_to_cube, ungrouped_scan, in_projection, cube_members
19671976
)
19681977
}
19691978

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

Lines changed: 66 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ use crate::{
1212
wrapped_select_window_expr_empty_tail, wrapper_pullup_replacer, wrapper_pushdown_replacer,
1313
AggregateFunctionExprDistinct, AggregateFunctionExprFun, AliasExprAlias, ColumnExprColumn,
1414
ListType, LogicalPlanLanguage, WrappedSelectPushToCube, WrapperPullupReplacerAliasToCube,
15-
WrapperPullupReplacerPushToCube, WrapperPushdownReplacerPushToCube,
15+
WrapperPullupReplacerPushToCube, WrapperPullupReplacerUngroupedScan,
16+
WrapperPushdownReplacerPushToCube, WrapperPushdownReplacerUngroupedScan,
1617
},
1718
copy_flag,
1819
transport::V1CubeMetaMeasureExt,
@@ -33,6 +34,7 @@ impl WrapperRules {
3334
"?cube_scan_input",
3435
"?alias_to_cube",
3536
"?push_to_cube",
37+
"?ungrouped_scan",
3638
"?in_projection",
3739
"?cube_members",
3840
),
@@ -49,41 +51,47 @@ impl WrapperRules {
4951
wrapped_select_projection_expr_empty_tail(),
5052
"?alias_to_cube",
5153
"?push_to_cube",
54+
"?ungrouped_scan",
5255
"WrapperPullupReplacerInProjection:false",
5356
"?cube_members",
5457
),
5558
wrapper_pullup_replacer(
5659
wrapped_select_subqueries_empty_tail(),
5760
"?alias_to_cube",
5861
"?push_to_cube",
62+
"?ungrouped_scan",
5963
"WrapperPullupReplacerInProjection:false",
6064
"?cube_members",
6165
),
6266
wrapper_pushdown_replacer(
6367
"?group_expr",
6468
"?alias_to_cube",
6569
"?pushdown_push_to_cube",
70+
"?pushdown_ungrouped_scan",
6671
"WrapperPullupReplacerInProjection:false",
6772
"?cube_members",
6873
),
6974
wrapper_pushdown_replacer(
7075
"?aggr_expr",
7176
"?alias_to_cube",
7277
"?pushdown_push_to_cube",
78+
"?pushdown_ungrouped_scan",
7379
"WrapperPullupReplacerInProjection:false",
7480
"?cube_members",
7581
),
7682
wrapper_pullup_replacer(
7783
wrapped_select_window_expr_empty_tail(),
7884
"?alias_to_cube",
7985
"?push_to_cube",
86+
"?ungrouped_scan",
8087
"WrapperPullupReplacerInProjection:false",
8188
"?cube_members",
8289
),
8390
wrapper_pullup_replacer(
8491
"?cube_scan_input",
8592
"?alias_to_cube",
8693
"?push_to_cube",
94+
"?ungrouped_scan",
8795
"WrapperPullupReplacerInProjection:false",
8896
"?cube_members",
8997
),
@@ -92,6 +100,7 @@ impl WrapperRules {
92100
wrapped_select_filter_expr_empty_tail(),
93101
"?alias_to_cube",
94102
"?push_to_cube",
103+
"?ungrouped_scan",
95104
"WrapperPullupReplacerInProjection:false",
96105
"?cube_members",
97106
),
@@ -102,6 +111,7 @@ impl WrapperRules {
102111
wrapped_select_order_expr_empty_tail(),
103112
"?alias_to_cube",
104113
"?push_to_cube",
114+
"?ungrouped_scan",
105115
"WrapperPullupReplacerInProjection:false",
106116
"?cube_members",
107117
),
@@ -117,6 +127,8 @@ impl WrapperRules {
117127
"?aggr_expr",
118128
"?push_to_cube",
119129
"?pushdown_push_to_cube",
130+
"?ungrouped_scan",
131+
"?pushdown_ungrouped_scan",
120132
"?select_push_to_cube",
121133
),
122134
),
@@ -126,6 +138,7 @@ impl WrapperRules {
126138
grouping_set_expr("?rollout_members", "?type"),
127139
"?alias_to_cube",
128140
"?push_to_cube",
141+
"?ungrouped_scan",
129142
"WrapperPullupReplacerInProjection:false",
130143
"?cube_members",
131144
),
@@ -134,6 +147,7 @@ impl WrapperRules {
134147
"?rollout_members",
135148
"?alias_to_cube",
136149
"?push_to_cube",
150+
"?ungrouped_scan",
137151
"WrapperPullupReplacerInProjection:false",
138152
"?cube_members",
139153
),
@@ -148,6 +162,7 @@ impl WrapperRules {
148162
"?rollout_members",
149163
"?alias_to_cube",
150164
"?push_to_cube",
165+
"?ungrouped_scan",
151166
"WrapperPullupReplacerInProjection:false",
152167
"?cube_members",
153168
),
@@ -157,6 +172,7 @@ impl WrapperRules {
157172
grouping_set_expr("?rollout_members", "?type"),
158173
"?alias_to_cube",
159174
"?push_to_cube",
175+
"?ungrouped_scan",
160176
"WrapperPullupReplacerInProjection:false",
161177
"?cube_members",
162178
),
@@ -179,6 +195,7 @@ impl WrapperRules {
179195
"?aggr_expr",
180196
"?alias_to_cube",
181197
"WrapperPushdownReplacerPushToCube:true",
198+
"WrapperPushdownReplacerUngroupedScan:true",
182199
"?in_projection",
183200
"?cube_members",
184201
),
@@ -187,6 +204,7 @@ impl WrapperRules {
187204
"?measure",
188205
"?alias_to_cube",
189206
"WrapperPullupReplacerPushToCube:true",
207+
"WrapperPullupReplacerUngroupedScan:true",
190208
"?in_projection",
191209
"?cube_members",
192210
),
@@ -256,6 +274,7 @@ impl WrapperRules {
256274
"?cube_scan_input",
257275
"?alias_to_cube",
258276
"?push_to_cube",
277+
"?ungrouped_scan",
259278
"?in_projection",
260279
"?cube_members",
261280
),
@@ -275,41 +294,47 @@ impl WrapperRules {
275294
wrapped_select_projection_expr_empty_tail(),
276295
"?alias_to_cube",
277296
"?push_to_cube",
297+
"?ungrouped_scan",
278298
"WrapperPullupReplacerInProjection:false",
279299
"?cube_members",
280300
),
281301
wrapper_pushdown_replacer(
282302
"?subqueries",
283303
"?alias_to_cube",
284304
"?pushdown_push_to_cube",
305+
"?pushdown_ungrouped_scan",
285306
"WrapperPullupReplacerInProjection:false",
286307
"?cube_members",
287308
),
288309
wrapper_pushdown_replacer(
289310
"?group_expr",
290311
"?alias_to_cube",
291312
"?pushdown_push_to_cube",
313+
"?pushdown_ungrouped_scan",
292314
"WrapperPullupReplacerInProjection:false",
293315
"?cube_members",
294316
),
295317
wrapper_pushdown_replacer(
296318
"?aggr_expr",
297319
"?alias_to_cube",
298320
"?pushdown_push_to_cube",
321+
"?pushdown_ungrouped_scan",
299322
"WrapperPullupReplacerInProjection:false",
300323
"?cube_members",
301324
),
302325
wrapper_pullup_replacer(
303326
wrapped_select_window_expr_empty_tail(),
304327
"?alias_to_cube",
305328
"?push_to_cube",
329+
"?ungrouped_scan",
306330
"WrapperPullupReplacerInProjection:false",
307331
"?cube_members",
308332
),
309333
wrapper_pullup_replacer(
310334
"?cube_scan_input",
311335
"?alias_to_cube",
312336
"?push_to_cube",
337+
"?ungrouped_scan",
313338
"WrapperPullupReplacerInProjection:false",
314339
"?cube_members",
315340
),
@@ -318,6 +343,7 @@ impl WrapperRules {
318343
wrapped_select_filter_expr_empty_tail(),
319344
"?alias_to_cube",
320345
"?push_to_cube",
346+
"?ungrouped_scan",
321347
"WrapperPullupReplacerInProjection:false",
322348
"?cube_members",
323349
),
@@ -328,6 +354,7 @@ impl WrapperRules {
328354
wrapped_select_order_expr_empty_tail(),
329355
"?alias_to_cube",
330356
"?push_to_cube",
357+
"?ungrouped_scan",
331358
"WrapperPullupReplacerInProjection:false",
332359
"?cube_members",
333360
),
@@ -344,6 +371,8 @@ impl WrapperRules {
344371
"?aggr_expr",
345372
"?push_to_cube",
346373
"?pushdown_push_to_cube",
374+
"?ungrouped_scan",
375+
"?pushdown_ungrouped_scan",
347376
"?select_push_to_cube",
348377
),
349378
)]);
@@ -355,12 +384,16 @@ impl WrapperRules {
355384
aggr_expr_var: &'static str,
356385
push_to_cube_var: &'static str,
357386
pushdown_push_to_cube_var: &'static str,
387+
ungrouped_scan_var: &'static str,
388+
pushdown_ungrouped_scan_var: &'static str,
358389
select_push_to_cube_var: &'static str,
359390
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
360391
let group_expr_var = var!(group_expr_var);
361392
let aggr_expr_var = var!(aggr_expr_var);
362393
let push_to_cube_var = var!(push_to_cube_var);
363394
let pushdown_push_to_cube_var = var!(pushdown_push_to_cube_var);
395+
let ungrouped_scan_var = var!(ungrouped_scan_var);
396+
let pushdown_ungrouped_scan_var = var!(pushdown_ungrouped_scan_var);
364397
let select_push_to_cube_var = var!(select_push_to_cube_var);
365398
move |egraph, subst| {
366399
Self::transform_aggregate_impl(
@@ -370,6 +403,8 @@ impl WrapperRules {
370403
aggr_expr_var,
371404
push_to_cube_var,
372405
pushdown_push_to_cube_var,
406+
ungrouped_scan_var,
407+
pushdown_ungrouped_scan_var,
373408
select_push_to_cube_var,
374409
)
375410
}
@@ -382,13 +417,17 @@ impl WrapperRules {
382417
aggr_expr_var: &'static str,
383418
push_to_cube_var: &'static str,
384419
pushdown_push_to_cube_var: &'static str,
420+
ungrouped_scan_var: &'static str,
421+
pushdown_ungrouped_scan_var: &'static str,
385422
select_push_to_cube_var: &'static str,
386423
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
387424
let alias_to_cube_var = var!(alias_to_cube_var);
388425
let group_expr_var = var!(group_expr_var);
389426
let aggr_expr_var = var!(aggr_expr_var);
390427
let push_to_cube_var = var!(push_to_cube_var);
391428
let pushdown_push_to_cube_var = var!(pushdown_push_to_cube_var);
429+
let ungrouped_scan_var = var!(ungrouped_scan_var);
430+
let pushdown_ungrouped_scan_var = var!(pushdown_ungrouped_scan_var);
392431
let select_push_to_cube_var = var!(select_push_to_cube_var);
393432
let meta = self.meta_context.clone();
394433
move |egraph, subst| {
@@ -405,6 +444,8 @@ impl WrapperRules {
405444
aggr_expr_var,
406445
push_to_cube_var,
407446
pushdown_push_to_cube_var,
447+
ungrouped_scan_var,
448+
pushdown_ungrouped_scan_var,
408449
select_push_to_cube_var,
409450
)
410451
} else {
@@ -420,6 +461,8 @@ impl WrapperRules {
420461
aggr_expr_var: Var,
421462
push_to_cube_var: Var,
422463
pushdown_push_to_cube_var: Var,
464+
ungrouped_scan_var: Var,
465+
pushdown_ungrouped_scan_var: Var,
423466
select_push_to_cube_var: Var,
424467
) -> bool {
425468
if egraph[subst[group_expr_var]].data.referenced_expr.is_none() {
@@ -440,21 +483,29 @@ impl WrapperRules {
440483
return false;
441484
}
442485

443-
for push_to_cube in var_iter!(
444-
egraph[subst[push_to_cube_var]],
445-
WrapperPullupReplacerPushToCube
446-
)
447-
.cloned()
448-
{
449-
subst.insert(
450-
select_push_to_cube_var,
451-
egraph.add(LogicalPlanLanguage::WrappedSelectPushToCube(
452-
WrappedSelectPushToCube(push_to_cube),
453-
)),
454-
);
455-
return true;
486+
if !copy_flag!(
487+
egraph,
488+
subst,
489+
push_to_cube_var,
490+
WrapperPullupReplacerPushToCube,
491+
select_push_to_cube_var,
492+
WrappedSelectPushToCube
493+
) {
494+
return false;
456495
}
457-
false
496+
497+
if !copy_flag!(
498+
egraph,
499+
subst,
500+
ungrouped_scan_var,
501+
WrapperPullupReplacerUngroupedScan,
502+
pushdown_ungrouped_scan_var,
503+
WrapperPushdownReplacerUngroupedScan
504+
) {
505+
return false;
506+
}
507+
508+
true
458509
}
459510

460511
fn check_rollup_allowed(

0 commit comments

Comments
 (0)