Added support for day and other datetime keywords in DF (#294)

DanCodedThis · web-flow · commit eeaddd3d7aee · 2025-02-27T15:15:18.000+02:00
* start here

* no macro, just simple vistor mut

* remvoe redundant trait creation

* integration test add, fixed logic, next unit tests

* activated impl and cargo fmt + clippy

* added unit tests + impl extension

* cargo fmt + clippy

* cargo clippy

* cargo clippy

* removed redundant check

* fix

* merged with Artem PR

* cargo fmt + clippy

* small fixes

* cargo fmt + clippy
diff --git a/crates/runtime/src/datafusion/execution.rs b/crates/runtime/src/datafusion/execution.rs
@@ -106,12 +106,11 @@ impl SqlExecutor {
         warehouse_name: &str,
     ) -> IceBucketSQLResult<Vec<RecordBatch>> {
         // Update query to use custom JSON functions
-        let query = self.preprocess_query(query);
+        let query = Self::preprocess_query(query);
         let mut statement = self
             .parse_query(query.as_str())
             .context(super::error::DataFusionSnafu)?;
         Self::postprocess_query_statement(&mut statement);
-
         // statement = self.update_statement_references(statement, warehouse_name);
         // query = statement.to_string();
 
@@ -244,18 +243,15 @@ impl SqlExecutor {
     /// Panics if .
     #[must_use]
     #[allow(clippy::unwrap_used)]
-    #[tracing::instrument(level = "trace", skip(self), ret)]
-    pub fn preprocess_query(&self, query: &str) -> String {
+    #[tracing::instrument(level = "trace", ret)]
+    pub fn preprocess_query(query: &str) -> String {
         // Replace field[0].subfield -> json_get(json_get(field, 0), 'subfield')
         // TODO: This regex should be a static allocation
         let re = regex::Regex::new(r"(\w+.\w+)\[(\d+)][:\.](\w+)").unwrap();
-        let date_add =
-            regex::Regex::new(r"(date|time|timestamp)(_?add|_?diff)\(\s*([a-zA-Z]+),").unwrap();
-
+        //date_add processing moved to `postprocess_query_statement`
         let mut query = re
             .replace_all(query, "json_get(json_get($1, $2), '$3')")
             .to_string();
-        query = date_add.replace_all(&query, "$1$2('$3',").to_string();
         let alter_iceberg_table = regex::Regex::new(r"alter\s+iceberg\s+table").unwrap();
         query = alter_iceberg_table
             .replace_all(&query, "alter table")
@@ -398,6 +394,7 @@ impl SqlExecutor {
                     .await
                     .context(ih_error::IcebergSnafu)?;
             };
+
             // Create new table
             rest_catalog
                 .create_table(
@@ -1336,8 +1333,144 @@ pub fn created_entity_response() -> Result<Vec<RecordBatch>, arrow::error::Arrow
 }
 
 #[cfg(test)]
-mod test {
-    use crate::datafusion::execution::SqlExecutor;
+mod tests {
+    use super::SqlExecutor;
+    use crate::datafusion::{session::SessionParams, type_planner::CustomTypePlanner};
+    use datafusion::sql::parser::Statement as DFStatement;
+    use datafusion::sql::sqlparser::ast::visit_expressions;
+    use datafusion::sql::sqlparser::ast::{Expr, ObjectName};
+    use datafusion::{
+        execution::SessionStateBuilder,
+        prelude::{SessionConfig, SessionContext},
+    };
+    use datafusion_iceberg::planner::IcebergQueryPlanner;
+    use sqlparser::ast::Value;
+    use sqlparser::ast::{
+        Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, FunctionArguments,
+    };
+    use std::ops::ControlFlow;
+    use std::sync::Arc;
+
+    struct Test<'a, T> {
+        input: &'a str,
+        expected: T,
+        should_work: bool,
+    }
+    impl<'a, T> Test<'a, T> {
+        pub const fn new(input: &'a str, expected: T, should_work: bool) -> Self {
+            Self {
+                input,
+                expected,
+                should_work,
+            }
+        }
+    }
+    #[test]
+    #[allow(
+        clippy::unwrap_used,
+        clippy::explicit_iter_loop,
+        clippy::collapsible_match
+    )]
+    fn test_timestamp_keywords_postprocess() {
+        let state = SessionStateBuilder::new()
+            .with_config(
+                SessionConfig::new()
+                    .with_information_schema(true)
+                    .with_option_extension(SessionParams::default())
+                    .set_str("datafusion.sql_parser.dialect", "SNOWFLAKE"),
+            )
+            .with_default_features()
+            .with_query_planner(Arc::new(IcebergQueryPlanner {}))
+            .with_type_planner(Arc::new(CustomTypePlanner {}))
+            .build();
+        let ctx = SessionContext::new_with_state(state);
+        let executor = SqlExecutor::new(ctx).unwrap();
+        let test = vec![
+            Test::new(
+                "SELECT dateadd(year, 5, '2025-06-01')",
+                Value::SingleQuotedString("year".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT dateadd(\"year\", 5, '2025-06-01')",
+                Value::SingleQuotedString("year".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT dateadd('year', 5, '2025-06-01')",
+                Value::SingleQuotedString("year".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT dateadd(\"'year'\", 5, '2025-06-01')",
+                Value::SingleQuotedString("year".to_owned()),
+                false,
+            ),
+            Test::new(
+                "SELECT dateadd(\'year\', 5, '2025-06-01')",
+                Value::SingleQuotedString("year".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT datediff(day, 5, '2025-06-01')",
+                Value::SingleQuotedString("day".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT datediff('week', 5, '2025-06-01')",
+                Value::SingleQuotedString("week".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT datediff(nsecond, 10000000, '2025-06-01')",
+                Value::SingleQuotedString("nsecond".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT date_diff(hour, 5, '2025-06-01')",
+                Value::SingleQuotedString("hour".to_owned()),
+                true,
+            ),
+            Test::new(
+                "SELECT date_add(us, 100000, '2025-06-01')",
+                Value::SingleQuotedString("us".to_owned()),
+                true,
+            ),
+        ];
+        for test in test.iter() {
+            let mut statement = executor.parse_query(test.input).unwrap();
+            SqlExecutor::postprocess_query_statement(&mut statement);
+            if let DFStatement::Statement(statement) = statement {
+                visit_expressions(&statement, |expr| {
+                    if let Expr::Function(Function {
+                        name: ObjectName(idents),
+                        args: FunctionArguments::List(FunctionArgumentList { args, .. }),
+                        ..
+                    }) = expr
+                    {
+                        match idents.first().unwrap().value.as_str() {
+                            "dateadd" | "date_add" | "datediff" | "date_diff" => {
+                                if let FunctionArg::Unnamed(FunctionArgExpr::Expr(ident)) =
+                                    args.iter().next().unwrap()
+                                {
+                                    if let Expr::Value(found) = ident {
+                                        if test.should_work {
+                                            assert_eq!(*found, test.expected);
+                                        } else {
+                                            assert_ne!(*found, test.expected);
+                                        }
+                                    }
+                                }
+                            }
+                            _ => {}
+                        }
+                    }
+                    ControlFlow::<()>::Continue(())
+                });
+            }
+        }
+    }
+
     use datafusion::sql::parser::DFParser;
 
     #[allow(clippy::unwrap_used)]
diff --git a/crates/runtime/src/datafusion/functions/convert_timezone.rs b/crates/runtime/src/datafusion/functions/convert_timezone.rs
@@ -341,11 +341,10 @@ mod tests {
                 );
                 assert_eq!(
                     result, expected,
-                    "convert_timezone created wrong value for {}",
-                    source_timestamp_tz_value
-                )
+                    "convert_timezone created wrong value for {source_timestamp_tz_value}"
+                );
             }
-            _ => panic!("Conversion of {} failed", source_timestamp_tz),
+            _ => panic!("Conversion of {source_timestamp_tz} failed"),
         }
     }
     #[test]
@@ -371,11 +370,10 @@ mod tests {
                 );
                 assert_eq!(
                     result, expected,
-                    "convert_timezone created wrong value for {}",
-                    source_timestamp_tz_value
-                )
+                    "convert_timezone created wrong value for {source_timestamp_tz_value}"
+                );
             }
-            _ => panic!("Conversion of {} failed", source_timestamp_tz_value),
+            _ => panic!("Conversion of {source_timestamp_tz_value} failed"),
         }
     }
     #[test]
@@ -403,11 +401,10 @@ mod tests {
                 );
                 assert_ne!(
                     result, expected,
-                    "convert_timezone created wrong value for {}",
-                    source_timestamp_tz_value
-                )
+                    "convert_timezone created wrong value for {source_timestamp_tz_value}"
+                );
             }
-            _ => panic!("Conversion of {} failed", source_timestamp_tz_value),
+            _ => panic!("Conversion of {source_timestamp_tz_value} failed"),
         }
     }
 }
