Address list-of-struct projection review

Author: godnight10061

vortex-array/src/expr/analysis/immediate_access.rs

@@ -11,6 +11,7 @@ use crate::expr::analysis::AnnotationFn;
 use crate::expr::analysis::Annotations;
 use crate::expr::descendent_annotations;
 use crate::expr::exprs::get_item::GetItem;
+use crate::expr::exprs::get_item_list::GetItemList;
 use crate::expr::exprs::root::Root;
 use crate::expr::exprs::select::Select;
 
@@ -28,7 +29,17 @@ pub fn annotate_scope_access(scope: &StructFields) -> impl AnnotationFn<Annotati
             if expr.child(0).is::<Root>() {
                 return vec![field_name.clone()];
             }
-        } else if expr.is::<Root>() {
+        }
+
+        if expr.is::<GetItemList>() {
+            if let Some(field_name) = expr.child(0).as_opt::<GetItem>() {
+                if expr.child(0).child(0).is::<Root>() {
+                    return vec![field_name.clone()];
+                }
+            }
+        }
+
+        if expr.is::<Root>() {
             return scope.names().iter().cloned().collect();
         }
 

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

@@ -92,15 +92,7 @@ impl VTable for GetItem {
                 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));
-        }
-
-        Ok(field_dtype)
+        Ok(field_dtype.union_nullability(struct_dtype.nullability()))
     }
 
     fn execute(

vortex-array/src/expr/vtable.rs

@@ -815,4 +815,16 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn get_item_list_is_not_serializable() {
+        let expr = get_item_list("field", root());
+        let err = expr
+            .serialize_proto()
+            .expect_err("get_item_list must not be serializable");
+        assert!(
+            err.to_string().contains("must not be serialized"),
+            "unexpected error: {err}"
+        );
+    }
 }

vortex-datafusion/src/persistent/opener.rs

@@ -204,8 +204,8 @@ impl FileOpener for VortexOpener {
 
             // The schema of the stream returned from the vortex scan.
             // We use the physical_file_schema as reference for types that don't roundtrip.
-            let scan_dtype = scan_projection.return_dtype(vxf.dtype()).map_err(|_e| {
-                exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan")
+            let scan_dtype = scan_projection.return_dtype(vxf.dtype()).map_err(|e| {
+                exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan: {e}")
             })?;
             let stream_schema = calculate_physical_schema(&scan_dtype, &projected_physical_schema)?;
 

vortex-file/tests/test_write_table.rs

@@ -13,6 +13,7 @@ use vortex_array::ToCanonical;
 use vortex_array::arrays::ListArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::arrays::StructArray;
+use vortex_array::assert_arrays_eq;
 use vortex_array::expr::get_item;
 use vortex_array::expr::root;
 use vortex_array::expr::session::ExprSession;
@@ -167,7 +168,7 @@ async fn test_list_of_struct_nested_projection() {
                 ],
             )]),
         ],
-        element_dtype,
+        element_dtype.clone(),
     )
     .unwrap();
 
@@ -178,7 +179,7 @@ async fn test_list_of_struct_nested_projection() {
 
     let data = StructArray::new(
         FieldNames::from(["id", "items"]),
-        vec![ids, items],
+        vec![ids, items.clone()],
         row_count,
         Validity::NonNullable,
     )
@@ -204,19 +205,33 @@ async fn test_list_of_struct_nested_projection() {
     let mut stream = vxf
         .scan()
         .expect("scan")
-        .with_projection(projection)
+        .with_projection(projection.clone())
         .into_stream()
         .expect("into_stream");
 
     let out = stream.next().await.expect("first batch").expect("batch");
 
-    // Smoke-check the projected shape; detailed value semantics are covered by unit tests in
-    // `vortex-array`.
-    assert_eq!(out.len(), row_count);
-    assert!(matches!(
-        out.dtype(),
-        vortex_dtype::DType::List(_, Nullability::Nullable)
-    ));
+    let expected = ListArray::from_iter_opt_slow::<u32, _, _>(
+        [
+            Some(vec![
+                Scalar::primitive(1i32, Nullability::NonNullable),
+                Scalar::primitive(2i32, Nullability::NonNullable),
+            ]),
+            Some(Vec::new()),
+            None,
+            Some(vec![Scalar::primitive(3i32, Nullability::NonNullable)]),
+        ],
+        Arc::new(vortex_dtype::DType::Primitive(
+            PType::I32,
+            Nullability::NonNullable,
+        )),
+    )
+    .unwrap()
+    .into_array();
+    let simplified = projection.optimize_recursive(data.dtype()).unwrap();
+    let expected_dtype = simplified.return_dtype(data.dtype()).unwrap();
+    assert_eq!(out.dtype(), &expected_dtype);
+    assert_arrays_eq!(expected, out);
 
     // Verify the list column is not stored as a single flat blob layout.
     // This is the root cause of poor nested support described in #4889.

vortex-layout/src/layouts/list/reader.rs

@@ -121,6 +121,9 @@ impl ListReader {
         self.lazy_children.get(idx)
     }
 
+    /// Creates a future that will produce a slice of this list array.
+    ///
+    /// The produced slice may have a projection applied to its elements.
     fn list_slice_future(
         &self,
         row_range: Range<u64>,