Skip to content

Commit ca7d298

Browse files
mcheshkovmcheshkov
committed
fix(cubesql): Avoid mixing flags from pushdown and pullup replacers
1 parent ff0b485 commit ca7d298

File tree

14 files changed

+436
-65
lines changed

14 files changed

+436
-65
lines changed

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -524,6 +524,27 @@ macro_rules! var {
524524
};
525525
}
526526

527+
#[macro_export]
528+
macro_rules! copy_flag {
529+
($egraph:expr, $subst:expr, $in_var:expr, $in_kind:ident, $out_var:expr, $out_kind:ident) => {{
530+
let mut found = false;
531+
for in_value in $crate::var_iter!($egraph[$subst[$in_var]], $in_kind) {
532+
// Typechecking for $in_kind, only booleans are supported for now
533+
let in_value: bool = *in_value;
534+
$subst.insert(
535+
$out_var,
536+
$egraph.add($crate::compile::rewrite::LogicalPlanLanguage::$out_kind(
537+
$out_kind(in_value),
538+
)),
539+
);
540+
found = true;
541+
// This is safe, because we expect only enode with one child, with boolena inside, and expect that they would never unify
542+
break;
543+
}
544+
found
545+
}};
546+
}
547+
527548
pub struct WithColumnRelation(Option<String>);
528549

