fix(cubesql): column does not exist in case of ORDER BY is called on CASE column

Author: paveltiunov

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

@@ -437,15 +437,28 @@ impl CubeScanWrapperNode {
                                 .chain(group_expr.iter())
                                 .chain(aggr_expr.iter())
                                 .map(|e| {
-                                    Ok((
-                                        Column {
-                                            relation: alias.clone(),
-                                            name: expr_name(&e, &schema)?,
-                                        },
-                                        e.clone(),
-                                    ))
+                                    let name = expr_name(&e, &schema)?;
+                                    Ok(vec![
+                                        (
+                                            Column {
+                                                relation: alias.clone(),
+                                                name: name.clone(),
+                                            },
+                                            e.clone(),
+                                        ),
+                                        (
+                                            Column {
+                                                relation: None,
+                                                name: name,
+                                            },
+                                            e.clone(),
+                                        ),
+                                    ])
                                 })
-                                .collect::<Result<HashMap<_, _>>>()?;
+                                .collect::<Result<Vec<_>>>()?
+                                .into_iter()
+                                .flatten()
+                                .collect::<HashMap<_, _>>();
 
                             let (order, mut sql) = Self::generate_column_expr(
                                 plan.clone(),

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

@@ -18570,6 +18570,35 @@ ORDER BY \"COUNT(count)\" DESC"
         );
     }
 
+    #[tokio::test]
+    async fn test_case_wrapper_alias_with_order() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            "SELECT CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END AS \"f822c516-3515-11c2-8464-5d4845a02f73\", AVG(avgPrice) mp FROM KibanaSampleDataEcommerce a GROUP BY CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END ORDER BY CASE WHEN customer_gender = 'female' THEN 'f' ELSE 'm' END NULLS FIRST LIMIT 500"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await;
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert!(logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql
+            .contains("ORDER BY \"case_when_a_cust\""));
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+    }
+
     #[tokio::test]
     async fn test_case_wrapper_ungrouped() {
         if !Rewriter::sql_push_down_enabled() {