diff --git a/crates/runtime/src/datafusion/functions/date_add.rs b/crates/runtime/src/datafusion/functions/date_add.rs
@@ -235,7 +235,7 @@ mod tests {
             )),
         ];
         let fn_args = ScalarFunctionArgs {
-            args: args,
+            args,
             number_rows: 0,
             return_type: &arrow_schema::DataType::Timestamp(
                 arrow_schema::TimeUnit::Microsecond,
@@ -248,7 +248,7 @@ mod tests {
                     Some(1736600400000000i64),
                     Some(Arc::from(String::from("+00").into_boxed_str())),
                 );
-                assert_eq!(&result, &expected, "date_add created a wrong value")
+                assert_eq!(&result, &expected, "date_add created a wrong value");
             }
             _ => panic!("Conversion failed"),
         }
@@ -268,7 +268,7 @@ mod tests {
             ),
         ];
         let fn_args = ScalarFunctionArgs {
-            args: args,
+            args,
             number_rows: 0,
             return_type: &arrow_schema::DataType::Timestamp(
                 arrow_schema::TimeUnit::Microsecond,
@@ -283,7 +283,7 @@ mod tests {
                 )
                 .to_array()
                 .unwrap();
-                assert_eq!(&result, &expected, "date_add created a wrong value")
+                assert_eq!(&result, &expected, "date_add created a wrong value");
             }
             _ => panic!("Conversion failed"),
         }
