fix: Rewrite date_trunc and from_unixtime for the SQLite unparser (apache#15630)

* feat: Support date_trunc and from_unixtime in SQLite unparser

* feat: Add tests

* refactor: Update datafusion/sql/src/unparser/expr.rs

Co-authored-by: Phillip LeBlanc <[email protected]>

* chore: Cargo fmt

* chore: Use import to format properly

---------

Co-authored-by: Phillip LeBlanc <[email protected]>

Author: 2 people

datafusion/sql/src/unparser/dialect.rs

@@ -17,7 +17,10 @@
 
 use std::{collections::HashMap, sync::Arc};
 
-use super::{utils::character_length_to_sql, utils::date_part_to_sql, Unparser};
+use super::{
+    utils::character_length_to_sql, utils::date_part_to_sql,
+    utils::sqlite_date_trunc_to_sql, utils::sqlite_from_unixtime_to_sql, Unparser,
+};
 use arrow::datatypes::TimeUnit;
 use datafusion_common::Result;
 use datafusion_expr::Expr;
@@ -490,6 +493,8 @@ impl Dialect for SqliteDialect {
             "character_length" => {
                 character_length_to_sql(unparser, self.character_length_style(), args)
             }
+            "from_unixtime" => sqlite_from_unixtime_to_sql(unparser, args),
+            "date_trunc" => sqlite_date_trunc_to_sql(unparser, args),
             _ => Ok(None),
         }
     }

datafusion/sql/src/unparser/expr.rs

