Skip to content

Commit 77d2848

Browse files
committed
feat(cubesql): Add projection flattening rule
Allow to flatten projection node into internal WrappedSelect. This should allow to execute plans like Projection(Filter(...)) as a single ungrouped wrapper with push-to-Cube
1 parent 7ef6282 commit 77d2848

File tree

3 files changed

+289
-6
lines changed

3 files changed

+289
-6
lines changed

rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ impl RewriteRules for WrapperRules {
6060
self.aggregate_merge_rules(&mut rules);
6161
self.projection_rules(&mut rules);
6262
self.projection_rules_subquery(&mut rules);
63+
self.projection_merge_rules(&mut rules);
6364
self.limit_rules(&mut rules);
6465
self.filter_rules(&mut rules);
6566
self.filter_rules_subquery(&mut rules);

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

Lines changed: 249 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
use crate::{
22
compile::rewrite::{
3-
cube_scan_wrapper, projection,
3+
cube_scan_wrapper, projection, rewrite,
44
rewriter::{CubeEGraph, CubeRewrite},
55
rules::wrapper::WrapperRules,
66
subquery, transforming_rewrite, wrapped_select, wrapped_select_aggr_expr_empty_tail,
77
wrapped_select_filter_expr_empty_tail, wrapped_select_group_expr_empty_tail,
88
wrapped_select_having_expr_empty_tail, wrapped_select_joins_empty_tail,
9-
wrapped_select_order_expr_empty_tail, wrapped_select_subqueries_empty_tail,
10-
wrapped_select_window_expr_empty_tail, wrapper_pullup_replacer, wrapper_pushdown_replacer,
11-
wrapper_replacer_context, ListType, LogicalPlanLanguage, ProjectionAlias,
12-
WrappedSelectAlias, WrappedSelectPushToCube, WrappedSelectUngroupedScan,
13-
WrapperReplacerContextPushToCube, WrapperReplacerContextUngroupedScan,
9+
wrapped_select_order_expr_empty_tail, wrapped_select_projection_expr_empty_tail,
10+
wrapped_select_subqueries_empty_tail, wrapped_select_window_expr_empty_tail,
11+
wrapper_pullup_replacer, wrapper_pushdown_replacer, wrapper_replacer_context, ListType,
12+
LogicalPlanLanguage, ProjectionAlias, WrappedSelectAlias, WrappedSelectPushToCube,
13+
WrappedSelectUngroupedScan, WrapperReplacerContextPushToCube,
14+
WrapperReplacerContextUngroupedScan,
1415
},
1516
copy_flag, var, var_iter,
1617
};
@@ -329,6 +330,248 @@ impl WrapperRules {
329330
),
330331
)]);
331332
}
333+
334+
pub fn projection_merge_rules(&self, rules: &mut Vec<CubeRewrite>) {
335+
rules.extend(vec![rewrite(
336+
"wrapper-merge-projection-with-inner-wrapped-select",
337+
// Input is not a finished wrapper_pullup_replacer, but WrappedSelect just before pullup
338+
// After pullup replacer would disable push to cube, because any node on top would have WrappedSelect in `from`
339+
// So there would be no CubeScan to push to
340+
// Instead, this rule tries to catch `from` before pulling up, and merge outer Filter into inner WrappedSelect
341+
projection(
342+
"?projection_expr",
343+
cube_scan_wrapper(
344+
wrapped_select(
345+
"WrappedSelectSelectType:Projection",
346+
wrapper_pullup_replacer(
347+
wrapped_select_projection_expr_empty_tail(),
348+
wrapper_replacer_context(
349+
"?alias_to_cube",
350+
"WrapperReplacerContextPushToCube:true",
351+
"WrapperReplacerContextInProjection:false",
352+
"?cube_members",
353+
"?grouped_subqueries",
354+
"?ungrouped_scan",
355+
),
356+
),
357+
wrapper_pullup_replacer(
358+
wrapped_select_subqueries_empty_tail(),
359+
wrapper_replacer_context(
360+
"?alias_to_cube",
361+
"WrapperReplacerContextPushToCube:true",
362+
"WrapperReplacerContextInProjection:false",
363+
"?cube_members",
364+
"?grouped_subqueries",
365+
"?ungrouped_scan",
366+
),
367+
),
368+
wrapper_pullup_replacer(
369+
wrapped_select_group_expr_empty_tail(),
370+
wrapper_replacer_context(
371+
"?alias_to_cube",
372+
"WrapperReplacerContextPushToCube:true",
373+
"WrapperReplacerContextInProjection:false",
374+
"?cube_members",
375+
"?grouped_subqueries",
376+
"?ungrouped_scan",
377+
),
378+
),
379+
wrapper_pullup_replacer(
380+
wrapped_select_aggr_expr_empty_tail(),
381+
wrapper_replacer_context(
382+
"?alias_to_cube",
383+
"WrapperReplacerContextPushToCube:true",
384+
"WrapperReplacerContextInProjection:false",
385+
"?cube_members",
386+
"?grouped_subqueries",
387+
"?ungrouped_scan",
388+
),
389+
),
390+
wrapper_pullup_replacer(
391+
wrapped_select_window_expr_empty_tail(),
392+
wrapper_replacer_context(
393+
"?alias_to_cube",
394+
"WrapperReplacerContextPushToCube:true",
395+
"WrapperReplacerContextInProjection:false",
396+
"?cube_members",
397+
"?grouped_subqueries",
398+
"?ungrouped_scan",
399+
),
400+
),
401+
wrapper_pullup_replacer(
402+
"?inner_from",
403+
wrapper_replacer_context(
404+
"?alias_to_cube",
405+
"WrapperReplacerContextPushToCube:true",
406+
"WrapperReplacerContextInProjection:false",
407+
"?cube_members",
408+
"?grouped_subqueries",
409+
"?ungrouped_scan",
410+
),
411+
),
412+
wrapper_pullup_replacer(
413+
"?inner_joins",
414+
wrapper_replacer_context(
415+
"?alias_to_cube",
416+
"WrapperReplacerContextPushToCube:true",
417+
"WrapperReplacerContextInProjection:false",
418+
"?cube_members",
419+
"?grouped_subqueries",
420+
"?ungrouped_scan",
421+
),
422+
),
423+
wrapper_pullup_replacer(
424+
"?inner_filter",
425+
wrapper_replacer_context(
426+
"?alias_to_cube",
427+
"WrapperReplacerContextPushToCube:true",
428+
"WrapperReplacerContextInProjection:false",
429+
"?cube_members",
430+
"?grouped_subqueries",
431+
"?ungrouped_scan",
432+
),
433+
),
434+
wrapped_select_having_expr_empty_tail(),
435+
// Inner must not have limit and offset, because they are not commutative with aggregation
436+
"WrappedSelectLimit:None",
437+
"WrappedSelectOffset:None",
438+
wrapper_pullup_replacer(
439+
wrapped_select_order_expr_empty_tail(),
440+
wrapper_replacer_context(
441+
"?alias_to_cube",
442+
"WrapperReplacerContextPushToCube:true",
443+
"WrapperReplacerContextInProjection:false",
444+
"?cube_members",
445+
"?grouped_subqueries",
446+
"?ungrouped_scan",
447+
),
448+
),
449+
"WrappedSelectAlias:None",
450+
"WrappedSelectDistinct:false",
451+
"WrappedSelectPushToCube:true",
452+
"WrappedSelectUngroupedScan:true",
453+
),
454+
"CubeScanWrapperFinalized:false",
455+
),
456+
// TODO support merging projection with aliases
457+
"ProjectionAlias:None",
458+
"ProjectionSplit:false",
459+
),
460+
cube_scan_wrapper(
461+
wrapped_select(
462+
"WrappedSelectSelectType:Projection",
463+
wrapper_pushdown_replacer(
464+
"?projection_expr",
465+
wrapper_replacer_context(
466+
"?alias_to_cube",
467+
"WrapperReplacerContextPushToCube:true",
468+
"WrapperReplacerContextInProjection:true",
469+
"?cube_members",
470+
"?grouped_subqueries",
471+
"WrapperReplacerContextUngroupedScan:true",
472+
),
473+
),
474+
wrapper_pullup_replacer(
475+
wrapped_select_subqueries_empty_tail(),
476+
wrapper_replacer_context(
477+
"?alias_to_cube",
478+
"WrapperReplacerContextPushToCube:true",
479+
"WrapperReplacerContextInProjection:true",
480+
"?cube_members",
481+
"?grouped_subqueries",
482+
"WrapperReplacerContextUngroupedScan:true",
483+
),
484+
),
485+
wrapper_pullup_replacer(
486+
wrapped_select_group_expr_empty_tail(),
487+
wrapper_replacer_context(
488+
"?alias_to_cube",
489+
"WrapperReplacerContextPushToCube:true",
490+
"WrapperReplacerContextInProjection:true",
491+
"?cube_members",
492+
"?grouped_subqueries",
493+
"WrapperReplacerContextUngroupedScan:true",
494+
),
495+
),
496+
wrapper_pullup_replacer(
497+
wrapped_select_aggr_expr_empty_tail(),
498+
wrapper_replacer_context(
499+
"?alias_to_cube",
500+
"WrapperReplacerContextPushToCube:true",
501+
"WrapperReplacerContextInProjection:true",
502+
"?cube_members",
503+
"?grouped_subqueries",
504+
"WrapperReplacerContextUngroupedScan:true",
505+
),
506+
),
507+
wrapper_pullup_replacer(
508+
wrapped_select_window_expr_empty_tail(),
509+
wrapper_replacer_context(
510+
"?alias_to_cube",
511+
"WrapperReplacerContextPushToCube:true",
512+
"WrapperReplacerContextInProjection:true",
513+
"?cube_members",
514+
"?grouped_subqueries",
515+
"WrapperReplacerContextUngroupedScan:true",
516+
),
517+
),
518+
wrapper_pullup_replacer(
519+
"?inner_from",
520+
wrapper_replacer_context(
521+
"?alias_to_cube",
522+
"WrapperReplacerContextPushToCube:true",
523+
"WrapperReplacerContextInProjection:true",
524+
"?cube_members",
525+
"?grouped_subqueries",
526+
"WrapperReplacerContextUngroupedScan:true",
527+
),
528+
),
529+
wrapper_pullup_replacer(
530+
"?inner_joins",
531+
wrapper_replacer_context(
532+
"?alias_to_cube",
533+
"WrapperReplacerContextPushToCube:true",
534+
"WrapperReplacerContextInProjection:true",
535+
"?cube_members",
536+
"?grouped_subqueries",
537+
"WrapperReplacerContextUngroupedScan:true",
538+
),
539+
),
540+
wrapper_pullup_replacer(
541+
"?inner_filter",
542+
wrapper_replacer_context(
543+
"?alias_to_cube",
544+
"WrapperReplacerContextPushToCube:true",
545+
"WrapperReplacerContextInProjection:true",
546+
"?cube_members",
547+
"?grouped_subqueries",
548+
"WrapperReplacerContextUngroupedScan:true",
549+
),
550+
),
551+
wrapped_select_having_expr_empty_tail(),
552+
"WrappedSelectLimit:None",
553+
"WrappedSelectOffset:None",
554+
wrapper_pullup_replacer(
555+
wrapped_select_order_expr_empty_tail(),
556+
wrapper_replacer_context(
557+
"?alias_to_cube",
558+
"WrapperReplacerContextPushToCube:true",
559+
"WrapperReplacerContextInProjection:true",
560+
"?cube_members",
561+
"?grouped_subqueries",
562+
"WrapperReplacerContextUngroupedScan:true",
563+
),
564+
),
565+
"WrappedSelectAlias:None",
566+
"WrappedSelectDistinct:false",
567+
"WrappedSelectPushToCube:true",
568+
"WrappedSelectUngroupedScan:true",
569+
),
570+
"CubeScanWrapperFinalized:false",
571+
),
572+
)]);
573+
}
574+
332575
fn transform_projection(
333576
&self,
334577
expr_var: &'static str,

rust/cubesql/cubesql/src/compile/test/test_wrapper.rs

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1417,3 +1417,42 @@ async fn wrapper_agg_dimension_over_limit() {
14171417
.sql
14181418
.contains("\"ungrouped\": true"));
14191419
}
1420+
1421+
// TODO allow number measures and add test for those
1422+
/// Projection(Filter(CubeScan(ungrouped))) should have projection expressions pushed down to Cube
1423+
#[tokio::test]
1424+
async fn wrapper_projection_flatten_simple_measure() {
1425+
if !Rewriter::sql_push_down_enabled() {
1426+
return;
1427+
}
1428+
init_testing_logger();
1429+
1430+
let query_plan = convert_select_to_query_plan(
1431+
// language=PostgreSQL
1432+
r#"
1433+
SELECT
1434+
maxPrice
1435+
FROM
1436+
MultiTypeCube
1437+
WHERE
1438+
LOWER(CAST(dim_num0 AS TEXT)) = 'all'
1439+
;
1440+
"#
1441+
.to_string(),
1442+
DatabaseProtocol::PostgreSQL,
1443+
)
1444+
.await;
1445+
1446+
let physical_plan = query_plan.as_physical_plan().await.unwrap();
1447+
println!(
1448+
"Physical plan: {}",
1449+
displayable(physical_plan.as_ref()).indent()
1450+
);
1451+
1452+
let request = query_plan
1453+
.as_logical_plan()
1454+
.find_cube_scan_wrapped_sql()
1455+
.request;
1456+
assert_eq!(request.measures.unwrap().len(), 1);
1457+
assert_eq!(request.dimensions.unwrap().len(), 0);
1458+
}

0 commit comments

Comments
 (0)