fix(schema-compiler): Fix BigQuery DATE_ADD push down template for years/quarters/months

packages/cubejs-schema-compiler/src/adapter/BigqueryQuery.ts

@@ -186,11 +186,21 @@
   }
 
   public subtractTimestampInterval(date, interval) {
-    return `TIMESTAMP_SUB(${date}, INTERVAL ${this.formatInterval(interval)[0]})`;
+    const [intervalFormatted, timeUnit] = this.formatInterval(interval);
+    if (['YEAR', 'MONTH', 'QUARTER'].includes(timeUnit)) {
+      return this.timeStampCast(`DATETIME_SUB(DATETIME(${date}), INTERVAL ${intervalFormatted})`);
+    }
+
+    return `TIMESTAMP_SUB(${date}, INTERVAL ${intervalFormatted})`;
   }
 
   public addTimestampInterval(date, interval) {
-    return `TIMESTAMP_ADD(${date}, INTERVAL ${this.formatInterval(interval)[0]})`;
+    const [intervalFormatted, timeUnit] = this.formatInterval(interval);
+    if (['YEAR', 'MONTH', 'QUARTER'].includes(timeUnit)) {
+      return this.timeStampCast(`DATETIME_ADD(DATETIME(${date}), INTERVAL ${intervalFormatted})`);
+    }
+
+    return `TIMESTAMP_ADD(${date}, INTERVAL ${intervalFormatted})`;
   }
 
   public nowTimestampSql() {
@@ -242,7 +252,7 @@
     templates.functions.STRPOS = 'STRPOS({{ args_concat }})';
     templates.functions.DATEDIFF = 'DATETIME_DIFF(CAST({{ args[2] }} AS DATETIME), CAST({{ args[1] }} AS DATETIME), {{ date_part }})';
     // DATEADD is being rewritten to DATE_ADD
-    // templates.functions.DATEADD = 'DATETIME_ADD(CAST({{ args[2] }} AS DATETTIME), INTERVAL {{ interval }} {{ date_part }})';
+    templates.functions.DATE_ADD = '{% if date_part|upper in [\'YEAR\', \'MONTH\', \'QUARTER\'] %}TIMESTAMP(DATETIME_ADD(DATETIME({{ args[0] }}), INTERVAL {{ interval }} {{ date_part }})){% else %}TIMESTAMP_ADD({{ args[0] }}, INTERVAL {{ interval }} {{ date_part }}){% endif %}';
     templates.functions.CURRENTDATE = 'CURRENT_DATE';
     delete templates.functions.TO_CHAR;
     templates.expressions.binary = '{% if op == \'%\' %}MOD({{ left }}, {{ right }}){% else %}({{ left }} {{ op }} {{ right }}){% endif %}';

rust/cubesql/cubesql/src/compile/engine/df/wrapper.rs

@@ -2491,10 +2491,47 @@
                                 if DATE_PART_REGEX.is_match(date_part) {
                                     Ok(Some(date_part.to_string()))
                                 } else {
-                                    Err(date_part_err(date_part))
+                                    Err(date_part_err(date_part.to_string()))
                                 }
                             }
-                            _ => Err(date_part_err(&args[0].to_string())),
+                            _ => Err(date_part_err(args[0].to_string())),
+                        },
+                        "date_add" => match &args[1] {
+                            Expr::Literal(ScalarValue::IntervalDayTime(Some(interval))) => {
+                                let days = (*interval >> 32) as i32;
+                                let ms = (*interval & 0xFFFF_FFFF) as i32;
+
+                                if days != 0 && ms == 0 {
+                                    Ok(Some("DAY".to_string()))
+                                } else if ms != 0 && days == 0 {
+                                    Ok(Some("MILLISECOND".to_string()))
+                                } else {
+                                    Err(DataFusionError::Internal(format!(
+                                        "Unsupported mixed IntervalDayTime: days = {days}, ms = {ms}"
+                                    )))
+                                }
+                            }
+                            Expr::Literal(ScalarValue::IntervalYearMonth(Some(_months))) => {
+                                Ok(Some("MONTH".to_string()))
+                            }
+                            Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(interval))) => {
+                                let months = (interval >> 96) as i32;
+                                let days = ((interval >> 64) & 0xFFFF_FFFF) as i32;
+                                let nanos = *interval as i64;
+
+                                if months != 0 && days == 0 && nanos == 0 {
+                                    Ok(Some("MONTH".to_string()))
+                                } else if days != 0 && months == 0 && nanos == 0 {
+                                    Ok(Some("DAY".to_string()))
+                                } else if nanos != 0 && months == 0 && days == 0 {
+                                    Ok(Some("NANOSECOND".to_string()))
+                                } else {
+                                    Err(DataFusionError::Internal(format!(
+                                        "Unsupported mixed IntervalMonthDayNano: months = {months}, days = {days}, nanos = {nanos}"
+                                    )))
+                                }
+                            }
+                            _ => Err(date_part_err(args[1].to_string())),
                         },
                         _ => Ok(None),
                     }?;
