feat(cubesql): starts_with, ends_with, LOWER(?column) = ?literal (cube-js#5310)

Author: ovr

rust/cubesql/cubesql/src/compile/engine/udf.rs

@@ -676,6 +676,35 @@ pub fn create_timediff_udf() -> ScalarUDF {
     )
 }
 
+pub fn create_ends_with_udf() -> ScalarUDF {
+    let fun = make_scalar_function(move |args: &[ArrayRef]| {
+        assert!(args.len() == 2);
+
+        let string_array = downcast_string_arg!(args[0], "string", i32);
+        let prefix_array = downcast_string_arg!(args[1], "prefix", i32);
+
+        let result = string_array
+            .iter()
+            .zip(prefix_array.iter())
+            .map(|(string, prefix)| match (string, prefix) {
+                (Some(string), Some(prefix)) => Some(string.ends_with(prefix)),
+                _ => None,
+            })
+            .collect::<BooleanArray>();
+
+        Ok(Arc::new(result) as ArrayRef)
+    });
+
+    let return_type: ReturnTypeFunction = Arc::new(move |_| Ok(Arc::new(DataType::Boolean)));
+
+    ScalarUDF::new(
+        "ends_with",
+        &Signature::exact(vec![DataType::Utf8, DataType::Utf8], Volatility::Immutable),
+        &return_type,
+        &fun,
+    )
+}
+
 // https://docs.aws.amazon.com/redshift/latest/dg/r_DATEDIFF_function.html
 pub fn create_datediff_udf() -> ScalarUDF {
     let fun = make_scalar_function(move |args: &[ArrayRef]| {

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

@@ -41,13 +41,13 @@ use self::{
             create_current_schemas_udf, create_current_timestamp_udf, create_current_user_udf,
             create_date_add_udf, create_date_sub_udf, create_date_udf, create_dateadd_udf,
             create_datediff_udf, create_dayofmonth_udf, create_dayofweek_udf, create_dayofyear_udf,
-            create_db_udf, create_format_type_udf, create_generate_series_udtf,
-            create_generate_subscripts_udtf, create_has_schema_privilege_udf, create_hour_udf,
-            create_if_udf, create_instr_udf, create_interval_mul_udf, create_isnull_udf,
-            create_json_build_object_udf, create_least_udf, create_locate_udf, create_makedate_udf,
-            create_measure_udaf, create_minute_udf, create_pg_backend_pid_udf,
-            create_pg_datetime_precision_udf, create_pg_expandarray_udtf,
-            create_pg_get_constraintdef_udf, create_pg_get_expr_udf,
+            create_db_udf, create_ends_with_udf, create_format_type_udf,
+            create_generate_series_udtf, create_generate_subscripts_udtf,
+            create_has_schema_privilege_udf, create_hour_udf, create_if_udf, create_instr_udf,
+            create_interval_mul_udf, create_isnull_udf, create_json_build_object_udf,
+            create_least_udf, create_locate_udf, create_makedate_udf, create_measure_udaf,
+            create_minute_udf, create_pg_backend_pid_udf, create_pg_datetime_precision_udf,
+            create_pg_expandarray_udtf, create_pg_get_constraintdef_udf, create_pg_get_expr_udf,
             create_pg_get_serial_sequence_udf, create_pg_get_userbyid_udf,
             create_pg_is_other_temp_schema, create_pg_my_temp_schema,
             create_pg_numeric_precision_udf, create_pg_numeric_scale_udf,
@@ -1128,6 +1128,7 @@ WHERE `TABLE_SCHEMA` = '{}'",
         ctx.register_udf(create_json_build_object_udf());
         ctx.register_udf(create_regexp_substr_udf());
         ctx.register_udf(create_interval_mul_udf());
+        ctx.register_udf(create_ends_with_udf());
 
         // udaf
         ctx.register_udaf(create_measure_udaf());
@@ -1766,6 +1767,112 @@ mod tests {
         )
     }
 
+    #[tokio::test]
+    async fn test_starts_with() {
+        let query_plan = convert_select_to_query_plan(
+            "SELECT COUNT(*) as cnt FROM KibanaSampleDataEcommerce WHERE starts_with(customer_gender, 'fe')"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await;
+
+        let cube_scan = query_plan.as_logical_plan().find_cube_scan();
+
+        assert_eq!(
+            cube_scan.request,
+            V1LoadRequestQuery {
+                measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string(),]),
+                segments: Some(vec![]),
+                dimensions: Some(vec![]),
+                time_dimensions: None,
+                order: None,
+                limit: None,
+                offset: None,
+                filters: Some(vec![V1LoadRequestQueryFilterItem {
+                    member: Some("KibanaSampleDataEcommerce.customer_gender".to_string()),
+                    operator: Some("startsWith".to_string()),
+                    values: Some(vec!["fe".to_string()]),
+                    or: None,
+                    and: None
+                }])
+            }
+        )
+    }
+
+    #[tokio::test]
+    async fn test_ends_with_query() {
+        let query_plan = convert_select_to_query_plan(
+            "SELECT COUNT(*) as cnt FROM KibanaSampleDataEcommerce WHERE ends_with(customer_gender, 'emale')"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await;
+
+        let cube_scan = query_plan.as_logical_plan().find_cube_scan();
+
+        assert_eq!(
+            cube_scan.request,
+            V1LoadRequestQuery {
+                measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string(),]),
+                segments: Some(vec![]),
+                dimensions: Some(vec![]),
+                time_dimensions: None,
+                order: None,
+                limit: None,
+                offset: None,
+                filters: Some(vec![V1LoadRequestQueryFilterItem {
+                    member: Some("KibanaSampleDataEcommerce.customer_gender".to_string()),
+                    operator: Some("endsWith".to_string()),
+                    values: Some(vec!["emale".to_string()]),
+                    or: None,
+                    and: None
+                }])
+            }
+        )
+    }
+
+    #[tokio::test]
+    async fn test_lower_equals_thoughtspot() {
+        let query_plan = convert_select_to_query_plan(
+            "SELECT COUNT(*) as cnt FROM KibanaSampleDataEcommerce WHERE LOWER(customer_gender) = 'female'"
+                .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+            .await;
+
+        let cube_scan = query_plan.as_logical_plan().find_cube_scan();
+
+        assert_eq!(
+            cube_scan.request,
+            V1LoadRequestQuery {
+                measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string(),]),
+                segments: Some(vec![]),
+                dimensions: Some(vec![]),
+                time_dimensions: None,
+                order: None,
+                limit: None,
+                offset: None,
+                // TODO: Migrate to equalsLower operator, when it will be available in Cube?
+                filters: Some(vec![
+                    V1LoadRequestQueryFilterItem {
+                        member: Some("KibanaSampleDataEcommerce.customer_gender".to_string()),
+                        operator: Some("startsWith".to_string()),
+                        values: Some(vec!["female".to_string()]),
+                        or: None,
+                        and: None
+                    },
+                    V1LoadRequestQueryFilterItem {
+                        member: Some("KibanaSampleDataEcommerce.customer_gender".to_string()),
+                        operator: Some("endsWith".to_string()),
+                        values: Some(vec!["female".to_string()]),
+                        or: None,
+                        and: None
+                    }
+                ])
+            }
+        )
+    }
+
     #[tokio::test]
     async fn test_change_user_via_in_filter_thoughtspot() {
         let query_plan = convert_select_to_query_plan(
@@ -5219,6 +5326,24 @@ ORDER BY \"COUNT(count)\" DESC"
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_ends_with() -> Result<(), CubeError> {
+        insta::assert_snapshot!(
+            "ends_with",
+            execute_query(
+                "select \
+                    ends_with('rust is killing me', 'me') as r1,
+                    ends_with('rust is killing me', 'no') as r2
+                "
+                .to_string(),
+                DatabaseProtocol::MySQL
+            )
+            .await?
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_locate() -> Result<(), CubeError> {
         assert_eq!(

rust/cubesql/cubesql/src/compile/rewrite/analysis.rs

@@ -276,6 +276,10 @@ impl LogicalPlanAnalysis {
                 push_referenced_columns(params[1], &mut vec)?;
                 Some(vec)
             }
+            LogicalPlanLanguage::ScalarUDFExpr(params) => {
+                push_referenced_columns(params[1], &mut vec)?;
+                Some(vec)
+            }
             LogicalPlanLanguage::AggregateFunctionExpr(params) => {
                 push_referenced_columns(params[1], &mut vec)?;
                 Some(vec)
@@ -297,7 +301,8 @@ impl LogicalPlanAnalysis {
             | LogicalPlanLanguage::CaseExprElseExpr(params)
             | LogicalPlanLanguage::CaseExprExpr(params)
             | LogicalPlanLanguage::AggregateFunctionExprArgs(params)
-            | LogicalPlanLanguage::ScalarFunctionExprArgs(params) => {
+            | LogicalPlanLanguage::ScalarFunctionExprArgs(params)
+            | LogicalPlanLanguage::ScalarUDFExprArgs(params) => {
                 for p in params.iter() {
                     vec.extend(referenced_columns(*p)?.into_iter());
                 }

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

@@ -508,11 +508,24 @@ fn list_expr(list_type: impl Display, list: Vec<impl Display>) -> String {
 }
 
 fn udf_expr(fun_name: impl Display, args: Vec<impl Display>) -> String {
-    format!(
-        "(ScalarUDFExpr ScalarUDFExprFun:{} {})",
-        fun_name,
-        list_expr("ScalarUDFExprArgs", args)
-    )
+    udf_expr_var_arg(fun_name, list_expr("ScalarUDFExprArgs", args))
+}
+
+fn udf_expr_var_arg(fun_name: impl Display, arg_list: impl Display) -> String {
+    let prefix = if fun_name.to_string().starts_with("?") {
+        ""
+    } else {
+        "ScalarUDFExprFun:"
+    };
+    format!("(ScalarUDFExpr {}{} {})", prefix, fun_name, arg_list)
+}
+
+fn udf_fun_expr_args(left: impl Display, right: impl Display) -> String {
+    format!("(ScalarUDFExprArgs {} {})", left, right)
+}
+
+fn udf_fun_expr_args_empty_tail() -> String {
+    "ScalarUDFExprArgs".to_string()
 }
 
 fn fun_expr(fun_name: impl Display, args: Vec<impl Display>) -> String {