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

vortex-data / vortex / 16945584710

13 Aug 2025 06:09PM UTC coverage: 87.832% (+0.2%) from 87.664%
16945584710

Pull #4226

github

web-flow
Merge 90f7ce6f4 into df21f4278
Pull Request #4226: Support converting TimestampTZ to and from duckdb

432 of 451 new or added lines in 15 files covered. (95.79%)

37 existing lines in 7 files now uncovered.

56593 of 64433 relevant lines covered (87.83%)

627980.13 hits per line

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

88.74
/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 {
32
    fn eq(&self, other: &Self) -> bool {
3✔
33
        self.fields == other.fields && self.child.eq(&other.child)
3✔
34
    }
3✔
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

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

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

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

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

66
        Some(ProstMetadata(SelectOpts { opts: Some(opts) }))
2✔
67
    }
2✔
68

69
    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
2,723✔
70
        vec![&expr.child]
2,723✔
71
    }
2,723✔
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

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

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

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

108
        Ok(SelectExpr { fields, child })
2✔
109
    }
2✔
110

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

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

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

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

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

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

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

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

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

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

174
    pub fn fields(&self) -> &SelectField {
1,365✔
175
        &self.fields
1,365✔
176
    }
1,365✔
177

178
    pub fn child(&self) -> &ExprRef {
1,361✔
179
        &self.child
1,361✔
180
    }
1,361✔
181

182
    /// Turn the select expression into an `include`, relative to a provided array of field names.
183
    ///
184
    /// For example:
185
    /// ```rust
186
    /// # use vortex_expr::root;
187
    /// # use vortex_expr::select::{SelectExpr, SelectField};
188
    /// # use vortex_dtype::FieldNames;
189
    /// let field_names = FieldNames::from(["a", "b", "c"]);
190
    /// let include = SelectExpr::new(SelectField::Include(["a"].into()), root());
191
    /// let exclude = SelectExpr::new(SelectField::Exclude(["b", "c"].into()), root());
192
    /// assert_eq!(
193
    ///     &include.as_include(&field_names).unwrap(),
194
    ///     &exclude.as_include(&field_names).unwrap()
195
    /// );
196
    /// ```
197
    pub fn as_include(&self, field_names: &FieldNames) -> VortexResult<ExprRef> {
2✔
198
        Ok(Self::new(
2✔
199
            SelectField::Include(self.fields.as_include_names(field_names)?),
2✔
200
            self.child.clone(),
2✔
201
        )
202
        .into_expr())
2✔
203
    }
2✔
204
}
205

206
impl SelectField {
207
    pub fn include(columns: FieldNames) -> Self {
×
208
        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
×
UNCOV
209
        Self::Include(columns)
×
UNCOV
210
    }
×
211

UNCOV
212
    pub fn exclude(columns: FieldNames) -> Self {
×
UNCOV
213
        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
×
UNCOV
214
        Self::Exclude(columns)
×
UNCOV
215
    }
×
216

217
    pub fn is_include(&self) -> bool {
2✔
218
        matches!(self, Self::Include(_))
2✔
219
    }
2✔
220

UNCOV
221
    pub fn is_exclude(&self) -> bool {
×
222
        matches!(self, Self::Exclude(_))
×
223
    }
×
224

225
    pub fn fields(&self) -> &FieldNames {
1,365✔
226
        let (SelectField::Include(fields) | SelectField::Exclude(fields)) = self;
1,365✔
227

228
        fields
1,365✔
229
    }
1,365✔
230

231
    pub fn as_include_names(&self, field_names: &FieldNames) -> VortexResult<FieldNames> {
1,363✔
232
        if self
1,363✔
233
            .fields()
1,363✔
234
            .iter()
1,363✔
235
            .any(|f| !field_names.iter().contains(f))
3,205✔
236
        {
UNCOV
237
            vortex_bail!(
×
UNCOV
238
                "Field {:?} in select not in field names {:?}",
×
239
                self,
240
                field_names
241
            );
242
        }
1,363✔
243
        match self {
1,363✔
244
            SelectField::Include(fields) => Ok(fields.clone()),
1,362✔
245
            SelectField::Exclude(exc_fields) => Ok(field_names
1✔
246
                .iter()
1✔
247
                .filter(|f| !exc_fields.iter().contains(f))
3✔
248
                .cloned()
1✔
249
                .collect()),
1✔
250
        }
251
    }
1,363✔
252
}
253

