Skip to content

Commit 351ac7a

Browse files
authored
fix(cubesql): Allow more filters in CubeScan before aggregation pushdown (#9409)
Without this queries with duplicated aggregations on top of `__user` filter can execute with unnecessary post-processing, or, in aggregate-sort-limit case, even can fall back to SQL pushdown. For now, only `__user` filters are enabled, others will be done later.
1 parent f5e69f7 commit 351ac7a

File tree

4 files changed

+94
-2
lines changed

4 files changed

+94
-2
lines changed

rust/cubesql/cubesql/src/compile/rewrite/analysis.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -737,6 +737,13 @@ impl LogicalPlanAnalysis {
737737
) -> Option<Vec<(String, String)>> {
738738
let filter_operators = |id| egraph.index(id).data.filter_operators.clone();
739739
match enode {
740+
LogicalPlanLanguage::CubeScanFilters(params) => {
741+
let mut map = Vec::new();
742+
for id in params.iter() {
743+
map.extend(filter_operators(*id)?.into_iter());
744+
}
745+
Some(map)
746+
}
740747
LogicalPlanLanguage::FilterOp(params) => filter_operators(params[0]),
741748
LogicalPlanLanguage::FilterOpFilters(params) => {
742749
let mut map = Vec::new();
@@ -763,6 +770,9 @@ impl LogicalPlanAnalysis {
763770
.to_string();
764771
Some(vec![(member, "equals".to_string())])
765772
}
773+
LogicalPlanLanguage::ChangeUserMember(_) => {
774+
Some(vec![("__user".to_string(), "equals".to_string())])
775+
}
766776
_ => None,
767777
}
768778
}

rust/cubesql/cubesql/src/compile/rewrite/rules/members.rs

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1625,11 +1625,16 @@ impl MemberRules {
16251625
if *ungrouped {
16261626
continue;
16271627
}
1628-
let Some(empty_filters) = &egraph.index(subst[filters_var]).data.is_empty_list
1628+
let Some(filter_operators) =
1629+
&egraph.index(subst[filters_var]).data.filter_operators
16291630
else {
16301631
return false;
16311632
};
1632-
if !empty_filters {
1633+
let only_allowed_filters = filter_operators.iter().all(|(member, _op)| {
1634+
// TODO this should allow even more, like dimensions and segments
1635+
member == "__user"
1636+
});
1637+
if !only_allowed_filters {
16331638
return false;
16341639
}
16351640
if referenced_aggr_expr.len() == 0 {

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

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,3 +275,66 @@ GROUP BY 1
275275
assert!(sql_query.sql.contains(r#""changeUser": "gopher""#));
276276
assert_eq!(load_calls[0].meta.change_user(), Some("gopher".to_string()));
277277
}
278+
279+
/// Repeated aggregation should be flattened even in presence of __user filter
280+
#[tokio::test]
281+
async fn flatten_aggregation_into_user_change() {
282+
init_testing_logger();
283+
284+
let query_plan = convert_select_to_query_plan(
285+
// language=PostgreSQL
286+
r#"
287+
SELECT
288+
dim_str0
289+
FROM
290+
(
291+
SELECT
292+
dim_str0
293+
FROM
294+
(
295+
SELECT
296+
dim_str0,
297+
AVG(avgPrice)
298+
FROM
299+
MultiTypeCube
300+
WHERE
301+
__user = 'gopher'
302+
GROUP BY
303+
1
304+
) t
305+
GROUP BY
306+
dim_str0
307+
) AS t
308+
GROUP BY
309+
dim_str0
310+
ORDER BY
311+
dim_str0 ASC
312+
LIMIT
313+
1
314+
"#
315+
.to_string(),
316+
DatabaseProtocol::PostgreSQL,
317+
)
318+
.await;
319+
320+
// This query should rewrite completely as CubeScan
321+
let logical_plan = query_plan.as_logical_plan();
322+
let cube_scan = logical_plan.expect_root_cube_scan();
323+
324+
assert_eq!(cube_scan.options.change_user, Some("gopher".to_string()));
325+
326+
assert_eq!(
327+
cube_scan.request,
328+
V1LoadRequestQuery {
329+
measures: Some(vec![]),
330+
segments: Some(vec![]),
331+
dimensions: Some(vec!["MultiTypeCube.dim_str0".to_string(),]),
332+
order: Some(vec![vec![
333+
"MultiTypeCube.dim_str0".to_string(),
334+
"asc".to_string(),
335+
],],),
336+
limit: Some(1),
337+
..Default::default()
338+
}
339+
)
340+
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,13 @@ use crate::{
1111
};
1212

1313
pub trait LogicalPlanTestUtils {
14+
fn try_expect_root_cube_scan(&self) -> Option<&CubeScanNode>;
15+
16+
fn expect_root_cube_scan(&self) -> &CubeScanNode {
17+
self.try_expect_root_cube_scan()
18+
.expect("Root node is not CubeScan")
19+
}
20+
1421
fn find_cube_scan(&self) -> CubeScanNode;
1522

1623
fn find_cube_scan_wrapped_sql(&self) -> CubeScanWrappedSqlNode;
@@ -21,6 +28,13 @@ pub trait LogicalPlanTestUtils {
2128
}
2229

2330
impl LogicalPlanTestUtils for LogicalPlan {
31+
fn try_expect_root_cube_scan(&self) -> Option<&CubeScanNode> {
32+
let LogicalPlan::Extension(ext) = self else {
33+
return None;
34+
};
35+
ext.node.as_any().downcast_ref::<CubeScanNode>()
36+
}
37+
2438
fn find_cube_scan(&self) -> CubeScanNode {
2539
let cube_scans = find_cube_scans_deep_search(Arc::new(self.clone()), true);
2640
if cube_scans.len() != 1 {

0 commit comments

Comments
 (0)