@@ -1689,6 +1689,7 @@ mod tests {
     use std::ops::{Add, Sub};
     use std::{any::Any, sync::Arc, vec};
 
+    use crate::unparser::dialect::SqliteDialect;
     use arrow::array::{LargeListArray, ListArray};
     use arrow::datatypes::{DataType::Int8, Field, Int32Type, Schema, TimeUnit};
     use ast::ObjectName;
@@ -1701,6 +1702,7 @@ mod tests {
         ScalarUDFImpl, Signature, Volatility, WindowFrame, WindowFunctionDefinition,
     };
     use datafusion_expr::{interval_month_day_nano_lit, ExprFunctionExt};
+    use datafusion_functions::datetime::from_unixtime::FromUnixtimeFunc;
     use datafusion_functions::expr_fn::{get_field, named_struct};
     use datafusion_functions_aggregate::count::count_udaf;
     use datafusion_functions_aggregate::expr_fn::sum;
@@ -1712,7 +1714,7 @@ mod tests {
 
     use crate::unparser::dialect::{
         CharacterLengthStyle, CustomDialect, CustomDialectBuilder, DateFieldExtractStyle,
-        Dialect, DuckDBDialect, PostgreSqlDialect, ScalarFnToSqlHandler,
+        DefaultDialect, Dialect, DuckDBDialect, PostgreSqlDialect, ScalarFnToSqlHandler,
     };
 
     use super::*;
@@ -2871,6 +2873,115 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_from_unixtime() -> Result<()> {
+        let default_dialect: Arc<dyn Dialect> = Arc::new(DefaultDialect {});
+        let sqlite_dialect: Arc<dyn Dialect> = Arc::new(SqliteDialect {});
+
+        for (dialect, expected) in [
+            (default_dialect, "from_unixtime(date_col)"),
+            (sqlite_dialect, "datetime(`date_col`, 'unixepoch')"),
+        ] {
+            let unparser = Unparser::new(dialect.as_ref());
+            let expr = Expr::ScalarFunction(ScalarFunction {
+                func: Arc::new(ScalarUDF::from(FromUnixtimeFunc::new())),
+                args: vec![col("date_col")],
+            });
+
+            let ast = unparser.expr_to_sql(&expr)?;
+
+            let actual = ast.to_string();
+            let expected = expected.to_string();
+
+            assert_eq!(actual, expected);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_date_trunc() -> Result<()> {
+        let default_dialect: Arc<dyn Dialect> = Arc::new(DefaultDialect {});
+        let sqlite_dialect: Arc<dyn Dialect> = Arc::new(SqliteDialect {});
+
+        for (dialect, precision, expected) in [
+            (
+                Arc::clone(&default_dialect),
+                "YEAR",
+                "date_trunc('YEAR', date_col)",
+            ),
+            (
+                Arc::clone(&sqlite_dialect),
+                "YEAR",
+                "strftime('%Y', `date_col`)",
+            ),
+            (
+                Arc::clone(&default_dialect),
+                "MONTH",
+                "date_trunc('MONTH', date_col)",
+            ),
+            (
+                Arc::clone(&sqlite_dialect),
+                "MONTH",
+                "strftime('%Y-%m', `date_col`)",
+            ),
+            (
+                Arc::clone(&default_dialect),
+                "DAY",
+                "date_trunc('DAY', date_col)",
+            ),
+            (
+                Arc::clone(&sqlite_dialect),
+                "DAY",
+                "strftime('%Y-%m-%d', `date_col`)",
+            ),
+            (
+                Arc::clone(&default_dialect),
+                "HOUR",
+                "date_trunc('HOUR', date_col)",
+            ),
+            (
+                Arc::clone(&sqlite_dialect),
+                "HOUR",
+                "strftime('%Y-%m-%d %H', `date_col`)",
+            ),
+            (
+                Arc::clone(&default_dialect),
+                "MINUTE",
+                "date_trunc('MINUTE', date_col)",
+            ),
+            (
+                Arc::clone(&sqlite_dialect),
+                "MINUTE",
+                "strftime('%Y-%m-%d %H:%M', `date_col`)",
+            ),
+            (default_dialect, "SECOND", "date_trunc('SECOND', date_col)"),
+            (
+                sqlite_dialect,
+                "SECOND",
+                "strftime('%Y-%m-%d %H:%M:%S', `date_col`)",
+            ),
+        ] {
+            let unparser = Unparser::new(dialect.as_ref());
+            let expr = Expr::ScalarFunction(ScalarFunction {
+                func: Arc::new(ScalarUDF::from(
+                    datafusion_functions::datetime::date_trunc::DateTruncFunc::new(),
+                )),
+                args: vec![
+                    Expr::Literal(ScalarValue::Utf8(Some(precision.to_string()))),
+                    col("date_col"),
+                ],
+            });
+
+            let ast = unparser.expr_to_sql(&expr)?;
+
+            let actual = ast.to_string();
+            let expected = expected.to_string();
+
+            assert_eq!(actual, expected);
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_dictionary_to_sql() -> Result<()> {
         let dialect = CustomDialectBuilder::new().build();

datafusion/sql/src/unparser/utils.rs

@@ -500,3 +500,72 @@ pub(crate) fn character_length_to_sql(
         character_length_args,
     )?))
 }
+
+/// SQLite does not support timestamp/date scalars like `to_timestamp`, `from_unixtime`, `date_trunc`, etc.
+/// This remaps `from_unixtime` to `datetime(expr, 'unixepoch')`, expecting the input to be in seconds.
+/// It supports no other arguments, so if any are supplied it will return an error.
+///
+/// # Errors
+///
+/// - If the number of arguments is not 1 - the column or expression to convert.
+/// - If the scalar function cannot be converted to SQL.
+pub(crate) fn sqlite_from_unixtime_to_sql(
+    unparser: &Unparser,
+    from_unixtime_args: &[Expr],
+) -> Result<Option<ast::Expr>> {
+    if from_unixtime_args.len() != 1 {
+        return internal_err!(
+            "from_unixtime for SQLite expects 1 argument, found {}",
+            from_unixtime_args.len()
+        );
+    }
+
+    Ok(Some(unparser.scalar_function_to_sql(
+        "datetime",
+        &[
+            from_unixtime_args[0].clone(),
+            Expr::Literal(ScalarValue::Utf8(Some("unixepoch".to_string()))),
+        ],
+    )?))
+}
+
+/// SQLite does not support timestamp/date scalars like `to_timestamp`, `from_unixtime`, `date_trunc`, etc.
+/// This uses the `strftime` function to format the timestamp as a string depending on the truncation unit.
+///
+/// # Errors
+///
+/// - If the number of arguments is not 2 - truncation unit and the column or expression to convert.
+/// - If the scalar function cannot be converted to SQL.
+pub(crate) fn sqlite_date_trunc_to_sql(
+    unparser: &Unparser,
+    date_trunc_args: &[Expr],
+) -> Result<Option<ast::Expr>> {
+    if date_trunc_args.len() != 2 {
+        return internal_err!(
+            "date_trunc for SQLite expects 2 arguments, found {}",
+            date_trunc_args.len()
+        );
+    }
+
+    if let Expr::Literal(ScalarValue::Utf8(Some(unit))) = &date_trunc_args[0] {
+        let format = match unit.to_lowercase().as_str() {
+            "year" => "%Y",
+            "month" => "%Y-%m",
+            "day" => "%Y-%m-%d",
+            "hour" => "%Y-%m-%d %H",
+            "minute" => "%Y-%m-%d %H:%M",
+            "second" => "%Y-%m-%d %H:%M:%S",
+            _ => return Ok(None),
+        };
+
+        return Ok(Some(unparser.scalar_function_to_sql(
+            "strftime",
+            &[
+                Expr::Literal(ScalarValue::Utf8(Some(format.to_string()))),
+                date_trunc_args[1].clone(),
+            ],
+        )?));
+    }
+
+    Ok(None)
+}