feat(cubesql): Flatten IN lists expressions to improve performance (#8235)

Nikita-str · web-flow · commit 66aa01d4f8f8 · 2024-05-14T16:39:26.000+06:00
diff --git a/rust/cubesql/cubesql/benches/benchmarks.rs b/rust/cubesql/cubesql/benches/benchmarks.rs
@@ -2,6 +2,7 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use cubesql::compile::test::rewrite_engine::{
     cube_context, query_to_logical_plan, rewrite_rules, rewrite_runner,
 };
+use itertools::Itertools;
 use std::sync::Arc;
 
 macro_rules! bench_func {
@@ -348,9 +349,83 @@ pub fn power_bi_sum_wrap(c: &mut Criterion) {
     bench_func!("power_bi_sum_wrap", get_power_bi_sum_wrap(), c);
 }
 
+fn get_simple_long_in_expr() -> String {
+    const N: usize = 50;
+    let set = (1..=N).join(", ");
+    format!("SELECT * FROM NumberCube WHERE someNumber IN ({set})")
+}
+
+pub fn long_simple_in_expr(c: &mut Criterion) {
+    std::env::set_var("CUBESQL_SQL_PUSH_DOWN", "true");
+    bench_func!("long_simple_in_expr", get_simple_long_in_expr(), c);
+}
+
+fn get_long_in_expr() -> String {
+    r#"
+    SELECT 
+        "WideCube"."dim1" as "column1", 
+        "WideCube"."dim2" as "column2", 
+        "WideCube"."dim3" as "column3", 
+        "WideCube"."dim4" as "column4", 
+        "WideCube"."dim5" as "column5", 
+        "WideCube"."dim6" as "column6", 
+        "WideCube"."dim7" as "column7", 
+        "WideCube"."dim8" as "column8", 
+        "WideCube"."dim9" as "column9", 
+        "WideCube"."dim10" as "column10", 
+        "WideCube"."dim11" as "column11", 
+        "WideCube"."dim12" as "column12", 
+        "WideCube"."dim13" as "column13", 
+        "WideCube"."dim14" as "column14", 
+        "WideCube"."dim15" as "column15", 
+        SUM("WideCube"."dim16") as "some_sum" 
+    FROM 
+        "WideCube" 
+    WHERE 
+        "WideCube"."dim1" = 1 
+        AND "WideCube"."dim2" = 2 
+        AND "WideCube"."dim3" = 3 
+        AND "WideCube"."dim4" = 4 
+        AND "WideCube"."dim5" = 5 
+        AND "WideCube"."dim6" = 6 
+        AND "WideCube"."dim7" = 7 
+        AND "WideCube"."dim8" = 8 
+        AND "WideCube"."dim9" = 9 
+        AND "WideCube"."dim10" = 10 
+        AND ("WideCube"."dim11" = 42 OR "WideCube"."dim11" IS NULL) 
+        AND (
+            "WideCube"."dim12" IN (
+                1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 
+                27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50
+            ) OR "WideCube"."dim12" IS NULL
+        ) AND "WideCube"."dim20" = 55 
+    GROUP BY 
+        "WideCube"."dim1", 
+        "WideCube"."dim2", 
+        "WideCube"."dim3", 
+        "WideCube"."dim4", 
+        "WideCube"."dim5", 
+        "WideCube"."dim6", 
+        "WideCube"."dim7", 
+        "WideCube"."dim8", 
+        "WideCube"."dim9", 
+        "WideCube"."dim10", 
+        "WideCube"."dim11", 
+        "WideCube"."dim12", 
+        "WideCube"."dim13", 
+        "WideCube"."dim14", 
+        "WideCube"."dim15"
+  "#.into()
+}
+
+pub fn long_in_expr(c: &mut Criterion) {
+    std::env::set_var("CUBESQL_SQL_PUSH_DOWN", "true");
+    bench_func!("long_in_expr", get_long_in_expr(), c);
+}
+
 criterion_group! {
     name = benches;
     config = Criterion::default().measurement_time(std::time::Duration::from_secs(15)).sample_size(10);
-    targets = split_query, split_query_count_distinct, wrapped_query, power_bi_wrap, power_bi_sum_wrap
+    targets = split_query, split_query_count_distinct, wrapped_query, power_bi_wrap, power_bi_sum_wrap, long_in_expr, long_simple_in_expr
 }
 criterion_main!(benches);
diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs
@@ -19422,6 +19422,9 @@ ORDER BY "source"."str0" ASC
 
     #[tokio::test]
     async fn test_thoughtspot_where_binary_in_true_false() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
         init_logger();
 
         let logical_plan = convert_select_to_query_plan(
@@ -19453,88 +19456,17 @@ ORDER BY "source"."str0" ASC
         .await
         .as_logical_plan();
 
-        assert_eq!(
-            logical_plan.find_cube_scan().request,
-            V1LoadRequestQuery {
-                measures: Some(vec!["KibanaSampleDataEcommerce.count".to_string()]),
-                dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
-                segments: Some(vec![]),
-                time_dimensions: None,
-                order: None,
-                limit: None,
-                offset: None,
-                filters: Some(vec![
-                    V1LoadRequestQueryFilterItem {
-                        member: None,
-                        operator: None,
-                        values: None,
-                        or: Some(vec![
-                            json!(V1LoadRequestQueryFilterItem {
-                                member: None,
-                                operator: None,
-                                values: None,
-                                or: None,
-                                and: Some(vec![
-                                    json!(V1LoadRequestQueryFilterItem {
-                                        member: Some(
-                                            "KibanaSampleDataEcommerce.customer_gender".to_string()
-                                        ),
-                                        operator: Some("startsWith".to_string()),
-                                        values: Some(vec!["female".to_string()]),
-                                        or: None,
-                                        and: None,
-                                    }),
-                                    json!(V1LoadRequestQueryFilterItem {
-                                        member: Some(
-                                            "KibanaSampleDataEcommerce.customer_gender".to_string()
-                                        ),
-                                        operator: Some("endsWith".to_string()),
-                                        values: Some(vec!["female".to_string()]),
-                                        or: None,
-                                        and: None,
-                                    }),
-                                ]),
-                            }),
-                            json!(V1LoadRequestQueryFilterItem {
-                                member: None,
-                                operator: None,
-                                values: None,
-                                or: None,
-                                and: Some(vec![
-                                    json!(V1LoadRequestQueryFilterItem {
-                                        member: Some(
-                                            "KibanaSampleDataEcommerce.customer_gender".to_string()
-                                        ),
-                                        operator: Some("startsWith".to_string()),
-                                        values: Some(vec!["male".to_string()]),
-                                        or: None,
-                                        and: None,
-                                    }),
-                                    json!(V1LoadRequestQueryFilterItem {
-                                        member: Some(
-                                            "KibanaSampleDataEcommerce.customer_gender".to_string()
-                                        ),
-                                        operator: Some("endsWith".to_string()),
-                                        values: Some(vec!["male".to_string()]),
-                                        or: None,
-                                        and: None,
-                                    }),
-                                ]),
-                            }),
-                        ]),
-                        and: None,
-                    },
-                    V1LoadRequestQueryFilterItem {
-                        member: Some("KibanaSampleDataEcommerce.customer_gender".to_string()),
-                        operator: Some("set".to_string()),
-                        values: None,
-                        or: None,
-                        and: None,
-                    },
-                ]),
-                ungrouped: None,
-            }
-        )
+        // check if contains `(LOWER(..) = .. OR ..LOWER(..) = ..) IN (TRUE, FALSE)`
+        let re = Regex::new(r"\(LOWER ?\(.+\) = .+ OR .+LOWER ?\(.+\) = .+\) IN \(TRUE, FALSE\)")
+            .unwrap();
+
+        let sql = logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql;
+
+        assert!(re.is_match(&sql));
     }
 
     #[tokio::test]
@@ -22901,4 +22833,39 @@ LIMIT {{ limit }}{% endif %}"#.to_string(),
             displayable(physical_plan.as_ref()).indent()
         );
     }
