16198217859

Committed 10 Jul 2025 02:41PM UTC coverage: 78.257% (+0.07%) from 78.188%

Build # 16198217859

Build Type

push

github

Committed by

web-flow

Commit Message

Row idx layout reader (#3819)

Fixes #3613 

* Implements a RowIdx layout reader without using scope vars (in
preparation of removing them).
* Works the same way, by partitioning the expression into the row index
and non-row index parts.
* Abstracts out partitioned evaluations (shared with struct layouts that
partition by field)
* Introduces a RowIdx expression ensuring identity doesn't resolve to
include this meta column.

---------

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

Run Details

347 of 442 new or added lines in 9 files covered. (78.51%)

19 existing lines in 2 files now uncovered.

44183 of 56459 relevant lines covered (78.26%)

53569.79 hits per line

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

58.15

/vortex-expr/src/scope.rs

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

use std::any::{Any, TypeId};
use std::fmt::{Debug, Display};
use std::str::FromStr;
use std::sync::Arc;

use itertools::Itertools;
use vortex_array::{Array, ArrayRef};
use vortex_dtype::{DType, FieldPathSet};
use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
use vortex_utils::aliases::hash_map::HashMap;

use crate::scope_vars::{ScopeVar, ScopeVars};

type ExprScope<T> = HashMap<Identifier, T>;

#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub enum Identifier {
    Identity,
    Other(Arc<str>),
}

impl FromStr for Identifier {
    type Err = VortexError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.is_empty() {
            vortex_bail!("Empty strings aren't allowed in identifiers")
        } else {
            Ok(Identifier::Other(s.into()))
        }
    }
}

impl PartialEq<str> for Identifier {
    fn eq(&self, other: &str) -> bool {
        match self {
            Identifier::Identity => other.is_empty(),
            Identifier::Other(str) => str.as_ref() == other,
        }
    }
}

impl From<&str> for Identifier {
    fn from(value: &str) -> Self {
        if value.is_empty() {
            Identifier::Identity
        } else {
            Identifier::Other(Arc::from(value))
        }
    }
}

impl Identifier {
    pub fn is_identity(&self) -> bool {
        matches!(self, Self::Identity)
    }
}

impl Display for Identifier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Identifier::Identity => write!(f, ""),
            Identifier::Other(v) => write!(f, "{v}"),
        }
    }
}

/// Scope define the evaluation context/scope that an expression uses when being evaluated.
/// There is a special `Identifier` (`Identity`) which is used to bind the initial array being evaluated
///
/// Other identifier can be bound with variables either before execution or while executing (see `Let`).
/// Values can be extracted from the scope using the `Var` expression.
///
/// ```code
/// <let x = lit(1) in var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3]> ->
/// <var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
/// <Primitive[1,2,3] + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
/// <Primitive[1,2,3] + ConstantArray(1), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
/// <Primitive[2,3,4], { Identity -> Primitive[1,2,3], x -> ConstantArray(1)>
/// ```
///
/// Other values can be bound before execution e.g.
///  `<var("x") + var("y") + var("z"), x -> ..., y -> ..., z -> ...>`
#[derive(Clone, Default)]
pub struct Scope {
    array_len: usize,
    root_scope: Option<ArrayRef>,
    /// A map from identifiers to arrays
    arrays: ExprScope<ArrayRef>,
    /// A map identifiers to opaque values used by expressions, but
    /// cannot affect the result type/shape.
    vars: ExprScope<Arc<dyn Any + Send + Sync>>,
    /// Variables that can be set on the scope during expression evaluation.
    scope_vars: ScopeVars,
}

pub type ScopeElement = (Identifier, ArrayRef);

impl Scope {
    pub fn new(arr: ArrayRef) -> Self {
        Self {
            array_len: arr.len(),
            root_scope: Some(arr),
            ..Default::default()
        }
    }

    pub fn empty(len: usize) -> Self {
        Self {
            array_len: len,
            ..Default::default()
        }
    }

    /// Get a value out of the scope by its [`Identifier`]
    pub fn array(&self, id: &Identifier) -> Option<&ArrayRef> {
        if id.is_identity() {
            return self.root_scope.as_ref();
        }
        self.arrays.get(id)
    }

    pub fn vars(&self, id: Identifier) -> VortexResult<&Arc<dyn Any + Send + Sync>> {
        self.vars
            .get(&id)
            .ok_or_else(|| vortex_err!("cannot find {} in var scope", id))
    }

    pub fn is_empty(&self) -> bool {
        self.array_len == 0
    }

    pub fn len(&self) -> usize {
        self.array_len
    }

