16262933935

Committed 14 Jul 2025 09:17AM UTC coverage: 81.432% (+0.3%) from 81.15%

Build # 16262933935

Build Type

Pull #3844

github

Committed by

web-flow

Commit Message

Merge a49b86c01 into 382a2c489

Pull Request Pull Request #3844: feat[compressor]: wire up sequence array into the compressor

Run Details

155 of 159 new or added lines in 8 files covered. (97.48%)

217 existing lines in 24 files now uncovered.

46015 of 56507 relevant lines covered (81.43%)

146742.82 hits per line

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

71.08

/vortex-expr/src/exprs/select.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::fmt::Display;

use itertools::Itertools;
use vortex_array::{ArrayRef, DeserializeMetadata, EmptyMetadata, IntoArray, ToCanonical};
use vortex_dtype::{DType, FieldNames};
use vortex_error::{VortexResult, vortex_bail, vortex_err};

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

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum SelectField {
    Include(FieldNames),
    Exclude(FieldNames),
}

vtable!(Select);

#[derive(Debug, Clone, Hash)]
#[allow(clippy::derived_hash_with_manual_eq)]
pub struct SelectExpr {
    fields: SelectField,
    child: ExprRef,
}

impl PartialEq for SelectExpr {
    fn eq(&self, other: &Self) -> bool {
        self.fields == other.fields && self.child.eq(&other.child)
    }
}

pub struct SelectExprEncoding;

impl VTable for SelectVTable {
    type Expr = SelectExpr;
    type Encoding = SelectExprEncoding;
    type Metadata = EmptyMetadata;

    fn id(_encoding: &Self::Encoding) -> ExprId {
        ExprId::new_ref("select")
    }

    fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {
        ExprEncodingRef::new_ref(SelectExprEncoding.as_ref())
    }

    fn metadata(_expr: &Self::Expr) -> Option<Self::Metadata> {
        // Select does not support serialization
        None
    }

    fn children(expr: &Self::Expr) -> Vec<&ExprRef> {
        vec![&expr.child]
    }

    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
        Ok(SelectExpr {
            fields: expr.fields.clone(),
            child: children[0].clone(),
        })
    }

    fn build(
        _encoding: &Self::Encoding,
        _metadata: &<Self::Metadata as DeserializeMetadata>::Output,
        _children: Vec<ExprRef>,
    ) -> VortexResult<Self::Expr> {
        vortex_bail!("Select does not support deserialization")
    }

    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
        let batch = expr.child.unchecked_evaluate(scope)?.to_struct()?;
        Ok(match &expr.fields {
            SelectField::Include(f) => batch.project(f.as_ref()),
            SelectField::Exclude(names) => {
                let included_names = batch
                    .names()
                    .iter()
                    .filter(|&f| !names.as_ref().contains(f))
                    .cloned()
                    .collect::<Vec<_>>();
                batch.project(included_names.as_slice())
            }
        }?
        .into_array())
    }

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

        let projected = match &expr.fields {
            SelectField::Include(fields) => child_struct_dtype.project(fields.as_ref())?,
            SelectField::Exclude(fields) => child_struct_dtype
                .names()
                .iter()
                .cloned()
                .zip_eq(child_struct_dtype.fields())
                .filter(|(name, _)| !fields.as_ref().contains(name))
                .collect(),
        };

        Ok(DType::Struct(projected, child_dtype.nullability()))
    }
}

pub fn select(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
    SelectExpr::include_expr(fields.into(), child)
}

pub fn select_exclude(fields: impl Into<FieldNames>, child: ExprRef) -> ExprRef {
    SelectExpr::exclude_expr(fields.into(), child)
}

impl SelectExpr {
    pub fn new(fields: SelectField, child: ExprRef) -> Self {
        Self { fields, child }
    }

    pub fn new_expr(fields: SelectField, child: ExprRef) -> ExprRef {
        Self::new(fields, child).into_expr()
    }

