vortex_array/scalar_fn/fns/
literal.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Formatter;
5
6use prost::Message;
7use vortex_error::VortexResult;
8use vortex_error::vortex_err;
9use vortex_proto::expr as pb;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::IntoArray;
16use crate::arrays::ConstantArray;
17use crate::dtype::DType;
18use crate::expr::Expression;
19use crate::expr::display::ExprDisplay;
20use crate::scalar::Scalar;
21use crate::scalar_fn::Arity;
22use crate::scalar_fn::ChildName;
23use crate::scalar_fn::ExecutionArgs;
24use crate::scalar_fn::ScalarFnId;
25use crate::scalar_fn::ScalarFnVTable;
26use crate::scalar_fn::ScalarFnVTableExt;
27
28fn lit(value: impl Into<Scalar>) -> Expression {
29    Literal.new_expr(value.into(), [])
30}
31
32/// Expression that represents a literal scalar value.
33#[derive(Clone)]
34pub struct Literal;
35
36impl ScalarFnVTable for Literal {
37    type Options = Scalar;
38
39    fn id(&self) -> ScalarFnId {
40        static ID: CachedId = CachedId::new("vortex.literal");
41        *ID
42    }
43
44    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
45        Ok(Some(
46            pb::LiteralOpts {
47                value: Some(instance.into()),
48            }
49            .encode_to_vec(),
50        ))
51    }
52
53    fn deserialize(
54        &self,
55        _metadata: &[u8],
56        session: &VortexSession,
57    ) -> VortexResult<Self::Options> {
58        let ops = pb::LiteralOpts::decode(_metadata)?;
59        Scalar::from_proto(
60            ops.value
61                .as_ref()
62                .ok_or_else(|| vortex_err!("Literal metadata missing value"))?,
63            session,
64        )
65    }
66
67    fn arity(&self, _options: &Self::Options) -> Arity {
68        Arity::Exact(0)
69    }
70
71    fn child_name(&self, _instance: &Self::Options, _child_idx: usize) -> ChildName {
72        unreachable!()
73    }
74
75    fn fmt_sql(
76        &self,
77        scalar: &Scalar,
78        _expr: &dyn ExprDisplay,
79        f: &mut Formatter<'_>,
80    ) -> std::fmt::Result {
81        write!(f, "{}", scalar)
82    }
83
84    fn return_dtype(&self, options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
85        Ok(options.dtype().clone())
86    }
87
88    fn execute(
89        &self,
90        scalar: &Scalar,
91        args: &dyn ExecutionArgs,
92        _ctx: &mut ExecutionCtx,
93    ) -> VortexResult<ArrayRef> {
94        Ok(ConstantArray::new(scalar.clone(), args.row_count()).into_array())
95    }
96
97    fn validity(
98        &self,
99        scalar: &Scalar,
100        _expression: &Expression,
101    ) -> VortexResult<Option<Expression>> {
102        Ok(Some(lit(scalar.is_valid())))
103    }
104
105    fn is_strict(&self, _instance: &Self::Options) -> bool {
106        true
107    }
108
109    fn is_fallible(&self, _instance: &Self::Options) -> bool {
110        false
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use crate::dtype::DType;
117    use crate::dtype::Nullability;
118    use crate::dtype::PType;
119    use crate::dtype::StructFields;
120    use crate::expr::lit;
121    use crate::expr::test_harness;
122    use crate::scalar::Scalar;
123
124    #[test]
125    fn dtype() {
126        let dtype = test_harness::struct_dtype();
127
128        assert_eq!(
129            lit(10).return_dtype(&dtype).unwrap(),
130            DType::Primitive(PType::I32, Nullability::NonNullable)
131        );
132        assert_eq!(
133            lit(i64::MAX).return_dtype(&dtype).unwrap(),
134            DType::Primitive(PType::I64, Nullability::NonNullable)
135        );
136        assert_eq!(
137            lit(true).return_dtype(&dtype).unwrap(),
138            DType::Bool(Nullability::NonNullable)
139        );
140        assert_eq!(
141            lit(Scalar::null(DType::Bool(Nullability::Nullable)))
142                .return_dtype(&dtype)
143                .unwrap(),
144            DType::Bool(Nullability::Nullable)
145        );
146
147        let sdtype = DType::Struct(
148            StructFields::new(
149                ["dog", "cat"].into(),
150                vec![
151                    DType::Primitive(PType::U32, Nullability::NonNullable),
152                    DType::Utf8(Nullability::NonNullable),
153                ],
154            ),
155            Nullability::NonNullable,
156        );
157        assert_eq!(
158            lit(Scalar::struct_(
159                sdtype.clone(),
160                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
161            ))
162            .return_dtype(&dtype)
163            .unwrap(),
164            sdtype
165        );
166    }
167}
vortex_array/scalar_fn/fns/literal.rs

vortex_array/scalar_fn/fns/
literal.rs
