vortex_expr/
get_item.rs

1use std::any::Any;
2use std::fmt::{Debug, Display, Formatter};
3use std::hash::Hash;
4use std::sync::Arc;
5
6use vortex_array::stats::Stat;
7use vortex_array::{ArrayRef, ToCanonical};
8use vortex_dtype::{DType, FieldName};
9use vortex_error::{VortexResult, vortex_err};
10
11use crate::{AccessPath, AnalysisExpr, ExprRef, Scope, ScopeDType, StatsCatalog, VortexExpr, root};
12
13#[derive(Debug, Clone, Eq, Hash)]
14#[allow(clippy::derived_hash_with_manual_eq)]
15pub struct GetItem {
16    field: FieldName,
17    child: ExprRef,
18}
19
20impl GetItem {
21    pub fn new_expr(field: impl Into<FieldName>, child: ExprRef) -> ExprRef {
22        Arc::new(Self {
23            field: field.into(),
24            child,
25        })
26    }
27
28    pub fn field(&self) -> &FieldName {
29        &self.field
30    }
31
32    pub fn child(&self) -> &ExprRef {
33        &self.child
34    }
35
36    pub fn is(expr: &ExprRef) -> bool {
37        expr.as_any().is::<Self>()
38    }
39}
40
41pub fn col(field: impl Into<FieldName>) -> ExprRef {
42    GetItem::new_expr(field, root())
43}
44
45pub fn get_item(field: impl Into<FieldName>, child: ExprRef) -> ExprRef {
46    GetItem::new_expr(field, child)
47}
48
49pub fn get_item_scope(field: impl Into<FieldName>) -> ExprRef {
50    GetItem::new_expr(field, root())
51}
52
53impl Display for GetItem {
54    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
55        write!(f, "{}.{}", self.child, &self.field)
56    }
57}
58
59#[cfg(feature = "proto")]
60pub(crate) mod proto {
61    use vortex_error::{VortexResult, vortex_bail};
62    use vortex_proto::expr::kind;
63    use vortex_proto::expr::kind::Kind;
64
65    use crate::{ExprDeserialize, ExprRef, ExprSerializable, GetItem, Id};
66
67    pub(crate) struct GetItemSerde;
68
69    impl Id for GetItemSerde {
70        fn id(&self) -> &'static str {
71            "get_item"
72        }
73    }
74
75    impl ExprDeserialize for GetItemSerde {
76        fn deserialize(&self, kind: &Kind, children: Vec<ExprRef>) -> VortexResult<ExprRef> {
77            let Kind::GetItem(kind::GetItem { path }) = kind else {
78                vortex_bail!("wrong kind {:?}, want get_item", kind)
79            };
80
81            Ok(GetItem::new_expr(path.to_string(), children[0].clone()))
82        }
83    }
84
85    impl ExprSerializable for GetItem {
86        fn id(&self) -> &'static str {
87            GetItemSerde.id()
88        }
89
90        fn serialize_kind(&self) -> VortexResult<Kind> {
91            Ok(Kind::GetItem(kind::GetItem {
92                path: self.field.to_string(),
93            }))
94        }
95    }
96}
97
98impl AnalysisExpr for GetItem {
99    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
100        catalog.stats_ref(&self.field_path()?, Stat::Max)
101    }
102
103    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
104        catalog.stats_ref(&self.field_path()?, Stat::Min)
105    }
106
107    fn field_path(&self) -> Option<AccessPath> {
108        self.child()
109            .field_path()
110            .map(|fp| AccessPath::new(fp.field_path.push(self.field.clone()), fp.identifier))
111    }
112}
113
114impl VortexExpr for GetItem {
115    fn as_any(&self) -> &dyn Any {
116        self
117    }
118
119    fn unchecked_evaluate(&self, scope: &Scope) -> VortexResult<ArrayRef> {
120        self.child
121            .unchecked_evaluate(scope)?
122            .to_struct()?
123            .field_by_name(self.field())
124            .cloned()
125    }
126
127    fn children(&self) -> Vec<&ExprRef> {
128        vec![self.child()]
129    }
130
131    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
132        assert_eq!(children.len(), 1);
133        Self::new_expr(self.field().clone(), children[0].clone())
134    }
135
136    fn return_dtype(&self, scope: &ScopeDType) -> VortexResult<DType> {
137        let input = self.child.return_dtype(scope)?;
138        input
139            .as_struct()
140            .ok_or_else(|| vortex_err!("GetItem: child dtype is not a struct"))?
141            .field(self.field())
142    }
143}
144
145impl PartialEq for GetItem {
146    fn eq(&self, other: &GetItem) -> bool {
147        self.field == other.field && self.child.eq(&other.child)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use vortex_array::IntoArray;
154    use vortex_array::arrays::StructArray;
155    use vortex_buffer::buffer;
156    use vortex_dtype::DType;
157    use vortex_dtype::PType::I32;
158
159    use crate::get_item::get_item;
160    use crate::{Scope, root};
161
162    fn test_array() -> StructArray {
163        StructArray::from_fields(&[
164            ("a", buffer![0i32, 1, 2].into_array()),
165            ("b", buffer![4i64, 5, 6].into_array()),
166        ])
167        .unwrap()
168    }
169
170    #[test]
171    pub fn get_item_by_name() {
172        let st = test_array();
173        let get_item = get_item("a", root());
174        let item = get_item.evaluate(&Scope::new(st.to_array())).unwrap();
175        assert_eq!(item.dtype(), &DType::from(I32))
176    }
177
178    #[test]
179    pub fn get_item_by_name_none() {
180        let st = test_array();
181        let get_item = get_item("c", root());
182        assert!(get_item.evaluate(&Scope::new(st.to_array())).is_err());
183    }
184}
vortex_expr/get_item.rs

vortex_expr/
get_item.rs
