use original order columns for suboptimal case

Author: RusovDmitriy

rust/cubestore/cubestore-sql-tests/src/tests.rs

@@ -8005,6 +8005,10 @@ async fn limit_pushdown_without_group(service: Box<dyn SqlClient>) {
         .exec_query("CREATE TABLE foo.pushdown_where_group2_with_alias (a_alias int, b_alias int, c_alias int) index ind1 (a_alias, b_alias, c_alias) index ind2 (c_alias, b_alias)")
         .await
         .unwrap();
+    service
+        .exec_query("CREATE TABLE foo.pushdown_where_group3_with_alias (a_alias int, b_alias int, c_alias int) index ind1 (c_alias, b_alias)")
+        .await
+        .unwrap();
     service
         .exec_query(
             "INSERT INTO foo.pushdown_where_group1
@@ -8318,6 +8322,38 @@ async fn limit_pushdown_without_group(service: Box<dyn SqlClient>) {
     )
     .await
     .unwrap();
+
+    // ====================================
+    assert_limit_pushdown(
+        &service,
+        "SELECT a, b, c FROM (
+                    SELECT a_alias a, b_alias b, c_alias c FROM foo.pushdown_where_group3_with_alias
+                ) as `tb`
+                WHERE a = 20
+                ORDER BY 3 DESC
+                LIMIT 3",
+        Some("ind1"),
+        true,
+        true,
+    )
+    .await
+    .unwrap();
+
+    // ====================================
+    assert_limit_pushdown(
+        &service,
+        "SELECT a, b, c FROM (
+                    SELECT a_alias a, b_alias b, c_alias c FROM foo.pushdown_where_group3_with_alias
+                ) as `tb`
+                WHERE a > 20
+                ORDER BY 3 DESC
+                LIMIT 3",
+        Some("ind1"),
+        true,
+        true,
+    )
+    .await
+    .unwrap();
 }
 async fn limit_pushdown_without_group_resort(service: Box<dyn SqlClient>) {
     service.exec_query("CREATE SCHEMA foo").await.unwrap();

rust/cubestore/cubestore/src/queryplanner/planning.rs

@@ -389,6 +389,7 @@ impl<'a> PlanIndexStore for &'a dyn MetaStore {
 #[derive(Clone)]
 struct SortColumns {
     sort_on: Vec<String>,
+    sort_on_order_col_only: Vec<String>,
     required: bool,
 }
 
@@ -461,12 +462,18 @@ impl PlanRewriter for CollectConstraints {
                                                 .map(|n| n.clone())
                                                 .unique()
                                                 .collect::<Vec<_>>(),
+                                            sort_on_order_col_only: order_col_names
+                                                .iter()
+                                                .map(|n| n.clone())
+                                                .unique()
+                                                .collect::<Vec<_>>(),
                                             required: s.required,
                                         })
                                     }
                                 }
                                 None => Some(SortColumns {
                                     sort_on: order_col_names.clone(),
+                                    sort_on_order_col_only: vec![],
                                     required: false,
                                 }),
                             }
@@ -508,6 +515,7 @@ impl PlanRewriter for CollectConstraints {
                 let sort_on = if !sort_on.is_empty() && sort_on.iter().all(|c| c.is_some()) {
                     Some(SortColumns {
                         sort_on: sort_on.into_iter().map(|c| c.unwrap()).collect(),
+                        sort_on_order_col_only: vec![],
                         required: false,
                     })
                 } else {
@@ -555,6 +563,7 @@ impl PlanRewriter for CollectConstraints {
                 if single_value_filter_columns(predicate, &mut sort_on) {
                     if !sort_on.is_empty() {
                         let sort_on = Some(SortColumns {
+                            sort_on_order_col_only: vec![],
                             sort_on: sort_on
                                 .into_iter()
                                 .map(|c| c.name.to_string())
@@ -591,6 +600,7 @@ impl PlanRewriter for CollectConstraints {
         }
         Some(ConstraintsContext {
             sort_on: Some(SortColumns {
+                sort_on_order_col_only: vec![],
                 sort_on: join_on.iter().map(|(l, _)| l.name.clone()).collect(),
                 required: true,
             }),
@@ -612,6 +622,7 @@ impl PlanRewriter for CollectConstraints {
         }
         Some(ConstraintsContext {
             sort_on: Some(SortColumns {
+                sort_on_order_col_only: vec![],
                 sort_on: join_on.iter().map(|(_, r)| r.name.clone()).collect(),
                 required: true,
             }),
@@ -1060,6 +1071,10 @@ async fn pick_index(
     indices: Vec<IdRow<Index>>,
 ) -> Result<IndexCandidate, DataFusionError> {
     let sort_on = c.sort_on.as_ref().map(|sc| (&sc.sort_on, sc.required));
+    let sort_on_order_col_only = c
+        .sort_on
+        .as_ref()
+        .map(|sc| (&sc.sort_on_order_col_only, sc.required));
 
     let aggr_index_allowed = check_aggregates_expr(&table, &c.aggregates);
 
@@ -1180,12 +1195,14 @@ async fn pick_index(
                     &projection_columns,
                     &filter_columns,
                 );
-
                 let index = optimal.unwrap_or(default_index);
                 (
                     Ok(index),
                     index.get_row().multi_index_id().map(|_| index),
-                    None,
+                    match sort_on_order_col_only {
+                        Some((columns, flag)) if !columns.is_empty() => Some((columns, flag)),
+                        _ => None,
+                    },
                 )
             }
         }