    pub fn copy_with_array(&self, ident: Identifier, value: ArrayRef) -> Self {
        self.clone().with_array(ident, value)
    }

    /// Register an array with an identifier in the scope, overriding any existing value stored in it.
    pub fn with_array(mut self, ident: Identifier, value: ArrayRef) -> Self {
        assert_eq!(value.len(), self.len());

        if ident.is_identity() {
            self.root_scope = Some(value);
        } else {
            self.arrays.insert(ident, value);
        }
        self
    }

    /// Register an array with an identifier in the scope, overriding any existing value stored in it.
    pub fn with_array_pair(self, (ident, value): ScopeElement) -> Self {
        self.with_array(ident, value)
    }

    pub fn with_var(mut self, ident: Identifier, var: Arc<dyn Any + Send + Sync>) -> Self {
        self.vars.insert(ident, var);
        self
    }

    /// Returns a new evaluation scope with the given variable applied.
    pub fn with_scope_var<V: ScopeVar>(mut self, var: V) -> Self {
        self.scope_vars.insert(TypeId::of::<V>(), Box::new(var));
        self
    }

    /// Returns the scope variable of type `V` if it exists.
    pub fn scope_var<V: ScopeVar>(&self) -> Option<&V> {
        self.scope_vars
            .get(&TypeId::of::<V>())
            .and_then(|boxed| (**boxed).as_any().downcast_ref::<V>())
    }

    /// Returns the mutable scope variable of type `V` if it exists.
    pub fn scope_var_mut<V: ScopeVar>(&mut self) -> Option<&mut V> {
        self.scope_vars
            .get_mut(&TypeId::of::<V>())
            .and_then(|boxed| (**boxed).as_any_mut().downcast_mut::<V>())
    }

    pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &ArrayRef)> {
        let values = self.arrays.iter();

        self.root_scope
            .iter()
            .map(|s| (&Identifier::Identity, s))
            .chain(values)
    }
}

impl From<ArrayRef> for Scope {
    fn from(value: ArrayRef) -> Self {
        Self::new(value)
    }
}

#[derive(Clone, Default, Debug)]
pub struct ScopeDType {
    root: Option<DType>,
    types: ExprScope<DType>,
}

impl Display for ScopeDType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(root) = self.root.as_ref() {
            write!(f, "$: {}", root)?;
        }
        if !self.types.is_empty() {
            write!(f, ". ")?;
            write!(
                f,
                "{}",
                self.types
                    .iter()
                    .format_with(",", |x, f| f(&format_args!("{}: {}", x.0, x.1)))
            )?;
        }
        Ok(())
    }
}

pub type ScopeDTypeElement = (Identifier, DType);

impl From<&Scope> for ScopeDType {
    fn from(ctx: &Scope) -> Self {
        Self {
            root: ctx.root_scope.as_ref().map(|s| s.dtype().clone()),
            types: HashMap::from_iter(
                ctx.arrays
                    .iter()
                    .map(|(k, v)| (k.clone(), v.dtype().clone())),
            ),
        }
    }
}

impl ScopeDType {
    pub fn new(dtype: DType) -> Self {
        Self {
            root: Some(dtype),
            ..Default::default()
        }
    }

    pub fn dtype(&self, id: &Identifier) -> Option<&DType> {
        if id.is_identity() {
            return self.root.as_ref();
        }
        self.types.get(id)
    }

    pub fn copy_with_dtype(&self, ident: Identifier, dtype: DType) -> Self {
        self.clone().with_dtype(ident, dtype)
    }

    pub fn with_dtype(mut self, ident: Identifier, dtype: DType) -> Self {
        if ident.is_identity() {
            self.root = Some(dtype);
        } else {
            self.types.insert(ident, dtype);
        }
        self
    }

    pub fn with_dtype_element(self, (ident, dtype): ScopeDTypeElement) -> Self {
        self.with_dtype(ident, dtype)
    }
}

#[derive(Default, Clone, Debug)]
pub struct ScopeFieldPathSet {
    root: Option<FieldPathSet>,
    sets: ExprScope<FieldPathSet>,
}

pub type ScopeFieldPathSetElement = (Identifier, FieldPathSet);

impl ScopeFieldPathSet {
    pub fn new(path_set: FieldPathSet) -> Self {
        Self {
            root: Some(path_set),
            ..Default::default()
        }
    }

    pub fn set(&self, id: &Identifier) -> Option<&FieldPathSet> {
        if id.is_identity() {
            return self.root.as_ref();
        }
        self.sets.get(id)
    }

    pub fn copy_with_set(&self, ident: Identifier, set: FieldPathSet) -> Self {
        self.clone().with_set(ident, set)
    }

