• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

vortex-data / vortex / 16935267080

13 Aug 2025 11:00AM UTC coverage: 24.312% (-63.3%) from 87.658%
16935267080

Pull #4226

github

web-flow
Merge 81b48c7fb into baa6ea202
Pull Request #4226: Support converting TimestampTZ to and from duckdb

0 of 2 new or added lines in 1 file covered. (0.0%)

20666 existing lines in 469 files now uncovered.

8726 of 35892 relevant lines covered (24.31%)

147.74 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0
/vortex-expr/src/exprs/select.rs
1
// SPDX-License-Identifier: Apache-2.0
2
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3

4
use std::fmt::Display;
5

6
use itertools::Itertools;
7
use vortex_array::{ArrayRef, DeserializeMetadata, IntoArray, ProstMetadata, ToCanonical};
8
use vortex_dtype::{DType, FieldNames};
9
use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
10
use vortex_proto::expr::select_opts::Opts;
11
use vortex_proto::expr::{FieldNames as ProtoFieldNames, SelectOpts};
12

13
use crate::field::DisplayFieldNames;
14
use crate::{AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, VTable, vtable};
15

16
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
17
pub enum SelectField {
18
    Include(FieldNames),
19
    Exclude(FieldNames),
20
}
21

22
vtable!(Select);
23

24
#[derive(Debug, Clone, Hash, Eq)]
25
#[allow(clippy::derived_hash_with_manual_eq)]
26
pub struct SelectExpr {
27
    fields: SelectField,
28
    child: ExprRef,
29
}
30

31
impl PartialEq for SelectExpr {
UNCOV
32
    fn eq(&self, other: &Self) -> bool {
×
UNCOV
33
        self.fields == other.fields && self.child.eq(&other.child)
×
UNCOV
34
    }
×
35
}
36

37
pub struct SelectExprEncoding;
38

39
impl VTable for SelectVTable {
40
    type Expr = SelectExpr;
41
    type Encoding = SelectExprEncoding;
42
    type Metadata = ProstMetadata<SelectOpts>;
43

UNCOV
44
    fn id(_encoding: &Self::Encoding) -> ExprId {
×
UNCOV
45
        ExprId::new_ref("select")
×
UNCOV
46
    }
×
47

UNCOV
48
    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
×
UNCOV
49
        ExprEncodingRef::new_ref(SelectExprEncoding.as_ref())
×
UNCOV
50
    }
×
51

UNCOV
52
    fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {
×
UNCOV
53
        let names = expr
×
UNCOV
54
            .fields()
×
UNCOV
55
            .fields()
×
UNCOV
56
            .iter()
×
UNCOV
57
            .map(|f| f.to_string())
×
UNCOV
58
            .collect_vec();
×
59

UNCOV
60
        let opts = if expr.fields().is_include() {
×
UNCOV
61
            Opts::Include(ProtoFieldNames { names })
×
62
        } else {
UNCOV
63
            Opts::Exclude(ProtoFieldNames { names })
×
64
        };
65

UNCOV
66
        Some(ProstMetadata(SelectOpts { opts: Some(opts) }))
×
UNCOV
67
    }
×
68

UNCOV
69
    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
×
UNCOV
70
        vec![&expr.child]
×
UNCOV
71
    }
×
72

73
    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
×
74
        Ok(SelectExpr {
×
75
            fields: expr.fields.clone(),
×
76
            child: children[0].clone(),
×
77
        })
×
78
    }
×
79

UNCOV
80
    fn build(
×
UNCOV
81
        _encoding: &Self::Encoding,
×
UNCOV
82
        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
×
UNCOV
83
        mut children: Vec<ExprRef>,
×
UNCOV
84
    ) -> VortexResult<Self::Expr> {
×
UNCOV
85
        if children.len() != 1 {
×
86
            vortex_bail!("Select expression must have exactly one child");
×
UNCOV
87
        }
×
88

UNCOV
89
        let fields = match metadata.opts.as_ref() {
×
UNCOV
90
            Some(opts) => match opts {
×
UNCOV
91
                Opts::Include(field_names) => SelectField::Include(FieldNames::from_iter(
×
UNCOV
92
                    field_names.names.iter().map(|s| s.as_str()),
×
93
                )),
UNCOV
94
                Opts::Exclude(field_names) => SelectField::Exclude(FieldNames::from_iter(
×
UNCOV
95
                    field_names.names.iter().map(|s| s.as_str()),
×
96
                )),
97
            },
98
            None => {
99
                vortex_bail!("Select expressions must be provided with fields to select or exclude")
×
100
            }
101
        };
102

UNCOV
103
        let child = children
×
UNCOV
104
            .drain(..)
×
UNCOV
105
            .next()
×
UNCOV
106
            .vortex_expect("number of children validated to be one");
×
107

UNCOV
108
        Ok(SelectExpr { fields, child })
×
UNCOV
109
    }
