Skip to content

Commit 195402f

Browse files
MazterQyouMazterQyou
authored
feat(cubesql): Support BETWEEN SQL push down (cube-js#9834)
1 parent b35a4e0 commit 195402f

File tree

7 files changed

+203
-2
lines changed

7 files changed

+203
-2
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4233,7 +4233,8 @@ export class BaseQuery {
42334233
within_group: '{{ fun_sql }} WITHIN GROUP (ORDER BY {{ within_group_concat }})',
42344234
concat_strings: '{{ strings | join(\' || \' ) }}',
42354235
rolling_window_expr_timestamp_cast: '{{ value }}',
4236-
timestamp_literal: '{{ value }}'
4236+
timestamp_literal: '{{ value }}',
4237+
between: '{{ expr }} {% if negated %}NOT {% endif %}BETWEEN {{ low }} AND {{ high }}',
42374238
},
42384239
tesseract: {
42394240
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

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2018,7 +2018,47 @@ impl WrappedSelectNode {
20182018
Ok((resulting_sql, sql_query))
20192019
}
20202020
// Expr::GetIndexedField { .. } => {}
2021-
// Expr::Between { .. } => {}
2021+
Expr::Between {
2022+
expr,
2023+
negated,
2024+
low,
2025+
high,
2026+
} => {
2027+
let (expr, sql_query) = Self::generate_sql_for_expr_rec(
2028+
sql_query,
2029+
sql_generator.clone(),
2030+
*expr,
2031+
push_to_cube_context,
2032+
subqueries,
2033+
)
2034+
.await?;
2035+
let (low, sql_query) = Self::generate_sql_for_expr_rec(
2036+
sql_query,
2037+
sql_generator.clone(),
2038+
*low,
2039+
push_to_cube_context,
2040+
subqueries,
2041+
)
2042+
.await?;
2043+
let (high, sql_query) = Self::generate_sql_for_expr_rec(
2044+
sql_query,
2045+
sql_generator.clone(),
2046+
*high,
2047+
push_to_cube_context,
2048+
subqueries,
2049+
)
2050+
.await?;
2051+
let resulting_sql = sql_generator
2052+
.get_sql_templates()
2053+
.between_expr(expr, negated, low, high)
2054+
.map_err(|e| {
2055+
DataFusionError::Internal(format!(
2056+
"Can't generate SQL for between expr: {}",
2057+
e
2058+
))
2059+
})?;
2060+
Ok((resulting_sql, sql_query))
2061+
}
20222062
Expr::Case {
20232063
expr,
20242064
when_then_expr,

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

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
use crate::{
2+
compile::rewrite::{
3+
between_expr, rewrite,
4+
rewriter::{CubeEGraph, CubeRewrite},
5+
rules::wrapper::WrapperRules,
6+
transforming_rewrite, wrapper_pullup_replacer, wrapper_pushdown_replacer,
7+
wrapper_replacer_context,
8+
},
9+
var,
10+
};
11+
use egg::Subst;
12+
13+
impl WrapperRules {
14+
pub fn between_expr_rules(&self, rules: &mut Vec<CubeRewrite>) {
15+
rules.extend(vec![
16+
rewrite(
17+
"wrapper-push-down-between-expr",
18+
wrapper_pushdown_replacer(
19+
between_expr("?expr", "?negated", "?low", "?high"),
20+
"?context",
21+
),
22+
between_expr(
23+
wrapper_pushdown_replacer("?expr", "?context"),
24+
"?negated",
25+
wrapper_pushdown_replacer("?low", "?context"),
26+
wrapper_pushdown_replacer("?high", "?context"),
27+
),
28+
),
29+
transforming_rewrite(
30+
"wrapper-pull-up-between-expr",
31+
between_expr(
32+
wrapper_pullup_replacer(
33+
"?expr",
34+
wrapper_replacer_context(
35+
"?alias_to_cube",
36+
"?push_to_cube",
37+
"?in_projection",
38+
"?cube_members",
39+
"?grouped_subqueries",
40+
"?ungrouped_scan",
41+
"?input_data_source",
42+
),
43+
),
44+
"?negated",
45+
wrapper_pullup_replacer(
46+
"?low",
47+
wrapper_replacer_context(
48+
"?alias_to_cube",
49+
"?push_to_cube",
50+
"?in_projection",
51+
"?cube_members",
52+
"?grouped_subqueries",
53+
"?ungrouped_scan",
54+
"?input_data_source",
55+
),
56+
),
57+
wrapper_pullup_replacer(
58+
"?high",
59+
wrapper_replacer_context(
60+
"?alias_to_cube",
61+
"?push_to_cube",
62+
"?in_projection",
63+
"?cube_members",
64+
"?grouped_subqueries",
65+
"?ungrouped_scan",
66+
"?input_data_source",
67+
),
68+
),
69+
),
70+
wrapper_pullup_replacer(
71+
between_expr("?expr", "?negated", "?low", "?high"),
72+
wrapper_replacer_context(
73+
"?alias_to_cube",
74+
"?push_to_cube",
75+
"?in_projection",
76+
"?cube_members",
77+
"?grouped_subqueries",
78+
"?ungrouped_scan",
79+
"?input_data_source",
80+
),
81+
),
82+
self.transform_between_expr("?input_data_source"),
83+
),
84+
]);
85+
}
86+
87+
fn transform_between_expr(
88+
&self,
89+
input_data_source_var: &'static str,
90+
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
91+
let input_data_source_var = var!(input_data_source_var);
92+
let meta = self.meta_context.clone();
93+
move |egraph, subst| {
94+
let Ok(data_source) = Self::get_data_source(egraph, subst, input_data_source_var)
95+
else {
96+
return false;
97+
};
98+
99+
Self::can_rewrite_template(&data_source, &meta, "expressions/between")
100+
}
101+
}
102+
}

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
mod aggregate;
22
mod aggregate_function;
33
mod alias;
4+
mod between_expr;
45
mod binary_expr;
56
mod case;
67
mod cast;
@@ -90,6 +91,7 @@ impl RewriteRules for WrapperRules {
9091
self.not_expr_rules(&mut rules);
9192
self.distinct_rules(&mut rules);
9293
self.like_expr_rules(&mut rules);
94+
self.between_expr_rules(&mut rules);
9395

9496
rules
9597
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -633,6 +633,7 @@ OFFSET {{ offset }}{% endif %}"#.to_string(),
633633
("expressions/ilike".to_string(), "{{ expr }} {% if negated %}NOT {% endif %}ILIKE {{ pattern }}".to_string()),
634634
("expressions/like_escape".to_string(), "{{ like_expr }} ESCAPE {{ escape_char }}".to_string()),
635635
("expressions/within_group".to_string(), "{{ fun_sql }} WITHIN GROUP (ORDER BY {{ within_group_concat }})".to_string()),
636+
("expressions/between".to_string(), "{{ expr }} {% if negated %}NOT {% endif %}BETWEEN {{ low }} AND {{ high }}".to_string()),
636637
("join_types/inner".to_string(), "INNER".to_string()),
637638
("join_types/left".to_string(), "LEFT".to_string()),
638639
("quotes/identifiers".to_string(), "\"".to_string()),

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

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1912,3 +1912,40 @@ GROUP BY
19121912
}
19131913
);
19141914
}
1915+
1916+
#[tokio::test]
1917+
async fn test_wrapper_between() {
1918+
if !Rewriter::sql_push_down_enabled() {
1919+
return;
1920+
}
1921+
init_testing_logger();
1922+
1923+
let query_plan = convert_select_to_query_plan(
1924+
// language=PostgreSQL
1925+
r#"
1926+
SELECT
1927+
customer_gender
1928+
FROM KibanaSampleDataEcommerce
1929+
WHERE
1930+
KibanaSampleDataEcommerce.customer_gender = customer_gender
1931+
AND order_date BETWEEN '2024-01-01' AND '2024-12-31'
1932+
GROUP BY 1
1933+
;"#
1934+
.to_string(),
1935+
DatabaseProtocol::PostgreSQL,
1936+
)
1937+
.await;
1938+
1939+
let physical_plan = query_plan.as_physical_plan().await.unwrap();
1940+
println!(
1941+
"Physical plan: {}",
1942+
displayable(physical_plan.as_ref()).indent()
1943+
);
1944+
1945+
assert!(query_plan
1946+
.as_logical_plan()
1947+
.find_cube_scan_wrapped_sql()
1948+
.wrapped_sql
1949+
.sql
1950+
.contains("BETWEEN $1 AND $2"));
1951+
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -880,6 +880,24 @@ impl SqlTemplates {
880880
)
881881
}
882882

883+
pub fn between_expr(
884+
&self,
885+
expr: String,
886+
negated: bool,
887+
low: String,
888+
high: String,
889+
) -> Result<String, CubeError> {
890+
self.render_template(
891+
"expressions/between",
892+
context! {
893+
expr => expr,
894+
negated => negated,
895+
low => low,
896+
high => high
897+
},
898+
)
899+
}
900+
883901
pub fn param(&self, param_index: usize) -> Result<String, CubeError> {
884902
self.render_template("params/param", context! { param_index => param_index })
885903
}

0 commit comments

Comments
 (0)