fix: Rewrite date_trunc and from_unixtime for the SQLite unparser (#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 <phillip@leblanc.tech> * chore: Cargo fmt * chore: Use import to format properly --------- Co-authored-by: Phillip LeBlanc <phillip@leblanc.tech>

Author: peasee

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)
+}