fix(cubesql): ORDER BY references outer alias instead of inner expression for SQL push down

Author: paveltiunov

rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs

@@ -13,8 +13,8 @@ use cubeclient::models::V1LoadRequestQuery;
 use datafusion::{
     error::{DataFusionError, Result},
     logical_plan::{
-        plan::Extension, replace_col, Column, DFSchema, DFSchemaRef, Expr, LogicalPlan,
-        UserDefinedLogicalNode,
+        plan::Extension, replace_col, replace_col_to_expr, Column, DFSchema, DFSchemaRef, Expr,
+        LogicalPlan, UserDefinedLogicalNode,
     },
     physical_plan::aggregates::AggregateFunction,
     scalar::ScalarValue,
@@ -405,10 +405,37 @@ impl CubeScanWrapperNode {
                                 ungrouped_scan_node.clone(),
                             )
                             .await?;
+                            // Sort node always comes on top and pushed down to select so we need to replace columns here by appropriate column definitions
+                            let order_replace_map = projection_expr
+                                .iter()
+                                .chain(group_expr.iter())
+                                .chain(aggr_expr.iter())
+                                .map(|e| {
+                                    Ok((
+                                        Column {
+                                            relation: alias.clone(),
+                                            name: expr_name(&e, &schema)?,
+                                        },
+                                        e.clone(),
+                                    ))
+                                })
+                                .collect::<Result<HashMap<_, _>>>()?;
+
                             let (order, mut sql) = Self::generate_column_expr(
                                 plan.clone(),
                                 schema.clone(),
-                                order_expr.clone(),
+                                order_expr
+                                    .iter()
+                                    .map(|o| {
+                                        replace_col_to_expr(
+                                            o.clone(),
+                                            &order_replace_map
+                                                .iter()
+                                                .map(|(k, v)| (k, v))
+                                                .collect(),
+                                        )
+                                    })
+                                    .collect::<Result<Vec<_>>>()?,
                                 sql,
                                 generator.clone(),
                                 &column_remapping,

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

@@ -18438,6 +18438,36 @@ ORDER BY \"COUNT(count)\" DESC"
             .contains("ORDER BY"));
     }
 
+    #[tokio::test]
+    async fn test_case_wrapper_ungrouped_sorted_aliased() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            "SELECT x FROM (SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END x, AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY 1 ORDER BY 1 DESC) b"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await;
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert!(logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql
+            // TODO test without depend on column name
+            .contains("ORDER BY \"case_when"));
+    }
+
     #[tokio::test]
     async fn test_case_wrapper_with_limit() {
         if !Rewriter::sql_push_down_enabled() {