datafusion_expr/
literal.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Literal module contains foundational types that are used to represent literals in DataFusion.
19
20use crate::Expr;
21use datafusion_common::{metadata::FieldMetadata, ScalarValue};
22
23/// Create a literal expression
24pub fn lit<T: Literal>(n: T) -> Expr {
25    n.lit()
26}
27
28pub fn lit_with_metadata<T: Literal>(n: T, metadata: Option<FieldMetadata>) -> Expr {
29    let Some(metadata) = metadata else {
30        return n.lit();
31    };
32
33    let Expr::Literal(sv, prior_metadata) = n.lit() else {
34        unreachable!();
35    };
36    let new_metadata = match prior_metadata {
37        Some(mut prior) => {
38            prior.extend(metadata);
39            prior
40        }
41        None => metadata,
42    };
43
44    Expr::Literal(sv, Some(new_metadata))
45}
46
47/// Create a literal timestamp expression
48pub fn lit_timestamp_nano<T: TimestampLiteral>(n: T) -> Expr {
49    n.lit_timestamp_nano()
50}
51
52/// Trait for converting a type to a [`Literal`] literal expression.
53pub trait Literal {
54    /// convert the value to a Literal expression
55    fn lit(&self) -> Expr;
56}
57
58/// Trait for converting a type to a literal timestamp
59pub trait TimestampLiteral {
60    fn lit_timestamp_nano(&self) -> Expr;
61}
62
63impl Literal for &str {
64    fn lit(&self) -> Expr {
65        Expr::Literal(ScalarValue::from(*self), None)
66    }
67}
68
69impl Literal for String {
70    fn lit(&self) -> Expr {
71        Expr::Literal(ScalarValue::from(self.as_ref()), None)
72    }
73}
74
75impl Literal for &String {
76    fn lit(&self) -> Expr {
77        Expr::Literal(ScalarValue::from(self.as_ref()), None)
78    }
79}
80
81impl Literal for Vec<u8> {
82    fn lit(&self) -> Expr {
83        Expr::Literal(ScalarValue::Binary(Some((*self).to_owned())), None)
84    }
85}
86
87impl Literal for &[u8] {
88    fn lit(&self) -> Expr {
89        Expr::Literal(ScalarValue::Binary(Some((*self).to_owned())), None)
90    }
91}
92
93impl Literal for ScalarValue {
94    fn lit(&self) -> Expr {
95        Expr::Literal(self.clone(), None)
96    }
97}
98
99macro_rules! make_literal {
100    ($TYPE:ty, $SCALAR:ident, $DOC: expr) => {
101        #[doc = $DOC]
102        impl Literal for $TYPE {
103            fn lit(&self) -> Expr {
104                Expr::Literal(ScalarValue::$SCALAR(Some(self.clone())), None)
105            }
106        }
107    };
108}
109
110macro_rules! make_nonzero_literal {
111    ($TYPE:ty, $SCALAR:ident, $DOC: expr) => {
112        #[doc = $DOC]
113        impl Literal for $TYPE {
114            fn lit(&self) -> Expr {
115                Expr::Literal(ScalarValue::$SCALAR(Some(self.get())), None)
116            }
117        }
118    };
119}
120
121macro_rules! make_timestamp_literal {
122    ($TYPE:ty, $SCALAR:ident, $DOC: expr) => {
123        #[doc = $DOC]
124        impl TimestampLiteral for $TYPE {
125            fn lit_timestamp_nano(&self) -> Expr {
126                Expr::Literal(
127                    ScalarValue::TimestampNanosecond(Some((self.clone()).into()), None),
128                    None,
129                )
130            }
131        }
132    };
133}
134
135make_literal!(bool, Boolean, "literal expression containing a bool");
136make_literal!(f32, Float32, "literal expression containing an f32");
137make_literal!(f64, Float64, "literal expression containing an f64");
138make_literal!(i8, Int8, "literal expression containing an i8");
139make_literal!(i16, Int16, "literal expression containing an i16");
140make_literal!(i32, Int32, "literal expression containing an i32");
141make_literal!(i64, Int64, "literal expression containing an i64");
142make_literal!(u8, UInt8, "literal expression containing a u8");
143make_literal!(u16, UInt16, "literal expression containing a u16");
144make_literal!(u32, UInt32, "literal expression containing a u32");
145make_literal!(u64, UInt64, "literal expression containing a u64");
146
147make_nonzero_literal!(
148    std::num::NonZeroI8,
149    Int8,
150    "literal expression containing an i8"
151);
152make_nonzero_literal!(
153    std::num::NonZeroI16,
154    Int16,
155    "literal expression containing an i16"
156);
157make_nonzero_literal!(
158    std::num::NonZeroI32,
159    Int32,
160    "literal expression containing an i32"
161);
162make_nonzero_literal!(
163    std::num::NonZeroI64,
164    Int64,
165    "literal expression containing an i64"
166);
167make_nonzero_literal!(
168    std::num::NonZeroU8,
169    UInt8,
170    "literal expression containing a u8"
171);
172make_nonzero_literal!(
173    std::num::NonZeroU16,
174    UInt16,
175    "literal expression containing a u16"
176);
177make_nonzero_literal!(
178    std::num::NonZeroU32,
179    UInt32,
180    "literal expression containing a u32"
181);
182make_nonzero_literal!(
183    std::num::NonZeroU64,
184    UInt64,
185    "literal expression containing a u64"
186);
187
188make_timestamp_literal!(i8, Int8, "literal expression containing an i8");
189make_timestamp_literal!(i16, Int16, "literal expression containing an i16");
190make_timestamp_literal!(i32, Int32, "literal expression containing an i32");
191make_timestamp_literal!(i64, Int64, "literal expression containing an i64");
192make_timestamp_literal!(u8, UInt8, "literal expression containing a u8");
193make_timestamp_literal!(u16, UInt16, "literal expression containing a u16");
194make_timestamp_literal!(u32, UInt32, "literal expression containing a u32");
195
196#[cfg(test)]
197mod test {
198    use std::num::NonZeroU32;
199
200    use super::*;
201    use crate::expr_fn::col;
202
203    #[test]
204    fn test_lit_nonzero() {
205        let expr = col("id").eq(lit(NonZeroU32::new(1).unwrap()));
206        let expected = col("id").eq(lit(ScalarValue::UInt32(Some(1))));
207        assert_eq!(expr, expected);
208    }
209
210    #[test]
211    fn test_lit_timestamp_nano() {
212        let expr = col("time").eq(lit_timestamp_nano(10)); // 10 is an implicit i32
213        let expected =
214            col("time").eq(lit(ScalarValue::TimestampNanosecond(Some(10), None)));
215        assert_eq!(expr, expected);
216
217        let i: i64 = 10;
218        let expr = col("time").eq(lit_timestamp_nano(i));
219        assert_eq!(expr, expected);
220
221        let i: u32 = 10;
222        let expr = col("time").eq(lit_timestamp_nano(i));
223        assert_eq!(expr, expected);
224    }
225}
datafusion_expr/literal.rs

datafusion_expr/
literal.rs
