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::Formatter;
5
6use prost::Message;
7use vortex_error::VortexResult;
8use vortex_error::vortex_err;
9use vortex_proto::expr as pb;
10use vortex_session::VortexSession;
11use vortex_session::registry::CachedId;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::arrays::StructArray;
16use crate::arrays::struct_::StructArrayExt;
17use crate::builtins::ArrayBuiltins;
18use crate::builtins::ExprBuiltins;
19use crate::dtype::DType;
20use crate::dtype::FieldName;
21use crate::dtype::Nullability;
22use crate::expr::Expression;
23use crate::expr::lit;
24use crate::scalar_fn::Arity;
25use crate::scalar_fn::ChildName;
26use crate::scalar_fn::EmptyOptions;
27use crate::scalar_fn::ExecutionArgs;
28use crate::scalar_fn::ReduceCtx;
29use crate::scalar_fn::ReduceNode;
30use crate::scalar_fn::ReduceNodeRef;
31use crate::scalar_fn::ScalarFnId;
32use crate::scalar_fn::ScalarFnVTable;
33use crate::scalar_fn::ScalarFnVTableExt;
34use crate::scalar_fn::fns::literal::Literal;
35use crate::scalar_fn::fns::mask::Mask;
36use crate::scalar_fn::fns::pack::Pack;
37
38#[derive(Clone)]
39pub struct GetItem;
40
41impl ScalarFnVTable for GetItem {
42    type Options = FieldName;
43
44    fn id(&self) -> ScalarFnId {
45        static ID: CachedId = CachedId::new("vortex.get_item");
46        *ID
47    }
48
49    fn serialize(&self, instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
50        Ok(Some(
51            pb::GetItemOpts {
52                path: instance.to_string(),
53            }
54            .encode_to_vec(),
55        ))
56    }
57
58    fn deserialize(
59        &self,
60        _metadata: &[u8],
61        _session: &VortexSession,
62    ) -> VortexResult<Self::Options> {
63        let opts = pb::GetItemOpts::decode(_metadata)?;
64        Ok(FieldName::from(opts.path))
65    }
66
67    fn arity(&self, _field_name: &FieldName) -> Arity {
68        Arity::Exact(1)
69    }
70
71    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
72        match child_idx {
73            0 => ChildName::from("input"),
74            _ => unreachable!("Invalid child index {} for GetItem expression", child_idx),
75        }
76    }
77
78    fn fmt_sql(
79        &self,
80        field_name: &FieldName,
81        expr: &Expression,
82        f: &mut Formatter<'_>,
83    ) -> std::fmt::Result {
84        expr.children()[0].fmt_sql(f)?;
85        write!(f, ".{}", field_name)
86    }
87
88    fn return_dtype(&self, field_name: &FieldName, arg_dtypes: &[DType]) -> VortexResult<DType> {
89        let struct_dtype = &arg_dtypes[0];
90        let field_dtype = struct_dtype
91            .as_struct_fields_opt()
92            .and_then(|st| st.field(field_name))
93            .ok_or_else(|| {
94                vortex_err!("Couldn't find the {} field in the input scope", field_name)
95            })?;
96
97        // Match here to avoid cloning the dtype if nullability doesn't need to change
98        if matches!(
99            (struct_dtype.nullability(), field_dtype.nullability()),
100            (Nullability::Nullable, Nullability::NonNullable)
101        ) {
102            return Ok(field_dtype.with_nullability(Nullability::Nullable));
103        }
104
105        Ok(field_dtype)
106    }
107
108    fn execute(
109        &self,
110        field_name: &FieldName,
111        args: &dyn ExecutionArgs,
112        ctx: &mut ExecutionCtx,
113    ) -> VortexResult<ArrayRef> {
114        let input = args.get(0)?.execute::<StructArray>(ctx)?;
115        let field = input.unmasked_field_by_name(field_name).cloned()?;
116
117        match input.dtype().nullability() {
118            Nullability::NonNullable => Ok(field),
119            Nullability::Nullable => field.mask(input.validity()?.to_array(input.len())),
120        }
121    }
122
123    fn reduce(
124        &self,
125        field_name: &FieldName,
126        node: &dyn ReduceNode,
127        ctx: &dyn ReduceCtx,
128    ) -> VortexResult<Option<ReduceNodeRef>> {
129        let child = node.child(0);
130        if let Some(child_fn) = child.scalar_fn()
131            && let Some(pack) = child_fn.as_opt::<Pack>()
132            && let Some(idx) = pack.names.find(field_name)
133        {
134            let mut field = child.child(idx);
135
136            // Possibly mask the field if the pack is nullable
137            if pack.nullability.is_nullable() {
138                field = ctx.new_node(
139                    Mask.bind(EmptyOptions),
140                    &[field, ctx.new_node(Literal.bind(true.into()), &[])?],
141                )?;
142            }
143
144            return Ok(Some(field));
145        }
146
147        Ok(None)
148    }
149
150    fn simplify_untyped(
151        &self,
152        field_name: &FieldName,
153        expr: &Expression,
154    ) -> VortexResult<Option<Expression>> {
155        let child = expr.child(0);
156
157        // If the child is a Pack expression, we can directly return the corresponding child.
158        if let Some(pack) = child.as_opt::<Pack>() {
159            let idx = pack
160                .names
161                .iter()
162                .position(|name| name == field_name)
163                .ok_or_else(|| {
164                    vortex_err!(
165                        "Cannot find field {} in pack fields {:?}",
166                        field_name,
167                        pack.names
168                    )
169                })?;
170
171            let mut field = child.child(idx).clone();
172
173            // It's useful to simplify this node without type info, but we need to make sure
174            // the nullability is correct. We cannot cast since we don't have the dtype info here,
175            // so instead we insert a Mask expression that we know converts a child's dtype to
176            // nullable.
177            if pack.nullability.is_nullable() {
178                // Mask with an all-true array to ensure the field DType is nullable.
179                field = field.mask(lit(true))?;
180            }
181
182            return Ok(Some(field));
183        }
184
185        Ok(None)
186    }
187
188    // This will apply struct nullability field. We could add a dtype??
189    fn is_null_sensitive(&self, _field_name: &FieldName) -> bool {
190        true
191    }
192
193    fn is_fallible(&self, _field_name: &FieldName) -> bool {
194        // If this type-checks its infallible.
195        false
196    }
197}
198
199#[cfg(test)]
200mod tests {
201    use vortex_buffer::buffer;
202
203    use crate::IntoArray;
204    use crate::dtype::DType;
205    use crate::dtype::FieldNames;
206    use crate::dtype::Nullability;
207    use crate::dtype::Nullability::NonNullable;
208    use crate::dtype::PType;
209    use crate::dtype::StructFields;
210    use crate::expr::checked_add;
211    use crate::expr::get_item;
212    use crate::expr::lit;
213    use crate::expr::pack;
214    use crate::expr::root;
215    use crate::scalar_fn::fns::get_item::StructArray;
216    use crate::validity::Validity;
217
218    fn test_array() -> StructArray {
219        StructArray::from_fields(&[
220            ("a", buffer![0i32, 1, 2].into_array()),
221            ("b", buffer![4i64, 5, 6].into_array()),
222        ])
223        .unwrap()
224    }
225
226    #[test]
227    fn get_item_by_name() {
228        let st = test_array();
229        let get_item = get_item("a", root());
230        let item = st.into_array().apply(&get_item).unwrap();
231        assert_eq!(item.dtype(), &DType::from(PType::I32))
232    }
233
234    #[test]
235    fn get_item_by_name_none() {
236        let st = test_array();
237        let get_item = get_item("c", root());
238        assert!(st.into_array().apply(&get_item).is_err());
239    }
240
241    #[test]
242    fn get_nullable_field() {
243        let st = StructArray::try_new(
244            FieldNames::from(["a"]),
245            vec![buffer![1i32].into_array()],
246            1,
247            Validity::AllInvalid,
248        )
249        .unwrap()
250        .into_array();
251
252        let get_item_expr = get_item("a", root());
253        let item = st.apply(&get_item_expr).unwrap();
254        // The dtype should be nullable since it inherits struct validity
255        assert_eq!(
256            item.dtype(),
257            &DType::Primitive(PType::I32, Nullability::Nullable)
258        );
259    }
260
261    #[test]
262    fn test_pack_get_item_rule() {
263        // Create: pack(a: lit(1), b: lit(2)).get_item("b")
264        let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
265        let get_item_expr = get_item("b", pack_expr);
266
267        let result = get_item_expr
268            .optimize_recursive(&DType::Struct(StructFields::empty(), NonNullable))
269            .unwrap();
270
271        assert_eq!(result, lit(2));
272    }
273
274    #[test]
275    fn test_multi_level_pack_get_item_simplify() {
276        let inner_pack = pack([("a", lit(1)), ("b", lit(2))], NonNullable);
277        let get_a = get_item("a", inner_pack);
278
279        let outer_pack = pack([("x", get_a), ("y", lit(3)), ("z", lit(4))], NonNullable);
280        let get_z = get_item("z", outer_pack);
281
282        let dtype = DType::Primitive(PType::I32, NonNullable);
283
284        let result = get_z.optimize_recursive(&dtype).unwrap();
285        assert_eq!(result, lit(4));
286    }
287
288    #[test]
289    fn test_deeply_nested_pack_get_item() {
290        let innermost = pack([("a", lit(42))], NonNullable);
291        let get_a = get_item("a", innermost);
292
293        let level2 = pack([("b", get_a)], NonNullable);
294        let get_b = get_item("b", level2);
295
296        let level3 = pack([("c", get_b)], NonNullable);
297        let get_c = get_item("c", level3);
298
299        let outermost = pack([("final", get_c)], NonNullable);
300        let get_final = get_item("final", outermost);
301
302        let dtype = DType::Primitive(PType::I32, NonNullable);
303
304        let result = get_final.optimize_recursive(&dtype).unwrap();
305        assert_eq!(result, lit(42));
306    }
307
308    #[test]
309    fn test_partial_pack_get_item_simplify() {
310        let inner_pack = pack([("x", lit(1)), ("y", lit(2))], NonNullable);
311        let get_x = get_item("x", inner_pack);
312        let add_expr = checked_add(get_x, lit(10));
313
314        let outer_pack = pack([("result", add_expr)], NonNullable);
315        let get_result = get_item("result", outer_pack);
316
317        let dtype = DType::Primitive(PType::I32, NonNullable);
318
319        let result = get_result.optimize_recursive(&dtype).unwrap();
320        let expected = checked_add(lit(1), lit(10));
321        assert_eq!(&result, &expected);
322    }
323
324    #[test]
325    fn get_item_filter_list_field() {
326        use vortex_mask::Mask;
327
328        use crate::arrays::BoolArray;
329        use crate::arrays::FilterArray;
330        use crate::arrays::ListArray;
331
332        let list = ListArray::try_new(
333            buffer![0f32, 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., 11.].into_array(),
334            buffer![2u64, 4, 6, 8, 10, 12].into_array(),
335            Validity::Array(BoolArray::from_iter([true, true, false, true, true]).into_array()),
336        )
337        .unwrap();
338
339        let filtered = FilterArray::try_new(
340            list.into_array(),
341            Mask::from_iter([true, true, false, false, false]),
342        )
343        .unwrap();
344
345        let st = StructArray::try_new(
346            FieldNames::from(["data"]),
347            vec![filtered.into_array()],
348            2,
349            Validity::AllValid,
350        )
351        .unwrap();
352
353        st.into_array().apply(&get_item("data", root())).unwrap();
354    }
355}
vortex_array/scalar_fn/fns/get_item.rs

vortex_array/scalar_fn/fns/
get_item.rs
