Support to_timestamp

dask_planner/src/dialect.rs

@@ -1,6 +1,12 @@
-use core::{iter::Peekable, str::Chars};
-
+use core::iter::Peekable;
+use core::str::Chars;
+use datafusion_sql::sqlparser::ast::{
+    Expr, Function, FunctionArg, FunctionArgExpr, Ident, ObjectName, Value,
+};
 use datafusion_sql::sqlparser::dialect::Dialect;
+use datafusion_sql::sqlparser::parser::{Parser, ParserError};
+use datafusion_sql::sqlparser::tokenizer::Token;
+use datafusion_sql::sqlparser::keywords::Keyword;
 
 #[derive(Debug)]
 pub struct DaskDialect {}
@@ -37,4 +43,93 @@ impl Dialect for DaskDialect {
     fn supports_filter_during_aggregation(&self) -> bool {
         true
     }
+
+    /// override expression parsing
+    fn parse_prefix(&self, parser: &mut Parser) -> Option<Result<Expr, ParserError>> {
+        fn parse_expr(parser: &mut Parser) -> Result<Option<Expr>, ParserError> {
+            match parser.peek_token() {
+                Token::Word(w) if w.value.to_lowercase() == "timestampadd" => {
+                    // TIMESTAMPADD(YEAR, 2, d)
+                    parser.next_token(); // skip timestampadd
+                    parser.expect_token(&Token::LParen)?;
+                    let time_unit = parser.next_token();
+                    parser.expect_token(&Token::Comma)?;
+                    let n = parser.parse_expr()?;
+                    parser.expect_token(&Token::Comma)?;
+                    let expr = parser.parse_expr()?;
+                    parser.expect_token(&Token::RParen)?;
+
+                    // convert to function args
+                    let args = vec![
+                        FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(
+                            Value::SingleQuotedString(time_unit.to_string()),
+                        ))),
+                        FunctionArg::Unnamed(FunctionArgExpr::Expr(n)),
+                        FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)),
+                    ];
+
+                    Ok(Some(Expr::Function(Function {
+                        name: ObjectName(vec![Ident::new("TIMESTAMPADD")]),
+                        args,
+                        over: None,
+                        distinct: false,
+                        special: false,
+                    })))
+                }
+                Token::Word(w) if w.value.to_lowercase() == "to_timestamp" => {
+                    parser.next_token(); // skip to_timestamp
+                    parser.expect_token(&Token::LParen)?;
+                    let expr = parser.parse_expr()?;
+                    let comma = parser.consume_token(&Token::Comma);
+                    if comma {
+                        // Parse TO_TIMESTAMP(d, "%Y-%m-%d %H:%M:%S")
+                        let time_format = parser.next_token();
+                        parser.expect_token(&Token::RParen)?;
+
+                        // convert to function args
+                        let args = vec![
+                            FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)),
+                            FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(
+                                Value::SingleQuotedString(time_format.to_string()),
+                            ))),
+                        ];
+
+                        Ok(Some(Expr::Function(Function {
+                            name: ObjectName(vec![Ident::new("dsql_totimestamp")]),
+                            args,
+                            over: None,
+                            distinct: false,
+                            special: false,
+                        })))
+                    } else {
+                        // Parse TO_TIMESTAMP(d)
+                        let time_format = "%Y-%m-%d %H:%M:%S";
+                        parser.expect_token(&Token::RParen)?;
+
+                        // convert to function args
+                        let args = vec![
+                            FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)),
+                            FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(
+                                Value::SingleQuotedString(time_format.to_string()),
+                            ))),
+                        ];
+
+                        Ok(Some(Expr::Function(Function {
+                            name: ObjectName(vec![Ident::new("dsql_totimestamp")]),
+                            args,
+                            over: None,
+                            distinct: false,
+                            special: false,
+                        })))
+                    }
+                }
+                _ => Ok(None)
+            }
+        }
+        match parse_expr(parser) {
+            Ok(Some(expr)) => Some(Ok(expr)),
+            Ok(None) => None,
+            Err(e) => Some(Err(e)),
+        }
+    }
 }

dask_planner/src/sql.rs

@@ -152,6 +152,14 @@ impl ContextProvider for DaskSQLContext {
                 let rtf: ReturnTypeFunction = Arc::new(|_| Ok(Arc::new(DataType::Int64)));
                 return Some(Arc::new(ScalarUDF::new(name, &sig, &rtf, &fun)));
             }
+            "dsql_totimestamp" => {
+                let sig = Signature::variadic(
+                    vec![DataType::Date64, DataType::Utf8],
+                    Volatility::Immutable,
+                );
+                let rtf: ReturnTypeFunction = Arc::new(|_| Ok(Arc::new(DataType::Utf8)));
+                return Some(Arc::new(ScalarUDF::new(name, &sig, &rtf, &fun)));
+            }
             "atan2" | "mod" => {
                 let sig = Signature::variadic(
                     vec![DataType::Float64, DataType::Float64],

dask_sql/physical/rex/core/call.py

@@ -600,6 +600,18 @@ def extract(self, what, df: SeriesOrScalar):
             raise NotImplementedError(f"Extraction of {what} is not (yet) implemented.")
 
 
+class ToTimestampOperation(Operation):
+    def __init__(self):
+        super().__init__(self.to_timestamp)
+
+    def to_timestamp(self, df, format):
+        df = pd.to_datetime(df, unit="s")
+        df = df.strftime(format)
+        result = [timestamp for timestamp in df]
+        result = pd.Series(result)
+        return result
+
+
 class YearOperation(Operation):
     def __init__(self):
         super().__init__(self.extract_year)
@@ -976,6 +988,7 @@ class RexCallPlugin(BaseRexPlugin):
             lambda x: x + pd.tseries.offsets.MonthEnd(1),
             lambda x: convert_to_datetime(x) + pd.tseries.offsets.MonthEnd(1),
         ),
+        "dsql_totimestamp": ToTimestampOperation(),
         # Temporary UDF functions that need to be moved after this POC
         "datepart": DatePartOperation(),
         "year": YearOperation(),

tests/integration/test_rex.py

@@ -656,3 +656,48 @@ def test_date_functions(c):
             FROM df
             """
         )
+
+
+def test_totimestamp(c):
+    df = pd.DataFrame(
+        {
+            "a": [1203073300, 1406073600, 2806073600],
+        }
+    )
+    c.create_table("df", df)
+
+    df = c.sql(
+        """
+        SELECT to_timestamp(a) AS date FROM df
+    """
+    )
+
+    expected_df = pd.DataFrame(
+        {
+            "date": [
+                datetime(2008, 2, 15, 11, 1, 40),
+                datetime(2014, 7, 23),
+                datetime(2058, 12, 2, 16, 53, 20),
+            ],
+        }
+    )
+
+    assert_eq(df, expected_df, check_dtype=False)
+
+    df = c.sql(
+        """
+        SELECT to_timestamp(a, "%d/%m/%Y") AS date FROM df
+    """
+    )
+
+    expected_df = pd.DataFrame(
+        {
+            "date": [
+                "15/02/2008",
+                "23/07/2014",
+                "02/12/2058",
+            ],
+        }
+    )
+
+    assert_eq(df, expected_df, check_dtype=False)