dev

Author: RusovDmitriy

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

@@ -470,13 +470,61 @@ impl KafkaPostProcessPlanner {
     }
 
     fn get_source_unique_column(&self, expr: &Expr) -> Result<Column, CubeError> {
+        fn find_column_name(expr: &Expr) -> Option<String> {
+            match expr {
+                Expr::Column(c) => Some(c.name.clone()),
+                Expr::Alias(e, _) => find_column_name(&**e),
+                Expr::ScalarUDF { args, .. } => {
+                    let mut column_name: Option<String> = None;
+                    for arg in args {
+                        if let Some(name) = find_column_name(arg) {
+                            if let Some(existing_name) = &column_name {
+                                if existing_name != &name {
+                                    return None;
+                                }
+                            } else {
+                                column_name = Some(name);
+                            }
+                        }
+                    }
+                    column_name
+                }
+                _ => None,
+            }
+        }
+
         let source_name = match expr {
             Expr::Column(c) => Ok(c.name.clone()),
             Expr::Alias(e, _) => match &**e {
                 Expr::Column(c) => Ok(c.name.clone()),
-                _ => Err(CubeError::user(format!(
-                    "Unique key can't be an expression in kafka streaming queries"
-                ))),
+                Expr::ScalarUDF { args, .. } => {
+                    let mut column_name: Option<String> = None;
+                    for arg in args {
+                        match find_column_name(arg) {
+                            Some(name) => {
+                                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(),
+                                        ));
+                                    }
+                                } else {
+                                    column_name = Some(name);
+                                }
+                            }
+                            _ => {},
+                        }
+                    }
+                    column_name.ok_or_else(|| {
+                        CubeError::user("Scalar function must contain at least one column".to_string())
+                    })
+                },
+                _ => {
+                    println!("unknown expr: {:?}", e);
+                    Err(CubeError::user(format!(
+                        "Unique key can't be an expression in kafka streaming queries"
+                    )))
+                },
             },
             _ => Err(CubeError::user(
                 "All expressions must have aliases in kafka streaming queries".to_string(),