feat(cubesql): Support NOT SQL push down (#7422)

Author: MazterQyou

packages/cubejs-schema-compiler/src/adapter/BaseQuery.js

@@ -2486,6 +2486,7 @@ class BaseQuery {
         window_function: '{{ fun_call }} OVER ({% if partition_by_concat %}PARTITION BY {{ partition_by_concat }}{% if order_by_concat %} {% endif %}{% endif %}{% if order_by_concat %}ORDER BY {{ order_by_concat }}{% endif %})',
         in_list: '{{ expr }} {% if negated %}NOT {% endif %}IN ({{ in_exprs_concat }})',
         negative: '-({{ expr }})',
+        not: 'NOT ({{ expr }})',
       },
       quotes: {
         identifiers: '"',

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

@@ -930,7 +930,27 @@ impl CubeScanWrapperNode {
                 // Expr::Like(_) => {}-=
                 // Expr::ILike(_) => {}
                 // Expr::SimilarTo(_) => {}
-                // Expr::Not(_) => {}
+                Expr::Not(expr) => {
+                    let (expr, sql_query) = Self::generate_sql_for_expr(
+                        plan.clone(),
+                        sql_query,
+                        sql_generator.clone(),
+                        *expr,
+                        ungrouped_scan_node.clone(),
+                    )
+                    .await?;
+                    let resulting_sql =
+                        sql_generator
+                            .get_sql_templates()
+                            .not_expr(expr)
+                            .map_err(|e| {
+                                DataFusionError::Internal(format!(
+                                    "Can't generate SQL for not expr: {}",
+                                    e
+                                ))
+                            })?;
+                    Ok((resulting_sql, sql_query))
+                }
                 Expr::IsNotNull(expr) => {
                     let (expr, sql_query) = Self::generate_sql_for_expr(
                         plan.clone(),

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

@@ -19767,4 +19767,38 @@ ORDER BY \"COUNT(count)\" DESC"
             .sql
             .contains("-("));
     }
+
+    #[tokio::test]
+    async fn test_not_expr() {
+        if !Rewriter::sql_push_down_enabled() {
+            return;
+        }
+        init_logger();
+
+        let query_plan = convert_select_to_query_plan(
+            "
+            SELECT NOT has_subscription AS has_no_subscription
+            FROM KibanaSampleDataEcommerce AS k
+            GROUP BY 1
+            ORDER BY 1 DESC
+            "
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+        )
+        .await;
+
+        let physical_plan = query_plan.as_physical_plan().await.unwrap();
+        println!(
+            "Physical plan: {}",
+            displayable(physical_plan.as_ref()).indent()
+        );
+
+        let logical_plan = query_plan.as_logical_plan();
+        assert!(logical_plan
+            .find_cube_scan_wrapper()
+            .wrapped_sql
+            .unwrap()
+            .sql
+            .contains("NOT ("));
+    }
 }

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

@@ -12,6 +12,7 @@ mod is_null_expr;
 mod limit;
 mod literal;
 mod negative_expr;
+mod not_expr;
 mod order;
 mod projection;
 mod scalar_function;
@@ -64,6 +65,7 @@ impl RewriteRules for WrapperRules {
         self.literal_rules(&mut rules);
         self.in_list_expr_rules(&mut rules);
         self.negative_expr_rules(&mut rules);
+        self.not_expr_rules(&mut rules);
 
         rules
     }

rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/not_expr.rs

@@ -0,0 +1,77 @@
+use crate::{
+    compile::rewrite::{
+        analysis::LogicalPlanAnalysis, not_expr, rewrite, rules::wrapper::WrapperRules,
+        transforming_rewrite, wrapper_pullup_replacer, wrapper_pushdown_replacer,
+        LogicalPlanLanguage, WrapperPullupReplacerAliasToCube,
+    },
+    var, var_iter,
+};
+use egg::{EGraph, Rewrite, Subst};
+
+impl WrapperRules {
+    pub fn not_expr_rules(
+        &self,
+        rules: &mut Vec<Rewrite<LogicalPlanLanguage, LogicalPlanAnalysis>>,
+    ) {
+        rules.extend(vec![
+            rewrite(
+                "wrapper-not-push-down",
+                wrapper_pushdown_replacer(
+                    not_expr("?expr"),
+                    "?alias_to_cube",
+                    "?ungrouped",
+                    "?cube_members",
+                ),
+                not_expr(wrapper_pushdown_replacer(
+                    "?expr",
+                    "?alias_to_cube",
+                    "?ungrouped",
+                    "?cube_members",
+                )),
+            ),
+            transforming_rewrite(
+                "wrapper-not-pull-up",
+                not_expr(wrapper_pullup_replacer(
+                    "?expr",
+                    "?alias_to_cube",
+                    "?ungrouped",
+                    "?cube_members",
+                )),
+                wrapper_pullup_replacer(
+                    not_expr("?expr"),
+                    "?alias_to_cube",
+                    "?ungrouped",
+                    "?cube_members",
+                ),
+                self.transform_not_expr("?alias_to_cube"),
+            ),
+        ]);
+    }
+
+    fn transform_not_expr(
+        &self,
+        alias_to_cube_var: &'static str,
+    ) -> impl Fn(&mut EGraph<LogicalPlanLanguage, LogicalPlanAnalysis>, &mut Subst) -> bool {
+        let alias_to_cube_var = var!(alias_to_cube_var);
+        let meta = self.cube_context.meta.clone();
+        move |egraph, subst| {
+            for alias_to_cube in var_iter!(
+                egraph[subst[alias_to_cube_var]],
+                WrapperPullupReplacerAliasToCube
+            )
+            .cloned()
+            {
+                if let Some(sql_generator) = meta.sql_generator_by_alias_to_cube(&alias_to_cube) {
+                    if sql_generator
+                        .get_sql_templates()
+                        .templates
+                        .contains_key("expressions/not")
+                    {
+                        return true;
+                    }
+                }
+            }
+            false
+        }
+    }
+}

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

@@ -236,6 +236,7 @@ pub fn get_test_tenant_ctx() -> Arc<MetaContext> {
                     ("expressions/window_function".to_string(), "{{ fun_call }} OVER ({% if partition_by %}PARTITION BY {{ partition_by }}{% if order_by %} {% endif %}{% endif %}{% if order_by %}ORDER BY {{ order_by }}{% endif %})".to_string()),
                     ("expressions/in_list".to_string(), "{{ expr }} {% if negated %}NOT {% endif %}IN ({{ in_exprs_concat }})".to_string()),
                     ("expressions/negative".to_string(), "-({{ expr }})".to_string()),
+                    ("expressions/not".to_string(), "NOT ({{ expr }})".to_string()),
                     ("quotes/identifiers".to_string(), "\"".to_string()),
                     ("quotes/escape".to_string(), "\"\"".to_string()),
                     ("params/param".to_string(), "${{ param_index + 1 }}".to_string())

rust/cubesql/cubesql/src/transport/service.rs

@@ -502,6 +502,10 @@ impl SqlTemplates {
         self.render_template("expressions/negative", context! { expr => expr })
     }
 
+    pub fn not_expr(&self, expr: String) -> Result<String, CubeError> {
+        self.render_template("expressions/not", context! { expr => expr })
+    }
+
     pub fn sort_expr(
         &self,
         expr: String,