×
110

UNCOV
111
    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
×
UNCOV
112
        let batch = expr.child.unchecked_evaluate(scope)?.to_struct()?;
×
UNCOV
113
        Ok(match &expr.fields {
×
UNCOV
114
            SelectField::Include(f) => batch.project(f.as_ref()),
×
UNCOV
115
            SelectField::Exclude(names) => {
×
UNCOV
116
                let included_names = batch
×
UNCOV
117
                    .names()
×
UNCOV
118
                    .iter()
×
UNCOV
119
                    .filter(|&f| !names.as_ref().contains(f))
×
UNCOV
120
                    .cloned()
×
UNCOV
121
                    .collect::<Vec<_>>();
×
UNCOV
122
                batch.project(included_names.as_slice())
×
123
            }
124
        }?
×
UNCOV
125
        .into_array())
×
UNCOV
126
    }
×
127

UNCOV
128
    fn return_dtype(expr: &Self::Expr, scope: &DType) -> VortexResult<DType> {
×
UNCOV
129
        let child_dtype = expr.child.return_dtype(scope)?;
×
UNCOV
130
        let child_struct_dtype = child_dtype
×
UNCOV
131
            .as_struct()
×
UNCOV
132
            .ok_or_else(|| vortex_err!("Select child not a struct dtype"))?;
×
133

UNCOV
134
        let projected = match &expr.fields {
×
UNCOV
135
            SelectField::Include(fields) => child_struct_dtype.project(fields.as_ref())?,
×
UNCOV
136
            SelectField::Exclude(fields) => child_struct_dtype
×
UNCOV
137
                .names()
×
UNCOV
138
                .iter()
×
UNCOV
139
                .cloned()
×
UNCOV
140
                .zip_eq(child_struct_dtype.fields())
×
UNCOV
141
                .filter(|(name, _)| !fields.as_ref().contains(name))
×
UNCOV
142
                .collect(),
×
143
        };
144

UNCOV
145
        Ok(DType::Struct(projected, child_dtype.nullability()))
×
UNCOV
146
    }
×
147
}
148

UNCOV
149
pub fn select(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
×
UNCOV
150
    SelectExpr::include_expr(fields.into(), child)
×
UNCOV
151
}
×
152

UNCOV
153
pub fn select_exclude(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
×
UNCOV
154
    SelectExpr::exclude_expr(fields.into(), child)
×
UNCOV
155
}
×
156

157
impl SelectExpr {
UNCOV
158
    pub fn new(fields: SelectField, child: ExprRef) -> Self {
×
UNCOV
159
        Self { fields, child }
×
UNCOV
160
    }
×
161

162
    pub fn new_expr(fields: SelectField, child: ExprRef) -> ExprRef {
×
163
        Self::new(fields, child).into_expr()
×
164
    }
×
165

UNCOV
166
    pub fn include_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
×
UNCOV
167
        Self::new(SelectField::Include(columns), child).into_expr()
×
UNCOV
168
    }
×
169

UNCOV
170
    pub fn exclude_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
×
UNCOV
171
        Self::new(SelectField::Exclude(columns), child).into_expr()
×
UNCOV
172
    }
×
173

UNCOV
174
    pub fn fields(&self) -> &SelectField {
×
UNCOV
175
        &self.fields
×
UNCOV
176
    }
×
177

UNCOV
178
    pub fn child(&self) -> &ExprRef {
×
UNCOV
179
        &self.child
×
UNCOV
180
    }
×
181

182
    pub fn as_include(&self, field_names: &FieldNames) -> VortexResult<ExprRef> {
×
183
        Ok(Self::new(
×
184
            SelectField::Include(self.fields.as_include_names(field_names)?),
×
185
            self.child.clone(),
×
186
        )
187
        .into_expr())
×
188
    }
×
189
}
190

