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_proto::expr as pb;
13use vortex_session::VortexSession;
14use vortex_session::registry::CachedId;
15
16use crate::ArrayRef;
17use crate::ExecutionCtx;
18use crate::IntoArray;
19use crate::arrays::StructArray;
20use crate::dtype::DType;
21use crate::dtype::FieldName;
22use crate::dtype::FieldNames;
23use crate::dtype::Nullability;
24use crate::dtype::StructFields;
25use crate::expr::Expression;
26use crate::expr::display::ExprDisplay;
27use crate::expr::lit;
28use crate::scalar_fn::Arity;
29use crate::scalar_fn::ChildName;
30use crate::scalar_fn::ExecutionArgs;
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(
131        &self,
132        _options: &Self::Options,
133        _expression: &Expression,
134    ) -> VortexResult<Option<Expression>> {
135        Ok(Some(lit(true)))
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_fallible(&self, _instance: &Self::Options) -> bool {
160        false
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 primitive_field(array: &ArrayRef, field_path: &[&str]) -> VortexResult<PrimitiveArray> {
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        let result = array.execute::<PrimitiveArray>(&mut ctx)?;
217        Ok(result)
218    }
219
220    #[test]
221    pub fn test_empty_pack() {
222        let mut ctx = array_session().create_execution_ctx();
223        let expr = Pack.new_expr(
224            PackOptions {
225                names: Default::default(),
226                nullability: Default::default(),
227            },
228            [],
229        );
230
231        let test_array = test_array();
232        let actual_array = test_array.clone().apply(&expr).unwrap();
233        assert_eq!(actual_array.len(), test_array.len());
234        let nfields = actual_array
235            .execute::<StructArray>(&mut ctx)
236            .unwrap()
237            .struct_fields()
238            .nfields();
239        assert_eq!(nfields, 0);
240    }
241
242    #[test]
243    pub fn test_simple_pack() {
244        let mut ctx = array_session().create_execution_ctx();
245        let expr = Pack.new_expr(
246            PackOptions {
247                names: ["one", "two", "three"].into(),
248                nullability: Nullability::NonNullable,
249            },
250            [col("a"), col("b"), col("a")],
251        );
252
253        let actual_array = test_array()
254            .apply(&expr)
255            .unwrap()
256            .execute::<StructArray>(&mut ctx)
257            .unwrap();
258
259        assert_eq!(actual_array.names(), ["one", "two", "three"]);
260        assert!(matches!(actual_array.validity(), Ok(Validity::NonNullable)));
261
262        assert_arrays_eq!(
263            primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(),
264            PrimitiveArray::from_iter([0i32, 1, 2]),
265            &mut ctx
266        );
267        assert_arrays_eq!(
268            primitive_field(&actual_array.clone().into_array(), &["two"]).unwrap(),
269            PrimitiveArray::from_iter([4i32, 5, 6]),
270            &mut ctx
271        );
272        assert_arrays_eq!(
273            primitive_field(&actual_array.into_array(), &["three"]).unwrap(),
274            PrimitiveArray::from_iter([0i32, 1, 2]),
275            &mut ctx
276        );
277    }
278
279    #[test]
280    pub fn test_nested_pack() {
281        let mut ctx = array_session().create_execution_ctx();
282        let expr = Pack.new_expr(
283            PackOptions {
284                names: ["one", "two", "three"].into(),
285                nullability: Nullability::NonNullable,
286            },
287            [
288                col("a"),
289                Pack.new_expr(
290                    PackOptions {
291                        names: ["two_one", "two_two"].into(),
292                        nullability: Nullability::NonNullable,
293                    },
294                    [col("b"), col("b")],
295                ),
296                col("a"),
297            ],
298        );
299
300        let actual_array = test_array()
301            .apply(&expr)
302            .unwrap()
303            .execute::<StructArray>(&mut ctx)
304            .unwrap();
305
306        assert_eq!(actual_array.names(), ["one", "two", "three"]);
307
308        assert_arrays_eq!(
309            primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(),
310            PrimitiveArray::from_iter([0i32, 1, 2]),
311            &mut ctx
312        );
313        assert_arrays_eq!(
314            primitive_field(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(),
315            PrimitiveArray::from_iter([4i32, 5, 6]),
316            &mut ctx
317        );
318        assert_arrays_eq!(
319            primitive_field(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(),
320            PrimitiveArray::from_iter([4i32, 5, 6]),
321            &mut ctx
322        );
323        assert_arrays_eq!(
324            primitive_field(&actual_array.into_array(), &["three"]).unwrap(),
325            PrimitiveArray::from_iter([0i32, 1, 2]),
326            &mut ctx
327        );
328    }
329
330    #[test]
331    pub fn test_pack_nullable() {
332        let mut ctx = array_session().create_execution_ctx();
333        let expr = Pack.new_expr(
334            PackOptions {
335                names: ["one", "two", "three"].into(),
336                nullability: Nullability::Nullable,
337            },
338            [col("a"), col("b"), col("a")],
339        );
340
341        let actual_array = test_array()
342            .apply(&expr)
343            .unwrap()
344            .execute::<StructArray>(&mut ctx)
345            .unwrap();
346
347        assert_eq!(actual_array.names(), ["one", "two", "three"]);
348        assert!(matches!(actual_array.validity(), Ok(Validity::AllValid)));
349    }
350
351    #[test]
352    pub fn test_display() {
353        let expr = pack(
354            [("id", col("user_id")), ("name", col("username"))],
355            Nullability::NonNullable,
356        );
357        assert_eq!(expr.to_string(), "pack(id: $.user_id, name: $.username)");
358
359        let expr2 = Pack.new_expr(
360            PackOptions {
361                names: ["x", "y"].into(),
362                nullability: Nullability::Nullable,
363            },
364            [col("a"), col("b")],
365        );
366        assert_eq!(expr2.to_string(), "pack(x: $.a, y: $.b)?");
367    }
368}
vortex_array/scalar_fn/fns/pack.rs

vortex_array/scalar_fn/fns/
pack.rs
