feat(cubesql): Support BETWEEN SQL push down (#9834)

Author: MazterQyou

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

@@ -4233,7 +4233,8 @@ export class BaseQuery {
         within_group: '{{ fun_sql }} WITHIN GROUP (ORDER BY {{ within_group_concat }})',
         concat_strings: '{{ strings | join(\' || \' ) }}',
         rolling_window_expr_timestamp_cast: '{{ value }}',
-        timestamp_literal: '{{ value }}'
+        timestamp_literal: '{{ value }}',
+        between: '{{ expr }} {% if negated %}NOT {% endif %}BETWEEN {{ low }} AND {{ high }}',
       },
       tesseract: {
         ilike: '{{ expr }} {% if negated %}NOT {% endif %}ILIKE {{ pattern }}', // May require different overloads in Tesseract than the ilike from expressions used in SQLAPI.

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

@@ -2018,7 +2018,47 @@ impl WrappedSelectNode {
                 Ok((resulting_sql, sql_query))
             }
             // Expr::GetIndexedField { .. } => {}
-            // Expr::Between { .. } => {}
+            Expr::Between {
+                expr,
+                negated,
+                low,
+                high,
+            } => {
+                let (expr, sql_query) = Self::generate_sql_for_expr_rec(
+                    sql_query,
+                    sql_generator.clone(),
+                    *expr,
+                    push_to_cube_context,
+                    subqueries,
+                )
+                .await?;
+                let (low, sql_query) = Self::generate_sql_for_expr_rec(
+                    sql_query,
+                    sql_generator.clone(),
+                    *low,
+                    push_to_cube_context,
+                    subqueries,
+                )
+                .await?;
+                let (high, sql_query) = Self::generate_sql_for_expr_rec(
+                    sql_query,
+                    sql_generator.clone(),
+                    *high,
+                    push_to_cube_context,
+                    subqueries,
+                )
+                .await?;
+                let resulting_sql = sql_generator
+                    .get_sql_templates()
+                    .between_expr(expr, negated, low, high)
+                    .map_err(|e| {
+                        DataFusionError::Internal(format!(
+                            "Can't generate SQL for between expr: {}",
+                            e
+                        ))
+                    })?;
+                Ok((resulting_sql, sql_query))
+            }
             Expr::Case {
                 expr,
                 when_then_expr,

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

@@ -0,0 +1,102 @@
+use crate::{
+    compile::rewrite::{
+        between_expr, rewrite,
+        rewriter::{CubeEGraph, CubeRewrite},
+        rules::wrapper::WrapperRules,
+        transforming_rewrite, wrapper_pullup_replacer, wrapper_pushdown_replacer,
+        wrapper_replacer_context,
+    },
+    var,
+};
+use egg::Subst;
+
+impl WrapperRules {
+    pub fn between_expr_rules(&self, rules: &mut Vec<CubeRewrite>) {
+        rules.extend(vec![
+            rewrite(
+                "wrapper-push-down-between-expr",
+                wrapper_pushdown_replacer(
+                    between_expr("?expr", "?negated", "?low", "?high"),
+                    "?context",
+                ),
+                between_expr(
+                    wrapper_pushdown_replacer("?expr", "?context"),
+                    "?negated",
+                    wrapper_pushdown_replacer("?low", "?context"),
+                    wrapper_pushdown_replacer("?high", "?context"),
+                ),
+            ),
+            transforming_rewrite(
+                "wrapper-pull-up-between-expr",
+                between_expr(
+                    wrapper_pullup_replacer(
+                        "?expr",
+                        wrapper_replacer_context(
+                            "?alias_to_cube",
+                            "?push_to_cube",
+                            "?in_projection",
+                            "?cube_members",
+                            "?grouped_subqueries",
+                            "?ungrouped_scan",
+                            "?input_data_source",
+                        ),
+                    ),
+                    "?negated",
+                    wrapper_pullup_replacer(
+                        "?low",
+                        wrapper_replacer_context(
+                            "?alias_to_cube",
+                            "?push_to_cube",
+                            "?in_projection",
+                            "?cube_members",
+                            "?grouped_subqueries",
+                            "?ungrouped_scan",
+                            "?input_data_source",
+                        ),
+                    ),
+                    wrapper_pullup_replacer(
+                        "?high",
+                        wrapper_replacer_context(
+                            "?alias_to_cube",
+                            "?push_to_cube",
+                            "?in_projection",
+                            "?cube_members",
+                            "?grouped_subqueries",
+                            "?ungrouped_scan",
+                            "?input_data_source",
+                        ),
+                    ),
+                ),
+                wrapper_pullup_replacer(
+                    between_expr("?expr", "?negated", "?low", "?high"),
+                    wrapper_replacer_context(
+                        "?alias_to_cube",
+                        "?push_to_cube",
+                        "?in_projection",
+                        "?cube_members",
+                        "?grouped_subqueries",
+                        "?ungrouped_scan",
+                        "?input_data_source",
+                    ),
+                ),
+                self.transform_between_expr("?input_data_source"),
+            ),
+        ]);
+    }
+
+    fn transform_between_expr(
+        &self,
+        input_data_source_var: &'static str,
+    ) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
+        let input_data_source_var = var!(input_data_source_var);
+        let meta = self.meta_context.clone();
+        move |egraph, subst| {
+            let Ok(data_source) = Self::get_data_source(egraph, subst, input_data_source_var)
+            else {
+                return false;
+            };
+
+            Self::can_rewrite_template(&data_source, &meta, "expressions/between")
+        }
+    }
+}

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

@@ -1,6 +1,7 @@
 mod aggregate;
 mod aggregate_function;
 mod alias;
+mod between_expr;
 mod binary_expr;
 mod case;
 mod cast;
@@ -90,6 +91,7 @@ impl RewriteRules for WrapperRules {
         self.not_expr_rules(&mut rules);
         self.distinct_rules(&mut rules);
         self.like_expr_rules(&mut rules);
+        self.between_expr_rules(&mut rules);
 
         rules
     }

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

@@ -633,6 +633,7 @@ OFFSET {{ offset }}{% endif %}"#.to_string(),
                     ("expressions/ilike".to_string(), "{{ expr }} {% if negated %}NOT {% endif %}ILIKE {{ pattern }}".to_string()),
                     ("expressions/like_escape".to_string(), "{{ like_expr }} ESCAPE {{ escape_char }}".to_string()),
                     ("expressions/within_group".to_string(), "{{ fun_sql }} WITHIN GROUP (ORDER BY {{ within_group_concat }})".to_string()),
+                    ("expressions/between".to_string(), "{{ expr }} {% if negated %}NOT {% endif %}BETWEEN {{ low }} AND {{ high }}".to_string()),
                     ("join_types/inner".to_string(), "INNER".to_string()),
                     ("join_types/left".to_string(), "LEFT".to_string()),
                     ("quotes/identifiers".to_string(), "\"".to_string()),

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

@@ -1912,3 +1912,40 @@ GROUP BY
         }
     );
 }