    pub fn include_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
        Self::new(SelectField::Include(columns), child).into_expr()
    }

    pub fn exclude_expr(columns: FieldNames, child: ExprRef) -> ExprRef {
        Self::new(SelectField::Exclude(columns), child).into_expr()
    }

    pub fn fields(&self) -> &SelectField {
        &self.fields
    }

    pub fn child(&self) -> &ExprRef {
        &self.child
    }

    pub fn as_include(&self, field_names: &FieldNames) -> VortexResult<ExprRef> {
        Ok(Self::new(
            SelectField::Include(self.fields.as_include_names(field_names)?),
            self.child.clone(),
        )
        .into_expr())
    }
}

impl SelectField {
    pub fn include(columns: FieldNames) -> Self {
        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
        Self::Include(columns)
    }

    pub fn exclude(columns: FieldNames) -> Self {
        assert_eq!(columns.iter().unique().collect_vec().len(), columns.len());
        Self::Exclude(columns)
    }

    pub fn is_include(&self) -> bool {
        matches!(self, Self::Include(_))
    }

    pub fn is_exclude(&self) -> bool {
        matches!(self, Self::Exclude(_))
    }

    pub fn fields(&self) -> &FieldNames {
        match self {
            SelectField::Include(fields) => fields,
            SelectField::Exclude(fields) => fields,
        }
    }

    pub fn as_include_names(&self, field_names: &FieldNames) -> VortexResult<FieldNames> {
        if self
            .fields()
            .iter()
            .any(|f| !field_names.iter().contains(f))
        {
            vortex_bail!(
                "Field {:?} in select not in field names {:?}",
                self,
                field_names
            );
        }
        match self {
            SelectField::Include(fields) => Ok(fields.clone()),
            SelectField::Exclude(exc_fields) => Ok(field_names
                .iter()
                .filter(|f| exc_fields.iter().contains(f))
                .cloned()
                .collect()),
        }
    }
}

impl Display for SelectField {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SelectField::Include(fields) => write!(f, "{{{}}}", DisplayFieldNames(fields)),
            SelectField::Exclude(fields) => write!(f, "~{{{}}}", DisplayFieldNames(fields)),
        }
    }
}

impl Display for SelectExpr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}{}", self.child, self.fields)
    }
}

impl AnalysisExpr for SelectExpr {}

#[cfg(test)]
mod tests {

    use vortex_array::arrays::StructArray;
    use vortex_array::{IntoArray, ToCanonical};
    use vortex_buffer::buffer;
    use vortex_dtype::{DType, FieldName, Nullability};

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

    fn test_array() -> StructArray {
        StructArray::from_fields(&[
            ("a", buffer![0, 1, 2].into_array()),
            ("b", buffer![4, 5, 6].into_array()),
        ])
        .unwrap()
    }

    #[test]
    pub fn include_columns() {
        let st = test_array();
        let select = select(vec![FieldName::from("a")], root());
        let selected = select
            .evaluate(&Scope::new(st.to_array()))
            .unwrap()
            .to_struct()
            .unwrap();
        let selected_names = selected.names().clone();
        assert_eq!(selected_names.as_ref(), &["a".into()]);
    }

    #[test]
    pub fn exclude_columns() {
        let st = test_array();
        let select = select_exclude(vec![FieldName::from("a")], root());
        let selected = select
            .evaluate(&Scope::new(st.to_array()))
            .unwrap()
            .to_struct()
            .unwrap();
        let selected_names = selected.names().clone();
        assert_eq!(selected_names.as_ref(), &["b".into()]);
    }

    #[test]
    fn dtype() {
        let dtype = test_harness::struct_dtype();

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

        let select_expr_exclude = select_exclude(
            vec![
                FieldName::from("col1"),
                FieldName::from("col2"),
                FieldName::from("bool1"),
                FieldName::from("bool2"),
            ],
            root(),
        );
        assert_eq!(
            select_expr_exclude.return_dtype(&dtype).unwrap(),
            expected_dtype
        );

        let select_expr_exclude = select_exclude(
            vec![FieldName::from("col1"), FieldName::from("col2")],
            root(),
        );
        assert_eq!(
            select_expr_exclude.return_dtype(&dtype).unwrap(),
            DType::Struct(
                dtype
                    .as_struct()
                    .unwrap()
                    .project(&["a".into(), "bool1".into(), "bool2".into()])
                    .unwrap(),
                Nullability::NonNullable
            )
        );
    }
}

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

vortex-data / vortex / 16262933935

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

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