From cd32f5532874437e28f8e816b3591da8707c9156 Mon Sep 17 00:00:00 2001 From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com> Date: Thu, 7 Aug 2025 18:28:21 +0400 Subject: [PATCH] fix(cubesql): Improve Trino SQL push down compatibility Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com> --- .../src/adapter/PrestodbQuery.ts | 1 + .../cubesql/src/compile/date_parser.rs | 4 +- rust/cubesql/cubesql/src/compile/mod.rs | 47 +++++++++++++++++++ rust/cubesql/pg-srv/src/values/timestamp.rs | 1 + 4 files changed, 52 insertions(+), 1 deletion(-) diff --git a/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts b/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts index 70505ef1a60e8..43bdd06da45b2 100644 --- a/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts +++ b/packages/cubejs-schema-compiler/src/adapter/PrestodbQuery.ts @@ -136,6 +136,7 @@ export class PrestodbQuery extends BaseQuery { const templates = super.sqlTemplates(); templates.functions.DATETRUNC = 'DATE_TRUNC({{ args_concat }})'; templates.functions.DATEPART = 'DATE_PART({{ args_concat }})'; + templates.functions.DATEDIFF = 'DATE_DIFF(\'{{ date_part }}\', {{ args[1] }}, {{ args[2] }})'; templates.functions.CURRENTDATE = 'CURRENT_DATE'; delete templates.functions.PERCENTILECONT; templates.statements.select = '{% if ctes %} WITH \n' + diff --git a/rust/cubesql/cubesql/src/compile/date_parser.rs b/rust/cubesql/cubesql/src/compile/date_parser.rs index 77900a3c0f59d..7fc46a694abf4 100644 --- a/rust/cubesql/cubesql/src/compile/date_parser.rs +++ b/rust/cubesql/cubesql/src/compile/date_parser.rs @@ -2,9 +2,11 @@ use crate::compile::engine::df::scan::DataFusionError; use chrono::{NaiveDate, NaiveDateTime}; pub fn parse_date_str(s: &str) -> Result { - let parsed = NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f") + let parsed = NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S") .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f")) .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S")) + .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S%.f UTC")) + .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f")) .or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.fZ")) .or_else(|_| { NaiveDate::parse_from_str(s, "%Y-%m-%d").map(|date| date.and_hms_opt(0, 0, 0).unwrap()) diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index a6df1bdfd217e..2b3066d80ecaa 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -17267,4 +17267,51 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; assert!(sql.contains("CAST(1 AS VARCHAR)")); } + + #[tokio::test] + async fn test_trino_datediff() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + let query_plan = convert_select_to_query_plan_customized( + r#" + SELECT + KibanaSampleDataEcommerce.id, + KibanaSampleDataEcommerce.order_date, + KibanaSampleDataEcommerce.last_mod, + DATEDIFF( + day, + KibanaSampleDataEcommerce.order_date, + KibanaSampleDataEcommerce.last_mod + ) as conv_date_diff, + COUNT(*) + FROM KibanaSampleDataEcommerce + WHERE ( + KibanaSampleDataEcommerce.order_date > cast('2025-01-01T00:00:00.000' as timestamp) + AND KibanaSampleDataEcommerce.order_date < cast('2025-01-01T23:59:59.999' as timestamp) + AND KibanaSampleDataEcommerce.customer_gender = 'test' + ) + GROUP BY 1, 2, 3, 4 + "# + .to_string(), + DatabaseProtocol::PostgreSQL, + vec![( + "functions/DATEDIFF".to_string(), + "DATE_DIFF('{{ date_part }}', {{ args[1] }}, {{ args[2] }})".to_string(), + )], + ) + .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(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!(sql.contains("DATE_DIFF('day', ")); + } } diff --git a/rust/cubesql/pg-srv/src/values/timestamp.rs b/rust/cubesql/pg-srv/src/values/timestamp.rs index ac495bb5d1178..9a744ab8215a3 100644 --- a/rust/cubesql/pg-srv/src/values/timestamp.rs +++ b/rust/cubesql/pg-srv/src/values/timestamp.rs @@ -157,6 +157,7 @@ impl FromProtocolValue for TimestampValue { // more formats, so let's align this with parse_date_str function from cubesql crate. let parsed_datetime = NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S") .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S%.f")) + .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%d %H:%M:%S%.f UTC")) .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S")) .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S%.f")) .or_else(|_| NaiveDateTime::parse_from_str(as_str, "%Y-%m-%dT%H:%M:%S%.fZ"))