191
impl SelectField {
192
    pub fn include(columns: FieldNames) -> Self {
×
193
        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
×
194
        Self::Include(columns)
×
195
    }
×
196

197
    pub fn exclude(columns: FieldNames) -> Self {
×
198
        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
×
199
        Self::Exclude(columns)
×
200
    }
×
201

UNCOV
202
    pub fn is_include(&self) -> bool {
×
UNCOV
203
        matches!(self, Self::Include(_))
×
UNCOV
204
    }
×
205

206
    pub fn is_exclude(&self) -> bool {
×
207
        matches!(self, Self::Exclude(_))
×
208
    }
×
209

UNCOV
210
    pub fn fields(&self) -> &FieldNames {
×
UNCOV
211
        let (SelectField::Include(fields) | SelectField::Exclude(fields)) = self;
×
212

UNCOV
213
        fields
×
UNCOV
214
    }
×
215

UNCOV
216
    pub fn as_include_names(&self, field_names: &FieldNames) -> VortexResult<FieldNames> {
×
UNCOV
217
        if self
×
UNCOV
218
            .fields()
×
UNCOV
219
            .iter()
×
UNCOV
220
            .any(|f| !field_names.iter().contains(f))
×
221
        {
222
            vortex_bail!(
×
223
                "Field {:?} in select not in field names {:?}",
×
224
                self,
225
                field_names
226
            );
UNCOV
227
        }
×
UNCOV
228
        match self {
×
UNCOV
229
            SelectField::Include(fields) => Ok(fields.clone()),
×
230
            SelectField::Exclude(exc_fields) => Ok(field_names
×
231
                .iter()
×
232
                .filter(|f| exc_fields.iter().contains(f))
×
233
                .cloned()
×
234
                .collect()),
×
235
        }
UNCOV
236
    }
×
237
}
238

239
impl Display for SelectField {
UNCOV
240
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
×
UNCOV
241
        match self {
×
UNCOV
242
            SelectField::Include(fields) => write!(f, "{{{}}}", DisplayFieldNames(fields)),
×
UNCOV
243
            SelectField::Exclude(fields) => write!(f, "~{{{}}}", DisplayFieldNames(fields)),
×
244
        }
UNCOV
245
    }
×
246
}
247

248
impl Display for SelectExpr {
UNCOV
249
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
×
UNCOV
250
        write!(f, "{}{}", self.child, self.fields)
×
UNCOV
251
    }
×
252
}
253

254
impl AnalysisExpr for SelectExpr {}
255

256
#[cfg(test)]
257
mod tests {
258

259
    use vortex_array::arrays::StructArray;
260
    use vortex_array::{IntoArray, ToCanonical};
261
    use vortex_buffer::buffer;
262
    use vortex_dtype::{DType, FieldName, Nullability};
263

264
    use crate::{Scope, root, select, select_exclude, test_harness};
265

266
    fn test_array() -> StructArray {
267
        StructArray::from_fields(&[
268
            ("a", buffer![0, 1, 2].into_array()),
269
            ("b", buffer![4, 5, 6].into_array()),
270
        ])
271
        .unwrap()
272
    }
273

274
    #[test]
275
    pub fn include_columns() {
276
        let st = test_array();
277
        let select = select(vec![FieldName::from("a")], root());
278
        let selected = select
279
            .evaluate(&Scope::new(st.to_array()))
280
            .unwrap()
281
            .to_struct()
282
            .unwrap();
283
        let selected_names = selected.names().clone();
284
        assert_eq!(selected_names.as_ref(), &["a".into()]);
285
    }
286

287
    #[test]
288
    pub fn exclude_columns() {
289
        let st = test_array();
290
        let select = select_exclude(vec![FieldName::from("a")], root());
291
        let selected = select
292
            .evaluate(&Scope::new(st.to_array()))
293
            .unwrap()
294
            .to_struct()
295
            .unwrap();
296
        let selected_names = selected.names().clone();
297
        assert_eq!(selected_names.as_ref(), &["b".into()]);
298
    }
299

300
    #[test]
301
    fn dtype() {
302
        let dtype = test_harness::struct_dtype();
303

304
        let select_expr = select(vec![FieldName::from("a")], root());
305
        let expected_dtype = DType::Struct(
306
            dtype.as_struct().unwrap().project(&["a".into()]).unwrap(),
307
            Nullability::NonNullable,
308
        );
309
        assert_eq!(select_expr.return_dtype(&dtype).unwrap(), expected_dtype);
310

311
        let select_expr_exclude = select_exclude(
312
            vec![
313
                FieldName::from("col1"),
314
                FieldName::from("col2"),
315
                FieldName::from("bool1"),
316
                FieldName::from("bool2"),
317
            ],
318
            root(),
319
        );
320
        assert_eq!(
321
            select_expr_exclude.return_dtype(&dtype).unwrap(),
322
            expected_dtype
323
        );
324

325
        let select_expr_exclude = select_exclude(
326
            vec![FieldName::from("col1"), FieldName::from("col2")],
327
            root(),
328
        );
329
        assert_eq!(
330
            select_expr_exclude.return_dtype(&dtype).unwrap(),
331
            DType::Struct(
332
                dtype
333
                    .as_struct()
334
                    .unwrap()
335
                    .project(&["a".into(), "bool1".into(), "bool2".into()])
336
                    .unwrap(),
337
                Nullability::NonNullable
338
            )
339
        );
340
    }
341
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc