vortex_array/scalar_fn/fns/
pack.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6use std::hash::Hash;
7use std::sync::Arc;
8
9use itertools::Itertools as _;
10use prost::Message;
11use vortex_error::VortexResult;
12use vortex_session::VortexSession;
13use vortex_session::registry::CachedId;
14
15use crate::ArrayRef;
16use crate::ExecutionCtx;
17use crate::IntoArray;
18use crate::arrays::StructArray;
19use crate::dtype::DType;
20use crate::dtype::FieldName;
21use crate::dtype::FieldNames;
22use crate::dtype::Nullability;
23use crate::dtype::StructFields;
24use crate::expr::display::ExprDisplay;
25use crate::proto::expr as pb;
26use crate::scalar_fn::Arity;
27use crate::scalar_fn::ChildName;
28use crate::scalar_fn::ExecutionArgs;
29use crate::scalar_fn::ReduceNode;
30use crate::scalar_fn::ReduceNodeValidity;
31use crate::scalar_fn::ScalarFnId;
32use crate::scalar_fn::ScalarFnVTable;
33use crate::validity::Validity;
34
35/// Pack zero or more expressions into a structure with named fields.
36#[derive(Clone)]
37pub struct Pack;
38
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
40pub struct PackOptions {
41    pub names: FieldNames,
42    pub nullability: Nullability,
43}
44
45impl Display for PackOptions {
46    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
47        write!(
48            f,
49            "names: [{}], nullability: {:#}",
50            self.names.iter().join(", "),
51            self.nullability
52        )
53    }
54}
55
56impl ScalarFnVTable for Pack {
57    type Options = PackOptions;
58
59    fn id(&self) -> ScalarFnId {
60        static ID: CachedId = CachedId::new("vortex.pack");
61        *ID
62    }
63
64    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
65        Ok(Some(
66            pb::PackOpts {
67                paths: instance.names.iter().map(|n| n.to_string()).collect(),
68                nullable: instance.nullability.into(),
69            }
70            .encode_to_vec(),
71        ))
72    }
73
74    fn deserialize(
75        &self,
76        _metadata: &[u8],
77        _session: &VortexSession,
78    ) -> VortexResult<Self::Options> {
79        let opts = pb::PackOpts::decode(_metadata)?;
80        let names: FieldNames = opts
81            .paths
82            .iter()
83            .map(|name| FieldName::from(name.as_str()))
84            .collect();
85        Ok(PackOptions {
86            names,
87            nullability: opts.nullable.into(),
88        })
89    }
90
91    fn arity(&self, options: &Self::Options) -> Arity {
92        Arity::Exact(options.names.len())
93    }
94
95    fn child_name(&self, instance: &Self::Options, child_idx: usize) -> ChildName {
96        match instance.names.get(child_idx) {
97            Some(name) => ChildName::from(Arc::clone(name.inner())),
98            None => unreachable!(
99                "Invalid child index {} for Pack expression with {} fields",
100                child_idx,
101                instance.names.len()
102            ),
103        }
104    }
105
106    fn fmt_sql(
107        &self,
108        options: &Self::Options,
109        expr: &dyn ExprDisplay,
110        f: &mut Formatter<'_>,
111    ) -> std::fmt::Result {
112        write!(f, "pack(")?;
113        for (i, name) in options.names.iter().enumerate() {
114            write!(f, "{}: ", name)?;
115            Display::fmt(expr.display_child(i), f)?;
116            if i + 1 < options.names.len() {
117                write!(f, ", ")?;
118            }
119        }
120        write!(f, "){}", options.nullability)
121    }
122
123    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
124        Ok(DType::Struct(
125            StructFields::new(options.names.clone(), arg_dtypes.to_vec()),
126            options.nullability,
127        ))
128    }
129
130    fn validity<T: ReduceNode>(
131        &self,
132        _options: &Self::Options,
133        node: &T,
134    ) -> VortexResult<ReduceNodeValidity<T>> {
135        Ok(ReduceNodeValidity::Reduced(node.new_constant(true.into())))
136    }
137
138    fn execute(
139        &self,
140        options: &Self::Options,
141        args: &dyn ExecutionArgs,
142        ctx: &mut ExecutionCtx,
143    ) -> VortexResult<ArrayRef> {
144        let len = args.row_count();
145        let value_arrays: Vec<ArrayRef> = (0..args.num_inputs())
146            .map(|i| args.get(i))
147            .collect::<VortexResult<_>>()?;
148        let validity: Validity = options.nullability.into();
149        StructArray::try_new(options.names.clone(), value_arrays, len, validity)?
150            .into_array()
151            .execute(ctx)
152    }
153
154    fn is_strict(&self, _instance: &Self::Options) -> bool {
155        // A null field value does not force the packed struct row to be null.
156        false
157    }
158
159    fn is_infallible(&self, _instance: &Self::Options) -> bool {
160        true
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use vortex_buffer::buffer;
167    use vortex_error::VortexResult;
168    use vortex_error::vortex_bail;
169
170    use super::Pack;
171    use super::PackOptions;
172    use crate::ArrayRef;
173    use crate::IntoArray;
174    use crate::VortexSessionExecute;
175    use crate::array_session;
176    use crate::arrays::PrimitiveArray;
177    use crate::arrays::struct_::StructArrayExt;
178    use crate::assert_arrays_eq;
179    use crate::dtype::Nullability;
180    use crate::expr::col;
181    use crate::expr::pack;
182    use crate::scalar_fn::ScalarFnVTableExt;
183    use crate::scalar_fn::fns::pack::StructArray;
184    use crate::validity::Validity;
185
186    fn test_array() -> ArrayRef {
187        StructArray::from_fields(&[
188            ("a", buffer![0, 1, 2].into_array()),
189            ("b", buffer![4, 5, 6].into_array()),
190        ])
191        .unwrap()
192        .into_array()
193    }
194
195    fn field_array(array: &ArrayRef, field_path: &[&str]) -> VortexResult<ArrayRef> {
196        let mut ctx = array_session().create_execution_ctx();
197        let mut field_path = field_path.iter();
198
199        let Some(field) = field_path.next() else {
200            vortex_bail!("empty field path");
201        };
202
203        let mut array = array
204            .clone()
205            .execute::<StructArray>(&mut ctx)?
206            .unmasked_field_by_name(field)?
207            .clone();
208        for field in field_path {
209            let next = array
210                .clone()
211                .execute::<StructArray>(&mut ctx)?
212                .unmasked_field_by_name(field)?
213                .clone();
214            array = next;
215        }
216        Ok(array)
217    }
218
219    #[test]
220    pub fn test_empty_pack() {
221        let mut ctx = array_session().create_execution_ctx();
222        let expr = Pack.new_expr(
223            PackOptions {
224                names: Default::default(),
225                nullability: Default::default(),
226            },
227            [],
228        );
229
230        let test_array = test_array();
231        let actual_array = test_array.clone().apply(&expr).unwrap();
232        assert_eq!(actual_array.len(), test_array.len());
233        let nfields = actual_array
234            .execute::<StructArray>(&mut ctx)
235            .unwrap()
236            .struct_fields()
237            .nfields();
238        assert_eq!(nfields, 0);
239    }
240
241    #[test]
242    pub fn test_simple_pack() {
243        let mut ctx = array_session().create_execution_ctx();
244        let expr = Pack.new_expr(
245            PackOptions {
246                names: ["one", "two", "three"].into(),
247                nullability: Nullability::NonNullable,
248            },
249            [col("a"), col("b"), col("a")],
250        );
251
252        let actual_array = test_array()
253            .apply(&expr)
254            .unwrap()
255            .execute::<StructArray>(&mut ctx)
256            .unwrap();
257
258        assert_eq!(actual_array.names(), ["one", "two", "three"]);
259        assert!(matches!(actual_array.validity(), Ok(Validity::NonNullable)));
260
261        assert_arrays_eq!(
262            field_array(&actual_array.clone().into_array(), &["one"]).unwrap(),
263            PrimitiveArray::from_iter([0i32, 1, 2]),
264            &mut ctx
265        );
266        assert_arrays_eq!(
267            field_array(&actual_array.clone().into_array(), &["two"]).unwrap(),
268            PrimitiveArray::from_iter([4i32, 5, 6]),
269            &mut ctx
270        );
271        assert_arrays_eq!(
272            field_array(&actual_array.into_array(), &["three"]).unwrap(),
273            PrimitiveArray::from_iter([0i32, 1, 2]),
274            &mut ctx
275        );
276    }
277
278    #[test]
279    pub fn test_nested_pack() {
280        let mut ctx = array_session().create_execution_ctx();
281        let expr = Pack.new_expr(
282            PackOptions {
283                names: ["one", "two", "three"].into(),
284                nullability: Nullability::NonNullable,
285            },
286            [
287                col("a"),
288                Pack.new_expr(
289                    PackOptions {
290                        names: ["two_one", "two_two"].into(),
291                        nullability: Nullability::NonNullable,
292                    },
293                    [col("b"), col("b")],
294                ),
295                col("a"),
296            ],
297        );
298
299        let actual_array = test_array()
300            .apply(&expr)
301            .unwrap()
302            .execute::<StructArray>(&mut ctx)
303            .unwrap();
304
305        assert_eq!(actual_array.names(), ["one", "two", "three"]);
306
307        assert_arrays_eq!(
308            field_array(&actual_array.clone().into_array(), &["one"]).unwrap(),
309            PrimitiveArray::from_iter([0i32, 1, 2]),
310            &mut ctx
311        );
312        assert_arrays_eq!(
313            field_array(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(),
314            PrimitiveArray::from_iter([4i32, 5, 6]),
315            &mut ctx
316        );
317        assert_arrays_eq!(
318            field_array(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(),
319            PrimitiveArray::from_iter([4i32, 5, 6]),
320            &mut ctx
321        );
322        assert_arrays_eq!(
323            field_array(&actual_array.into_array(), &["three"]).unwrap(),
324            PrimitiveArray::from_iter([0i32, 1, 2]),
325            &mut ctx
326        );
327    }
328
329    #[test]
330    pub fn test_pack_nullable() {
331        let mut ctx = array_session().create_execution_ctx();
332        let expr = Pack.new_expr(
333            PackOptions {
334                names: ["one", "two", "three"].into(),
335                nullability: Nullability::Nullable,
336            },
337            [col("a"), col("b"), col("a")],
338        );
339
340        let actual_array = test_array()
341            .apply(&expr)
342            .unwrap()
343            .execute::<StructArray>(&mut ctx)
344            .unwrap();
345
346        assert_eq!(actual_array.names(), ["one", "two", "three"]);
347        assert!(matches!(actual_array.validity(), Ok(Validity::AllValid)));
348        let expected = StructArray::try_new(
349            ["one", "two", "three"].into(),
350            vec![
351                buffer![0i32, 1, 2].into_array(),
352                buffer![4i32, 5, 6].into_array(),
353                buffer![0i32, 1, 2].into_array(),
354            ],
355            3,
356            Validity::AllValid,
357        )
358        .unwrap();
359        assert_arrays_eq!(actual_array, expected, &mut ctx);
360    }
361
362    #[test]
363    pub fn test_display() {
364        let expr = pack(
365            [("id", col("user_id")), ("name", col("username"))],
366            Nullability::NonNullable,
367        );
368        assert_eq!(expr.to_string(), "pack(id: $.user_id, name: $.username)");
369
370        let expr2 = Pack.new_expr(
371            PackOptions {
372                names: ["x", "y"].into(),
373                nullability: Nullability::Nullable,
374            },
375            [col("a"), col("b")],
376        );
377        assert_eq!(expr2.to_string(), "pack(x: $.a, y: $.b)?");
378    }
379}
vortex_array/scalar_fn/fns/pack.rs

vortex_array/scalar_fn/fns/
pack.rs