+
+#[tokio::test]
+async fn test_wrapper_between() {
+    if !Rewriter::sql_push_down_enabled() {
+        return;
+    }
+    init_testing_logger();
+
+    let query_plan = convert_select_to_query_plan(
+        // language=PostgreSQL
+        r#"
+        SELECT
+            customer_gender
+        FROM KibanaSampleDataEcommerce
+        WHERE
+            KibanaSampleDataEcommerce.customer_gender = customer_gender
+            AND order_date BETWEEN '2024-01-01' AND '2024-12-31'
+        GROUP BY 1
+        ;"#
+        .to_string(),
+        DatabaseProtocol::PostgreSQL,
+    )
+    .await;
+
+    let physical_plan = query_plan.as_physical_plan().await.unwrap();
+    println!(
+        "Physical plan: {}",
+        displayable(physical_plan.as_ref()).indent()
+    );
+
+    assert!(query_plan
+        .as_logical_plan()
+        .find_cube_scan_wrapped_sql()
+        .wrapped_sql
+        .sql
+        .contains("BETWEEN $1 AND $2"));
+}

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

@@ -880,6 +880,24 @@ impl SqlTemplates {
         )
     }
 
+    pub fn between_expr(
+        &self,
+        expr: String,
+        negated: bool,
+        low: String,
+        high: String,
+    ) -> Result<String, CubeError> {
+        self.render_template(
+            "expressions/between",
+            context! {
+                expr => expr,
+                negated => negated,
+                low => low,
+                high => high
+            },
+        )
+    }
+
     pub fn param(&self, param_index: usize) -> Result<String, CubeError> {
         self.render_template("params/param", context! { param_index => param_index })
     }