From ecaaf1ca5566ff7cc27289468440e6a902a09609 Mon Sep 17 00:00:00 2001 From: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com> Date: Mon, 29 Jan 2024 21:24:22 +0400 Subject: [PATCH] feat(cubesql): Extend `DATEDIFF` push down support --- .../src/DatabricksQuery.ts | 3 +- .../src/adapter/PostgresQuery.ts | 3 +- .../src/adapter/RedshiftQuery.ts | 1 + rust/cubesql/cubesql/src/compile/mod.rs | 63 +++++++++++++++++++ 4 files changed, 67 insertions(+), 3 deletions(-) diff --git a/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts b/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts index 5128cd80880b6..0a38daac050c5 100644 --- a/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts +++ b/packages/cubejs-databricks-jdbc-driver/src/DatabricksQuery.ts @@ -103,8 +103,7 @@ export class DatabricksQuery extends BaseQuery { templates.functions.BTRIM = 'TRIM({% if args[1] is defined %}{{ args[1] }} FROM {% endif %}{{ args[0] }})'; templates.functions.LTRIM = 'LTRIM({{ args|reverse|join(", ") }})'; templates.functions.RTRIM = 'RTRIM({{ args|reverse|join(", ") }})'; - // Databricks has a DATEDIFF function but produces values different from Redshift - delete templates.functions.DATEDIFF; + templates.functions.DATEDIFF = 'DATEDIFF({{ date_part }}, DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}))'; return templates; } } diff --git a/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.ts b/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.ts index a8773a0ff5c46..e3d8f069bda86 100644 --- a/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.ts +++ b/packages/cubejs-schema-compiler/src/adapter/PostgresQuery.ts @@ -54,7 +54,8 @@ export class PostgresQuery extends BaseQuery { templates.functions.NOW = 'NOW({{ args_concat }})'; // DATEADD is being rewritten to DATE_ADD // templates.functions.DATEADD = '({{ args[2] }} + \'{{ interval }} {{ date_part }}\'::interval)'; - delete templates.functions.DATEDIFF; + // TODO: is DATEDIFF expr worth documenting? + templates.functions.DATEDIFF = 'CASE WHEN LOWER(\'{{ date_part }}\') IN (\'year\', \'quarter\', \'month\') THEN (EXTRACT(YEAR FROM AGE(DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }}))) * 12 + EXTRACT(MONTH FROM AGE(DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }})))) / CASE LOWER(\'{{ date_part }}\') WHEN \'year\' THEN 12 WHEN \'quarter\' THEN 3 WHEN \'month\' THEN 1 END ELSE EXTRACT(EPOCH FROM DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}) - DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }})) / EXTRACT(EPOCH FROM \'1 {{ date_part }}\'::interval) END::bigint'; templates.expressions.interval = 'INTERVAL \'{{ interval }}\''; templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})'; diff --git a/packages/cubejs-schema-compiler/src/adapter/RedshiftQuery.ts b/packages/cubejs-schema-compiler/src/adapter/RedshiftQuery.ts index a74f48fe8f8a0..12dc05e6a92bb 100644 --- a/packages/cubejs-schema-compiler/src/adapter/RedshiftQuery.ts +++ b/packages/cubejs-schema-compiler/src/adapter/RedshiftQuery.ts @@ -15,6 +15,7 @@ export class RedshiftQuery extends PostgresQuery { public sqlTemplates() { const templates = super.sqlTemplates(); templates.functions.DLOG10 = 'LOG(10, {{ args_concat }})'; + templates.functions.DATEDIFF = 'DATEDIFF({{ date_part }}, {{ args[1] }}, {{ args[2] }})'; delete templates.functions.COVAR_POP; delete templates.functions.COVAR_SAMP; return templates; diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index 956fbbb3336e5..b3f76b4549b14 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -20803,6 +20803,69 @@ limit .sql; assert!(sql.contains("DATETIME_DIFF(CAST(")); assert!(sql.contains("day)")); + + // Databricks + let query_plan = convert_select_to_query_plan_customized( + " + SELECT DATEDIFF(DAY, order_date, last_mod) AS d + FROM KibanaSampleDataEcommerce AS k + GROUP BY 1 + ORDER BY 1 DESC + " + .to_string(), + DatabaseProtocol::PostgreSQL, + vec![ + ("functions/DATEDIFF".to_string(), "DATEDIFF({{ date_part }}, DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }}), DATE_TRUNC(\'{{ date_part }}\', {{ 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_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("DATEDIFF(day,")); + assert!(sql.contains("DATE_TRUNC('day',")); + + // PostgreSQL + let query_plan = convert_select_to_query_plan_customized( + " + SELECT DATEDIFF(DAY, order_date, last_mod) AS d + FROM KibanaSampleDataEcommerce AS k + GROUP BY 1 + ORDER BY 1 DESC + " + .to_string(), + DatabaseProtocol::PostgreSQL, + vec![ + ("functions/DATEDIFF".to_string(), "CASE WHEN LOWER(\'{{ date_part }}\') IN (\'year\', \'quarter\', \'month\') THEN (EXTRACT(YEAR FROM AGE(DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }}))) * 12 + EXTRACT(MONTH FROM AGE(DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}), DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }})))) / CASE LOWER(\'{{ date_part }}\') WHEN \'year\' THEN 12 WHEN \'quarter\' THEN 3 WHEN \'month\' THEN 1 END ELSE EXTRACT(EPOCH FROM DATE_TRUNC(\'{{ date_part }}\', {{ args[2] }}) - DATE_TRUNC(\'{{ date_part }}\', {{ args[1] }})) / EXTRACT(EPOCH FROM \'1 {{ date_part }}\'::interval) END::bigint".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_wrapper() + .wrapped_sql + .unwrap() + .sql; + assert!(sql.contains("CASE WHEN LOWER('day')")); + assert!(sql.contains("WHEN 'year' THEN 12 WHEN 'quarter' THEN 3 WHEN 'month' THEN 1 END")); + assert!(sql.contains("EXTRACT(EPOCH FROM")); } // redshift-dateadd-[literal-date32-]to-interval rewrites DATEADD to DATE_ADD