feat: add lit_timestamp_nanosecond (apache#1030)

* feat: add lit_timestamp_nanosecond

* refactor: make a traie for timestamp literal

* refactor: cleanup

* refactor: let nonly support appropriate numbers for now

Author: Nga Tran

datafusion/src/logical_plan/expr.rs

@@ -1362,6 +1362,11 @@ pub trait Literal {
     fn lit(&self) -> Expr;
 }
 
+/// Trait for converting a type to a literal timestamp
+pub trait TimestampLiteral {
+    fn lit_timestamp_nano(&self) -> Expr;
+}
+
 impl Literal for &str {
     fn lit(&self) -> Expr {
         Expr::Literal(ScalarValue::Utf8(Some((*self).to_owned())))
@@ -1391,6 +1396,19 @@ macro_rules! make_literal {
     };
 }
 
+macro_rules! make_timestamp_literal {
+    ($TYPE:ty, $SCALAR:ident, $DOC: expr) => {
+        #[doc = $DOC]
+        impl TimestampLiteral for $TYPE {
+            fn lit_timestamp_nano(&self) -> Expr {
+                Expr::Literal(ScalarValue::TimestampNanosecond(Some(
+                    (self.clone()).into(),
+                )))
+            }
+        }
+    };
+}
+
 make_literal!(bool, Boolean, "literal expression containing a bool");
 make_literal!(f32, Float32, "literal expression containing an f32");
 make_literal!(f64, Float64, "literal expression containing an f64");
@@ -1403,11 +1421,24 @@ make_literal!(u16, UInt16, "literal expression containing a u16");
 make_literal!(u32, UInt32, "literal expression containing a u32");
 make_literal!(u64, UInt64, "literal expression containing a u64");
 
+make_timestamp_literal!(i8, Int8, "literal expression containing an i8");
+make_timestamp_literal!(i16, Int16, "literal expression containing an i16");
+make_timestamp_literal!(i32, Int32, "literal expression containing an i32");
+make_timestamp_literal!(i64, Int64, "literal expression containing an i64");
+make_timestamp_literal!(u8, UInt8, "literal expression containing a u8");
+make_timestamp_literal!(u16, UInt16, "literal expression containing a u16");
+make_timestamp_literal!(u32, UInt32, "literal expression containing a u32");
+
 /// Create a literal expression
 pub fn lit<T: Literal>(n: T) -> Expr {
     n.lit()
 }
 
+/// Create a literal timestamp expression
+pub fn lit_timestamp_nano<T: TimestampLiteral>(n: T) -> Expr {
+    n.lit_timestamp_nano()
+}
+
 /// Concatenates the text representations of all the arguments. NULL arguments are ignored.
 pub fn concat(args: &[Expr]) -> Expr {
     Expr::ScalarFunction {
@@ -1871,6 +1902,21 @@ mod tests {
         assert!(maybe_expr.is_err());
     }
 
+    #[test]
+    fn test_lit_timestamp_nano() {
+        let expr = col("time").eq(lit_timestamp_nano(10)); // 10 is an implicit i32
+        let expected = col("time").eq(lit(ScalarValue::TimestampNanosecond(Some(10))));
+        assert_eq!(expr, expected);
+
+        let i: i64 = 10;
+        let expr = col("time").eq(lit_timestamp_nano(i));
+        assert_eq!(expr, expected);
+
+        let i: u32 = 10;
+        let expr = col("time").eq(lit_timestamp_nano(i));
+        assert_eq!(expr, expected);
+    }
+
     #[test]
     fn rewriter_visit() {
         let mut rewriter = RecordingRewriter::default();

datafusion/src/logical_plan/mod.rs

@@ -39,13 +39,13 @@ pub use expr::{
     abs, acos, and, array, ascii, asin, atan, avg, binary_expr, bit_length, btrim, case,
     ceil, character_length, chr, col, columnize_expr, combine_filters, concat, concat_ws,
     cos, count, count_distinct, create_udaf, create_udf, date_part, date_trunc, exp,
-    exprlist_to_fields, floor, in_list, initcap, left, length, lit, ln, log10, log2,
-    lower, lpad, ltrim, max, md5, min, normalize_col, normalize_cols, now, octet_length,
-    or, random, regexp_match, regexp_replace, repeat, replace, replace_col, reverse,
-    right, round, rpad, rtrim, sha224, sha256, sha384, sha512, signum, sin, split_part,
-    sqrt, starts_with, strpos, substr, sum, tan, to_hex, translate, trim, trunc,
-    unnormalize_col, unnormalize_cols, upper, when, Column, Expr, ExprRewriter,
-    ExpressionVisitor, Literal, Recursion, RewriteRecursion,
+    exprlist_to_fields, floor, in_list, initcap, left, length, lit, lit_timestamp_nano,
+    ln, log10, log2, lower, lpad, ltrim, max, md5, min, normalize_col, normalize_cols,
+    now, octet_length, or, random, regexp_match, regexp_replace, repeat, replace,
+    replace_col, reverse, right, round, rpad, rtrim, sha224, sha256, sha384, sha512,
+    signum, sin, split_part, sqrt, starts_with, strpos, substr, sum, tan, to_hex,
+    translate, trim, trunc, unnormalize_col, unnormalize_cols, upper, when, Column, Expr,
+    ExprRewriter, ExpressionVisitor, Literal, Recursion, RewriteRecursion,
 };
 pub use extension::UserDefinedLogicalNode;
 pub use operators::Operator;