File tree Expand file tree Collapse file tree 4 files changed +94
-2
lines changed
rust/cubesql/cubesql/src/compile Expand file tree Collapse file tree 4 files changed +94
-2
lines changed Original file line number Diff line number Diff 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 }
Original file line number Diff line number Diff 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 {
Original file line number Diff line number Diff 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+ }
Original file line number Diff line number Diff line change @@ -11,6 +11,13 @@ use crate::{
1111} ;
1212
1313pub 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
2330impl 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 {
You can’t perform that action at this time.
0 commit comments