vortex_expr/
var.rs

1use std::any::Any;
2use std::fmt::Display;
3use std::sync::Arc;
4
5use vortex_array::ArrayRef;
6use vortex_array::stats::Stat;
7use vortex_dtype::{DType, FieldPath};
8use vortex_error::{VortexResult, vortex_err};
9
10use crate::{
11    AccessPath, AnalysisExpr, ExprRef, Identifier, Scope, ScopeDType, StatsCatalog, VortexExpr,
12};
13
14#[derive(Debug, PartialEq, Eq, Hash)]
15pub struct Var {
16    var: Identifier,
17}
18
19/// Used to extract values (Arrays from the Scope).
20/// see `Scope`.
21impl Var {
22    pub fn new_expr(var: Identifier) -> ExprRef {
23        Arc::new(Self { var })
24    }
25
26    pub fn var(&self) -> &Identifier {
27        &self.var
28    }
29}
30
31#[cfg(feature = "proto")]
32pub(crate) mod proto {
33    use vortex_error::{VortexResult, vortex_bail};
34    use vortex_proto::expr::kind::{Kind, Var as ProtoVar};
35
36    use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id, Var};
37
38    pub(crate) struct VarSerde;
39
40    impl Id for VarSerde {
41        fn id(&self) -> &'static str {
42            "var"
43        }
44    }
45
46    impl ExprDeserialize for VarSerde {
47        fn deserialize(&self, kind: &Kind, _children: Vec<ExprRef>) -> VortexResult<ExprRef> {
48            let Kind::Var(op) = kind else {
49                vortex_bail!("wrong kind {:?}, wanted var", kind)
50            };
51
52            match op.var.as_str() {
53                "" => Ok(Var::new_expr(crate::Identifier::Identity)),
54                other => Ok(Var::new_expr(other.parse()?)),
55            }
56        }
57    }
58
59    impl ExprSerializable for Var {
60        fn id(&self) -> &'static str {
61            VarSerde.id()
62        }
63
64        fn serialize_kind(&self) -> VortexResult<Kind> {
65            Ok(Kind::Var(ProtoVar {
66                var: self.var.to_string(),
67            }))
68        }
69    }
70}
71
72impl Display for Var {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        write!(f, "${}", self.var)
75    }
76}
77
78impl AnalysisExpr for Var {
79    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
80        catalog.stats_ref(&self.field_path()?, Stat::Max)
81    }
82
83    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
84        catalog.stats_ref(&self.field_path()?, Stat::Min)
85    }
86
87    fn field_path(&self) -> Option<AccessPath> {
88        Some(AccessPath::new(FieldPath::root(), self.var.clone()))
89    }
90}
91
92impl VortexExpr for Var {
93    fn as_any(&self) -> &dyn Any {
94        self
95    }
96
97    fn unchecked_evaluate(&self, ctx: &Scope) -> VortexResult<ArrayRef> {
98        ctx.array(&self.var)
99            .cloned()
100            .ok_or_else(|| vortex_err!("cannot find '{}' in arrays scope", self.var))
101    }
102
103    fn children(&self) -> Vec<&ExprRef> {
104        vec![]
105    }
106
107    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
108        assert_eq!(children.len(), 0);
109        Var::new_expr(self.var.clone())
110    }
111
112    fn return_dtype(&self, dt_ctx: &ScopeDType) -> VortexResult<DType> {
113        dt_ctx
114            .dtype(&self.var)
115            .cloned()
116            .ok_or_else(|| vortex_err!("cannot find '{}' in dtype scope", self.var))
117    }
118}
119
120pub fn var(ident: Identifier) -> ExprRef {
121    Var::new_expr(ident)
122}
123
124/// Return a global pointer to the identity token.
125/// This is the name of the data found in a vortex array or file.
126pub fn root() -> ExprRef {
127    Var::new_expr(Identifier::Identity)
128}
129
130pub fn is_root(expr: &ExprRef) -> bool {
131    expr.as_any()
132        .downcast_ref::<Var>()
133        .is_some_and(|v| v.var().is_identity())
134}
135
136#[cfg(test)]
137mod tests {
138    use std::str::FromStr;
139
140    use itertools::Itertools;
141    use vortex_array::ToCanonical;
142    use vortex_array::arrays::PrimitiveArray;
143    use vortex_array::validity::Validity;
144    use vortex_buffer::buffer;
145
146    use crate::{Identifier, Scope, eq, var};
147
148    #[test]
149    fn test_two_vars() {
150        let a1 = PrimitiveArray::new(buffer![5, 4, 3, 2, 1, 0], Validity::AllValid).to_array();
151        let a2 = PrimitiveArray::from_iter(1..=6).to_array();
152
153        let expr = eq(var(Identifier::Identity), var("row".parse().unwrap()));
154        let res = expr
155            .evaluate(&Scope::new(a1).with_array("row".parse().unwrap(), a2))
156            .unwrap();
157        let res = res.to_bool().unwrap().boolean_buffer().iter().collect_vec();
158
159        assert_eq!(res, vec![false, false, true, false, false, false])
160    }
161
162    #[test]
163    fn test_empty_string_ident_not_allowed() {
164        assert!(Identifier::from_str("").is_err());
165    }
166}
vortex_expr/var.rs

vortex_expr/
var.rs
