Skip to content

Commit

Permalink
feat: add lit_timestamp_nanosecond (#1030)
Browse files Browse the repository at this point in the history
* feat: add lit_timestamp_nanosecond

* refactor: make a traie for timestamp literal

* refactor: cleanup

* refactor: let nonly support appropriate numbers for now
  • Loading branch information
NGA-TRAN committed Sep 22, 2021
1 parent 2392788 commit 1c858ce
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 7 deletions.
46 changes: 46 additions & 0 deletions datafusion/src/logical_plan/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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())))
Expand Down Expand Up @@ -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");
Expand All @@ -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 {
Expand Down Expand Up @@ -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();
Expand Down
14 changes: 7 additions & 7 deletions datafusion/src/logical_plan/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit 1c858ce

Please sign in to comment.