vortex_expr/exprs/
get_item.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::{Debug, Display, Formatter};
5use std::hash::Hash;
6
7use vortex_array::stats::Stat;
8use vortex_array::{ArrayRef, DeserializeMetadata, ProstMetadata, ToCanonical};
9use vortex_dtype::{DType, FieldName, FieldPath};
10use vortex_error::{VortexResult, vortex_bail, vortex_err};
11use vortex_proto::expr as pb;
12
13use crate::{
14    AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, StatsCatalog, VTable, root,
15    vtable,
16};
17
18vtable!(GetItem);
19
20#[allow(clippy::derived_hash_with_manual_eq)]
21#[derive(Debug, Clone, Hash)]
22pub struct GetItemExpr {
23    field: FieldName,
24    child: ExprRef,
25}
26
27impl PartialEq for GetItemExpr {
28    fn eq(&self, other: &Self) -> bool {
29        self.field == other.field && self.child.eq(&other.child)
30    }
31}
32
33pub struct GetItemExprEncoding;
34
35impl VTable for GetItemVTable {
36    type Expr = GetItemExpr;
37    type Encoding = GetItemExprEncoding;
38    type Metadata = ProstMetadata<pb::GetItemOpts>;
39
40    fn id(_encoding: &Self::Encoding) -> ExprId {
41        ExprId::new_ref("get_item")
42    }
43
44    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
45        ExprEncodingRef::new_ref(GetItemExprEncoding.as_ref())
46    }
47
48    fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {
49        Some(ProstMetadata(pb::GetItemOpts {
50            path: expr.field.to_string(),
51        }))
52    }
53
54    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
55        vec![&expr.child]
56    }
57
58    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
59        Ok(GetItemExpr {
60            field: expr.field.clone(),
61            child: children[0].clone(),
62        })
63    }
64
65    fn build(
66        _encoding: &Self::Encoding,
67        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
68        children: Vec<ExprRef>,
69    ) -> VortexResult<Self::Expr> {
70        if children.len() != 1 {
71            vortex_bail!(
72                "GetItem expression must have exactly 1 child, got {}",
73                children.len()
74            );
75        }
76
77        let field = FieldName::from(metadata.path.clone());
78        Ok(GetItemExpr {
79            field,
80            child: children[0].clone(),
81        })
82    }
83
84    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
85        expr.child
86            .unchecked_evaluate(scope)?
87            .to_struct()?
88            .field_by_name(expr.field())
89            .cloned()
90    }
91
92    fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {
93        let input = expr.child.return_dtype(scope)?;
94        input
95            .as_struct()
96            .and_then(|st| st.field(expr.field()))
97            .ok_or_else(|| {
98                vortex_err!(
99                    "Couldn't find the {} field in the input scope",
100                    expr.field()
101                )
102            })
103    }
104}
105
106impl GetItemExpr {
107    pub fn new(field: impl Into<FieldName>, child: ExprRef) -> Self {
108        Self {
109            field: field.into(),
110            child,
111        }
112    }
113
114    pub fn new_expr(field: impl Into<FieldName>, child: ExprRef) -> ExprRef {
115        Self::new(field, child).into_expr()
116    }
117
118    pub fn field(&self) -> &FieldName {
119        &self.field
120    }
121
122    pub fn child(&self) -> &ExprRef {
123        &self.child
124    }
125
126    pub fn is(expr: &ExprRef) -> bool {
127        expr.is::<GetItemVTable>()
128    }
129}
130
131pub fn col(field: impl Into<FieldName>) -> ExprRef {
132    GetItemExpr::new(field, root()).into_expr()
133}
134
135pub fn get_item(field: impl Into<FieldName>, child: ExprRef) -> ExprRef {
136    GetItemExpr::new(field, child).into_expr()
137}
138
139impl Display for GetItemExpr {
140    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
141        write!(f, "{}.{}", self.child, &self.field)
142    }
143}
144
145impl AnalysisExpr for GetItemExpr {
146    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
147        catalog.stats_ref(&self.field_path()?, Stat::Max)
148    }
149
150    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
151        catalog.stats_ref(&self.field_path()?, Stat::Min)
152    }
153
154    fn nan_count(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
155        catalog.stats_ref(&self.field_path()?, Stat::NaNCount)
156    }
157
158    fn field_path(&self) -> Option<FieldPath> {
159        self.child()
160            .field_path()
161            .map(|fp| fp.push(self.field.clone()))
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use vortex_array::IntoArray;
168    use vortex_array::arrays::StructArray;
169    use vortex_buffer::buffer;
170    use vortex_dtype::DType;
171    use vortex_dtype::PType::I32;
172
173    use crate::get_item::get_item;
174    use crate::{Scope, root};
175
176    fn test_array() -> StructArray {
177        StructArray::from_fields(&[
178            ("a", buffer![0i32, 1, 2].into_array()),
179            ("b", buffer![4i64, 5, 6].into_array()),
180        ])
181        .unwrap()
182    }
183
184    #[test]
185    pub fn get_item_by_name() {
186        let st = test_array();
187        let get_item = get_item("a", root());
188        let item = get_item.evaluate(&Scope::new(st.to_array())).unwrap();
189        assert_eq!(item.dtype(), &DType::from(I32))
190    }
191
192    #[test]
193    pub fn get_item_by_name_none() {
194        let st = test_array();
195        let get_item = get_item("c", root());
196        assert!(get_item.evaluate(&Scope::new(st.to_array())).is_err());
197    }
198}
vortex_expr/exprs/get_item.rs

vortex_expr/exprs/
get_item.rs
