feat(cubesql): Push Limit-Sort down Projection (#9776)

Author: MazterQyou

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

@@ -17168,4 +17168,59 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
             }
         )
     }
+
+    #[tokio::test]
+    async fn test_push_down_limit_sort_projection() {
+        init_testing_logger();
+
+        let logical_plan = convert_select_to_query_plan(
+            r#"
+            SELECT
+                "ta_1"."customer_gender" AS "ca_1",
+                DATE_TRUNC('MONTH', CAST("ta_1"."order_date" AS date)) AS "ca_2",
+                COALESCE(sum("ta_1"."sumPrice"), 0) AS "ca_3"
+            FROM
+                "db"."public"."KibanaSampleDataEcommerce" AS "ta_1"
+            WHERE
+                (
+                    "ta_1"."order_date" >= TIMESTAMP '2024-01-01 00:00:00.0'
+                    AND "ta_1"."order_date" < TIMESTAMP '2025-01-01 00:00:00.0'
+                )
+            GROUP BY
+                "ca_1",
+                "ca_2"
+            ORDER BY
+                "ca_2" ASC NULLS LAST
+            LIMIT
+                5000
+            ;"#
+            .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![]),
+                time_dimensions: Some(vec![V1LoadRequestQueryTimeDimension {
+                    dimension: "KibanaSampleDataEcommerce.order_date".to_string(),
+                    granularity: Some("month".to_string()),
+                    date_range: Some(json!(vec![
+                        "2024-01-01T00:00:00.000Z".to_string(),
+                        "2024-12-31T23:59:59.999Z".to_string()
+                    ])),
+                },]),
+                order: Some(vec![vec![
+                    "KibanaSampleDataEcommerce.order_date".to_string(),
+                    "asc".to_string(),
+                ]]),
+                limit: Some(5000),
+                ..Default::default()
+            }
+        )
+    }
 }

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

@@ -51,11 +51,6 @@ impl BestCubePlan {
             _ => 0,
         };
 
-        let non_pushed_down_limit_sort = match enode {
-            LogicalPlanLanguage::Sort(_) => 1,
-            _ => 0,
-        };
-
         let ast_size_inside_wrapper = match enode {
             LogicalPlanLanguage::WrappedSelect(_) => 1,
             _ => 0,
@@ -130,6 +125,8 @@ impl BestCubePlan {
             LogicalPlanLanguage::JoinCheckStage(_) => 1,
             LogicalPlanLanguage::JoinCheckPushDown(_) => 1,
             LogicalPlanLanguage::JoinCheckPullUp(_) => 1,
+            LogicalPlanLanguage::SortProjectionPushdownReplacer(_) => 1,
+            LogicalPlanLanguage::SortProjectionPullupReplacer(_) => 1,
             // Not really replacers but those should be deemed as mandatory rewrites and as soon as
             // there's always rewrite rule it's fine to have replacer cost.
             // Needs to be added as alias rewrite always more expensive than original function.
@@ -220,7 +217,8 @@ impl BestCubePlan {
             member_errors,
             non_pushed_down_window,
             non_pushed_down_grouping_sets,
-            non_pushed_down_limit_sort,
+            // Will be filled in finalize
+            non_pushed_down_limit_sort: 0,
             zero_members_wrapper,
             cube_members,
             errors: this_errors,
@@ -405,9 +403,8 @@ impl CubePlanCost {
                 CubePlanState::Wrapper => 0,
             },
             non_pushed_down_limit_sort: match sort_state {
-                SortState::DirectChild => self.non_pushed_down_limit_sort,
-                SortState::Current => self.non_pushed_down_limit_sort,
-                _ => 0,
+                SortState::Current => self.non_pushed_down_limit_sort + 1,
+                _ => self.non_pushed_down_limit_sort,
             },
             // Don't track state here: we want representation that have fewer wrappers with zero members _in total_
             zero_members_wrapper: self.zero_members_wrapper,

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

@@ -508,6 +508,13 @@ crate::plan_to_language! {
             members: Vec<LogicalPlan>,
             alias_to_cube: Vec<(String, String)>,
         },
+        SortProjectionPushdownReplacer {
+            expr: Arc<Expr>,
+            column_to_expr: Vec<(Column, Expr)>,
+        },
+        SortProjectionPullupReplacer {
+            expr: Arc<Expr>,
+        },
         EventNotification {
             name: String,
             members: Vec<LogicalPlan>,
@@ -2236,6 +2243,17 @@ fn join_check_pull_up(expr: impl Display, left: impl Display, right: impl Displa
     format!("(JoinCheckPullUp {expr} {left} {right})")
 }
 
+fn sort_projection_pushdown_replacer(expr: impl Display, column_to_expr: impl Display) -> String {
+    format!(
+        "(SortProjectionPushdownReplacer {} {})",
+        expr, column_to_expr
+    )
+}
+
+fn sort_projection_pullup_replacer(expr: impl Display) -> String {
+    format!("(SortProjectionPullupReplacer {})", expr)
+}
+
 pub fn original_expr_name(egraph: &CubeEGraph, id: Id) -> Option<String> {
     egraph[id]
         .data

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

@@ -477,7 +477,7 @@ impl Rewriter {
                 eval_stable_functions,
             ),
             &DateRules::new(config_obj.clone()),
-            &OrderRules::new(),
+            &OrderRules::new(config_obj.clone()),
             &CommonRules::new(config_obj.clone()),
         ];
         let mut rewrites = Vec::new();