vortex_expr/
select.rs

1use std::any::Any;
2use std::fmt::Display;
3use std::sync::Arc;
4
5use itertools::Itertools;
6use vortex_array::{Array, ArrayRef, ArrayVariants};
7use vortex_dtype::{DType, FieldNames};
8use vortex_error::{VortexResult, vortex_bail, vortex_err};
9
10use crate::field::DisplayFieldNames;
11use crate::{ExprRef, VortexExpr};
12
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub enum SelectField {
15    Include(FieldNames),
16    Exclude(FieldNames),
17}
18
19#[derive(Debug, Clone, Eq, Hash)]
20#[allow(clippy::derived_hash_with_manual_eq)]
21pub struct Select {
22    fields: SelectField,
23    child: ExprRef,
24}
25
26pub fn select(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
27    Select::include_expr(fields.into(), child)
28}
29
30pub fn select_exclude(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
31    Select::exclude_expr(fields.into(), child)
32}
33
34impl Select {
35    pub fn new_expr(fields: SelectField, child: ExprRef) -> ExprRef {
36        Arc::new(Self { fields, child })
37    }
38
39    pub fn include_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
40        Self::new_expr(SelectField::Include(columns), child)
41    }
42
43    pub fn exclude_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
44        Self::new_expr(SelectField::Exclude(columns), child)
45    }
46
47    pub fn fields(&self) -> &SelectField {
48        &self.fields
49    }
50
51    pub fn child(&self) -> &ExprRef {
52        &self.child
53    }
54
55    pub fn as_include(&self, field_names: &FieldNames) -> VortexResult<ExprRef> {
56        Ok(Self::new_expr(
57            SelectField::Include(self.fields.as_include_names(field_names)?),
58            self.child.clone(),
59        ))
60    }
61}
62
63impl SelectField {
64    pub fn include(columns: FieldNames) -> Self {
65        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
66        Self::Include(columns)
67    }
68
69    pub fn exclude(columns: FieldNames) -> Self {
70        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
71        Self::Exclude(columns)
72    }
73
74    pub fn is_include(&self) -> bool {
75        matches!(self, Self::Include(_))
76    }
77
78    pub fn is_exclude(&self) -> bool {
79        matches!(self, Self::Exclude(_))
80    }
81
82    pub fn fields(&self) -> &FieldNames {
83        match self {
84            SelectField::Include(fields) => fields,
85            SelectField::Exclude(fields) => fields,
86        }
87    }
88
89    pub fn as_include_names(&self, field_names: &FieldNames) -> VortexResult<FieldNames> {
90        if self.fields().iter().any(|f| !field_names.contains(f)) {
91            vortex_bail!(
92                "Field {:?} in select not in field names {:?}",
93                self,
94                field_names
95            );
96        }
97        match self {
98            SelectField::Include(fields) => Ok(fields.clone()),
99            SelectField::Exclude(exc_fields) => Ok(field_names
100                .iter()
101                .filter(|f| exc_fields.contains(f))
102                .cloned()
103                .collect()),
104        }
105    }
106}
107
108impl Display for SelectField {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        match self {
111            SelectField::Include(fields) => write!(f, "{{{}}}", DisplayFieldNames(fields)),
112            SelectField::Exclude(fields) => write!(f, "~{{{}}}", DisplayFieldNames(fields)),
113        }
114    }
115}
116
117impl Display for Select {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        write!(f, "{}{}", self.child, self.fields)
120    }
121}
122
123#[cfg(feature = "proto")]
124pub(crate) mod proto {
125    use vortex_error::{VortexResult, vortex_bail};
126    use vortex_proto::expr::kind::Kind;
127
128    use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id, Select};
129
130    pub struct SelectSerde;
131
132    impl Id for SelectSerde {
133        fn id(&self) -> &'static str {
134            "select"
135        }
136    }
137
138    impl ExprDeserialize for SelectSerde {
139        fn deserialize(&self, _kind: &Kind, _children: Vec<ExprRef>) -> VortexResult<ExprRef> {
140            vortex_bail!(NotImplemented: "", self.id())
141        }
142    }
143
144    impl ExprSerializable for Select {
145        fn id(&self) -> &'static str {
146            SelectSerde.id()
147        }
148
149        fn serialize_kind(&self) -> VortexResult<Kind> {
150            vortex_bail!(NotImplemented: "", self.id())
151        }
152    }
153}
154
155impl VortexExpr for Select {
156    fn as_any(&self) -> &dyn Any {
157        self
158    }
159
160    fn unchecked_evaluate(&self, batch: &dyn Array) -> VortexResult<ArrayRef> {
161        let batch = self.child.evaluate(batch)?;
162        let st = batch
163            .as_struct_typed()
164            .ok_or_else(|| vortex_err!("Not a struct array"))?;
165        match &self.fields {
166            SelectField::Include(f) => st.project(f),
167            SelectField::Exclude(names) => {
168                let included_names = st
169                    .names()
170                    .iter()
171                    .filter(|&f| !names.contains(f))
172                    .cloned()
173                    .collect::<Vec<_>>();
174                st.project(included_names.as_slice())
175            }
176        }
177    }
178
179    fn children(&self) -> Vec<&ExprRef> {
180        vec![&self.child]
181    }
182
183    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
184        assert_eq!(children.len(), 1);
185        Self::new_expr(self.fields.clone(), children[0].clone())
186    }
187
188    fn return_dtype(&self, scope_dtype: &DType) -> VortexResult<DType> {
189        let child_dtype = self.child.return_dtype(scope_dtype)?;
190        let child_struct_dtype = child_dtype
191            .as_struct()
192            .ok_or_else(|| vortex_err!("Select child not a struct dtype"))?;
193
194        let projected = match &self.fields {
195            SelectField::Include(fields) => child_struct_dtype.project(fields)?,
196            SelectField::Exclude(fields) => child_struct_dtype
197                .names()
198                .iter()
199                .cloned()
200                .zip_eq(child_struct_dtype.fields())
201                .filter(|(name, _)| !fields.contains(name))
202                .collect(),
203        };
204
205        Ok(DType::Struct(
206            Arc::new(projected),
207            child_dtype.nullability(),
208        ))
209    }
210}
211
212impl PartialEq for Select {
213    fn eq(&self, other: &Select) -> bool {
214        self.fields == other.fields && self.child.eq(&other.child)
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use std::sync::Arc;
221
222    use vortex_array::arrays::StructArray;
223    use vortex_array::{ArrayVariants, IntoArray};
224    use vortex_buffer::buffer;
225    use vortex_dtype::{DType, FieldName, Nullability};
226
227    use crate::{ident, select, select_exclude, test_harness};
228
229    fn test_array() -> StructArray {
230        StructArray::from_fields(&[
231            ("a", buffer![0, 1, 2].into_array()),
232            ("b", buffer![4, 5, 6].into_array()),
233        ])
234        .unwrap()
235    }
236
237    #[test]
238    pub fn include_columns() {
239        let st = test_array();
240        let select = select(vec![FieldName::from("a")], ident());
241        let selected = select.evaluate(&st).unwrap();
242        let selected_names = selected.as_struct_typed().unwrap().names().clone();
243        assert_eq!(selected_names.as_ref(), &["a".into()]);
244    }
245
246    #[test]
247    pub fn exclude_columns() {
248        let st = test_array();
249        let select = select_exclude(vec![FieldName::from("a")], ident());
250        let selected = select.evaluate(&st).unwrap();
251        let selected_names = selected.as_struct_typed().unwrap().names().clone();
252        assert_eq!(selected_names.as_ref(), &["b".into()]);
253    }
254
255    #[test]
256    fn dtype() {
257        let dtype = test_harness::struct_dtype();
258
259        let select_expr = select(vec![FieldName::from("a")], ident());
260        let expected_dtype = DType::Struct(
261            Arc::new(dtype.as_struct().unwrap().project(&["a".into()]).unwrap()),
262            Nullability::NonNullable,
263        );
264        assert_eq!(select_expr.return_dtype(&dtype).unwrap(), expected_dtype);
265
266        let select_expr_exclude = select_exclude(
267            vec![
268                FieldName::from("col1"),
269                FieldName::from("col2"),
270                FieldName::from("bool1"),
271                FieldName::from("bool2"),
272            ],
273            ident(),
274        );
275        assert_eq!(
276            select_expr_exclude.return_dtype(&dtype).unwrap(),
277            expected_dtype
278        );
279
280        let select_expr_exclude = select_exclude(
281            vec![FieldName::from("col1"), FieldName::from("col2")],
282            ident(),
283        );
284        assert_eq!(
285            select_expr_exclude.return_dtype(&dtype).unwrap(),
286            DType::Struct(
287                Arc::new(
288                    dtype
289                        .as_struct()
290                        .unwrap()
291                        .project(&["a".into(), "bool1".into(), "bool2".into()])
292                        .unwrap()
293                ),
294                Nullability::NonNullable
295            )
296        );
297    }
298}
vortex_expr/select.rs

vortex_expr/
select.rs
