diff --git a/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts b/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts index c341f025cbb3c..b945fbb838bb1 100644 --- a/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts +++ b/packages/cubejs-schema-compiler/src/adapter/SnowflakeQuery.ts @@ -114,6 +114,12 @@ export class SnowflakeQuery extends BaseQuery { templates.functions.BTRIM = 'TRIM({{ args_concat }})'; templates.functions.STRING_AGG = 'LISTAGG({% if distinct %}DISTINCT {% endif %}{{ args_concat }})'; templates.expressions.extract = 'EXTRACT({{ date_part }} FROM {{ expr }})'; + // Snowflake can't EXTRACT(EPOCH FROM ), so the epoch of a timestamp + // difference (left - right) is rendered as fractional seconds between them. + // TIMESTAMPDIFF is measured once at microsecond granularity (no per-second + // boundary rounding) and divided to seconds, matching Postgres' fractional + // EXTRACT(EPOCH FROM interval). + templates.expressions.extract_epoch_diff = 'TIMESTAMPDIFF(MICROSECOND, {{ right }}, {{ left }}) / 1000000'; templates.expressions.interval = 'INTERVAL \'{{ interval }}\''; templates.expressions.timestamp_literal = '\'{{ value }}\'::timestamp_tz'; templates.expressions.like = '{{ expr }} {% if negated %}NOT {% endif %}LIKE {{ pattern }}{% if default_escape %} ESCAPE \'\\\\\'{% endif %}'; diff --git a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs index cdd5142376396..63686b47d596d 100644 --- a/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs +++ b/rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs @@ -3063,6 +3063,22 @@ impl WrappedSelectNode { )) } + /// Returns the operands of a timestamp/date subtraction `left - right`, + /// peeling any cast wrapping the subtraction. Used to rewrite + /// `EXTRACT(EPOCH FROM (left - right))` for dialects that don't support + /// taking the epoch of an interval. + fn timestamp_diff_operands(expr: &Expr) -> Option<(&Expr, &Expr)> { + match expr { + Expr::BinaryExpr { + left, + op: Operator::Minus, + right, + } => Some((left, right)), + Expr::Cast { expr, .. } => Self::timestamp_diff_operands(expr), + _ => None, + } + } + #[inline(never)] fn generate_sql_for_scalar_function<'ctx>( mut sql_query: SqlQuery, @@ -3124,6 +3140,42 @@ impl WrappedSelectNode { date_part ))); } + // Some dialects (e.g. Snowflake) can't EXTRACT(EPOCH FROM ), + // i.e. take the epoch of a timestamp difference `a - b`. When the + // dialect provides a dedicated template, render the difference as a + // diff in seconds instead. + let sql_templates = sql_generator.get_sql_templates(); + if date_part.eq_ignore_ascii_case("epoch") + && sql_templates.contains_template("expressions/extract_epoch_diff") + { + if let Some((left, right)) = Self::timestamp_diff_operands(&args[1]) { + let (left_sql, query) = Self::generate_sql_for_expr( + sql_query, + sql_generator.clone(), + left.clone(), + push_to_cube_context, + subqueries, + )?; + let (right_sql, query) = Self::generate_sql_for_expr( + query, + sql_generator.clone(), + right.clone(), + push_to_cube_context, + subqueries, + )?; + return Ok(( + sql_templates + .extract_epoch_diff_expr(left_sql, right_sql) + .map_err(|e| { + DataFusionError::Internal(format!( + "Can't generate SQL for scalar function: {}", + e + )) + })?, + query, + )); + } + } let (arg_sql, query) = Self::generate_sql_for_expr( sql_query, sql_generator.clone(), @@ -3132,8 +3184,7 @@ impl WrappedSelectNode { subqueries, )?; return Ok(( - sql_generator - .get_sql_templates() + sql_templates .extract_expr(date_part.to_string(), arg_sql) .map_err(|e| { DataFusionError::Internal(format!( diff --git a/rust/cubesql/cubesql/src/compile/mod.rs b/rust/cubesql/cubesql/src/compile/mod.rs index a05b89d12004d..ebbf489cef84e 100644 --- a/rust/cubesql/cubesql/src/compile/mod.rs +++ b/rust/cubesql/cubesql/src/compile/mod.rs @@ -16930,6 +16930,51 @@ LIMIT {{ limit }}{% endif %}"#.to_string(), assert!(sql.contains("unix_timestamp")); } + #[tokio::test] + async fn test_extract_epoch_diff_pushdown() { + if !Rewriter::sql_push_down_enabled() { + return; + } + init_testing_logger(); + + // EXTRACT(EPOCH FROM (a - b)) — epoch of a timestamp difference. + let query = " + SELECT customer_gender, + AVG(EXTRACT(EPOCH FROM (order_date - last_mod)) / 86400) AS avg_days + FROM KibanaSampleDataEcommerce + GROUP BY 1 + "; + + // Generic (no dedicated template) keeps EXTRACT(EPOCH FROM (a - b)). + let query_plan = + convert_select_to_query_plan(query.to_string(), DatabaseProtocol::PostgreSQL).await; + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!(sql.contains("EXTRACT(epoch")); + + // Snowflake-style: epoch of a difference is rendered as a seconds diff. + let query_plan = convert_select_to_query_plan_customized( + query.to_string(), + DatabaseProtocol::PostgreSQL, + vec![ + ( + "expressions/extract".to_string(), + "EXTRACT({{ date_part }} FROM {{ expr }})".to_string(), + ), + ( + "expressions/extract_epoch_diff".to_string(), + "TIMESTAMPDIFF(MICROSECOND, {{ right }}, {{ left }}) / 1000000".to_string(), + ), + ], + ) + .await; + + let logical_plan = query_plan.as_logical_plan(); + let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql; + assert!(!sql.to_uppercase().contains("EXTRACT(EPOCH")); + assert!(sql.contains("TIMESTAMPDIFF(MICROSECOND,")); + } + #[tokio::test] async fn test_push_down_to_grouped_query_with_filters() { if !Rewriter::sql_push_down_enabled() { diff --git a/rust/cubesql/cubesql/src/transport/service.rs b/rust/cubesql/cubesql/src/transport/service.rs index b2b4b75bf6246..5367f488ea9b2 100644 --- a/rust/cubesql/cubesql/src/transport/service.rs +++ b/rust/cubesql/cubesql/src/transport/service.rs @@ -777,6 +777,20 @@ impl SqlTemplates { ) } + /// Renders the epoch (in seconds) of a timestamp difference `left - right`. + /// Used for dialects (e.g. Snowflake) where `EXTRACT(EPOCH FROM (left - right))` + /// is invalid because EPOCH can't be extracted from an interval. + pub fn extract_epoch_diff_expr( + &self, + left: String, + right: String, + ) -> Result { + self.render_template( + "expressions/extract_epoch_diff", + context! { left => left, right => right }, + ) + } + pub fn interval_any_expr( &self, interval: String,