16049945385

Committed 03 Jul 2025 12:06PM UTC coverage: 77.865% (-0.02%) from 77.885%

Build # 16049945385

Build Type

push

github

Committed by

web-flow

Commit Message

feat: Hide as many of FieldNames' Arcs as possible (#3747)

The main purpose of this PR is to hide the implementation details of
`FieldNames` from the public API, and pushing all the required
`from/into` calls internally.

The change in API can be seen in all the test changes, which I think are
nicer.

---------

Signed-off-by: Adam Gutglick <adam@spiraldb.com>

Coverage Stats

138 of 173 new or added lines in 30 files covered. (79.77%)

43324 of 55640 relevant lines covered (77.86%)

55702.55 hits per line

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

76.96

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

use std::any::Any;
use std::fmt::Display;
use std::sync::Arc;

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

use crate::field::DisplayFieldNames;
use crate::{AnalysisExpr, ExprRef, Scope, ScopeDType, VortexExpr};

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

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

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

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

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

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

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

    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_expr(
            SelectField::Include(self.fields.as_include_names(field_names)?),
            self.child.clone(),
        ))
    }
}

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 Select {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}{}", self.child, self.fields)
    }
}

#[cfg(feature = "proto")]
pub(crate) mod proto {
    use vortex_error::{VortexResult, vortex_bail};
    use vortex_proto::expr::kind::Kind;

    use crate::{ExprDeserialize, ExprRef, ExprSerializable, Id, Select};

    pub struct SelectSerde;

    impl Id for SelectSerde {
        fn id(&self) -> &'static str {
            "select"
        }
    }

    impl ExprDeserialize for SelectSerde {
        fn deserialize(&self, _kind: &Kind, _children: Vec<ExprRef>) -> VortexResult<ExprRef> {
            vortex_bail!(NotImplemented: "", self.id())
        }
    }

    impl ExprSerializable for Select {
        fn id(&self) -> &'static str {
            SelectSerde.id()
        }

        fn serialize_kind(&self) -> VortexResult<Kind> {
            vortex_bail!(NotImplemented: "", self.id())
        }
    }
}

impl AnalysisExpr for Select {}

impl VortexExpr for Select {
    fn as_any(&self) -> &dyn Any {
        self
    }

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

    fn children(&self) -> Vec<&ExprRef> {
        vec![&self.child]
    }

    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
        assert_eq!(children.len(), 1);
        Self::new_expr(self.fields.clone(), children[0].clone())
    }

    fn return_dtype(&self, scope: &ScopeDType) -> VortexResult<DType> {
        let child_dtype = self.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 &self.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.iter().contains(name))
                .collect(),
        };

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

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

#[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, ScopeDType, 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(&ScopeDType::new(dtype.clone()))
                .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(&ScopeDType::new(dtype.clone()))
                .unwrap(),
            expected_dtype
        );

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

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

vortex-data / vortex / 16049945385

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