@@ -2507,6 +2544,43 @@
                                 "Can't generate SQL for scalar function: interval must be Int64"
                             ))),
                         },
+                        "date_add" => match &args[1] {
+                            Expr::Literal(ScalarValue::IntervalDayTime(Some(interval))) => {
+                                let days = (*interval >> 32) as i32;
+                                let ms = (*interval & 0xFFFF_FFFF) as i32;
+
+                                if days != 0 && ms == 0 {
+                                    Ok(Some(days.to_string()))
+                                } else if ms != 0 && days == 0 {
+                                    Ok(Some(ms.to_string()))
+                                } else {
+                                    Err(DataFusionError::Internal(format!(
+                                        "Unsupported mixed IntervalDayTime: days = {days}, ms = {ms}"
+                                    )))
+                                }
+                            }
+                            Expr::Literal(ScalarValue::IntervalYearMonth(Some(months))) => {
+                                Ok(Some(months.to_string()))
+                            }
+                            Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(interval))) => {
+                                let months = (interval >> 96) as i32;
+                                let days = ((interval >> 64) & 0xFFFF_FFFF) as i32;
+                                let nanos = *interval as i64;
+
+                                if months != 0 && days == 0 && nanos == 0 {
+                                    Ok(Some(months.to_string()))
+                                } else if days != 0 && months == 0 && nanos == 0 {
+                                    Ok(Some(days.to_string()))
+                                } else if nanos != 0 && months == 0 && days == 0 {
+                                    Ok(Some(nanos.to_string()))
+                                } else {
+                                    Err(DataFusionError::Internal(format!(
+                                        "Unsupported mixed IntervalMonthDayNano: months = {months}, days = {days}, nanos = {nanos}"
+                                    )))
+                                }
+                            }
+                            _ => Err(date_part_err(args[1].to_string())),
+                        },
                         _ => Ok(None),
                     }?;
                     let mut sql_args = Vec::new();

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

@@ -14624,19 +14624,19 @@ ORDER BY "source"."str0" ASC
         assert!(sql.contains("EXTRACT(EPOCH FROM"));
     }
 
-    // redshift-dateadd-[literal-date32-]to-interval rewrites DATEADD to DATE_ADD
     #[tokio::test]
