16139349253

Committed 08 Jul 2025 09:26AM UTC coverage: 78.057% (-0.2%) from 78.253%

Build # 16139349253

Build Type

push

github

Committed by

web-flow

Commit Message

VortexExpr VTables (#3713)

Adds the same vtable machinery as arrays and layouts already use. It
uses the "Encoding" naming scheme from arrays and layouts. I don't
particularly like it, but it's consistent. Open to renames later.

Further, adds an expression registry to the Vortex session that will be
used for deserialization.

Expressions only decide their "options" serialization. So in theory, can
support many container formats, not just proto, provided each expression
can deserialize their own options format.

---------

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Run Details

800 of 1190 new or added lines in 38 files covered. (67.23%)

40 existing lines in 13 files now uncovered.

44100 of 56497 relevant lines covered (78.06%)

54989.55 hits per line

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

90.24

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

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

use std::fmt::{Debug, Display, Formatter};
use std::hash::Hash;

use vortex_array::stats::Stat;
use vortex_array::{ArrayRef, DeserializeMetadata, ProstMetadata, ToCanonical};
use vortex_dtype::{DType, FieldName};
use vortex_error::{VortexResult, vortex_bail, vortex_err};
use vortex_proto::expr as pb;

use crate::{
    AccessPath, AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, ScopeDType,
    StatsCatalog, VTable, root, vtable,
};

vtable!(GetItem);

#[allow(clippy::derived_hash_with_manual_eq)]
#[derive(Debug, Clone, Hash)]
pub struct GetItemExpr {
    field: FieldName,
    child: ExprRef,
}

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

pub struct GetItemExprEncoding;

impl VTable for GetItemVTable {
    type Expr = GetItemExpr;
    type Encoding = GetItemExprEncoding;
    type Metadata = ProstMetadata<pb::GetItemOpts>;

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

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

    fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {
        Some(ProstMetadata(pb::GetItemOpts {
            path: expr.field.to_string(),
        }))
    }

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

    fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {
        if children.len() != 1 {
            vortex_bail!(
                "GetItem expression must have exactly 1 child, got {}",
                children.len()
            );
        }

        Ok(GetItemExpr {
            field: expr.field.clone(),
            child: children[0].clone(),
        })
    }

    fn build(
        _encoding: &Self::Encoding,
        metadata: &<Self::Metadata as DeserializeMetadata>::Output,
        children: Vec<ExprRef>,
    ) -> VortexResult<Self::Expr> {
        if children.len() != 1 {
            vortex_bail!(
                "GetItem expression must have exactly 1 child, got {}",
                children.len()
            );
        }

        let field = FieldName::from(metadata.path.clone());
        Ok(GetItemExpr {
            field,
            child: children[0].clone(),
        })
    }

    fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {
        expr.child
            .unchecked_evaluate(scope)?
            .to_struct()?
            .field_by_name(expr.field())
            .cloned()
    }

    fn return_dtype(expr: &Self::Expr, scope: &ScopeDType) -> VortexResult<DType> {
        let input = expr.child.return_dtype(scope)?;
        input
            .as_struct()
            .and_then(|st| st.field(expr.field()))
            .ok_or_else(|| {
                vortex_err!(
                    "Couldn't find the {} field in the input scope",
                    expr.field()
                )
            })
    }
}

impl GetItemExpr {
    pub fn new(field: impl Into<FieldName>, child: ExprRef) -> Self {
        Self {
            field: field.into(),
            child,
        }
    }

    pub fn field(&self) -> &FieldName {
        &self.field
    }

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

    pub fn is(expr: &ExprRef) -> bool {
        expr.is::<GetItemVTable>()
    }
}

pub fn col(field: impl Into<FieldName>) -> ExprRef {
    GetItemExpr::new(field, root()).into_expr()
}

pub fn get_item(field: impl Into<FieldName>, child: ExprRef) -> ExprRef {
    GetItemExpr::new(field, child).into_expr()
}

pub fn get_item_scope(field: impl Into<FieldName>) -> ExprRef {
    GetItemExpr::new(field, root()).into_expr()
}

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

impl AnalysisExpr for GetItemExpr {
    fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
        catalog.stats_ref(&self.field_path()?, Stat::Max)
    }

    fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {
        catalog.stats_ref(&self.field_path()?, Stat::Min)
    }

    fn field_path(&self) -> Option<AccessPath> {
        self.child()
            .field_path()
            .map(|fp| AccessPath::new(fp.field_path.push(self.field.clone()), fp.identifier))
    }
}

