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, ArrayVariants};
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
58impl VortexExpr for GetItem {
59    fn as_any(&self) -> &dyn Any {
60        self
61    }
62
63    fn unchecked_evaluate(&self, batch: &dyn Array) -> VortexResult<ArrayRef> {
64        let child = self.child.evaluate(batch)?;
65        child
66            .as_struct_typed()
67            .ok_or_else(|| vortex_err!("GetItem: child array into struct"))?
68            // TODO(joe): apply struct validity
69            .maybe_null_field_by_name(self.field())
70    }
71
72    fn children(&self) -> Vec<&ExprRef> {
73        vec![self.child()]
74    }
75
76    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
77        assert_eq!(children.len(), 1);
78        Self::new_expr(self.field().clone(), children[0].clone())
79    }
80
81    fn return_dtype(&self, scope_dtype: &DType) -> VortexResult<DType> {
82        let input = self.child.return_dtype(scope_dtype)?;
83        input
84            .as_struct()
85            .ok_or_else(|| vortex_err!("GetItem: child dtype is not a struct"))?
86            .field(self.field())
87    }
88}
89
90impl PartialEq for GetItem {
91    fn eq(&self, other: &GetItem) -> bool {
92        self.field == other.field && self.child.eq(&other.child)
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use vortex_array::IntoArray;
99    use vortex_array::arrays::StructArray;
100    use vortex_buffer::buffer;
101    use vortex_dtype::DType;
102    use vortex_dtype::PType::I32;
103
104    use crate::get_item::get_item;
105    use crate::ident;
106
107    fn test_array() -> StructArray {
108        StructArray::from_fields(&[
109            ("a", buffer![0i32, 1, 2].into_array()),
110            ("b", buffer![4i64, 5, 6].into_array()),
111        ])
112        .unwrap()
113    }
114
115    #[test]
116    pub fn get_item_by_name() {
117        let st = test_array();
118        let get_item = get_item("a", ident());
119        let item = get_item.evaluate(&st).unwrap();
120        assert_eq!(item.dtype(), &DType::from(I32))
121    }
122
123    #[test]
124    pub fn get_item_by_name_none() {
125        let st = test_array();
126        let get_item = get_item("c", ident());
127        assert!(get_item.evaluate(&st).is_err());
128    }
129}
vortex_expr/get_item.rs

vortex_expr/
get_item.rs
