diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index ab03a4624fa..88cc768689e 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -22,6 +22,7 @@ use vortex_array::dtype::StructFields; use vortex_array::expr::ExactExpr; use vortex_array::expr::Expression; use vortex_array::expr::col; +use vortex_array::expr::is_root; use vortex_array::expr::make_free_field_annotator; use vortex_array::expr::root; use vortex_array::expr::transform::PartitionedExpr; @@ -319,6 +320,7 @@ impl LayoutReader for StructReader { expr: &Expression, mask_fut: MaskFuture, ) -> VortexResult { + let is_root_expr = is_root(expr); let validity_fut = self .validity()? .map(|reader| reader.projection_evaluation(row_range, &root(), mask_fut.clone())) @@ -348,14 +350,17 @@ impl LayoutReader for StructReader { partitioned.root.is::() || partitioned.root.is::(), ), }; + let push_validity_into_fields = is_pack_merge && !is_root_expr; let session = self.session.clone(); Ok(Box::pin(async move { if let Some(validity_fut) = validity_fut { let (array, validity) = try_join!(projected, validity_fut)?; - // If root expression was a pack, then we apply the validity to each child field - if is_pack_merge { + // A user-authored pack or merge creates a new struct, so the source struct's + // validity applies to its fields. A root expression is also expanded to a pack for + // partitioning, but that synthetic pack must retain the source struct's validity. + if push_validity_into_fields { let mut ctx = session.create_execution_ctx(); let struct_array = array.execute::(&mut ctx)?; let masked_fields: Vec = struct_array @@ -421,6 +426,7 @@ mod tests { use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; @@ -748,6 +754,64 @@ mod tests { assert_nth_scalar!(result, 2, 3, &mut ctx); } + #[rstest] + fn test_struct_layout_root_preserves_validity( + #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; + let project = reader.projection_evaluation(&(0..3), &root(), MaskFuture::new_true(3))?; + + let result = block_on(move |_| project)?; + assert_eq!(result.dtype(), reader.dtype()); + assert!( + result + .execute_scalar(0, &mut SESSION.create_execution_ctx())? + .is_null() + ); + Ok(()) + } + + #[rstest] + fn test_struct_layout_pack_pushes_validity_into_fields( + #[from(null_struct_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; + let expr = pack([("a", get_item("a", root()))], Nullability::NonNullable); + let project = reader.projection_evaluation(&(0..3), &expr, MaskFuture::new_true(3))?; + + let mut ctx = SESSION.create_execution_ctx(); + let result = block_on(move |_| project)?.execute::(&mut ctx)?; + assert!(!result.dtype().is_nullable()); + assert!( + result + .unmasked_field_by_name("a")? + .execute_scalar(0, &mut ctx)? + .is_null() + ); + Ok(()) + } + + #[rstest] + fn test_nested_struct_layout_preserves_struct_validity( + #[from(nested_struct_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let reader = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; + let project = reader.projection_evaluation( + &(0..3), + &get_item("a", root()), + MaskFuture::new_true(3), + )?; + + let result = block_on(move |_| project)?; + assert!(result.dtype().is_nullable()); + assert!( + result + .execute_scalar(1, &mut SESSION.create_execution_ctx())? + .is_null() + ); + Ok(()) + } + #[rstest] fn test_struct_layout_nested( #[from(nested_struct_layout)] (segments, layout): (Arc, LayoutRef),