#[cfg(test)]
mod tests {
    use vortex_array::IntoArray;
    use vortex_array::arrays::StructArray;
    use vortex_buffer::buffer;
    use vortex_dtype::DType;
    use vortex_dtype::PType::I32;

    use crate::get_item::get_item;
    use crate::{Scope, root};

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

    #[test]
    pub fn get_item_by_name() {
        let st = test_array();
        let get_item = get_item("a", root());
        let item = get_item.evaluate(&Scope::new(st.to_array())).unwrap();
        assert_eq!(item.dtype(), &DType::from(I32))
    }

    #[test]
    pub fn get_item_by_name_none() {
        let st = test_array();
        let get_item = get_item("c", root());
        assert!(get_item.evaluate(&Scope::new(st.to_array())).is_err());
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::fmt::{Debug, Display, Formatter};
5	use std::hash::Hash;
6
7	use vortex_array::stats::Stat;
8	use vortex_array::{ArrayRef, DeserializeMetadata, ProstMetadata, ToCanonical};
9	use vortex_dtype::{DType, FieldName};
10	use vortex_error::{VortexResult, vortex_bail, vortex_err};
11	use vortex_proto::expr as pb;
12
13	use crate::{
14	AccessPath, AnalysisExpr, ExprEncodingRef, ExprId, ExprRef, IntoExpr, Scope, ScopeDType,
15	StatsCatalog, VTable, root, vtable,
16	};
17
18	vtable!(GetItem);
19
20	#[allow(clippy::derived_hash_with_manual_eq)]
21	#[derive(Debug, Clone, Hash)]
22	pub struct GetItemExpr {
23	field: FieldName,
24	child: ExprRef,
25	}
26
27	impl PartialEq for GetItemExpr {
28	fn eq(&self, other: &Self) -> bool {	9,241✔
29	self.field == other.field && self.child.eq(&other.child)	9,241✔
30	}	9,241✔
31	}
32
33	pub struct GetItemExprEncoding;
34
35	impl VTable for GetItemVTable {
36	type Expr = GetItemExpr;
37	type Encoding = GetItemExprEncoding;
38	type Metadata = ProstMetadata<pb::GetItemOpts>;
39
40	fn id(_encoding: &Self::Encoding) -> ExprId {	107✔
41	ExprId::new_ref("get_item")	107✔
42	}	107✔
43
44	fn encoding(_expr: &Self::Expr) -> ExprEncodingRef {	1✔
45	ExprEncodingRef::new_ref(GetItemExprEncoding.as_ref())	1✔
46	}	1✔
47
48	fn metadata(expr: &Self::Expr) -> Option<Self::Metadata> {	1✔
49	Some(ProstMetadata(pb::GetItemOpts {	1✔
50	path: expr.field.to_string(),	1✔
51	}))	1✔
52	}	1✔
53
54	fn children(expr: &Self::Expr) -> Vec<&ExprRef> {	38,986✔
55	vec![&expr.child]	38,986✔
56	}	38,986✔
57
58	fn with_children(expr: &Self::Expr, children: Vec<ExprRef>) -> VortexResult<Self::Expr> {	990✔
59	if children.len() != 1 {	990✔
NEW 60	vortex_bail!(	×
NEW 61	"GetItem expression must have exactly 1 child, got {}",	×
NEW 62	children.len()	×
NEW 63	);	×
64	}	990✔
65		990✔
66	Ok(GetItemExpr {	990✔
67	field: expr.field.clone(),	990✔
68	child: children[0].clone(),	990✔
69	})	990✔
70	}	990✔
71
72	fn build(	1✔
73	_encoding: &Self::Encoding,	1✔
74	metadata: &<Self::Metadata as DeserializeMetadata>::Output,	1✔
75	children: Vec<ExprRef>,	1✔
76	) -> VortexResult<Self::Expr> {	1✔
77	if children.len() != 1 {	1✔
NEW 78	vortex_bail!(	×
NEW 79	"GetItem expression must have exactly 1 child, got {}",	×
NEW 80	children.len()	×
NEW 81	);	×
82	}	1✔
83		1✔
84	let field = FieldName::from(metadata.path.clone());	1✔
85	Ok(GetItemExpr {	1✔
86	field,	1✔
87	child: children[0].clone(),	1✔
88	})	1✔
89	}	1✔
90
91	fn evaluate(expr: &Self::Expr, scope: &Scope) -> VortexResult<ArrayRef> {	1,555✔
92	expr.child	1,555✔
93	.unchecked_evaluate(scope)?	1,555✔
94	.to_struct()?	1,555✔
95	.field_by_name(expr.field())	1,555✔
96	.cloned()	1,555✔
97	}	1,555✔
98
99	fn return_dtype(expr: &Self::Expr, scope: &ScopeDType) -> VortexResult<DType> {	1,619✔
100	let input = expr.child.return_dtype(scope)?;	1,619✔
101	input	1,619✔
102	.as_struct()	1,619✔
103	.and_then(\|st\| st.field(expr.field()))	1,619✔
104	.ok_or_else(\|\| {	1,619✔
NEW 105	vortex_err!(	×
NEW 106	"Couldn't find the {} field in the input scope",	×
NEW 107	expr.field()	×
NEW 108	)	×
109	})	1,619✔
110	}	1,619✔
111	}
112
113	impl GetItemExpr {
114	pub fn new(field: impl Into<FieldName>, child: ExprRef) -> Self {	4,298✔
115	Self {	4,298✔
116	field: field.into(),	4,298✔
117	child,	4,298✔
118	}	4,298✔
119	}	4,298✔
120
121	pub fn field(&self) -> &FieldName {	9,226✔
122	&self.field	9,226✔
123	}	9,226✔
124
125	pub fn child(&self) -> &ExprRef {	8,907✔
126	&self.child	8,907✔
127	}	8,907✔
128
129	pub fn is(expr: &ExprRef) -> bool {	388✔
130	expr.is::<GetItemVTable>()	388✔
131	}	388✔
132	}
133
134	pub fn col(field: impl Into<FieldName>) -> ExprRef {	93✔
135	GetItemExpr::new(field, root()).into_expr()	93✔
136	}	93✔
137
138	pub fn get_item(field: impl Into<FieldName>, child: ExprRef) -> ExprRef {	4,155✔
139	GetItemExpr::new(field, child).into_expr()	4,155✔
140	}	4,155✔
141
142	pub fn get_item_scope(field: impl Into<FieldName>) -> ExprRef {	50✔
143	GetItemExpr::new(field, root()).into_expr()	50✔
144	}	50✔
145
146	impl Display for GetItemExpr {
147	fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {	22✔
148	write!(f, "{}.{}", self.child, &self.field)	22✔
149	}	22✔
150	}
151
152	impl AnalysisExpr for GetItemExpr {
153	fn max(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {	11✔
154	catalog.stats_ref(&self.field_path()?, Stat::Max)	11✔
155	}	11✔
156
157	fn min(&self, catalog: &mut dyn StatsCatalog) -> Option<ExprRef> {	10✔
158	catalog.stats_ref(&self.field_path()?, Stat::Min)	10✔
159	}	10✔
160
161	fn field_path(&self) -> Option<AccessPath> {	21✔
162	self.child()	21✔
163	.field_path()	21✔
164	.map(\|fp\| AccessPath::new(fp.field_path.push(self.field.clone()), fp.identifier))	21✔
165	}	21✔
166	}
167
168	#[cfg(test)]
169	mod tests {
170	use vortex_array::IntoArray;
171	use vortex_array::arrays::StructArray;
172	use vortex_buffer::buffer;
173	use vortex_dtype::DType;
174	use vortex_dtype::PType::I32;
175
176	use crate::get_item::get_item;
177	use crate::{Scope, root};
178
179	fn test_array() -> StructArray {	2✔
180	StructArray::from_fields(&[	2✔
181	("a", buffer![0i32, 1, 2].into_array()),	2✔
182	("b", buffer![4i64, 5, 6].into_array()),	2✔
183	])	2✔
184	.unwrap()	2✔
185	}	2✔
186
187	#[test]
188	pub fn get_item_by_name() {	1✔
189	let st = test_array();	1✔
190	let get_item = get_item("a", root());	1✔
191	let item = get_item.evaluate(&Scope::new(st.to_array())).unwrap();	1✔
192	assert_eq!(item.dtype(), &DType::from(I32))	1✔
193	}	1✔
194
195	#[test]
196	pub fn get_item_by_name_none() {	1✔
197	let st = test_array();	1✔
198	let get_item = get_item("c", root());	1✔
199	assert!(get_item.evaluate(&Scope::new(st.to_array())).is_err());	1✔
200	}	1✔
201	}

vortex-data / vortex / 16139349253

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