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;
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::Expression;
25use crate::expr::lit;
26use crate::scalar_fn::Arity;
27use crate::scalar_fn::ChildName;
28use crate::scalar_fn::ExecutionArgs;
29use crate::scalar_fn::ScalarFnId;
30use crate::scalar_fn::ScalarFnVTable;
31use crate::validity::Validity;
32
33/// Pack zero or more expressions into a structure with named fields.
34#[derive(Clone)]
35pub struct Pack;
36
37#[derive(Debug, Clone, PartialEq, Eq, Hash)]
38pub struct PackOptions {
39    pub names: FieldNames,
40    pub nullability: Nullability,
41}
42
43impl Display for PackOptions {
44    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
45        write!(
46            f,
47            "names: [{}], nullability: {:#}",
48            self.names.iter().join(", "),
49            self.nullability
50        )
51    }
52}
53
54impl ScalarFnVTable for Pack {
55    type Options = PackOptions;
56
57    fn id(&self) -> ScalarFnId {
58        ScalarFnId::from("vortex.pack")
59    }
60
61    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
62        Ok(Some(
63            pb::PackOpts {
64                paths: instance.names.iter().map(|n| n.to_string()).collect(),
65                nullable: instance.nullability.into(),
66            }
67            .encode_to_vec(),
68        ))
69    }
70
71    fn deserialize(
72        &self,
73        _metadata: &[u8],
74        _session: &VortexSession,
75    ) -> VortexResult<Self::Options> {
76        let opts = pb::PackOpts::decode(_metadata)?;
77        let names: FieldNames = opts
78            .paths
79            .iter()
80            .map(|name| FieldName::from(name.as_str()))
81            .collect();
82        Ok(PackOptions {
83            names,
84            nullability: opts.nullable.into(),
85        })
86    }
87
88    fn arity(&self, options: &Self::Options) -> Arity {
89        Arity::Exact(options.names.len())
90    }
91
92    fn child_name(&self, instance: &Self::Options, child_idx: usize) -> ChildName {
93        match instance.names.get(child_idx) {
94            Some(name) => ChildName::from(Arc::clone(name.inner())),
95            None => unreachable!(
96                "Invalid child index {} for Pack expression with {} fields",
97                child_idx,
98                instance.names.len()
99            ),
100        }
101    }
102
103    fn fmt_sql(
104        &self,
105        options: &Self::Options,
106        expr: &Expression,
107        f: &mut Formatter<'_>,
108    ) -> std::fmt::Result {
109        write!(f, "pack(")?;
110        for (i, (name, child)) in options.names.iter().zip(expr.children().iter()).enumerate() {
111            write!(f, "{}: ", name)?;
112            child.fmt_sql(f)?;
113            if i + 1 < options.names.len() {
114                write!(f, ", ")?;
115            }
116        }
117        write!(f, "){}", options.nullability)
118    }
119
120    fn return_dtype(&self, options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult<DType> {
121        Ok(DType::Struct(
122            StructFields::new(options.names.clone(), arg_dtypes.to_vec()),
123            options.nullability,
124        ))
125    }
126
127    fn validity(
128        &self,
129        _options: &Self::Options,
130        _expression: &Expression,
131    ) -> VortexResult<Option<Expression>> {
132        Ok(Some(lit(true)))
133    }
134
135    fn execute(
136        &self,
137        options: &Self::Options,
138        args: &dyn ExecutionArgs,
139        ctx: &mut ExecutionCtx,
140    ) -> VortexResult<ArrayRef> {
141        let len = args.row_count();
142        let value_arrays: Vec<ArrayRef> = (0..args.num_inputs())
143            .map(|i| args.get(i))
144            .collect::<VortexResult<_>>()?;
145        let validity: Validity = options.nullability.into();
146        StructArray::try_new(options.names.clone(), value_arrays, len, validity)?
147            .into_array()
148            .execute(ctx)
149    }
150
151    // This applies a nullability
152    fn is_null_sensitive(&self, _instance: &Self::Options) -> bool {
153        true
154    }
155
156    fn is_fallible(&self, _instance: &Self::Options) -> bool {
157        false
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use vortex_buffer::buffer;
164    use vortex_error::VortexResult;
165    use vortex_error::vortex_bail;
166
167    use super::Pack;
168    use super::PackOptions;
169    use crate::ArrayRef;
170    use crate::IntoArray;
171    use crate::ToCanonical;
172    use crate::arrays::PrimitiveArray;
173    use crate::arrays::struct_::StructArrayExt;
174    use crate::assert_arrays_eq;
175    use crate::dtype::Nullability;
176    use crate::expr::col;
177    use crate::expr::pack;
178    use crate::scalar_fn::ScalarFnVTableExt;
179    use crate::scalar_fn::fns::pack::StructArray;
180    use crate::validity::Validity;
181
182    fn test_array() -> ArrayRef {
183        StructArray::from_fields(&[
184            ("a", buffer![0, 1, 2].into_array()),
185            ("b", buffer![4, 5, 6].into_array()),
186        ])
187        .unwrap()
188        .into_array()
189    }
190
191    fn primitive_field(array: &ArrayRef, field_path: &[&str]) -> VortexResult<PrimitiveArray> {
192        let mut field_path = field_path.iter();
193
194        let Some(field) = field_path.next() else {
195            vortex_bail!("empty field path");
196        };
197
198        let mut array = array.to_struct().unmasked_field_by_name(field)?.clone();
199        for field in field_path {
200            array = array.to_struct().unmasked_field_by_name(field)?.clone();
201        }
202        Ok(array.to_primitive())
203    }
204
205    #[test]
206    pub fn test_empty_pack() {
207        let expr = Pack.new_expr(
208            PackOptions {
209                names: Default::default(),
210                nullability: Default::default(),
211            },
212            [],
213        );
214
215        let test_array = test_array();
216        let actual_array = test_array.clone().apply(&expr).unwrap();
217        assert_eq!(actual_array.len(), test_array.len());
218        assert_eq!(actual_array.to_struct().struct_fields().nfields(), 0);
219    }
220
221    #[test]
222    pub fn test_simple_pack() {
223        let expr = Pack.new_expr(
224            PackOptions {
225                names: ["one", "two", "three"].into(),
226                nullability: Nullability::NonNullable,
227            },
228            [col("a"), col("b"), col("a")],
229        );
230
231        let actual_array = test_array().apply(&expr).unwrap().to_struct();
232
233        assert_eq!(actual_array.names(), ["one", "two", "three"]);
234        assert!(matches!(actual_array.validity(), Ok(Validity::NonNullable)));
235
236        assert_arrays_eq!(
237            primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(),
238            PrimitiveArray::from_iter([0i32, 1, 2])
239        );
240        assert_arrays_eq!(
241            primitive_field(&actual_array.clone().into_array(), &["two"]).unwrap(),
242            PrimitiveArray::from_iter([4i32, 5, 6])
243        );
244        assert_arrays_eq!(
245            primitive_field(&actual_array.into_array(), &["three"]).unwrap(),
246            PrimitiveArray::from_iter([0i32, 1, 2])
247        );
248    }
249
250    #[test]
251    pub fn test_nested_pack() {
252        let expr = Pack.new_expr(
253            PackOptions {
254                names: ["one", "two", "three"].into(),
255                nullability: Nullability::NonNullable,
256            },
257            [
258                col("a"),
259                Pack.new_expr(
260                    PackOptions {
261                        names: ["two_one", "two_two"].into(),
262                        nullability: Nullability::NonNullable,
263                    },
264                    [col("b"), col("b")],
265                ),
266                col("a"),
267            ],
268        );
269
270        let actual_array = test_array().apply(&expr).unwrap().to_struct();
271
272        assert_eq!(actual_array.names(), ["one", "two", "three"]);
273
274        assert_arrays_eq!(
275            primitive_field(&actual_array.clone().into_array(), &["one"]).unwrap(),
276            PrimitiveArray::from_iter([0i32, 1, 2])
277        );
278        assert_arrays_eq!(
279            primitive_field(&actual_array.clone().into_array(), &["two", "two_one"]).unwrap(),
280            PrimitiveArray::from_iter([4i32, 5, 6])
281        );
282        assert_arrays_eq!(
283            primitive_field(&actual_array.clone().into_array(), &["two", "two_two"]).unwrap(),
284            PrimitiveArray::from_iter([4i32, 5, 6])
285        );
286        assert_arrays_eq!(
287            primitive_field(&actual_array.into_array(), &["three"]).unwrap(),
288            PrimitiveArray::from_iter([0i32, 1, 2])
289        );
290    }
291
292    #[test]
293    pub fn test_pack_nullable() {
294        let expr = Pack.new_expr(
295            PackOptions {
296                names: ["one", "two", "three"].into(),
297                nullability: Nullability::Nullable,
298            },
299            [col("a"), col("b"), col("a")],
300        );
301
302        let actual_array = test_array().apply(&expr).unwrap().to_struct();
303
304        assert_eq!(actual_array.names(), ["one", "two", "three"]);
305        assert!(matches!(actual_array.validity(), Ok(Validity::AllValid)));
306    }
307
308    #[test]
309    pub fn test_display() {
310        let expr = pack(
311            [("id", col("user_id")), ("name", col("username"))],
312            Nullability::NonNullable,
313        );
314        assert_eq!(expr.to_string(), "pack(id: $.user_id, name: $.username)");
315
316        let expr2 = Pack.new_expr(
317            PackOptions {
318                names: ["x", "y"].into(),
319                nullability: Nullability::Nullable,
320            },
321            [col("a"), col("b")],
322        );
323        assert_eq!(expr2.to_string(), "pack(x: $.a, y: $.b)?");
324    }
325}
vortex_array/scalar_fn/fns/pack.rs

vortex_array/scalar_fn/fns/
pack.rs
