diff --git a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js b/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js index c5de58b620094..7e03b61dc5bbd 100644 --- a/packages/cubejs-schema-compiler/src/adapter/BaseQuery.js +++ b/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. diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs index 4473834c3ce7a..34862a24eb1e1 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs +++ b/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, diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/between_expr.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/between_expr.rs new file mode 100644 index 0000000000000..0bc61a6bc3bce --- /dev/null +++ b/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) { + 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") + } + } +} diff --git a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/mod.rs b/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/mod.rs index 980e68c14ea47..e90b3558d3e14 100644 --- a/rust/cubesql/cubesql/src/compile/rewrite/rules/wrapper/mod.rs +++ b/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 } diff --git a/rust/cubesql/cubesql/src/compile/test/mod.rs b/rust/cubesql/cubesql/src/compile/test/mod.rs index d0a280fcca389..e9c15ab4fc0b5 100644 --- a/rust/cubesql/cubesql/src/compile/test/mod.rs +++ b/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()), diff --git a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs b/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs index ba029cc9299c7..94c28faf88220 100644 --- a/rust/cubesql/cubesql/src/compile/test/test_wrapper.rs +++ b/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")); +} diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs index e148971e44dda..e79e8a3f14f76 100644 --- a/rust/cubesql/cubesql/src/transport/service.rs +++ b/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 { + self.render_template( + "expressions/between", + context! { + expr => expr, + negated => negated, + low => low, + high => high + }, + ) + } + pub fn param(&self, param_index: usize) -> Result { self.render_template("params/param", context! { param_index => param_index }) }