diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index 6deda6ae7ba1a..638254f20250c 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -17737,4 +17737,45 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), ); Ok(()) } + + #[tokio::test] + async fn test_push_down_limit_sort_projection_recursion() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let logical_plan = convert_select_to_query_plan( + r#" + SELECT + customer_gender AS "customer_gender", + SUM(sumPrice) AS "SUM(KibanaSampleDataEcommerce.sumPrice)" + FROM KibanaSampleDataEcommerce + GROUP BY 1 + ORDER BY 2 DESC + LIMIT 3 + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + ) + .await + .as_logical_plan(); + + assert_eq!( + logical_plan.find_cube_scan().request, + V1LoadRequestQuery { + measures: Some(vec!["KibanaSampleDataEcommerce.sumPrice".to_string()]), + dimensions: Some(vec![ + "KibanaSampleDataEcommerce.customer_gender".to_string(), + ]), + segments: Some(vec![]), + order: Some(vec![vec![ + "KibanaSampleDataEcommerce.sumPrice".to_string(), + "desc".to_string(), + ]]), + limit: Some(3), + ..Default::default() + } + ) + } } diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/order.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/order.rs index a55f5fc9e0b5e..8ecafc383aff6 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/rules/order.rs +++ b/rust/cubesql/cubesql/src/compile/rewrite/rules/order.rs @@ -16,9 +16,10 @@ use crate::{ order_replacer, projection, referenced_columns, rewrite, rewriter::{CubeEGraph, CubeRewrite, RewriteRules}, sort, sort_exp, sort_exp_empty_tail, sort_expr, sort_projection_pullup_replacer, - sort_projection_pushdown_replacer, transforming_rewrite, ColumnExprColumn, - LogicalPlanLanguage, OrderAsc, OrderMember, OrderReplacerColumnNameToMember, - ProjectionAlias, SortExprAsc, SortProjectionPushdownReplacerColumnToExpr, + sort_projection_pushdown_replacer, transforming_chain_rewrite, transforming_rewrite, + ColumnExprColumn, LogicalPlanLanguage, OrderAsc, OrderMember, + OrderReplacerColumnNameToMember, ProjectionAlias, SortExprAsc, + SortProjectionPushdownReplacerColumnToExpr, }, }, config::ConfigObj, @@ -87,21 +88,19 @@ impl RewriteRules for OrderRules { ), // TODO: refactor this rule to `push-down-sort-projection`, // possibly adjust cost function to penalize Limit-...-Sort plan - transforming_rewrite( + transforming_chain_rewrite( "push-down-limit-sort-projection", - limit( - "?skip", - "?fetch", - sort( - "?sort_expr", - projection( - "?projection_expr", - "?input", - "?projection_alias", - "?projection_split", - ), + limit("?skip", "?fetch", sort("?sort_expr", "?projection")), + vec![( + // Avoid recursion when projection input is self + "?projection", + projection( + "?projection_expr", + "?input", + "?projection_alias", + "?projection_split", ), - ), + )], projection( "?projection_expr", limit(