+
+    #[tokio::test]
+    async fn test_long_in_expr() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+
+        const N: usize = 50;
+        let set = (1..=N).join(", ");
+
+        let query = format!("SELECT * FROM NumberCube WHERE someNumber IN ({set})");
+        let query_plan = convert_select_to_query_plan(query, DatabaseProtocol::PostgreSQL).await;
+        let logical_plan = query_plan.as_logical_plan();
+
+        assert_eq!(
+            logical_plan.find_cube_scan().request,
+            V1LoadRequestQuery {
+                measures: Some(vec!["NumberCube.someNumber".into()]),
+                dimensions: Some(vec![]),
+                segments: Some(vec![]),
+                time_dimensions: None,
+                order: None,
+                limit: None,
+                offset: None,
+                filters: Some(vec![V1LoadRequestQueryFilterItem {
+                    member: Some("NumberCube.someNumber".into()),
+                    operator: Some("equals".into()),
+                    values: Some((1..=N).map(|x| x.to_string()).collect()),
+                    or: None,
+                    and: None
+                }]),
+                ungrouped: Some(true),
+            }
+        );
+    }
 }
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/converter.rs b/rust/cubesql/cubesql/src/compile/rewrite/converter.rs
@@ -96,6 +96,16 @@ macro_rules! add_expr_list_node {
     }};
 }
 
