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

vortex_array/scalar_fn/fns/
literal.rs
