dev

Author: RusovDmitriy

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

@@ -438,6 +438,28 @@ impl PlanRewriter for CollectConstraints {
         c: &Self::Context,
     ) -> Result<LogicalPlan, DataFusionError> {
         match &n {
+            LogicalPlan::Projection {
+                expr,
+                input,
+                schema,
+            } => {
+                let mut alias_to_column = HashMap::new();
+                expr.iter().for_each(|e| {
+                    if let Expr::Alias(box Expr::Column(c), alias) = e {
+                        alias_to_column.insert(alias.clone(), c.clone());
+                    }
+                });
+
+                self.constraints.iter_mut().for_each(|c| {
+                    c.sort_on.iter_mut().for_each(|sort_columns| {
+                        sort_columns.sort_on.iter_mut().for_each(|sort_column| {
+                            if let Some(column) = alias_to_column.get(sort_column) {
+                                *sort_column = column.name.clone();
+                            }
+                        });
+                    });
+                });
+            }
             LogicalPlan::TableScan {
                 projection,
                 filters,
@@ -755,6 +777,32 @@ impl PlanRewriter for ChooseIndex<'_> {
 
     fn enter_node(&mut self, n: &LogicalPlan, context: &Self::Context) -> Option<Self::Context> {
         match n {
+            LogicalPlan::Projection { expr, .. } => {
+                let mut alias_to_column = HashMap::new();
+                expr.iter().for_each(|e| {
+                    if let Expr::Alias(box Expr::Column(c), alias) = e {
+                        alias_to_column.insert(alias.clone(), c.clone());
+                    }
+                });
+
+                let names: Vec<String> = context
+                    .sort
+                    .clone()
+                    .unwrap_or_default()
+                    .iter()
+                    .map(|k| {
+                        alias_to_column
+                            .get(k)
+                            .map_or_else(|| k.clone(), |v| v.name.clone())
+                    })
+                    .collect();
+
+                if !names.is_empty() {
+                    Some(context.update_sort(names, context.sort_is_asc))
+                } else {
+                    None
+                }
+            }
             LogicalPlan::Limit { n, .. } => Some(context.update_limit(Some(*n))),
             LogicalPlan::Skip { n, .. } => {
                 if let Some(limit) = context.limit {

rust/cubestore/cubestore/src/streaming/kafka_post_processing.rs

@@ -481,7 +481,7 @@ impl KafkaPostProcessPlanner {
                             if let Some(existing_name) = &column_name {
                                 if existing_name != &name {
                                     return Err(CubeError::user(
-                                        "Scalar function can only use a single column".to_string(),
+                                        format!("Scalar function can only use a single column, expression: {:?}", expr),
                                     ));
                                 }
                             } else {
@@ -500,17 +500,15 @@ impl KafkaPostProcessPlanner {
             Expr::Alias(e, _) => match &**e {
                 Expr::Column(c) => Ok(c.name.clone()),
                 Expr::ScalarUDF { .. } => find_column_name(expr)?.ok_or_else(|| {
-                    CubeError::user("Scalar function must contain at least one column".to_string())
+                    CubeError::user(format!("Scalar function must contain at least one column, expression: {:?}", expr))
                 }),
-                _ => {
-                    println!("unknown expr: {:?}", e);
-                    Err(CubeError::user(format!(
-                        "Unique key can't be an expression in kafka streaming queries"
-                    )))
-                }
+                _ => Err(CubeError::user(format!(
+                    "Unique key can't be an expression in kafka streaming queries, expression: {:?}",
+                    expr
+                ))),
             },
             _ => Err(CubeError::user(
-                "All expressions must have aliases in kafka streaming queries".to_string(),
+                format!("All expressions must have aliases in kafka streaming queries, expression: {:?}", expr),
             )),
         }?;