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::lit;
27use crate::scalar_fn::Arity;
28use crate::scalar_fn::ChildName;
29use crate::scalar_fn::ExecutionArgs;
30use crate::scalar_fn::ScalarFnId;
31use crate::scalar_fn::ScalarFnVTable;
32use crate::validity::Validity;
33
34/// Pack zero or more expressions into a structure with named fields.
35#[derive(Clone)]
36pub struct Pack;
37
38#[derive(Debug, Clone, PartialEq, Eq, Hash)]
39pub struct PackOptions {
40    pub names: FieldNames,
41    pub nullability: Nullability,
42}
43
44impl Display for PackOptions {
45    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
46        write!(
47            f,
48            "names: [{}], nullability: {:#}",
49            self.names.iter().join(", "),
50            self.nullability
51        )
52    }
53}
54
55impl ScalarFnVTable for Pack {
56    type Options = PackOptions;
57
58    fn id(&self) -> ScalarFnId {
59        static ID: CachedId = CachedId::new("vortex.pack");
60        *ID
61    }
62
63    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
64        Ok(Some(
65            pb::PackOpts {
66                paths: instance.names.iter().map(|n| n.to_string()).collect(),
67                nullable: instance.nullability.into(),
68            }
69            .encode_to_vec(),
70        ))
71    }
72
73    fn deserialize(
74        &self,
75        _metadata: &[u8],
76        _session: &VortexSession,
77    ) -> VortexResult<Self::Options> {
78        let opts = pb::PackOpts::decode(_metadata)?;
79        let names: FieldNames = opts
80            .paths
81            .iter()
82            .map(|name| FieldName::from(name.as_str()))
83            .collect();
84        Ok(PackOptions {
85            names,
86            nullability: opts.nullable.into(),
87        })
88    }
89
90    fn arity(&self, options: &Self::Options) -> Arity {
91        Arity::Exact(options.names.len())
92    }
93
94    fn child_name(&self, instance: &Self::Options, child_idx: usize) -> ChildName {
95        match instance.names.get(child_idx) {
96            Some(name) => ChildName::from(Arc::clone(name.inner())),
97            None => unreachable!(
98                "Invalid child index {} for Pack expression with {} fields",
99                child_idx,
100                instance.names.len()
101            ),
102        }
103    }
104
105    fn fmt_sql(
106        &self,
107        options: &Self::Options,
108        expr: &Expression,
109        f: &mut Formatter<'_>,
110    ) -> std::fmt::Result {
111        write!(f, "pack(")?;
112        for (i, (name, child)) in options.names.iter().zip(expr.children().iter()).enumerate() {
113            write!(f, "{}: ", name)?;
114            child.fmt_sql(f)?;
115            if i + 1 < options.names.len() {
116                write!(f, ", ")?;
117            }
118        }
119        write!(f, "){}", options.nullability)
120    }
121
122    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
123        Ok(DType::Struct(
124            StructFields::new(options.names.clone(), arg_dtypes.to_vec()),
125            options.nullability,
126        ))
127    }
128
129    fn validity(
130        &self,
131        _options: &Self::Options,
132        _expression: &Expression,
133    ) -> VortexResult<Option<Expression>> {
134        Ok(Some(lit(true)))
135    }
136
137    fn execute(
138        &self,
139        options: &Self::Options,
140        args: &dyn ExecutionArgs,
141        ctx: &mut ExecutionCtx,
142    ) -> VortexResult<ArrayRef> {
143        let len = args.row_count();
144        let value_arrays: Vec<ArrayRef> = (0..args.num_inputs())
145            .map(|i| args.get(i))
146            .collect::<VortexResult<_>>()?;
147        let validity: Validity = options.nullability.into();
148        StructArray::try_new(options.names.clone(), value_arrays, len, validity)?
149            .into_array()
150            .execute(ctx)
151    }
152
153    // This applies a nullability
154    fn is_null_sensitive(&self, _instance: &Self::Options) -> bool {
155        true
156    }
157
158    fn is_fallible(&self, _instance: &Self::Options) -> bool {
159        false
160    }
161}
162
163#[cfg(test)]
164mod tests {
165    use vortex_buffer::buffer;
166    use vortex_error::VortexResult;
167    use vortex_error::vortex_bail;
168
169    use super::Pack;
170    use super::PackOptions;
171    use crate::ArrayRef;
172    use crate::IntoArray;
173    use crate::VortexSessionExecute;
174    use crate::array_session;
175    use crate::arrays::PrimitiveArray;
176    use crate::arrays::struct_::StructArrayExt;
177    use crate::assert_arrays_eq;
178    use crate::dtype::Nullability;
179    use crate::expr::col;
180    use crate::expr::pack;
181    use crate::scalar_fn::ScalarFnVTableExt;
182    use crate::scalar_fn::fns::pack::StructArray;
183    use crate::validity::Validity;
184
185    fn test_array() -> ArrayRef {
186        StructArray::from_fields(&[
187            ("a", buffer![0, 1, 2].into_array()),
188            ("b", buffer![4, 5, 6].into_array()),
189        ])
190        .unwrap()
191        .into_array()
192    }
193
194    fn primitive_field(array: &ArrayRef, field_path: &[&str]) -> VortexResult<PrimitiveArray> {
195        let mut ctx = array_session().create_execution_ctx();
196        let mut field_path = field_path.iter();
197
198        let Some(field) = field_path.next() else {
199            vortex_bail!("empty field path");
200        };
201
202        let mut array = array
203            .clone()
204            .execute::<StructArray>(&mut ctx)?
205            .unmasked_field_by_name(field)?
206            .clone();
207        for field in field_path {
208            let next = array
209                .clone()
210                .execute::<StructArray>(&mut ctx)?
211                .unmasked_field_by_name(field)?
212                .clone();
213            array = next;
214        }
215        let result = array.execute::<PrimitiveArray>(&mut ctx)?;
216        Ok(result)
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            primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(),
263            PrimitiveArray::from_iter([0i32, 1, 2]),
264            &mut ctx
265        );
266        assert_arrays_eq!(
267            primitive_field(&actual_array.clone().into_array(), &["two"]).unwrap(),
268            PrimitiveArray::from_iter([4i32, 5, 6]),
269            &mut ctx
270        );
271        assert_arrays_eq!(
272            primitive_field(&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            primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(),
309            PrimitiveArray::from_iter([0i32, 1, 2]),
310            &mut ctx
311        );
312        assert_arrays_eq!(
313            primitive_field(&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            primitive_field(&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            primitive_field(&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    }
349
350    #[test]
351    pub fn test_display() {
352        let expr = pack(
353            [("id", col("user_id")), ("name", col("username"))],
354            Nullability::NonNullable,
355        );
356        assert_eq!(expr.to_string(), "pack(id: $.user_id, name: $.username)");
357
358        let expr2 = Pack.new_expr(
359            PackOptions {
360                names: ["x", "y"].into(),
361                nullability: Nullability::Nullable,
362            },
363            [col("a"), col("b")],
364        );
365        assert_eq!(expr2.to_string(), "pack(x: $.a, y: $.b)?");
366    }
367}
vortex_array/scalar_fn/fns/pack.rs

vortex_array/scalar_fn/fns/
pack.rs