-    #[ignore]
     async fn test_dateadd_push_down() {
         if !Rewriter::sql_push_down_enabled() {
             return;
         }
         init_testing_logger();
 
+        // Redshift function DATEADD
         let query_plan = convert_select_to_query_plan(
             "
             SELECT DATEADD(DAY, 7, order_date) AS d
             FROM KibanaSampleDataEcommerce AS k
+            WHERE LOWER(customer_gender) = 'test'
             GROUP BY 1
             ORDER BY 1 DESC
             "
@@ -14652,25 +14652,24 @@ ORDER BY "source"."str0" ASC
         );
 
         let logical_plan = query_plan.as_logical_plan();
-        assert!(logical_plan
-            .find_cube_scan_wrapped_sql()
-            .wrapped_sql
-            .sql
-            .contains("DATEADD(day, 7,"));
+        let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
+        // redshift-dateadd-[literal-date32-]to-interval rewrites DATEADD to DATE_ADD
+        assert!(sql.contains("DATE_ADD("));
+        assert!(sql.contains("INTERVAL '7 DAY')"));
 
-        // BigQuery
+        // BigQuery + Postgres DATE_ADD + DAYS
+        let bq_templates = vec![("functions/DATE_ADD".to_string(), "{% if date_part|upper in ['YEAR', 'MONTH', 'QUARTER'] %}TIMESTAMP(DATETIME_ADD(DATETIME({{ args[0] }}), INTERVAL {{ interval }} {{ date_part }})){% else %}TIMESTAMP_ADD({{ args[0] }}, INTERVAL {{ interval }} {{ date_part }}){% endif %}".to_string())];
         let query_plan = convert_select_to_query_plan_customized(
             "
-            SELECT DATEADD(DAY, 7, order_date) AS d
+            SELECT DATE_ADD(order_date, INTERVAL '7 DAYS') AS d
             FROM KibanaSampleDataEcommerce AS k
+            WHERE LOWER(customer_gender) = 'test'
             GROUP BY 1
             ORDER BY 1 DESC
             "
             .to_string(),
             DatabaseProtocol::PostgreSQL,
-            vec![
-                ("functions/DATEADD".to_string(), "DATETIME_ADD(CAST({{ args[2] }} AS DATETTIME), INTERVAL {{ interval }} {{ date_part }})".to_string()),
-            ],
+            bq_templates.clone(),
         )
         .await;
 
@@ -14682,23 +14681,101 @@ ORDER BY "source"."str0" ASC
 
         let logical_plan = query_plan.as_logical_plan();
         let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
-        assert!(sql.contains("DATETIME_ADD(CAST("));
-        assert!(sql.contains("INTERVAL 7 day)"));
+        assert!(sql.contains("TIMESTAMP_ADD("));
+        assert!(sql.contains("INTERVAL 7 DAY)"));
 
-        // Postgres
+        // BigQuery + Redshift DATEADD + DAYS
+        let bq_templates = vec![("functions/DATE_ADD".to_string(), "{% if date_part|upper in ['YEAR', 'MONTH', 'QUARTER'] %}TIMESTAMP(DATETIME_ADD(DATETIME({{ args[0] }}), INTERVAL {{ interval }} {{ date_part }})){% else %}TIMESTAMP_ADD({{ args[0] }}, INTERVAL {{ interval }} {{ date_part }}){% endif %}".to_string())];
         let query_plan = convert_select_to_query_plan_customized(
             "
             SELECT DATEADD(DAY, 7, order_date) AS d
             FROM KibanaSampleDataEcommerce AS k
+            WHERE LOWER(customer_gender) = 'test'
+            GROUP BY 1
+            ORDER BY 1 DESC
+            "
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+            bq_templates.clone(),
+        )
+        .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("TIMESTAMP_ADD("));
+        assert!(sql.contains("INTERVAL 7 DAY)"));
+
+        // BigQuery + Postgres DATE_ADD + MONTHS
+        let query_plan = convert_select_to_query_plan_customized(
+            "
+            SELECT DATE_ADD(order_date, INTERVAL '7 MONTHS') AS d
+            FROM KibanaSampleDataEcommerce AS k
+            WHERE LOWER(customer_gender) = 'test'
+            GROUP BY 1
+            ORDER BY 1 DESC
+            "
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+            bq_templates,
+        )
+        .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("TIMESTAMP(DATETIME_ADD(DATETIME("));
+        assert!(sql.contains("INTERVAL 7 MONTH)"));
+
+        // BigQuery + Redshift DATEADD + MONTHS
+        let bq_templates = vec![("functions/DATE_ADD".to_string(), "{% if date_part|upper in ['YEAR', 'MONTH', 'QUARTER'] %}TIMESTAMP(DATETIME_ADD(DATETIME({{ args[0] }}), INTERVAL {{ interval }} {{ date_part }})){% else %}TIMESTAMP_ADD({{ args[0] }}, INTERVAL {{ interval }} {{ date_part }}){% endif %}".to_string())];
+        let query_plan = convert_select_to_query_plan_customized(
+            "
+            SELECT DATEADD(MONTH, 7, order_date) AS d
+            FROM KibanaSampleDataEcommerce AS k
+            WHERE LOWER(customer_gender) = 'test'
+            GROUP BY 1
+            ORDER BY 1 DESC
+            "
+            .to_string(),
+            DatabaseProtocol::PostgreSQL,
+            bq_templates.clone(),
+        )
+        .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("TIMESTAMP(DATETIME_ADD(DATETIME("));
+        assert!(sql.contains("INTERVAL 7 MONTH)"));
+
+        // Postgres DATE_ADD
+        let query_plan = convert_select_to_query_plan_customized(
+            "
+            SELECT DATE_ADD(order_date, INTERVAL '7 DAYS') AS d
+            FROM KibanaSampleDataEcommerce AS k
+            WHERE LOWER(customer_gender) = 'test'
             GROUP BY 1
             ORDER BY 1 DESC
             "
             .to_string(),
             DatabaseProtocol::PostgreSQL,
-            vec![(
-                "functions/DATEADD".to_string(),
-                "({{ args[2] }} + \'{{ interval }} {{ date_part }}\'::interval)".to_string(),
-            )],
+            vec![],
         )
         .await;
 
@@ -14710,7 +14787,8 @@ ORDER BY "source"."str0" ASC
 
         let logical_plan = query_plan.as_logical_plan();
         let sql = logical_plan.find_cube_scan_wrapped_sql().wrapped_sql.sql;
-        assert!(sql.contains("+ '7 day'::interval"));
+        assert!(sql.contains("DATE_ADD("));
+        assert!(sql.contains("INTERVAL '7 DAY'"));
     }
 
     #[tokio::test]