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

vortex_expr/
get_item.rs
