vortex_expr/
pack.rs

1use std::any::Any;
2use std::fmt::Display;
3use std::hash::Hash;
4use std::sync::Arc;
5
6use itertools::Itertools as _;
7use vortex_array::arrays::StructArray;
8use vortex_array::validity::Validity;
9use vortex_array::{Array, ArrayRef};
10use vortex_dtype::{DType, FieldName, FieldNames, Nullability, StructDType};
11use vortex_error::{VortexExpect as _, VortexResult, vortex_bail, vortex_err};
12
13use crate::{ExprRef, VortexExpr};
14
15/// Pack zero or more expressions into a structure with named fields.
16///
17/// # Examples
18///
19/// ```
20/// use vortex_array::IntoArray;
21/// use vortex_array::compute::scalar_at;
22/// use vortex_buffer::buffer;
23/// use vortex_expr::{Pack, Identity, VortexExpr};
24/// use vortex_scalar::Scalar;
25///
26/// let example = Pack::try_new_expr(
27///     ["x".into(), "x copy".into(), "second x copy".into()].into(),
28///     vec![Identity::new_expr(), Identity::new_expr(), Identity::new_expr()],
29/// ).unwrap();
30/// let packed = example.evaluate(&buffer![100, 110, 200].into_array()).unwrap();
31/// let x_copy = packed
32///     .as_struct_typed()
33///     .unwrap()
34///     .maybe_null_field_by_name("x copy")
35///     .unwrap();
36/// assert_eq!(scalar_at(&x_copy, 0).unwrap(), Scalar::from(100));
37/// assert_eq!(scalar_at(&x_copy, 1).unwrap(), Scalar::from(110));
38/// assert_eq!(scalar_at(&x_copy, 2).unwrap(), Scalar::from(200));
39/// ```
40///
41#[derive(Debug, Clone, PartialEq, Eq, Hash)]
42pub struct Pack {
43    names: FieldNames,
44    values: Vec<ExprRef>,
45}
46
47impl Pack {
48    pub fn try_new_expr(names: FieldNames, values: Vec<ExprRef>) -> VortexResult<Arc<Self>> {
49        if names.len() != values.len() {
50            vortex_bail!("length mismatch {} {}", names.len(), values.len());
51        }
52        Ok(Arc::new(Pack { names, values }))
53    }
54
55    pub fn names(&self) -> &FieldNames {
56        &self.names
57    }
58
59    pub fn field(&self, field_name: &FieldName) -> VortexResult<ExprRef> {
60        let idx = self
61            .names
62            .iter()
63            .position(|name| name == field_name)
64            .ok_or_else(|| {
65                vortex_err!(
66                    "Cannot find field {} in pack fields {:?}",
67                    field_name,
68                    self.names
69                )
70            })?;
71
72        self.values
73            .get(idx)
74            .cloned()
75            .ok_or_else(|| vortex_err!("field index out of bounds: {}", idx))
76    }
77}
78
79pub fn pack(elements: impl IntoIterator<Item = (impl Into<FieldName>, ExprRef)>) -> ExprRef {
80    let (names, values): (Vec<_>, Vec<_>) = elements
81        .into_iter()
82        .map(|(name, value)| (name.into(), value))
83        .unzip();
84    Pack::try_new_expr(names.into(), values)
85        .vortex_expect("pack names and values have the same length")
86}
87
88impl Display for Pack {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        f.write_str("{")?;
91        self.names
92            .iter()
93            .zip(&self.values)
94            .format_with(", ", |(name, expr), f| f(&format_args!("{name}: {expr}")))
95            .fmt(f)?;
96        f.write_str("}")
97    }
98}
99
100impl VortexExpr for Pack {
101    fn as_any(&self) -> &dyn Any {
102        self
103    }
104
105    fn unchecked_evaluate(&self, batch: &dyn Array) -> VortexResult<ArrayRef> {
106        let len = batch.len();
107        let value_arrays = self
108            .values
109            .iter()
110            .map(|value_expr| value_expr.evaluate(batch))
111            .process_results(|it| it.collect::<Vec<_>>())?;
112        Ok(
113            StructArray::try_new(self.names.clone(), value_arrays, len, Validity::NonNullable)?
114                .into_array(),
115        )
116    }
117
118    fn children(&self) -> Vec<&ExprRef> {
119        self.values.iter().collect()
120    }
121
122    fn replacing_children(self: Arc<Self>, children: Vec<ExprRef>) -> ExprRef {
123        assert_eq!(children.len(), self.values.len());
124        Self::try_new_expr(self.names.clone(), children)
125            .vortex_expect("children are known to have the same length as names")
126    }
127
128    fn return_dtype(&self, scope_dtype: &DType) -> VortexResult<DType> {
129        let value_dtypes = self
130            .values
131            .iter()
132            .map(|value_expr| value_expr.return_dtype(scope_dtype))
133            .process_results(|it| it.collect())?;
134        Ok(DType::Struct(
135            Arc::new(StructDType::new(self.names.clone(), value_dtypes)),
136            Nullability::NonNullable,
137        ))
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use std::sync::Arc;
144
145    use vortex_array::arrays::{PrimitiveArray, StructArray};
146    use vortex_array::{Array, IntoArray, ToCanonical};
147    use vortex_buffer::buffer;
148    use vortex_dtype::FieldNames;
149    use vortex_error::{VortexResult, vortex_bail, vortex_err};
150
151    use crate::{Pack, VortexExpr, col};
152
153    fn test_array() -> StructArray {
154        StructArray::from_fields(&[
155            ("a", buffer![0, 1, 2].into_array()),
156            ("b", buffer![4, 5, 6].into_array()),
157        ])
158        .unwrap()
159    }
160
161    fn primitive_field(array: &dyn Array, field_path: &[&str]) -> VortexResult<PrimitiveArray> {
162        let mut field_path = field_path.iter();
163
164        let Some(field) = field_path.next() else {
165            vortex_bail!("empty field path");
166        };
167
168        let mut array = array
169            .as_struct_typed()
170            .ok_or_else(|| vortex_err!("expected a struct"))?
171            .maybe_null_field_by_name(field)?;
172
173        for field in field_path {
174            array = array
175                .as_struct_typed()
176                .ok_or_else(|| vortex_err!("expected a struct"))?
177                .maybe_null_field_by_name(field)?;
178        }
179        Ok(array.to_primitive().unwrap())
180    }
181
182    #[test]
183    pub fn test_empty_pack() {
184        let expr = Pack::try_new_expr(Arc::new([]), Vec::new()).unwrap();
185
186        let test_array = test_array().into_array();
187        let actual_array = expr.evaluate(&test_array).unwrap();
188        assert_eq!(actual_array.len(), test_array.len());
189        assert!(actual_array.as_struct_typed().unwrap().nfields() == 0);
190    }
191
192    #[test]
193    pub fn test_simple_pack() {
194        let expr = Pack::try_new_expr(
195            ["one".into(), "two".into(), "three".into()].into(),
196            vec![col("a"), col("b"), col("a")],
197        )
198        .unwrap();
199
200        let actual_array = expr.evaluate(&test_array()).unwrap();
201        let expected_names: FieldNames = ["one".into(), "two".into(), "three".into()].into();
202        assert_eq!(
203            actual_array.as_struct_typed().unwrap().names(),
204            &expected_names
205        );
206
207        assert_eq!(
208            primitive_field(&actual_array, &["one"])
209                .unwrap()
210                .as_slice::<i32>(),
211            [0, 1, 2]
212        );
213        assert_eq!(
214            primitive_field(&actual_array, &["two"])
215                .unwrap()
216                .as_slice::<i32>(),
217            [4, 5, 6]
218        );
219        assert_eq!(
220            primitive_field(&actual_array, &["three"])
221                .unwrap()
222                .as_slice::<i32>(),
223            [0, 1, 2]
224        );
225    }
226
227    #[test]
228    pub fn test_nested_pack() {
229        let expr = Pack::try_new_expr(
230            ["one".into(), "two".into(), "three".into()].into(),
231            vec![
232                col("a"),
233                Pack::try_new_expr(
234                    ["two_one".into(), "two_two".into()].into(),
235                    vec![col("b"), col("b")],
236                )
237                .unwrap(),
238                col("a"),
239            ],
240        )
241        .unwrap();
242
243        let actual_array = expr.evaluate(&test_array()).unwrap();
244        let expected_names: FieldNames = ["one".into(), "two".into(), "three".into()].into();
245        assert_eq!(
246            actual_array.as_struct_typed().unwrap().names(),
247            &expected_names
248        );
249
250        assert_eq!(
251            primitive_field(&actual_array, &["one"])
252                .unwrap()
253                .as_slice::<i32>(),
254            [0, 1, 2]
255        );
256        assert_eq!(
257            primitive_field(&actual_array, &["two", "two_one"])
258                .unwrap()
259                .as_slice::<i32>(),
260            [4, 5, 6]
261        );
262        assert_eq!(
263            primitive_field(&actual_array, &["two", "two_two"])
264                .unwrap()
265                .as_slice::<i32>(),
266            [4, 5, 6]
267        );
268        assert_eq!(
269            primitive_field(&actual_array, &["three"])
270                .unwrap()
271                .as_slice::<i32>(),
272            [0, 1, 2]
273        );
274    }
275}
vortex_expr/pack.rs

vortex_expr/
pack.rs