    pub fn with_set(mut self, ident: Identifier, set: FieldPathSet) -> Self {
        if ident.is_identity() {
            self.root = Some(set);
        } else {
            self.sets.insert(ident, set);
        }
        self
    }

    pub fn with_set_element(self, (ident, set): ScopeFieldPathSetElement) -> Self {
        self.with_set(ident, set)
    }
}

#[cfg(test)]
mod test {
    #[test]
    fn test_scope_var() {
        use super::*;

        #[derive(Clone, PartialEq, Eq, Debug)]
        struct TestVar {
            value: i32,
        }

        let scope = Scope::empty(100);
        assert!(scope.scope_var::<TestVar>().is_none());

        let var = TestVar { value: 42 };
        let mut scope = scope.with_scope_var(var.clone());
        assert_eq!(scope.scope_var::<TestVar>(), Some(&var));

        scope.scope_var_mut::<TestVar>().unwrap().value = 43;
        assert_eq!(scope.scope_var::<TestVar>(), Some(&TestVar { value: 43 }));
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::any::{Any, TypeId};
5	use std::fmt::{Debug, Display};
6	use std::str::FromStr;
7	use std::sync::Arc;
8
9	use itertools::Itertools;
10	use vortex_array::{Array, ArrayRef};
11	use vortex_dtype::{DType, FieldPathSet};
12	use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
13	use vortex_utils::aliases::hash_map::HashMap;
14
15	use crate::scope_vars::{ScopeVar, ScopeVars};
16
17	type ExprScope<T> = HashMap<Identifier, T>;
18
19	#[derive(Clone, Debug, Eq, PartialEq, Hash)]
20	pub enum Identifier {
21	Identity,
22	Other(Arc<str>),
23	}
24
25	impl FromStr for Identifier {
26	type Err = VortexError;
27
28	fn from_str(s: &str) -> Result<Self, Self::Err> {	2✔
29	if s.is_empty() {	2✔
30	vortex_bail!("Empty strings aren't allowed in identifiers")	1✔
31	} else {
32	Ok(Identifier::Other(s.into()))	1✔
33	}
34	}	2✔
35	}
36
37	impl PartialEq<str> for Identifier {
38	fn eq(&self, other: &str) -> bool {	3✔
39	match self {	3✔
40	Identifier::Identity => other.is_empty(),	1✔
41	Identifier::Other(str) => str.as_ref() == other,	2✔
42	}
43	}	3✔
44	}
45
46	impl From<&str> for Identifier {
47	fn from(value: &str) -> Self {	17✔
48	if value.is_empty() {	17✔
49	Identifier::Identity	7✔
50	} else {
51	Identifier::Other(Arc::from(value))	10✔
52	}
53	}	17✔
54	}
55
56	impl Identifier {
57	pub fn is_identity(&self) -> bool {	45,831✔
58	matches!(self, Self::Identity)	45,831✔
59	}	45,831✔
60	}
61
62	impl Display for Identifier {
63	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	26✔
64	match self {	26✔
65	Identifier::Identity => write!(f, ""),	26✔
66	Identifier::Other(v) => write!(f, "{v}"),	×
67	}
68	}	26✔
69	}
70
71	/// Scope define the evaluation context/scope that an expression uses when being evaluated.
72	/// There is a special `Identifier` (`Identity`) which is used to bind the initial array being evaluated
73	///
74	/// Other identifier can be bound with variables either before execution or while executing (see `Let`).
75	/// Values can be extracted from the scope using the `Var` expression.
76	///
77	/// ```code
78	/// <let x = lit(1) in var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3]> ->
79	/// <var(Identifier::Identity) + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
80	/// <Primitive[1,2,3] + var(x), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
81	/// <Primitive[1,2,3] + ConstantArray(1), { Identity -> Primitive[1,2,3], x -> ConstantArray(1)> ->
82	/// <Primitive[2,3,4], { Identity -> Primitive[1,2,3], x -> ConstantArray(1)>
83	/// ```
84	///
85	/// Other values can be bound before execution e.g.
86	/// `<var("x") + var("y") + var("z"), x -> ..., y -> ..., z -> ...>`
87	#[derive(Clone, Default)]
88	pub struct Scope {
89	array_len: usize,
90	root_scope: Option<ArrayRef>,
91	/// A map from identifiers to arrays
92	arrays: ExprScope<ArrayRef>,
93	/// A map identifiers to opaque values used by expressions, but
94	/// cannot affect the result type/shape.
95	vars: ExprScope<Arc<dyn Any + Send + Sync>>,
96	/// Variables that can be set on the scope during expression evaluation.
97	scope_vars: ScopeVars,
98	}
99
100	pub type ScopeElement = (Identifier, ArrayRef);
101
102	impl Scope {
103	pub fn new(arr: ArrayRef) -> Self {	2,526✔
104	Self {	2,526✔
105	array_len: arr.len(),	2,526✔
106	root_scope: Some(arr),	2,526✔
107	..Default::default()	2,526✔
108	}	2,526✔
109	}	2,526✔
110
111	pub fn empty(len: usize) -> Self {	1✔
112	Self {	1✔
113	array_len: len,	1✔
114	..Default::default()	1✔
115	}	1✔
116	}	1✔
117
118	/// Get a value out of the scope by its [`Identifier`]
119	pub fn array(&self, id: &Identifier) -> Option<&ArrayRef> {	3,296✔
120	if id.is_identity() {	3,296✔
121	return self.root_scope.as_ref();	3,295✔
122	}	1✔
123	self.arrays.get(id)	1✔
124	}	3,296✔
125
126	pub fn vars(&self, id: Identifier) -> VortexResult<&Arc<dyn Any + Send + Sync>> {	×
127	self.vars	×
128	.get(&id)	×
129	.ok_or_else(\|\| vortex_err!("cannot find {} in var scope", id))	×
130	}	×
131
132	pub fn is_empty(&self) -> bool {	×
133	self.array_len == 0	×
134	}	×
135
136	pub fn len(&self) -> usize {	2,656✔
137	self.array_len	2,656✔
138	}	2,656✔
139
140	pub fn copy_with_array(&self, ident: Identifier, value: ArrayRef) -> Self {	×
141	self.clone().with_array(ident, value)	×
142	}	×
143
144	/// Register an array with an identifier in the scope, overriding any existing value stored in it.
145	pub fn with_array(mut self, ident: Identifier, value: ArrayRef) -> Self {	1✔
146	assert_eq!(value.len(), self.len());	1✔
147
148	if ident.is_identity() {	1✔
149	self.root_scope = Some(value);	×
150	} else {	1✔
151	self.arrays.insert(ident, value);	1✔
152	}	1✔
153	self	1✔
154	}	1✔
155
156	/// Register an array with an identifier in the scope, overriding any existing value stored in it.
UNCOV 157	pub fn with_array_pair(self, (ident, value): ScopeElement) -> Self {	×
UNCOV 158	self.with_array(ident, value)	×
UNCOV 159	}	×
160
161	pub fn with_var(mut self, ident: Identifier, var: Arc<dyn Any + Send + Sync>) -> Self {	×
162	self.vars.insert(ident, var);	×
163	self	×
164	}	×
165
166	/// Returns a new evaluation scope with the given variable applied.
167	pub fn with_scope_var<V: ScopeVar>(mut self, var: V) -> Self {	1✔
168	self.scope_vars.insert(TypeId::of::<V>(), Box::new(var));	1✔
169	self	1✔
170	}	1✔
171
172	/// Returns the scope variable of type `V` if it exists.
173	pub fn scope_var<V: ScopeVar>(&self) -> Option<&V> {	3✔
174	self.scope_vars	3✔
175	.get(&TypeId::of::<V>())	3✔
176	.and_then(\|boxed\| (**boxed).as_any().downcast_ref::<V>())	3✔
177	}	3✔
178
179	/// Returns the mutable scope variable of type `V` if it exists.
180	pub fn scope_var_mut<V: ScopeVar>(&mut self) -> Option<&mut V> {	1✔
181	self.scope_vars	1✔
182	.get_mut(&TypeId::of::<V>())	1✔
183	.and_then(\|boxed\| (**boxed).as_any_mut().downcast_mut::<V>())	1✔
184	}	1✔
185
186	pub fn iter(&self) -> impl Iterator<Item = (&Identifier, &ArrayRef)> {	×
187	let values = self.arrays.iter();	×
188		×
189	self.root_scope	×
190	.iter()	×
191	.map(\|s\| (&Identifier::Identity, s))	×
192	.chain(values)	×
193	}	×
194	}
195
196	impl From<ArrayRef> for Scope {
197	fn from(value: ArrayRef) -> Self {	×
198	Self::new(value)	×
199	}	×
200	}
201
202	#[derive(Clone, Default, Debug)]
203	pub struct ScopeDType {
204	root: Option<DType>,
205	types: ExprScope<DType>,
206	}
207
208	impl Display for ScopeDType {
209	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {	×
210	if let Some(root) = self.root.as_ref() {	×
211	write!(f, "$: {}", root)?;	×
212	}	×
213	if !self.types.is_empty() {	×
214	write!(f, ". ")?;	×
215	write!(	×
216	f,	×
217	"{}",	×
218	self.types	×
219	.iter()	×
220	.format_with(",", \|x, f\| f(&format_args!("{}: {}", x.0, x.1)))	×
221	)?;	×
222	}	×
223	Ok(())	×
224	}	×
225	}
226
227	pub type ScopeDTypeElement = (Identifier, DType);
228
229	impl From<&Scope> for ScopeDType {
230	fn from(ctx: &Scope) -> Self {	2,576✔
231	Self {	2,576✔
232	root: ctx.root_scope.as_ref().map(\|s\| s.dtype().clone()),	2,576✔
233	types: HashMap::from_iter(	2,576✔
234	ctx.arrays	2,576✔
235	.iter()	2,576✔
236	.map(\|(k, v)\| (k.clone(), v.dtype().clone())),	2,576✔
237	),	2,576✔
238	}	2,576✔
239	}	2,576✔
240	}
241
242	impl ScopeDType {
243	pub fn new(dtype: DType) -> Self {	12,566✔
244	Self {	12,566✔
245	root: Some(dtype),	12,566✔
246	..Default::default()	12,566✔
247	}	12,566✔
248	}	12,566✔
249
250	pub fn dtype(&self, id: &Identifier) -> Option<&DType> {	37,512✔
251	if id.is_identity() {	37,512✔
252	return self.root.as_ref();	37,511✔
253	}	1✔
254	self.types.get(id)	1✔
255	}	37,512✔
256
257	pub fn copy_with_dtype(&self, ident: Identifier, dtype: DType) -> Self {	×
258	self.clone().with_dtype(ident, dtype)	×
259	}	×
260
UNCOV 261	pub fn with_dtype(mut self, ident: Identifier, dtype: DType) -> Self {	×
UNCOV 262	if ident.is_identity() {	×
263	self.root = Some(dtype);	×
UNCOV 264	} else {	×
UNCOV 265	self.types.insert(ident, dtype);	×
UNCOV 266	}	×
UNCOV 267	self	×
UNCOV 268	}	×
269
UNCOV 270	pub fn with_dtype_element(self, (ident, dtype): ScopeDTypeElement) -> Self {	×
UNCOV 271	self.with_dtype(ident, dtype)	×
UNCOV 272	}	×
273	}
274
275	#[derive(Default, Clone, Debug)]
276	pub struct ScopeFieldPathSet {
277	root: Option<FieldPathSet>,
278	sets: ExprScope<FieldPathSet>,
279	}
280
281	pub type ScopeFieldPathSetElement = (Identifier, FieldPathSet);
282
283	impl ScopeFieldPathSet {
284	pub fn new(path_set: FieldPathSet) -> Self {	301✔
285	Self {	301✔
286	root: Some(path_set),	301✔
287	..Default::default()	301✔
288	}	301✔
289	}	301✔
290
291	pub fn set(&self, id: &Identifier) -> Option<&FieldPathSet> {	362✔
292	if id.is_identity() {	362✔
293	return self.root.as_ref();	362✔
294	}	×
295	self.sets.get(id)	×
296	}	362✔
297
298	pub fn copy_with_set(&self, ident: Identifier, set: FieldPathSet) -> Self {	×
299	self.clone().with_set(ident, set)	×
300	}	×
301
302	pub fn with_set(mut self, ident: Identifier, set: FieldPathSet) -> Self {	×
303	if ident.is_identity() {	×
304	self.root = Some(set);	×
305	} else {	×
306	self.sets.insert(ident, set);	×
307	}	×
308	self	×
309	}	×
310
311	pub fn with_set_element(self, (ident, set): ScopeFieldPathSetElement) -> Self {	×
312	self.with_set(ident, set)	×
313	}	×
314	}
315
316	#[cfg(test)]
317	mod test {
318	#[test]
319	fn test_scope_var() {	1✔
320	use super::*;
321
322	#[derive(Clone, PartialEq, Eq, Debug)]
323	struct TestVar {
324	value: i32,
325	}
326
327	let scope = Scope::empty(100);	1✔
328	assert!(scope.scope_var::<TestVar>().is_none());	1✔
329
330	let var = TestVar { value: 42 };	1✔
331	let mut scope = scope.with_scope_var(var.clone());	1✔
332	assert_eq!(scope.scope_var::<TestVar>(), Some(&var));	1✔
333
334	scope.scope_var_mut::<TestVar>().unwrap().value = 43;	1✔
335	assert_eq!(scope.scope_var::<TestVar>(), Some(&TestVar { value: 43 }));	1✔
336	}	1✔
337	}

vortex-data / vortex / 16198217859

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