529550
impl ExprRewriter for WithColumnRelation {

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

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,9 @@ 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-
WrapperPullupReplacerUngrouped,
15+
WrapperPullupReplacerUngrouped, WrapperPushdownReplacerUngrouped,
1616
},
17+
copy_flag,
1718
transport::V1CubeMetaMeasureExt,
1819
var, var_iter,
1920
};
@@ -61,14 +62,14 @@ impl WrapperRules {
6162
wrapper_pushdown_replacer(
6263
"?group_expr",
6364
"?alias_to_cube",
64-
"?ungrouped",
65+
"?pushdown_ungrouped",
6566
"WrapperPullupReplacerInProjection:false",
6667
"?cube_members",
6768
),
6869
wrapper_pushdown_replacer(
6970
"?aggr_expr",
7071
"?alias_to_cube",
71-
"?ungrouped",
72+
"?pushdown_ungrouped",
7273
"WrapperPullupReplacerInProjection:false",
7374
"?cube_members",
7475
),
@@ -115,6 +116,7 @@ impl WrapperRules {
115116
"?group_expr",
116117
"?aggr_expr",
117118
"?ungrouped",
119+
"?pushdown_ungrouped",
118120
"?select_push_to_cube",
119121
),
120122
),
@@ -176,7 +178,7 @@ impl WrapperRules {
176178
wrapper_pushdown_replacer(
177179
"?aggr_expr",
178180
"?alias_to_cube",
179-
"WrapperPullupReplacerUngrouped:true",
181+
"WrapperPushdownReplacerUngrouped:true",
180182
"?in_projection",
181183
"?cube_members",
182184
),
@@ -279,21 +281,21 @@ impl WrapperRules {
279281
wrapper_pushdown_replacer(
280282
"?subqueries",
281283
"?alias_to_cube",
282-
"?ungrouped",
284+
"?pushdown_ungrouped",
283285
"WrapperPullupReplacerInProjection:false",
284286
"?cube_members",
285287
),
286288
wrapper_pushdown_replacer(
287289
"?group_expr",
288290
"?alias_to_cube",
289-
"?ungrouped",
291+
"?pushdown_ungrouped",
290292
"WrapperPullupReplacerInProjection:false",
291293
"?cube_members",
292294
),
293295
wrapper_pushdown_replacer(
294296
"?aggr_expr",
295297
"?alias_to_cube",
296-
"?ungrouped",
298+
"?pushdown_ungrouped",
297299
"WrapperPullupReplacerInProjection:false",
298300
"?cube_members",
299301
),
@@ -341,6 +343,7 @@ impl WrapperRules {
341343
"?group_expr",
342344
"?aggr_expr",
343345
"?ungrouped",
346+
"?pushdown_ungrouped",
344347
"?select_push_to_cube",
345348
),
346349
)]);
@@ -351,11 +354,13 @@ impl WrapperRules {
351354
group_expr_var: &'static str,
352355
aggr_expr_var: &'static str,
353356
ungrouped_var: &'static str,
357+
pushdown_ungrouped_var: &'static str,
354358
select_push_to_cube_var: &'static str,
355359
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
356360
let group_expr_var = var!(group_expr_var);
357361
let aggr_expr_var = var!(aggr_expr_var);
358362
let ungrouped_var = var!(ungrouped_var);
363+
let pushdown_ungrouped_var = var!(pushdown_ungrouped_var);
359364
let select_push_to_cube_var = var!(select_push_to_cube_var);
360365
move |egraph, subst| {
361366
Self::transform_aggregate_impl(
@@ -364,6 +369,7 @@ impl WrapperRules {
364369
group_expr_var,
365370
aggr_expr_var,
366371
ungrouped_var,
372+
pushdown_ungrouped_var,
367373
select_push_to_cube_var,
368374
)
369375
}
@@ -375,12 +381,14 @@ impl WrapperRules {
375381
group_expr_var: &'static str,
376382
aggr_expr_var: &'static str,
377383
ungrouped_var: &'static str,
384+
pushdown_ungrouped_var: &'static str,
378385
select_push_to_cube_var: &'static str,
379386
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
380387
let alias_to_cube_var = var!(alias_to_cube_var);
381388
let group_expr_var = var!(group_expr_var);
382389
let aggr_expr_var = var!(aggr_expr_var);
383390
let ungrouped_var = var!(ungrouped_var);
391+
let pushdown_ungrouped_var = var!(pushdown_ungrouped_var);
384392
let select_push_to_cube_var = var!(select_push_to_cube_var);
385393
let meta = self.meta_context.clone();
386394
move |egraph, subst| {
@@ -396,6 +404,7 @@ impl WrapperRules {
396404
group_expr_var,
397405
aggr_expr_var,
398406
ungrouped_var,
407+
pushdown_ungrouped_var,
399408
select_push_to_cube_var,
400409
)
401410
} else {
@@ -410,6 +419,7 @@ impl WrapperRules {
410419
group_expr_var: Var,
411420
aggr_expr_var: Var,
412421
ungrouped_var: Var,
422+
pushdown_ungrouped_var: Var,
413423
select_push_to_cube_var: Var,
414424
) -> bool {
415425
if egraph[subst[group_expr_var]].data.referenced_expr.is_none() {
@@ -418,6 +428,18 @@ impl WrapperRules {
418428
if egraph[subst[aggr_expr_var]].data.referenced_expr.is_none() {
419429
return false;
420430
}
431+
432+
if !copy_flag!(
433+
egraph,
434+
subst,
435+
ungrouped_var,
436+
WrapperPullupReplacerUngrouped,
437+
pushdown_ungrouped_var,
438+
WrapperPushdownReplacerUngrouped
439+
) {
440+
return false;
441+
}
442+
421443
for ungrouped in
422444
var_iter!(egraph[subst[ungrouped_var]], WrapperPullupReplacerUngrouped).cloned()
423445
{

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ impl WrapperRules {
1919
wrapper_pushdown_replacer(
2020
column_expr("?name"),
2121
"?alias_to_cube",
22-
"WrapperPullupReplacerUngrouped:false",
22+
"WrapperPushdownReplacerUngrouped:false",
2323
"?in_projection",
2424
"?cube_members",
2525
),
@@ -38,7 +38,7 @@ impl WrapperRules {
3838
wrapper_pushdown_replacer(
3939
column_expr("?name"),
4040
"?alias_to_cube",
41-
"WrapperPullupReplacerUngrouped:true",
41+
"WrapperPushdownReplacerUngrouped:true",
4242
"WrapperPullupReplacerInProjection:true",
4343
"?cube_members",
4444
),
@@ -57,7 +57,7 @@ impl WrapperRules {
5757
wrapper_pushdown_replacer(
5858
column_expr("?name"),
5959
"?alias_to_cube",
60-
"WrapperPullupReplacerUngrouped:true",
60+
"WrapperPushdownReplacerUngrouped:true",
6161
"?in_projection",
6262
"?cube_members",
6363
),

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

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@ use crate::{
1010
wrapped_select_projection_expr_empty_tail, wrapped_select_subqueries_empty_tail,
1111
wrapped_select_window_expr_empty_tail, wrapper_pullup_replacer, wrapper_pushdown_replacer,
1212
LogicalPlanLanguage, WrappedSelectPushToCube, WrappedSelectUngroupedScan,
13-
WrapperPullupReplacerUngrouped,
13+
WrapperPullupReplacerUngrouped, WrapperPushdownReplacerUngrouped,
1414
},
15-
var, var_iter,
15+
copy_flag, var, var_iter,
1616
};
1717
use egg::{Subst, Var};
1818

@@ -177,7 +177,7 @@ impl WrapperRules {
177177
wrapper_pushdown_replacer(
178178
"?filter_expr",
179179
"?alias_to_cube",
180-
"?ungrouped",
180+
"?pushdown_ungrouped",
181181
"?in_projection",
182182
"?cube_members",
183183
),
@@ -208,6 +208,7 @@ impl WrapperRules {
208208
),
209209
self.transform_filter(
210210
"?ungrouped",
211+
"?pushdown_ungrouped",
211212
"?select_push_to_cube",
212213
"?select_ungrouped_scan",
213214
),
@@ -254,7 +255,7 @@ impl WrapperRules {
254255
wrapper_pushdown_replacer(
255256
"?subqueries",
256257
"?alias_to_cube",
257-
"?ungrouped",
258+
"?pushdown_ungrouped",
258259
"?in_projection",
259260
"?cube_members",
260261
),
@@ -291,7 +292,7 @@ impl WrapperRules {
291292
wrapper_pushdown_replacer(
292293
"?filter_expr",
293294
"?alias_to_cube",
294-
"?ungrouped",
295+
"?pushdown_ungrouped",
295296
"?in_projection",
296297
"?cube_members",
297298
),
@@ -323,6 +324,7 @@ impl WrapperRules {
323324
self.transform_filter_subquery(
324325
"?alias_to_cube",
325326
"?ungrouped",
327+
"?pushdown_ungrouped",
326328
"?select_push_to_cube",
327329
"?select_ungrouped_scan",
328330
),
@@ -332,17 +334,20 @@ impl WrapperRules {
332334
fn transform_filter(
333335
&self,
334336
ungrouped_var: &'static str,
337+
pushdown_ungrouped_var: &'static str,
335338
select_push_to_cube_var: &'static str,
336339
select_ungrouped_scan_var: &'static str,
337340
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
338341
let ungrouped_var = var!(ungrouped_var);
342+
let pushdown_ungrouped_var = var!(pushdown_ungrouped_var);
339343
let select_push_to_cube_var = var!(select_push_to_cube_var);
340344
let select_ungrouped_scan_var = var!(select_ungrouped_scan_var);
341345
move |egraph, subst| {
342346
Self::transform_filter_impl(
343347
egraph,
344348
subst,
345349
ungrouped_var,
350+
pushdown_ungrouped_var,
346351
select_push_to_cube_var,
347352
select_ungrouped_scan_var,
348353
)
@@ -353,11 +358,13 @@ impl WrapperRules {
353358
&self,
354359
alias_to_cube_var: &'static str,
355360
ungrouped_var: &'static str,
361+
pushdown_ungrouped_var: &'static str,
356362
select_push_to_cube_var: &'static str,
357363
select_ungrouped_scan_var: &'static str,
358364
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
359365
let alias_to_cube_var = var!(alias_to_cube_var);
360366
let ungrouped_var = var!(ungrouped_var);
367+
let pushdown_ungrouped_var = var!(pushdown_ungrouped_var);
361368
let select_push_to_cube_var = var!(select_push_to_cube_var);
362369
let select_ungrouped_scan_var = var!(select_ungrouped_scan_var);
363370
let meta = self.meta_context.clone();
@@ -372,6 +379,7 @@ impl WrapperRules {
372379
egraph,
373380
subst,
374381
ungrouped_var,
382+
pushdown_ungrouped_var,
375383
select_push_to_cube_var,
376384
select_ungrouped_scan_var,
377385
)
@@ -385,9 +393,21 @@ impl WrapperRules {
385393
egraph: &mut CubeEGraph,
386394
subst: &mut Subst,
387395
ungrouped_var: Var,
396+
pushdown_ungrouped_var: Var,
388397
select_push_to_cube_var: Var,
389398
select_ungrouped_scan_var: Var,
390399
) -> bool {
400+
if !copy_flag!(
401+
egraph,
402+
subst,
403+
ungrouped_var,
404+
WrapperPullupReplacerUngrouped,
405+
pushdown_ungrouped_var,
406+
WrapperPushdownReplacerUngrouped
407+
) {
408+
return false;
409+
}
410+
391411
for ungrouped in
392412
var_iter!(egraph[subst[ungrouped_var]], WrapperPullupReplacerUngrouped).cloned()
393413
{

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

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@ use crate::{
44
rewriter::{CubeEGraph, CubeRewrite},
55
rules::wrapper::WrapperRules,
66
transforming_rewrite, wrapper_pullup_replacer, wrapper_pushdown_replacer,
7-
WrapperPullupReplacerAliasToCube,
7+
WrapperPullupReplacerAliasToCube, WrapperPullupReplacerUngrouped,
8+
WrapperPushdownReplacerUngrouped,
89
},
9-
var, var_iter,
10+
copy_flag, var, var_iter,
1011
};
1112
use egg::Subst;
1213

@@ -33,13 +34,13 @@ impl WrapperRules {
3334
wrapper_pullup_replacer(
3435
"?list",
3536
"?alias_to_cube",
36-
"?ungrouped",
37+
"?pullup_ungrouped",
3738
"?in_projection",
3839
"?cube_members",
3940
),
4041
"?negated",
4142
),
42-
self.transform_in_list_only_consts("?list"),
43+
self.transform_in_list_only_consts("?list", "?ungrouped", "?pullup_ungrouped"),
4344
),
4445
rewrite(
4546
"wrapper-in-list-push-down",
@@ -132,9 +133,23 @@ impl WrapperRules {
132133
fn transform_in_list_only_consts(
133134
&self,
134135
list_var: &'static str,
136+
ungrouped_var: &'static str,
137+
pullup_ungrouped_var: &'static str,
135138
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
136139
let list_var = var!(list_var);
140+
let ungrouped_var = var!(ungrouped_var);
141+
let pullup_ungrouped_var = var!(pullup_ungrouped_var);
137142
move |egraph: &mut CubeEGraph, subst| {
143+
if !copy_flag!(
144+
egraph,
145+
subst,
146+
ungrouped_var,
147+
WrapperPushdownReplacerUngrouped,
148+
pullup_ungrouped_var,
149+
WrapperPullupReplacerUngrouped
150+
) {
151+
return false;
152+
}
138153
return egraph[subst[list_var]].data.constant_in_list.is_some();
139154
}
140155
}

0 commit comments

Comments
 (0)