@@ -303,7 +303,7 @@ mod tests {
             ),
         ];
         let fn_args = ScalarFunctionArgs {
-            args: args,
+            args,
             number_rows: 0,
             return_type: &arrow_schema::DataType::Timestamp(
                 arrow_schema::TimeUnit::Microsecond,
@@ -318,7 +318,7 @@ mod tests {
                 ))
                 .to_array(2)
                 .unwrap();
-                assert_eq!(&result, &expected, "date_add created a wrong value")
+                assert_eq!(&result, &expected, "date_add created a wrong value");
             }
             _ => panic!("Conversion failed"),
         }
diff --git a/crates/runtime/src/datafusion/functions/mod.rs b/crates/runtime/src/datafusion/functions/mod.rs
@@ -18,8 +18,10 @@
 use std::sync::Arc;
 
 use datafusion::{common::Result, execution::FunctionRegistry, logical_expr::ScalarUDF};
-use sqlparser::ast::Value::SingleQuotedString;
-use sqlparser::ast::{Expr, Function, FunctionArg, FunctionArgExpr, FunctionArguments};
+use sqlparser::ast::Value::{self, SingleQuotedString};
+use sqlparser::ast::{
+    Expr, Function, FunctionArg, FunctionArgExpr, FunctionArgumentList, FunctionArguments, Ident,
+};
 
 mod convert_timezone;
 mod date_add;
