feat(cubesql): Push Limit-Sort down Projection

Signed-off-by: Alex Qyoun-ae <[email protected]>

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,
@@ -220,7 +215,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 +401,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/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();

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

@@ -1,18 +1,33 @@
+use std::{
+    collections::HashMap,
+    ops::{Index, IndexMut},
+    sync::Arc,
+};
+
+use egg::{Id, Subst};
+
 use crate::{
-    compile::rewrite::{
-        analysis::OriginalExpr,
-        column_name_to_member_vec, cube_scan, cube_scan_order, cube_scan_order_empty_tail,
-        expr_column_name, order, order_replacer, referenced_columns, rewrite,
-        rewriter::{CubeEGraph, CubeRewrite, RewriteRules},
-        sort, sort_exp, sort_exp_empty_tail, sort_expr, transforming_rewrite, LogicalPlanLanguage,
-        OrderAsc, OrderMember, OrderReplacerColumnNameToMember, SortExprAsc,
+    compile::{
+        datafusion::logical_plan::{replace_col_to_expr, Column, Expr},
+        rewrite::{
+            analysis::OriginalExpr,
+            column_name_to_member_vec,
+            converter::LogicalPlanToLanguageConverter,
+            cube_scan, cube_scan_order, cube_scan_order_empty_tail, expr_column_name, limit, order,
+            order_replacer, projection, referenced_columns, rewrite,
+            rewriter::{CubeEGraph, CubeRewrite, RewriteRules},
+            sort, sort_exp, sort_exp_empty_tail, sort_expr, transforming_rewrite,
+            LogicalPlanLanguage, OrderAsc, OrderMember, OrderReplacerColumnNameToMember,
+            ProjectionAlias, SortExprAsc,
+        },
     },
-    var, var_iter,
+    config::ConfigObj,
+    var, var_iter, CubeError,
 };
-use egg::Subst;
-use std::ops::{Index, IndexMut};
 
-pub struct OrderRules {}
+pub struct OrderRules {
+    config_obj: Arc<dyn ConfigObj>,
+}
 
 impl RewriteRules for OrderRules {
     fn rewrite_rules(&self) -> Vec<CubeRewrite> {
@@ -70,13 +85,42 @@ impl RewriteRules for OrderRules {
                 order_replacer(sort_exp_empty_tail(), "?aliases"),
                 cube_scan_order_empty_tail(),
             ),
+            transforming_rewrite(
+                "push-down-limit-sort-projection",
+                limit(
+                    "?skip",
+                    "?fetch",
+                    sort(
+                        "?sort_expr",
+                        projection(
+                            "?projection_expr",
+                            "?input",
+                            "?projection_alias",
+                            "?projection_split",
+                        ),
+                    ),
+                ),
+                projection(
+                    "?projection_expr",
+                    limit("?skip", "?fetch", sort("?new_sort_expr", "?input")),
+                    "?projection_alias",
+                    "?projection_split",
+                ),
+                self.push_down_limit_sort_projection(
+                    "?input",
+                    "?projection_expr",
+                    "?projection_alias",
+                    "?sort_expr",
+                    "?new_sort_expr",
+                ),
+            ),
         ]
     }
 }
 
 impl OrderRules {
-    pub fn new() -> Self {
-        Self {}
+    pub fn new(config_obj: Arc<dyn ConfigObj>) -> Self {
+        Self { config_obj }
     }
 
     fn push_down_sort(
@@ -173,4 +217,147 @@ impl OrderRules {
             false
         }
     }
+
+    fn push_down_limit_sort_projection(
+        &self,
+        input_var: &'static str,
+        projection_expr_var: &'static str,
+        projection_alias_var: &'static str,
+        sort_expr_var: &'static str,
+        new_sort_expr_var: &'static str,
+    ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
+        let input_var = var!(input_var);
+        let projection_expr_var = var!(projection_expr_var);
+        let projection_alias_var = var!(projection_alias_var);
+        let sort_expr_var = var!(sort_expr_var);
+        let new_sort_expr_var = var!(new_sort_expr_var);
+        let flat_list = self.config_obj.push_down_pull_up_split();
+        move |egraph, subst| {
+            let input_is_sort_or_limit = egraph[subst[input_var]].nodes.iter().any(|node| {
+                matches!(
+                    node,
+                    LogicalPlanLanguage::Sort(_) | LogicalPlanLanguage::Limit(_)
+                )
+            });
+            if input_is_sort_or_limit {
+                return false;
+            }
+
+            let Some(expr_to_alias) = egraph[subst[projection_expr_var]]
+                .data
+                .expr_to_alias
+                .as_deref()
+            else {
+                return false;
+            };
+
+            for projection_alias in var_iter!(egraph[subst[projection_alias_var]], ProjectionAlias)
+            {
+                let mut column_with_expr = vec![];
+                for (expr, alias, _) in expr_to_alias {
+                    let column = Column::from_name(alias);
+                    column_with_expr.push((column, expr));
+                    if let Some(projection_alias) = projection_alias.as_deref() {
+                        let column = Column {
+                            relation: Some(projection_alias.to_string()),
+                            name: alias.to_string(),
+                        };
+                        column_with_expr.push((column, expr));
+                    }
+                }
+
+                let column_to_expr = column_with_expr
+                    .iter()
+                    .map(|(column, expr)| (column, *expr))
+                    .collect::<HashMap<_, _>>();
+
+                let Ok(new_sort_expr) =
+                    rewrite_sort_expr(egraph, subst[sort_expr_var], &column_to_expr)
+                else {
+                    continue;
+                };
+                let Ok(new_sort_expr_id) = insert_sort_expr(new_sort_expr, egraph, flat_list)
+                else {
+                    // Insertion failure should never happen as it can be partial,
+                    // so fail right away.
+                    return false;
+                };
+
+                subst.insert(new_sort_expr_var, new_sort_expr_id);
+                return true;
+            }
+            false
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+enum SortExprNode {
+    List(Vec<SortExprNode>),
+    Expr(Expr, Id, Id),
+}
+
+fn rewrite_sort_expr(
+    egraph: &CubeEGraph,
+    id: Id,
+    column_to_expr: &HashMap<&Column, &Expr>,
+) -> Result<SortExprNode, CubeError> {
+    for node in &egraph[id].nodes {
+        match node {
+            LogicalPlanLanguage::SortExp(sort_exp) => {
+                if let Ok(exprs) = sort_exp
+                    .iter()
+                    .map(|id| rewrite_sort_expr(egraph, *id, column_to_expr))
+                    .collect::<Result<Vec<_>, _>>()
+                {
+                    return Ok(SortExprNode::List(exprs));
+                }
+            }
+            LogicalPlanLanguage::SortExpr(sort_expr) => {
+                let Some(OriginalExpr::Expr(original_expr)) =
+                    &egraph[sort_expr[0]].data.original_expr
+                else {
+                    continue;
+                };
+                let Ok(rewritten_expr) = replace_col_to_expr(original_expr.clone(), column_to_expr)
+                else {
+                    continue;
+                };
+                return Ok(SortExprNode::Expr(
+                    rewritten_expr,
+                    sort_expr[1],
+                    sort_expr[2],
+                ));
+            }
+            _ => continue,
+        }
+    }
+    Err(CubeError::internal(
+        "Unable to find replacable SortExp or SortExpr while pushing sort down the projection"
+            .to_string(),
+    ))
+}
+
+fn insert_sort_expr(
+    sort_expr: SortExprNode,
+    egraph: &mut CubeEGraph,
+    flat_list: bool,
+) -> Result<Id, CubeError> {
+    match sort_expr {
+        SortExprNode::List(list) => {
+            let ids = list
+                .into_iter()
+                .map(|sort_node| insert_sort_expr(sort_node, egraph, flat_list))
+                .collect::<Result<Vec<_>, _>>()?;
+            Ok(egraph.add(LogicalPlanLanguage::SortExp(ids)))
+        }
+        SortExprNode::Expr(expr, asc_id, nulls_first_id) => {
+            let expr_id = LogicalPlanToLanguageConverter::add_expr(egraph, &expr, flat_list)?;
+            Ok(egraph.add(LogicalPlanLanguage::SortExpr([
+                expr_id,
+                asc_id,
+                nulls_first_id,
+            ])))
+        }
+    }
 }