+macro_rules! add_expr_flat_list_node {
+    ($graph:expr, $value_expr:expr, $query_params:expr, $field_variant:ident) => {{
+        let list = $value_expr
+            .iter()
+            .map(|expr| Self::add_expr_replace_params($graph, expr, $query_params))
+            .collect::<Result<Vec<_>, _>>()?;
+        $graph.add(LogicalPlanLanguage::$field_variant(list))
+    }};
+}
+
 macro_rules! add_binary_expr_list_node {
     ($graph:expr, $value_expr:expr, $query_params:expr, $field_variant:ident) => {{
         fn to_binary_tree(
@@ -421,7 +431,7 @@ impl LogicalPlanToLanguageConverter {
                 negated,
             } => {
                 let expr = Self::add_expr_replace_params(graph, expr, query_params)?;
-                let list = add_expr_list_node!(graph, list, query_params, InListExprList);
+                let list = add_expr_flat_list_node!(graph, list, query_params, InListExprList);
                 let negated = add_expr_data_node!(graph, negated, InListExprNegated);
                 graph.add(LogicalPlanLanguage::InListExpr([expr, list, negated]))
             }
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/mod.rs b/rust/cubesql/cubesql/src/compile/rewrite/mod.rs
@@ -735,6 +735,13 @@ fn list_expr(list_type: impl Display, list: Vec<impl Display>) -> String {
     current
 }
 
+#[allow(unused)]
+fn flat_list_expr(list_type: impl Display, list: Vec<impl Display>) -> String {
+    use itertools::Itertools;
+    let list = list.iter().join(" ");
+    format!("({list_type} {list})")
+}
+
 fn udf_expr(fun_name: impl Display, args: Vec<impl Display>) -> String {
     udf_expr_var_arg(fun_name, list_expr("ScalarUDFExprArgs", args))
 }
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/filters.rs
@@ -7,7 +7,7 @@ use crate::{
         column_expr, cube_scan, cube_scan_filters, cube_scan_filters_empty_tail, cube_scan_members,
         dimension_expr, expr_column_name, filter, filter_member, filter_op, filter_op_filters,
         filter_op_filters_empty_tail, filter_replacer, filter_simplify_replacer, fun_expr,
-        fun_expr_var_arg, inlist_expr, is_not_null_expr, is_null_expr, like_expr, limit, list_expr,
+        fun_expr_var_arg, inlist_expr, is_not_null_expr, is_null_expr, like_expr, limit,
         literal_bool, literal_expr, literal_int, literal_string, measure_expr,
         member_name_by_alias, negative_expr, not_expr, projection, rewrite,
         rewriter::RewriteRules,
@@ -1674,56 +1674,6 @@ impl RewriteRules for FilterRules {
                     "?filter_aliases",
                 ),
             ),
-            rewrite(
-                "filter-thoughtspot-lower-in-true-false",
-                filter_replacer(
-                    inlist_expr(
-                        binary_expr(
-                            binary_expr(
-                                fun_expr("Lower", vec![column_expr("?column")]),
-                                "=",
-                                literal_expr("?left_literal"),
-                            ),
-                            "OR",
-                            binary_expr(
-                                fun_expr("Lower", vec![column_expr("?column")]),
-                                "=",
-                                literal_expr("?right_literal"),
-                            ),
-                        ),
-                        list_expr(
-                            "InListExprList",
-                            vec![literal_bool(true), literal_bool(false)],
-                        ),
-                        "InListExprNegated:false",
-                    ),
-                    "?alias_to_cube",
-                    "?members",
-                    "?filter_aliases",
-                ),
-                filter_replacer(
-                    binary_expr(
-                        binary_expr(
-                            binary_expr(
-                                fun_expr("Lower", vec![column_expr("?column")]),
-                                "=",
-                                literal_expr("?left_literal"),
-                            ),
-                            "OR",
-                            binary_expr(
-                                fun_expr("Lower", vec![column_expr("?column")]),
-                                "=",
-                                literal_expr("?right_literal"),
-                            ),
-                        ),
-                        "AND",
-                        is_not_null_expr(column_expr("?column")),
-                    ),
-                    "?alias_to_cube",
-                    "?members",
-                    "?filter_aliases",
-                ),
-            ),
             transforming_rewrite(
                 "extract-year-equals",
                 filter_replacer(
diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/in_list_expr.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/in_list_expr.rs
diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -735,6 +735,13 @@ fn list_expr(list_type: impl Display, list: Vec<impl Display>) -> String {`
`735`	`735`	`current`
`736`	`736`	`}`
`737`	`737`
	`738`	`+#[allow(unused)]`
	`739`	`+fn flat_list_expr(list_type: impl Display, list: Vec<impl Display>) -> String {`
	`740`	`+ use itertools::Itertools;`
	`741`	`+ let list = list.iter().join(" ");`
	`742`	`+ format!("({list_type} {list})")`
	`743`	`+}`
	`744`	`+`
`738`	`745`	`fn udf_expr(fun_name: impl Display, args: Vec<impl Display>) -> String {`
`739`	`746`	`udf_expr_var_arg(fun_name, list_expr("ScalarUDFExprArgs", args))`
`740`	`747`	`}`