feat: Add DateFieldExtractStyle::Strftime support for SqliteDialect u…

…nparser (#12161) * feat: Add DateFieldExtractStyle::Strftime support for SqliteDialect (#26) * feat: Add DateFieldExtractStyle::Strftime support for SqliteDialect * refactor: Refactor DateFieldExtractStyle if checks into if/match * Fix merge issue --------- Co-authored-by: Andrew Lamb <[email protected]>
apache · Sep 3, 2024 · 0cd7c25 · 0cd7c25
1 parent 7ca9810
commit 0cd7c25
Show file tree

Hide file tree

Showing 2 changed files with 81 additions and 21 deletions.
diff --git a/datafusion/sql/src/unparser/dialect.rs b/datafusion/sql/src/unparser/dialect.rs
@@ -130,6 +130,7 @@ pub enum IntervalStyle {
 pub enum DateFieldExtractStyle {
     DatePart,
     Extract,
+    Strftime,
 }
 
 pub struct DefaultDialect {}
@@ -213,6 +214,10 @@ impl Dialect for SqliteDialect {
         Some('`')
     }
 
+    fn date_field_extract_style(&self) -> DateFieldExtractStyle {
+        DateFieldExtractStyle::Strftime
+    }
+
     fn date32_cast_dtype(&self) -> sqlparser::ast::DataType {
         sqlparser::ast::DataType::Text
     }

diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs
@@ -524,30 +524,78 @@ impl Unparser<'_> {
         _func: &Arc<ScalarUDF>,
         args: &[Expr],
     ) -> Option<ast::Expr> {
-        if func_name.to_lowercase() == "date_part"
-            && self.dialect.date_field_extract_style() == DateFieldExtractStyle::Extract
-            && args.len() == 2
-        {
-            let date_expr = self.expr_to_sql(&args[1]).ok()?;
-
-            if let Expr::Literal(ScalarValue::Utf8(Some(field))) = &args[0] {
-                let field = match field.to_lowercase().as_str() {
-                    "year" => ast::DateTimeField::Year,
-                    "month" => ast::DateTimeField::Month,
-                    "day" => ast::DateTimeField::Day,
-                    "hour" => ast::DateTimeField::Hour,
-                    "minute" => ast::DateTimeField::Minute,
-                    "second" => ast::DateTimeField::Second,
-                    _ => return None,
-                };
+        if func_name.to_lowercase() == "date_part" {
+            match (self.dialect.date_field_extract_style(), args.len()) {
+                (DateFieldExtractStyle::Extract, 2) => {
+                    let date_expr = self.expr_to_sql(&args[1]).ok()?;
+
+                    if let Expr::Literal(ScalarValue::Utf8(Some(field))) = &args[0] {
+                        let field = match field.to_lowercase().as_str() {
+                            "year" => ast::DateTimeField::Year,
+                            "month" => ast::DateTimeField::Month,
+                            "day" => ast::DateTimeField::Day,
+                            "hour" => ast::DateTimeField::Hour,
+                            "minute" => ast::DateTimeField::Minute,
+                            "second" => ast::DateTimeField::Second,
+                            _ => return None,
+                        };
+
+                        return Some(ast::Expr::Extract {
+                            field,
+                            expr: Box::new(date_expr),
+                            syntax: ast::ExtractSyntax::From,
+                        });
+                    }
+                }
+                (DateFieldExtractStyle::Strftime, 2) => {
+                    let column = self.expr_to_sql(&args[1]).ok()?;
+
+                    if let Expr::Literal(ScalarValue::Utf8(Some(field))) = &args[0] {
+                        let field = match field.to_lowercase().as_str() {
+                            "year" => "%Y",
+                            "month" => "%m",
+                            "day" => "%d",
+                            "hour" => "%H",
+                            "minute" => "%M",
+                            "second" => "%S",
+                            _ => return None,
+                        };
 
-                return Some(ast::Expr::Extract {
-                    field,
-                    expr: Box::new(date_expr),
-                    syntax: ast::ExtractSyntax::From,
-                });
+                        return Some(ast::Expr::Function(ast::Function {
+                            name: ast::ObjectName(vec![ast::Ident {
+                                value: "strftime".to_string(),
+                                quote_style: None,
+                            }]),
+                            args: ast::FunctionArguments::List(
+                                ast::FunctionArgumentList {
+                                    duplicate_treatment: None,
+                                    args: vec![
+                                        ast::FunctionArg::Unnamed(
+                                            ast::FunctionArgExpr::Expr(ast::Expr::Value(
+                                                ast::Value::SingleQuotedString(
+                                                    field.to_string(),
+                                                ),
+                                            )),
+                                        ),
+                                        ast::FunctionArg::Unnamed(
+                                            ast::FunctionArgExpr::Expr(column),
+                                        ),
+                                    ],
+                                    clauses: vec![],
+                                },
+                            ),
+                            filter: None,
+                            null_treatment: None,
+                            over: None,
+                            within_group: vec![],
+                            parameters: ast::FunctionArguments::None,
+                        }));
+                    }
+                }
+                _ => {} // no overrides for DateFieldExtractStyle::DatePart, because it's already a date_part
             }
         }
+
         None
     }
 
@@ -2178,6 +2226,7 @@ mod tests {
                 "YEAR",
                 "EXTRACT(YEAR FROM x)",
             ),
+            (DateFieldExtractStyle::Strftime, "YEAR", "strftime('%Y', x)"),
             (
                 DateFieldExtractStyle::DatePart,
                 "MONTH",
@@ -2188,11 +2237,17 @@ mod tests {
                 "MONTH",
                 "EXTRACT(MONTH FROM x)",
             ),
+            (
+                DateFieldExtractStyle::Strftime,
+                "MONTH",
+                "strftime('%m', x)",
+            ),
             (
                 DateFieldExtractStyle::DatePart,
                 "DAY",
                 "date_part('DAY', x)",
             ),
+            (DateFieldExtractStyle::Strftime, "DAY", "strftime('%d', x)"),
             (DateFieldExtractStyle::Extract, "DAY", "EXTRACT(DAY FROM x)"),
         ] {
             let dialect = CustomDialectBuilder::new()