Fix list-of-struct nested projection

vortex-array/src/expr/exprs/get_item.rs

@@ -33,15 +33,28 @@ use crate::expr::Pack;
 use crate::expr::ReduceCtx;
 use crate::expr::ReduceNode;
 use crate::expr::ReduceNodeRef;
+use crate::expr::SimplifyCtx;
 use crate::expr::StatsCatalog;
 use crate::expr::VTable;
 use crate::expr::VTableExt;
 use crate::expr::exprs::root::root;
+use crate::expr::get_item_list;
 use crate::expr::lit;
 use crate::expr::stats::Stat;
 
 pub struct GetItem;
 
+fn propagate_nullability(parent_nullability: Nullability, field_dtype: DType) -> DType {
+    if matches!(
+        (parent_nullability, field_dtype.nullability()),
+        (Nullability::Nullable, Nullability::NonNullable)
+    ) {
+        return field_dtype.with_nullability(Nullability::Nullable);
+    }
+
+    field_dtype
+}
+
 impl VTable for GetItem {
     type Options = FieldName;
 
@@ -85,23 +98,24 @@ impl VTable for GetItem {
     }
 
     fn return_dtype(&self, field_name: &FieldName, arg_dtypes: &[DType]) -> VortexResult<DType> {
-        let struct_dtype = &arg_dtypes[0];
-        let field_dtype = struct_dtype
-            .as_struct_fields_opt()
-            .and_then(|st| st.field(field_name))
-            .ok_or_else(|| {
+        let input_dtype = &arg_dtypes[0];
+
+        // Struct field access: `$.a`.
+        if let Some(struct_fields) = input_dtype.as_struct_fields_opt() {
+            let field_dtype = struct_fields.field(field_name).ok_or_else(|| {
                 vortex_err!("Couldn't find the {} field in the input scope", field_name)
             })?;
 
-        // Match here to avoid cloning the dtype if nullability doesn't need to change
-        if matches!(
-            (struct_dtype.nullability(), field_dtype.nullability()),
-            (Nullability::Nullable, Nullability::NonNullable)
-        ) {
-            return Ok(field_dtype.with_nullability(Nullability::Nullable));
+            return Ok(propagate_nullability(
+                input_dtype.nullability(),
+                field_dtype,
+            ));
         }
 
-        Ok(field_dtype)
+        Err(vortex_err!(
+            "Expected struct dtype for child of GetItem expression, got {}",
+            input_dtype
+        ))
     }
 
     fn evaluate(
@@ -110,34 +124,75 @@ impl VTable for GetItem {
         expr: &Expression,
         scope: &ArrayRef,
     ) -> VortexResult<ArrayRef> {
-        let input = expr.children()[0].evaluate(scope)?.to_struct();
-        let field = input.field_by_name(field_name).cloned()?;
+        let input = expr.children()[0].evaluate(scope)?;
+
+        // Struct field access: `$.a`.
+        if input.dtype().is_struct() {
+            let input = input.to_struct();
+            let field = input.field_by_name(field_name).cloned()?;
 
-        match input.dtype().nullability() {
-            Nullability::NonNullable => Ok(field),
-            Nullability::Nullable => mask(&field, &input.validity_mask().not()),
+            return match input.dtype().nullability() {
+                Nullability::NonNullable => Ok(field),
+                Nullability::Nullable => mask(&field, &input.validity_mask().not()),
+            };
         }
+
+        Err(vortex_err!(
+            "Expected struct scope for GetItem evaluation, got {}",
+            input.dtype()
+        ))
     }
 
     fn execute(&self, field_name: &FieldName, mut args: ExecutionArgs) -> VortexResult<Datum> {
-        let struct_dtype = args.dtypes[0]
-            .as_struct_fields_opt()
-            .ok_or_else(|| vortex_err!("Expected struct dtype for child of GetItem expression"))?;
-        let field_idx = struct_dtype
-            .find(field_name)
-            .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", field_name))?;
-
-        match args.datums.pop().vortex_expect("missing input") {
-            Datum::Scalar(s) => {
-                let mut field = s.as_struct().field(field_idx);
-                field.mask_validity(s.is_valid());
-                Ok(Datum::Scalar(field))
-            }
-            Datum::Vector(v) => {
-                let mut field = v.as_struct().fields()[field_idx].clone();
-                field.mask_validity(v.validity());
-                Ok(Datum::Vector(field))
-            }
+        let input_dtype = &args.dtypes[0];
+
+        // Struct field access: `$.a`.
+        if let Some(struct_dtype) = input_dtype.as_struct_fields_opt() {
+            let field_idx = struct_dtype
+                .find(field_name)
+                .ok_or_else(|| vortex_err!("Field {} not found in struct dtype", field_name))?;
+
+            return match args.datums.pop().vortex_expect("missing input") {
+                Datum::Scalar(s) => {
+                    let mut field = s.as_struct().field(field_idx);
+                    field.mask_validity(s.is_valid());
+                    Ok(Datum::Scalar(field))
+                }
+                Datum::Vector(v) => {
+                    let mut field = v.as_struct().fields()[field_idx].clone();
+                    field.mask_validity(v.validity());
+                    Ok(Datum::Vector(field))
+                }
+            };
+        }
+
+        Err(vortex_err!(
+            "Expected struct dtype for child of GetItem expression, got {}",
+            input_dtype
+        ))
+    }
+
+    fn simplify(
+        &self,
+        field_name: &FieldName,
+        expr: &Expression,
+        ctx: &dyn SimplifyCtx,
+    ) -> VortexResult<Option<Expression>> {
+        let child = expr.child(0);
+        let child_dtype = ctx.return_dtype(child)?;
+
+        let element_dtype = match child_dtype {
+            DType::List(element_dtype, _) => Some(element_dtype),
+            DType::FixedSizeList(element_dtype, ..) => Some(element_dtype),
+            _ => None,
+        };
+
+        if let Some(element_dtype) = element_dtype
+            && element_dtype.as_struct_fields_opt().is_some()
+        {
+            Ok(Some(get_item_list(field_name.clone(), child.clone())))
+        } else {
+            Ok(None)
         }
     }
 
@@ -261,6 +316,8 @@ pub fn get_item(field: impl Into<FieldName>, child: Expression) -> Expression {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use vortex_buffer::buffer;
     use vortex_dtype::DType;
     use vortex_dtype::FieldNames;
@@ -272,9 +329,12 @@ mod tests {
 
     use crate::Array;
     use crate::IntoArray;
+    use crate::arrays::FixedSizeListArray;
+    use crate::arrays::ListArray;
     use crate::arrays::StructArray;
     use crate::expr::exprs::binary::checked_add;
     use crate::expr::exprs::get_item::get_item;
+    use crate::expr::exprs::get_item_list::get_item_list;
     use crate::expr::exprs::literal::lit;
     use crate::expr::exprs::pack::pack;
     use crate::expr::exprs::root::root;
@@ -322,6 +382,182 @@ mod tests {
         );
     }
 
+    #[test]
+    fn get_item_list_of_struct() {
+        let element_dtype = Arc::new(DType::Struct(
+            [
+                ("a", DType::Primitive(PType::I32, NonNullable)),
+                ("b", DType::Utf8(NonNullable)),
+            ]
+            .into_iter()
+            .collect(),
+            NonNullable,
+        ));
+
+        let row_count = 4;
+        let items = ListArray::from_iter_opt_slow::<u32, _, _>(
+            [
+                Some(vec![
+                    Scalar::struct_(
+                        (*element_dtype).clone(),
+                        vec![
+                            Scalar::primitive(1i32, NonNullable),
+                            Scalar::utf8("x", NonNullable),
+                        ],
+                    ),
+                    Scalar::struct_(
+                        (*element_dtype).clone(),
+                        vec![
+                            Scalar::primitive(2i32, NonNullable),
+                            Scalar::utf8("y", NonNullable),
+                        ],
+                    ),
+                ]),
+                Some(Vec::new()),
+                None,
+                Some(vec![Scalar::struct_(
+                    (*element_dtype).clone(),
+                    vec![
+                        Scalar::primitive(3i32, NonNullable),
+                        Scalar::utf8("z", NonNullable),
+                    ],
+                )]),
+            ],
+            element_dtype,
+        )
+        .unwrap();
+
+        let ids = buffer![0i32, 1, 2, 3].into_array();
+
+        let data = StructArray::new(
+            FieldNames::from(["id", "items"]),
+            vec![ids, items],
+            row_count,
+            Validity::NonNullable,
+        )
+        .into_array();
+
+        // Regression for nested field projection on list-of-struct: `items.a`.
+        let projection = get_item_list("a", get_item("items", root()));
+        let out = data.apply(&projection).expect("apply");
+
+        assert_eq!(
+            out.dtype(),
+            &DType::List(
+                Arc::new(DType::Primitive(PType::I32, NonNullable)),
+                Nullability::Nullable
+            )
+        );
+
+        assert_eq!(
+            out.scalar_at(0).as_list().elements().unwrap().to_vec(),
+            vec![
+                Scalar::primitive(1i32, NonNullable),
+                Scalar::primitive(2i32, NonNullable),
+            ]
+        );
+        assert!(out.scalar_at(1).as_list().elements().unwrap().is_empty());
+        assert!(out.scalar_at(2).is_null());
+        assert_eq!(
+            out.scalar_at(3).as_list().elements().unwrap().to_vec(),
+            vec![Scalar::primitive(3i32, NonNullable)]
+        );
+    }
+
+    #[test]
+    fn get_item_fixed_size_list_of_struct() {
+        let n_lists: usize = 3;
+        let list_size: u32 = 2;
+        let n_elements = n_lists * list_size as usize;
+
+        let struct_elems = StructArray::try_new(
+            FieldNames::from(["a", "b"]),
+            vec![
+                buffer![1i32, 2, 3, 4, 5, 6].into_array(),
+                buffer![10i64, 20, 30, 40, 50, 60].into_array(),
+            ],
+            n_elements,
+            Validity::from_iter([true, false, true, true, false, true]),
+        )
+        .unwrap()
+        .into_array();
+
+        let items = FixedSizeListArray::try_new(
+            struct_elems,
+            list_size,
+            Validity::from_iter([true, false, true]),
+            n_lists,
+        )
+        .unwrap()
+        .into_array();
+
+        let ids = buffer![0i32, 1, 2].into_array();
+
+        let data = StructArray::new(
+            FieldNames::from(["id", "items"]),
+            vec![ids, items],
+            n_lists,
+            Validity::NonNullable,
+        )
+        .into_array();
+
+        // FixedSizeList-of-struct projection: `items.a`, including struct-level nulls inside the list.
+        let projection = get_item_list("a", get_item("items", root()));
+        let out = data.apply(&projection).expect("apply");
+
+        assert_eq!(
+            out.dtype(),
+            &DType::FixedSizeList(
+                Arc::new(DType::Primitive(PType::I32, Nullability::Nullable)),
+                list_size,
+                Nullability::Nullable
+            )
+        );
+
+        assert_eq!(
+            out.scalar_at(0).as_list().elements().unwrap().to_vec(),
+            vec![
+                Scalar::primitive(1i32, Nullability::Nullable),
+                Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
+            ]
+        );
+        assert!(out.scalar_at(1).is_null());
+        assert_eq!(
+            out.scalar_at(2).as_list().elements().unwrap().to_vec(),
+            vec![
+                Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)),
+                Scalar::primitive(6i32, Nullability::Nullable),
+            ]
+        );
+    }
+
+    #[test]
+    fn get_item_list_of_struct_desugars_to_get_item_list_on_optimize() {
+        let scope_dtype = DType::Struct(
+            [
+                ("id", DType::from(PType::I32)),
+                (
+                    "items",
+                    DType::List(
+                        Arc::new(DType::Struct(
+                            [("a", DType::from(PType::I32))].into_iter().collect(),
+                            NonNullable,
+                        )),
+                        Nullability::Nullable,
+                    ),
+                ),
+            ]
+            .into_iter()
+            .collect(),
+            NonNullable,
+        );
+
+        let expr = get_item("a", get_item("items", root()));
+        let optimized = expr.optimize_recursive(&scope_dtype).unwrap();
+
+        assert!(optimized.is::<crate::expr::exprs::get_item_list::GetItemList>());
+    }
+
     #[test]
     fn test_pack_get_item_rule() {
         // Create: pack(a: lit(1), b: lit(2)).get_item("b")