vortex_array/expr/exprs/
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_dtype::{DType, match_each_float_ptype};
8use vortex_error::{VortexResult, vortex_bail, vortex_err};
9use vortex_proto::expr as pb;
10use vortex_scalar::Scalar;
11
12use crate::arrays::ConstantArray;
13use crate::expr::{ChildName, ExprId, Expression, ExpressionView, StatsCatalog, VTable, VTableExt};
14use crate::{Array, ArrayRef, IntoArray};
15
16/// Expression that represents a literal scalar value.
17pub struct Literal;
18
19impl VTable for Literal {
20    type Instance = Scalar;
21
22    fn id(&self) -> ExprId {
23        ExprId::new_ref("vortex.literal")
24    }
25
26    fn serialize(&self, instance: &Self::Instance) -> VortexResult<Option<Vec<u8>>> {
27        Ok(Some(
28            pb::LiteralOpts {
29                value: Some(instance.as_ref().into()),
30            }
31            .encode_to_vec(),
32        ))
33    }
34
35    fn deserialize(&self, metadata: &[u8]) -> VortexResult<Option<Self::Instance>> {
36        let ops = pb::LiteralOpts::decode(metadata)?;
37        Ok(Some(
38            ops.value
39                .as_ref()
40                .ok_or_else(|| vortex_err!("Literal metadata missing value"))?
41                .try_into()?,
42        ))
43    }
44
45    fn validate(&self, expr: &ExpressionView<Self>) -> VortexResult<()> {
46        if !expr.children().is_empty() {
47            vortex_bail!(
48                "Literal expression does not have children, got: {:?}",
49                expr.children()
50            );
51        }
52        Ok(())
53    }
54
55    fn child_name(&self, _instance: &Self::Instance, _child_idx: usize) -> ChildName {
56        unreachable!()
57    }
58
59    fn fmt_sql(&self, expr: &ExpressionView<Self>, f: &mut Formatter<'_>) -> std::fmt::Result {
60        write!(f, "{}", expr.data())
61    }
62
63    fn fmt_data(&self, instance: &Self::Instance, f: &mut Formatter<'_>) -> std::fmt::Result {
64        write!(f, "{}", instance)
65    }
66
67    fn return_dtype(&self, expr: &ExpressionView<Self>, _scope: &DType) -> VortexResult<DType> {
68        Ok(expr.data().dtype().clone())
69    }
70
71    fn evaluate(&self, expr: &ExpressionView<Self>, scope: &ArrayRef) -> VortexResult<ArrayRef> {
72        Ok(ConstantArray::new(expr.data().clone(), scope.len()).into_array())
73    }
74
75    fn stat_max(
76        &self,
77        expr: &ExpressionView<Self>,
78        _catalog: &mut dyn StatsCatalog,
79    ) -> Option<Expression> {
80        Some(lit(expr.data().clone()))
81    }
82
83    fn stat_min(
84        &self,
85        expr: &ExpressionView<Self>,
86        _catalog: &mut dyn StatsCatalog,
87    ) -> Option<Expression> {
88        Some(lit(expr.data().clone()))
89    }
90
91    fn stat_nan_count(
92        &self,
93        expr: &ExpressionView<Self>,
94        _catalog: &mut dyn StatsCatalog,
95    ) -> Option<Expression> {
96        // The NaNCount for a non-float literal is not defined.
97        // For floating point types, the NaNCount is 1 for lit(NaN), and 0 otherwise.
98        let value = expr.data().as_primitive_opt()?;
99        if !value.ptype().is_float() {
100            return None;
101        }
102
103        match_each_float_ptype!(value.ptype(), |T| {
104            match value.typed_value::<T>() {
105                None => Some(lit(0u64)),
106                Some(value) if value.is_nan() => Some(lit(1u64)),
107                _ => Some(lit(0u64)),
108            }
109        })
110    }
111}
112
113/// Create a new `Literal` expression from a type that coerces to `Scalar`.
114///
115///
116/// ## Example usage
117///
118/// ```
119/// use vortex_array::arrays::PrimitiveArray;
120/// use vortex_dtype::Nullability;
121/// use vortex_array::expr::{lit, Literal};
122/// use vortex_scalar::Scalar;
123///
124/// let number = lit(34i32);
125///
126/// let literal = number.as_::<Literal>();
127/// assert_eq!(literal.data(), &Scalar::primitive(34i32, Nullability::NonNullable));
128/// ```
129pub fn lit(value: impl Into<Scalar>) -> Expression {
130    Literal.new_expr(value.into(), [])
131}
132
133#[cfg(test)]
134mod tests {
135    use vortex_dtype::{DType, Nullability, PType, StructFields};
136    use vortex_scalar::Scalar;
137
138    use super::lit;
139    use crate::expr::test_harness;
140
141    #[test]
142    fn dtype() {
143        let dtype = test_harness::struct_dtype();
144
145        assert_eq!(
146            lit(10).return_dtype(&dtype).unwrap(),
147            DType::Primitive(PType::I32, Nullability::NonNullable)
148        );
149        assert_eq!(
150            lit(i64::MAX).return_dtype(&dtype).unwrap(),
151            DType::Primitive(PType::I64, Nullability::NonNullable)
152        );
153        assert_eq!(
154            lit(true).return_dtype(&dtype).unwrap(),
155            DType::Bool(Nullability::NonNullable)
156        );
157        assert_eq!(
158            lit(Scalar::null(DType::Bool(Nullability::Nullable)))
159                .return_dtype(&dtype)
160                .unwrap(),
161            DType::Bool(Nullability::Nullable)
162        );
163
164        let sdtype = DType::Struct(
165            StructFields::new(
166                ["dog", "cat"].into(),
167                vec![
168                    DType::Primitive(PType::U32, Nullability::NonNullable),
169                    DType::Utf8(Nullability::NonNullable),
170                ],
171            ),
172            Nullability::NonNullable,
173        );
174        assert_eq!(
175            lit(Scalar::struct_(
176                sdtype.clone(),
177                vec![Scalar::from(32_u32), Scalar::from("rufus".to_string())]
178            ))
179            .return_dtype(&dtype)
180            .unwrap(),
181            sdtype
182        );
183    }
184}