[arrow-datafusion] branch master updated: feat: add lit_timestamp_nanosecond (#1030)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 1c858ce  feat: add lit_timestamp_nanosecond (#1030)
1c858ce is described below

commit 1c858ce7baab1929cfdba97051ef4e5e4d0a866b
Author: Nga Tran <nt...@influxdata.com>
AuthorDate: Wed Sep 22 10:58:41 2021 -0400

    feat: add lit_timestamp_nanosecond (#1030)
    
    * feat: add lit_timestamp_nanosecond
    
    * refactor: make a traie for timestamp literal
    
    * refactor: cleanup
    
    * refactor: let nonly support appropriate numbers for now
---
 datafusion/src/logical_plan/expr.rs | 46 +++++++++++++++++++++++++++++++++++++
 datafusion/src/logical_plan/mod.rs  | 14 +++++------
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/datafusion/src/logical_plan/expr.rs b/datafusion/src/logical_plan/expr.rs
index d808b6b..0ec445a 100644
--- a/datafusion/src/logical_plan/expr.rs
+++ b/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 {
@@ -1872,6 +1903,21 @@ mod tests {
     }
 
     #[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();
         col("state").eq(lit("CO")).rewrite(&mut rewriter).unwrap();
diff --git a/datafusion/src/logical_plan/mod.rs b/datafusion/src/logical_plan/mod.rs
index 0b6cf60..84ee9e5 100644
--- a/datafusion/src/logical_plan/mod.rs
+++ b/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;