@@ -99,4 +101,18 @@ pub fn visit_functions_expressions(func: &mut Function) {
         _ => func_name,
     };
     func.name = sqlparser::ast::ObjectName(vec![sqlparser::ast::Ident::new(name)]);
+    if let FunctionArguments::List(FunctionArgumentList { args, .. }) = &mut func.args {
+        match func_name {
+            "dateadd" | "date_add" | "datediff" | "date_diff" => {
+                if let Some(FunctionArg::Unnamed(FunctionArgExpr::Expr(ident))) =
+                    args.iter_mut().next()
+                {
+                    if let Expr::Identifier(Ident { value, .. }) = ident {
+                        *ident = Expr::Value(Value::SingleQuotedString(value.clone()));
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
 }
diff --git a/crates/runtime/src/tests/queries.rs b/crates/runtime/src/tests/queries.rs
@@ -15,8 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-// use crate::tests::utils::macros::test_query;
+use crate::tests::utils::macros::test_query;
 
+test_query!(select_date_add_diff, "SELECT dateadd(day, 5, '2025-06-01')");
 // // SELECT
 // test_query!(select_star, "SELECT * FROM employee_table");
 // test_query!(select_ilike, "SELECT * ILIKE '%id%' FROM employee_table;");
diff --git a/crates/runtime/src/tests/snapshots/query_select_date_add_diff.snap b/crates/runtime/src/tests/snapshots/query_select_date_add_diff.snap
diff --git a/crates/runtime/src/tests/utils.rs b/crates/runtime/src/tests/utils.rs

Original file line number	Diff line number	Diff line change
`@@ -341,11 +341,10 @@ mod tests {`
`341`	`341`	`);`
`342`	`342`	`assert_eq!(`
`343`	`343`	`result, expected,`
`344`		`- "convert_timezone created wrong value for {}",`
`345`		`- source_timestamp_tz_value`
`346`		`- )`
	`344`	`+ "convert_timezone created wrong value for {source_timestamp_tz_value}"`
	`345`	`+ );`
`347`	`346`	`}`
`348`		`- _ => panic!("Conversion of {} failed", source_timestamp_tz),`
	`347`	`+ _ => panic!("Conversion of {source_timestamp_tz} failed"),`
`349`	`348`	`}`
`350`	`349`	`}`
`351`	`350`	`#[test]`
`@@ -371,11 +370,10 @@ mod tests {`
`371`	`370`	`);`
`372`	`371`	`assert_eq!(`
`373`	`372`	`result, expected,`
`374`		`- "convert_timezone created wrong value for {}",`
`375`		`- source_timestamp_tz_value`
`376`		`- )`
	`373`	`+ "convert_timezone created wrong value for {source_timestamp_tz_value}"`
	`374`	`+ );`
`377`	`375`	`}`
`378`		`- _ => panic!("Conversion of {} failed", source_timestamp_tz_value),`
	`376`	`+ _ => panic!("Conversion of {source_timestamp_tz_value} failed"),`
`379`	`377`	`}`
`380`	`378`	`}`
`381`	`379`	`#[test]`
`@@ -403,11 +401,10 @@ mod tests {`
`403`	`401`	`);`
`404`	`402`	`assert_ne!(`
`405`	`403`	`result, expected,`
`406`		`- "convert_timezone created wrong value for {}",`
`407`		`- source_timestamp_tz_value`
`408`		`- )`
	`404`	`+ "convert_timezone created wrong value for {source_timestamp_tz_value}"`
	`405`	`+ );`
`409`	`406`	`}`
`410`		`- _ => panic!("Conversion of {} failed", source_timestamp_tz_value),`
	`407`	`+ _ => panic!("Conversion of {source_timestamp_tz_value} failed"),`
`411`	`408`	`}`
`412`	`409`	`}`
`413`	`410`	`}`