Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add lit_timestamp_nanosecond #1030

Merged
merged 4 commits into from
Sep 22, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here are a few examples: is time column has timestamp data type, value 10 needs to also be a timestamp (in nanosecond in this case) to have the expression analysis works

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