254
impl Display for SelectField {
255
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3✔
256
        match self {
3✔
257
            SelectField::Include(fields) => write!(f, "{{{}}}", DisplayFieldNames(fields)),
2✔
258
            SelectField::Exclude(fields) => write!(f, "~{{{}}}", DisplayFieldNames(fields)),
1✔
259
        }
260
    }
3✔
261
}
262

263
impl Display for SelectExpr {
264
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
3✔
265
        write!(f, "{}{}", self.child, self.fields)
3✔
266
    }
3✔
267
}
268

269
impl AnalysisExpr for SelectExpr {}
270

271
#[cfg(test)]
272
mod tests {
273

274
    use vortex_array::arrays::StructArray;
275
    use vortex_array::{IntoArray, ToCanonical};
276
    use vortex_buffer::buffer;
277
    use vortex_dtype::{DType, FieldName, FieldNames, Nullability};
278

279
    use crate::{Scope, SelectExpr, SelectField, root, select, select_exclude, test_harness};
280

281
    fn test_array() -> StructArray {
2✔
282
        StructArray::from_fields(&[
2✔
283
            ("a", buffer![0, 1, 2].into_array()),
2✔
284
            ("b", buffer![4, 5, 6].into_array()),
2✔
285
        ])
2✔
286
        .unwrap()
2✔
287
    }
2✔
288

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

302
    #[test]
303
    pub fn exclude_columns() {
1✔
304
        let st = test_array();
1✔
305
        let select = select_exclude(vec![FieldName::from("a")], root());
1✔
306
        let selected = select
1✔
307
            .evaluate(&Scope::new(st.to_array()))
1✔
308
            .unwrap()
1✔
309
            .to_struct()
1✔
310
            .unwrap();
1✔
311
        let selected_names = selected.names().clone();
1✔
312
        assert_eq!(selected_names.as_ref(), &["b".into()]);
1✔
313
    }
1✔
314

315
    #[test]
316
    fn dtype() {
1✔
317
        let dtype = test_harness::struct_dtype();
1✔
318

319
        let select_expr = select(vec![FieldName::from("a")], root());
1✔
320
        let expected_dtype = DType::Struct(
1✔
321
            dtype.as_struct().unwrap().project(&["a".into()]).unwrap(),
1✔
322
            Nullability::NonNullable,
1✔
323
        );
1✔
324
        assert_eq!(select_expr.return_dtype(&dtype).unwrap(), expected_dtype);
1✔
325

326
        let select_expr_exclude = select_exclude(
1✔
327
            vec![
1✔
328
                FieldName::from("col1"),
1✔
329
                FieldName::from("col2"),
1✔
330
                FieldName::from("bool1"),
1✔
331
                FieldName::from("bool2"),
1✔
332
            ],
333
            root(),
1✔
334
        );
335
        assert_eq!(
1✔
336
            select_expr_exclude.return_dtype(&dtype).unwrap(),
1✔
337
            expected_dtype
338
        );
339

340
        let select_expr_exclude = select_exclude(
1✔
341
            vec![FieldName::from("col1"), FieldName::from("col2")],
1✔
342
            root(),
1✔
343
        );
344
        assert_eq!(
1✔
345
            select_expr_exclude.return_dtype(&dtype).unwrap(),
1✔
346
            DType::Struct(
1✔
347
                dtype
1✔
348
                    .as_struct()
1✔
349
                    .unwrap()
1✔
350
                    .project(&["a".into(), "bool1".into(), "bool2".into()])
1✔
351
                    .unwrap(),
1✔
352
                Nullability::NonNullable
1✔
353
            )
1✔
354
        );
355
    }
1✔
356

357
    #[test]
358
    fn test_as_include_names() {
1✔
359
        let field_names = FieldNames::from(["a", "b", "c"]);
1✔
360
        let include = SelectExpr::new(SelectField::Include(["a"].into()), root());
1✔
361
        let exclude = SelectExpr::new(SelectField::Exclude(["b", "c"].into()), root());
1✔
362
        assert_eq!(
1✔
363
            &include.as_include(&field_names).unwrap(),
1✔
364
            &exclude.as_include(&field_names).unwrap()
1✔
365
        );
366
    }
1✔
367
}
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