vortex_array/scalar_fn/fns/
get_item.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use prost::Message;
8use vortex_error::VortexResult;
9use vortex_error::vortex_err;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::IntoArray;
16use crate::arrays::Constant;
17use crate::arrays::ConstantArray;
18use crate::arrays::ScalarFnArray;
19use crate::arrays::StructArray;
20use crate::arrays::struct_::StructArrayExt;
21use crate::builtins::ArrayBuiltins;
22use crate::dtype::DType;
23use crate::dtype::FieldName;
24use crate::dtype::Nullability;
25use crate::expr::display::ExprDisplay;
26use crate::proto::expr as pb;
27use crate::scalar::Scalar;
28use crate::scalar_fn::Arity;
29use crate::scalar_fn::ChildName;
30use crate::scalar_fn::EmptyOptions;
31use crate::scalar_fn::ExecutionArgs;
32use crate::scalar_fn::ReduceNode;
33use crate::scalar_fn::ScalarFnId;
34use crate::scalar_fn::ScalarFnVTable;
35use crate::scalar_fn::ScalarFnVTableExt;
36use crate::scalar_fn::fns::literal::Literal;
37use crate::scalar_fn::fns::mask::Mask;
38use crate::scalar_fn::fns::pack::Pack;
39
40#[derive(Clone)]
41pub struct GetItem;
42
43impl GetItem {
44    /// Creates a lazy projection of `field_name` from `input`.
45    ///
46    /// # Errors
47    ///
48    /// Returns an error if `input` is not a struct or does not contain `field_name`.
49    pub fn try_new(
50        input: ArrayRef,
51        field_name: impl Into<FieldName>,
52    ) -> VortexResult<ScalarFnArray> {
53        ScalarFnArray::try_new(GetItem.bind(field_name.into()), vec![input])
54    }
55}
56
57impl ScalarFnVTable for GetItem {
58    type Options = FieldName;
59
60    fn id(&self) -> ScalarFnId {
61        static ID: CachedId = CachedId::new("vortex.get_item");
62        *ID
63    }
64
65    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
66        Ok(Some(
67            pb::GetItemOpts {
68                path: instance.to_string(),
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::GetItemOpts::decode(_metadata)?;
80        Ok(FieldName::from(opts.path))
81    }
82
83    fn arity(&self, _field_name: &FieldName) -> Arity {
84        Arity::Exact(1)
85    }
86
87    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
88        match child_idx {
89            0 => ChildName::from("input"),
90            _ => unreachable!("Invalid child index {} for GetItem expression", child_idx),
91        }
92    }
93
94    fn fmt_sql(
95        &self,
96        field_name: &FieldName,
97        expr: &dyn ExprDisplay,
98        f: &mut Formatter<'_>,
99    ) -> std::fmt::Result {
100        Display::fmt(expr.display_child(0), f)?;
101        write!(f, ".{}", field_name)
102    }
103
104    fn return_dtype(&self, field_name: &FieldName, arg_dtypes: &[DType]) -> VortexResult<DType> {
105        let struct_dtype = &arg_dtypes[0];
106        let field_dtype = struct_dtype
107            .as_struct_fields_opt()
108            .and_then(|st| st.field(field_name))
109            .ok_or_else(|| {
110                vortex_err!("Couldn't find the {} field in the input scope", field_name)
111            })?;
112
113        // Match here to avoid cloning the dtype if nullability doesn't need to change
114        if matches!(
115            (struct_dtype.nullability(), field_dtype.nullability()),
116            (Nullability::Nullable, Nullability::NonNullable)
117        ) {
118            return Ok(field_dtype.with_nullability(Nullability::Nullable));
119        }
120
121        Ok(field_dtype)
122    }
123
124    fn execute(
125        &self,
126        field_name: &FieldName,
127        args: &dyn ExecutionArgs,
128        ctx: &mut ExecutionCtx,
129    ) -> VortexResult<ArrayRef> {
130        let input = args.get(0)?;
131        if let Some(constant) = input.as_opt::<Constant>() {
132            let dtype = self.return_dtype(field_name, &[input.dtype().clone()])?;
133            let scalar = if constant.scalar().is_null() {
134                Scalar::null(dtype)
135            } else {
136                constant
137                    .scalar()
138                    .as_struct()
139                    .field(field_name)
140                    .ok_or_else(|| {
141                        vortex_err!(
142                            "Field '{}' missing from constant struct array {}",
143                            field_name,
144                            input.dtype()
145                        )
146                    })?
147                    .cast(&dtype)?
148            };
149
150            return Ok(ConstantArray::new(scalar, input.len()).into_array());
151        }
152
153        let input = input.execute::<StructArray>(ctx)?;
154        let field = input.unmasked_field_by_name(field_name).cloned()?;
155
156        match input.dtype().nullability() {
157            Nullability::NonNullable => Ok(field),
158            Nullability::Nullable => field.mask(input.validity()?.to_array(input.len())),
159        }
160    }
161
162    fn reduce<T: ReduceNode>(&self, field_name: &FieldName, node: &T) -> VortexResult<Option<T>> {
163        let child = node.child(0);
164        if let Some(child_fn) = child.scalar_fn()
165            && let Some(pack) = child_fn.as_opt::<Pack>()
166            && let Some(idx) = pack.names.find(field_name)
167        {
168            let mut field = child.child(idx);
169
170            // Possibly mask the field if the pack is nullable
171            if pack.nullability.is_nullable() {
172                field = node.new_node(
173                    Mask.bind(EmptyOptions),
174                    &[field, node.new_node(Literal.bind(true.into()), &[])?],
175                )?;
176            }
177
178            return Ok(Some(field));
179        }
180
181        Ok(None)
182    }
183
184    fn is_strict(&self, _field_name: &FieldName) -> bool {
185        true
186    }
187
188    fn is_infallible(&self, _field_name: &FieldName) -> bool {
189        // If this type-checks, it is infallible.
190        true
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use vortex_buffer::buffer;
197    use vortex_error::VortexResult;
198
199    use super::GetItem;
200    use crate::ArrayRef;
201    use crate::IntoArray;
202    use crate::VortexSessionExecute;
203    use crate::arrays::Constant;
204    use crate::arrays::ConstantArray;
205    use crate::arrays::Filter;
206    use crate::arrays::Primitive;
207    use crate::arrays::ScalarFn;
208    use crate::arrays::Slice;
209    use crate::builtins::ArrayBuiltins;
210    use crate::dtype::DType;
211    use crate::dtype::FieldName;
212    use crate::dtype::FieldNames;
213    use crate::dtype::Nullability;
214    use crate::dtype::Nullability::NonNullable;
215    use crate::dtype::PType;
216    use crate::dtype::StructFields;
217    use crate::expr::checked_add;
218    use crate::expr::get_item;
219    use crate::expr::lit;
220    use crate::expr::pack;
221    use crate::expr::root;
222    use crate::scalar::Scalar;
223    use crate::scalar_fn::fns::get_item::StructArray;
224    use crate::validity::Validity;
225
226    fn test_array() -> StructArray {
227        StructArray::from_fields(&[
228            ("a", buffer![0i32, 1, 2].into_array()),
229            ("b", buffer![4i64, 5, 6].into_array()),
230        ])
231        .unwrap()
232    }
233
234    #[test]
235    fn get_item_by_name() {
236        let st = test_array();
237        let get_item = get_item("a", root());
238        assert!(
239            get_item
240                .as_scalar()
241                .is_some_and(|f| f.signature().is_strict())
242        );
243        let item = st.into_array().apply(&get_item).unwrap();
244        assert_eq!(item.dtype(), &DType::from(PType::I32))
245    }
246
247    #[test]
248    fn get_item_by_name_none() {
249        let st = test_array();
250        let get_item = get_item("c", root());
251        assert!(st.into_array().apply(&get_item).is_err());
252    }
253
254    #[test]
255    fn get_nullable_field() {
256        let st = StructArray::try_new(
257            FieldNames::from(["a"]),
258            vec![buffer![1i32].into_array()],
259            1,
260            Validity::AllInvalid,
261        )
262        .unwrap()
263        .into_array();
264
265        let get_item_expr = get_item("a", root());
266        let item = st.apply(&get_item_expr).unwrap();
267        // The dtype should be nullable since it inherits struct validity
268        assert_eq!(
269            item.dtype(),
270            &DType::Primitive(PType::I32, Nullability::Nullable)
271        );
272    }
273
274    #[test]
275    fn get_non_nullable_field_from_all_valid_nullable_struct() -> VortexResult<()> {
276        let st = StructArray::try_new(
277            FieldNames::from(["a"]),
278            vec![buffer![1i32].into_array()],
279            1,
280            Validity::AllValid,
281        )?
282        .into_array();
283
284        let item = st.apply(&get_item("a", root()))?;
285        assert_eq!(
286            item.dtype(),
287            &DType::Primitive(PType::I32, Nullability::Nullable)
288        );
289        Ok(())
290    }
291
292    #[test]
293    fn execute_constant_struct_stays_constant() -> VortexResult<()> {
294        let field_count = 128usize;
295        let selected_idx = 97i32;
296        let names = (0..field_count)
297            .map(|idx| FieldName::from(format!("f{idx}")))
298            .collect::<Vec<_>>();
299        let dtypes = vec![DType::Primitive(PType::I32, NonNullable); field_count];
300        let fields = StructFields::new(FieldNames::from(names), dtypes);
301        let scalar = Scalar::struct_(
302            DType::Struct(fields, NonNullable),
303            (0..128).map(|idx| Scalar::primitive(idx, NonNullable)),
304        );
305        let array = ConstantArray::new(scalar, 1_000_000).into_array();
306        let mut ctx = crate::array_session().create_execution_ctx();
307
308        let item: ArrayRef = GetItem::try_new(array, "f97")?
309            .into_array()
310            .execute(&mut ctx)?;
311
312        let constant = item.as_::<Constant>();
313        assert_eq!(constant.len(), 1_000_000);
314        assert_eq!(
315            constant.scalar(),
316            &Scalar::primitive(selected_idx, NonNullable)
317        );
318        Ok(())
319    }
320
321    #[test]
322    fn get_item_pushes_through_filter() -> VortexResult<()> {
323        use vortex_mask::Mask;
324
325        let filtered = test_array()
326            .into_array()
327            .filter(Mask::from_iter([true, false, true]))?;
328
329        let item = filtered.get_item("a")?;
330
331        assert!(item.is::<Filter>());
332        assert!(!item.is::<ScalarFn>());
333        Ok(())
334    }
335
336    #[test]
337    fn get_item_pushes_through_slice() -> VortexResult<()> {
338        let sliced = test_array().into_array().slice(1..3)?;
339
340        let item = sliced.get_item("a")?;
341
342        assert!(item.is::<Primitive>());
343        assert!(!item.is::<Slice>());
344        assert!(!item.is::<ScalarFn>());
345        Ok(())
346    }
347
348    #[test]
349    fn get_item_pushes_through_masked() -> VortexResult<()> {
350        let masked = test_array()
351            .into_array()
352            .mask(Validity::from_iter([true, false, true]).to_array(3))?;
353
354        let item = masked.get_item("a")?;
355
356        assert!(item.is::<Primitive>());
357        assert_eq!(
358            item.dtype(),
359            &DType::Primitive(PType::I32, Nullability::Nullable)
360        );
361        assert!(!item.is::<ScalarFn>());
362        Ok(())
363    }
364
365    #[test]
366    fn test_pack_get_item_rule() {
367        // Create: pack(a: lit(1), b: lit(2)).get_item("b")
368        let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
369        let get_item_expr = get_item("b", pack_expr);
370
371        let dtype = DType::Struct(StructFields::empty(), NonNullable);
372        let result = get_item_expr
373            .bind(&dtype)
374            .and_then(|expr| expr.optimize_recursive())
375            .unwrap();
376
377        assert_eq!(result, lit(2).bind(&dtype).unwrap());
378    }
379
380    #[test]
381    fn test_multi_level_pack_get_item_simplify() {
382        let inner_pack = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
383        let get_a = get_item("a", inner_pack);
384
385        let outer_pack = pack([("x", get_a), ("y", lit(3)), ("z", lit(4))], NonNullable);
386        let get_z = get_item("z", outer_pack);
387
388        let dtype = DType::Primitive(PType::I32, NonNullable);
389
390        let result = get_z.bind(&dtype).unwrap().optimize_recursive().unwrap();
391        assert_eq!(result, lit(4).bind(&dtype).unwrap());
392    }
393
394    #[test]
395    fn test_deeply_nested_pack_get_item() {
396        let innermost = pack([("a", lit(42))], NonNullable);
397        let get_a = get_item("a", innermost);
398
399        let level2 = pack([("b", get_a)], NonNullable);
400        let get_b = get_item("b", level2);
401
402        let level3 = pack([("c", get_b)], NonNullable);
403        let get_c = get_item("c", level3);
404
405        let outermost = pack([("final", get_c)], NonNullable);
406        let get_final = get_item("final", outermost);
407
408        let dtype = DType::Primitive(PType::I32, NonNullable);
409
410        let result = get_final
411            .bind(&dtype)
412            .unwrap()
413            .optimize_recursive()
414            .unwrap();
415        assert_eq!(result, lit(42).bind(&dtype).unwrap());
416    }
417
418    #[test]
419    fn test_partial_pack_get_item_simplify() {
420        let inner_pack = pack([("x", lit(1)), ("y", lit(2))], NonNullable);
421        let get_x = get_item("x", inner_pack);
422        let add_expr = checked_add(get_x, lit(10));
423
424        let outer_pack = pack([("result", add_expr)], NonNullable);
425        let get_result = get_item("result", outer_pack);
426
427        let dtype = DType::Primitive(PType::I32, NonNullable);
428
429        let result = get_result
430            .bind(&dtype)
431            .unwrap()
432            .optimize_recursive()
433            .unwrap();
434        let expected = checked_add(lit(1), lit(10)).bind(&dtype).unwrap();
435        assert_eq!(&result, &expected);
436    }
437
438    #[test]
439    fn get_item_filter_list_field() {
440        use vortex_mask::Mask;
441
442        use crate::arrays::BoolArray;
443        use crate::arrays::FilterArray;
444        use crate::arrays::ListArray;
445
446        let list = ListArray::try_new(
447            buffer![0f32, 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11.].into_array(),
448            buffer![2u64, 4, 6, 8, 10, 12].into_array(),
449            Validity::Array(BoolArray::from_iter([true, true, false, true, true]).into_array()),
450        )
451        .unwrap();
452
453        let filtered = FilterArray::try_new(
454            list.into_array(),
455            Mask::from_iter([true, true, false, false, false]),
456        )
457        .unwrap();
458
459        let st = StructArray::try_new(
460            FieldNames::from(["data"]),
461            vec![filtered.into_array()],
462            2,
463            Validity::AllValid,
464        )
465        .unwrap();
466
467        st.into_array().apply(&get_item("data", root())).unwrap();
468    }
469}
vortex_array/scalar_fn/fns/get_item.rs

vortex_array/scalar_fn/fns/
get_item.rs
