From daada40b9bc77139553cc1af399c09593cfdd03f Mon Sep 17 00:00:00 2001 From: Liam Abourousse <75449798+Bardakor@users.noreply.github.com> Date: Thu, 20 Aug 2026 19:59:20 +0300 Subject: [PATCH] Handle inherited validity for nullable struct fields Signed-off-by: Liam Abourousse <75449798+Bardakor@users.noreply.github.com> --- vortex-file/tests/test_nullable_struct.rs | 109 ++++ vortex-layout/src/layouts/struct_/mod.rs | 26 +- .../src/layouts/struct_/nullable_reader.rs | 523 ++++++++++++++++++ 3 files changed, 651 insertions(+), 7 deletions(-) create mode 100644 vortex-file/tests/test_nullable_struct.rs create mode 100644 vortex-layout/src/layouts/struct_/nullable_reader.rs diff --git a/vortex-file/tests/test_nullable_struct.rs b/vortex-file/tests/test_nullable_struct.rs new file mode 100644 index 00000000000..5f665026e1f --- /dev/null +++ b/vortex-file/tests/test_nullable_struct.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::tests_outside_test_module)] + +use std::sync::LazyLock; + +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::FieldNames; +use vortex_array::expr::get_item; +use vortex_array::expr::gt; +use vortex_array::expr::is_not_null; +use vortex_array::expr::is_null; +use vortex_array::expr::lit; +use vortex_array::expr::root; +use vortex_array::stream::ArrayStreamExt; +use vortex_array::validity::Validity; +use vortex_buffer::ByteBuffer; +use vortex_buffer::buffer; +use vortex_error::VortexResult; +use vortex_file::OpenOptionsSessionExt; +use vortex_file::WriteOptionsSessionExt; +use vortex_io::session::RuntimeSession; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; + +mod common; + +use common::enable_all_registered_array_encodings; + +static SESSION: LazyLock = LazyLock::new(|| { + let session = vortex_array::array_session() + .with::() + .with::(); + + vortex_file::register_default_encodings(&session); + enable_all_registered_array_encodings(&session); + + session +}); + +#[tokio::test] +async fn nullable_struct_child_inherits_parent_validity() -> VortexResult<()> { + // The second child value is deliberately valid and would match `a > 1`, but its parent struct + // row is null. This is the representation produced by older writers and is valid Arrow data. + let nullable_struct = StructArray::try_new( + FieldNames::from(["a"]), + vec![buffer![1i32, 2].into_array()], + 2, + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + )?; + let data = StructArray::try_new( + FieldNames::from(["s"]), + vec![nullable_struct.into_array()], + 2, + Validity::NonNullable, + )? + .into_array(); + + let mut bytes = Vec::new(); + SESSION + .write_options() + .write(&mut bytes, data.to_array_stream()) + .await?; + let file = SESSION + .open_options() + .open_buffer(ByteBuffer::from(bytes))?; + + let field = get_item("a", get_item("s", root())); + let projected = file + .scan()? + .with_projection(field.clone()) + .into_array_stream()? + .read_all() + .await?; + assert_eq!( + projected.invalid_count(&mut SESSION.create_execution_ctx())?, + 1 + ); + + let nulls = file + .scan()? + .with_filter(is_null(field.clone())) + .into_array_stream()? + .read_all() + .await?; + assert_eq!(nulls.len(), 1); + + let non_nulls = file + .scan()? + .with_filter(is_not_null(field.clone())) + .into_array_stream()? + .read_all() + .await?; + assert_eq!(non_nulls.len(), 1); + + let raw_child_match = file + .scan()? + .with_filter(gt(field, lit(1i32))) + .into_array_stream()? + .read_all() + .await?; + assert_eq!(raw_child_match.len(), 0); + + Ok(()) +} diff --git a/vortex-layout/src/layouts/struct_/mod.rs b/vortex-layout/src/layouts/struct_/mod.rs index 862b789a349..7adabedf948 100644 --- a/vortex-layout/src/layouts/struct_/mod.rs +++ b/vortex-layout/src/layouts/struct_/mod.rs @@ -1,11 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod nullable_reader; mod reader; pub mod writer; use std::sync::Arc; +use nullable_reader::NullableStructReader; use reader::StructReader; use vortex_array::EmptyMetadata; use vortex_array::dtype::DType; @@ -114,13 +116,23 @@ impl VTable for Struct { session: &VortexSession, ctx: &LayoutReaderContext, ) -> VortexResult { - Ok(Arc::new(StructReader::try_new( - layout.clone(), - name, - segment_source, - session.session(), - ctx.clone(), - )?)) + if layout.dtype().is_nullable() { + Ok(Arc::new(NullableStructReader::try_new( + layout.clone(), + name, + segment_source, + session.session(), + ctx.clone(), + )?)) + } else { + Ok(Arc::new(StructReader::try_new( + layout.clone(), + name, + segment_source, + session.session(), + ctx.clone(), + )?)) + } } } diff --git a/vortex-layout/src/layouts/struct_/nullable_reader.rs b/vortex-layout/src/layouts/struct_/nullable_reader.rs new file mode 100644 index 00000000000..8bede76ed0c --- /dev/null +++ b/vortex-layout/src/layouts/struct_/nullable_reader.rs @@ -0,0 +1,523 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ops::Range; +use std::sync::Arc; +use std::sync::OnceLock; + +use itertools::Itertools; +use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldMask; +use vortex_array::dtype::FieldName; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::bound::get_item; +use vortex_array::expr::bound::pack; +use vortex_array::expr::make_bound_free_field_annotator; +use vortex_array::expr::root; +use vortex_array::expr::transform::BoundPartitionedExpr; +use vortex_array::expr::transform::partition_bound; +use vortex_array::expr::traversal::NodeExt; +use vortex_array::expr::traversal::Transformed; +use vortex_array::expr::traversal::TraversalOrder; +use vortex_array::scalar_fn::fns::get_item::GetItem; +use vortex_array::scalar_fn::fns::select::Select; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_utils::aliases::dash_map::DashMap; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::ArrayFuture; +use crate::LayoutReader; +use crate::LayoutReaderRef; +use crate::LazyReaderChildren; +use crate::RowSplits; +use crate::SplitRange; +use crate::layouts::partitioned::BoundPartitionedExprEval; +use crate::layouts::struct_::StructLayout; +use crate::segments::SegmentSource; + +/// Reader for nullable structs. +/// +/// The serialized child layouts retain their declared dtypes. Parent validity is applied at +/// evaluation time so existing files remain readable and older readers can still read newly +/// written files. Strict predicates may continue to use child pruning because parent nulls only +/// remove logical rows. Non-strict predicates must see inherited validity before evaluation. +pub(super) struct NullableStructReader { + layout: StructLayout, + name: Arc, + lazy_children: LazyReaderChildren, + session: VortexSession, + expanded_root_expr: BoundExpression, + field_lookup: Option>, + partitioned_expr_cache: DashMap>>, +} + +impl NullableStructReader { + pub(super) fn try_new( + layout: StructLayout, + name: Arc, + segment_source: Arc, + session: VortexSession, + ctx: crate::LayoutReaderContext, + ) -> VortexResult { + let struct_dt = layout.struct_fields(); + let field_lookup = (struct_dt.nfields() > 80).then(|| { + struct_dt + .names() + .iter() + .enumerate() + .map(|(i, name)| (name.clone(), i)) + .collect() + }); + let mut dtypes = Vec::with_capacity(struct_dt.nfields() + 1); + let mut names = Vec::with_capacity(struct_dt.nfields() + 1); + + dtypes.push(DType::Bool(Nullability::NonNullable)); + names.push(Arc::from("validity")); + dtypes.extend(struct_dt.fields()); + names.extend(struct_dt.names().iter().map(|name| Arc::clone(name.inner()))); + + let lazy_children = LazyReaderChildren::new( + Arc::clone(layout.children()), + dtypes, + names, + Arc::clone(&segment_source), + session.clone(), + ctx, + ); + let expanded_root_expr = expanded_struct_root(layout.dtype(), struct_dt)?; + + Ok(Self { + layout, + name, + lazy_children, + session, + expanded_root_expr, + field_lookup, + partitioned_expr_cache: Default::default(), + }) + } + + fn struct_fields(&self) -> &StructFields { + self.layout.struct_fields() + } + + fn field_reader(&self, name: &FieldName) -> VortexResult<&LayoutReaderRef> { + let idx = self + .field_lookup + .as_ref() + .and_then(|lookup| lookup.get(name).copied()) + .or_else(|| self.struct_fields().find(name)) + .ok_or_else(|| vortex_err!("Field {} not found in struct layout", name))?; + self.field_reader_by_index(idx) + } + + fn field_reader_by_index(&self, idx: usize) -> VortexResult<&LayoutReaderRef> { + let child_index = self + .layout + .slot_to_child(idx + 1) + .vortex_expect("struct field slot is always present"); + self.lazy_children.get(child_index) + } + + fn validity(&self) -> VortexResult<&LayoutReaderRef> { + let child_index = self + .layout + .slot_to_child(0) + .vortex_expect("nullable struct validity slot is always present"); + self.lazy_children.get(child_index) + } + + fn logical_field_dtype(&self, name: &FieldName) -> VortexResult { + Ok(self.field_reader(name)?.dtype().as_nullable()) + } + + fn partition_expr(&self, expr: &BoundExpression) -> VortexResult { + let key = ExactBoundExpr(expr.clone()); + let cell = match self.partitioned_expr_cache.get(&key) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.partitioned_expr_cache + .entry(key) + .or_insert_with(|| Arc::new(OnceLock::new())) + .value(), + ), + }; + + if let Some(value) = cell.get() { + return Ok(value.clone()); + } + + let expr = expand_struct_root( + expr.clone(), + &self.expanded_root_expr, + self.struct_fields(), + )?; + let mut partitioned = partition_bound( + expr.clone(), + make_bound_free_field_annotator(self.struct_fields()), + )?; + + let result = if partitioned.partitions.len() == 1 { + let name = partitioned.partition_names[0].clone(); + Partitioned::Single( + name.clone(), + step_into_struct_field(expr, &name, self.logical_field_dtype(&name)?)?, + ) + } else { + let partitions = partitioned + .partitions + .iter() + .zip_eq(partitioned.partition_names.iter()) + .map(|(expr, name)| { + step_into_struct_field(expr.clone(), name, self.logical_field_dtype(name)?) + }) + .try_collect::<_, Vec<_>, _>()? + .into_boxed_slice(); + partitioned.replace_partitions(partitions)?; + Partitioned::Multi(Arc::new(partitioned)) + }; + + Ok(cell.get_or_init(|| result).clone()) + } + + fn validity_projection_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let reader = self.validity()?; + let validity_root = root().bind(reader.dtype())?; + reader.projection_evaluation(row_range, &validity_root, mask) + } + + fn validity_filter_evaluation( + &self, + row_range: &Range, + mask: MaskFuture, + ) -> VortexResult { + let reader = self.validity()?; + let validity_root = root().bind(reader.dtype())?; + reader.filter_evaluation(row_range, &validity_root, mask) + } + + fn field_projection_evaluation( + &self, + row_range: &Range, + name: &FieldName, + expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + let reader = self.field_reader(name)?; + let validity = self.validity_projection_evaluation(row_range, mask.clone())?; + + if expression_is_strict(expr) { + let physical_expr = retype_roots(expr.clone(), reader.dtype().clone())?; + let field = reader.projection_evaluation(row_range, &physical_expr, mask)?; + + return Ok(Box::pin(async move { + let (field, validity) = futures::try_join!(field, validity)?; + field.mask(validity) + })); + } + + let field_root = root().bind(reader.dtype())?; + let field = reader.projection_evaluation(row_range, &field_root, mask)?; + let expr = expr.clone(); + + Ok(Box::pin(async move { + let (field, validity) = futures::try_join!(field, validity)?; + field.mask(validity)?.apply_bound(&expr) + })) + } + + fn field_filter_evaluation( + &self, + row_range: &Range, + name: &FieldName, + expr: &BoundExpression, + mask: MaskFuture, + ) -> VortexResult { + if expression_is_strict(expr) { + let reader = self.field_reader(name)?; + let physical_expr = retype_roots(expr.clone(), reader.dtype().clone())?; + let valid_mask = self.validity_filter_evaluation(row_range, mask)?; + return reader.filter_evaluation(row_range, &physical_expr, valid_mask); + } + + let input_mask = mask.clone(); + let result = self.field_projection_evaluation(row_range, name, expr, mask)?; + let session = self.session.clone(); + let len = input_mask.len(); + + Ok(MaskFuture::new(len, async move { + let (input_mask, result) = futures::try_join!(input_mask, result)?; + let mut ctx = session.create_execution_ctx(); + let result_mask = result.null_as_false().execute(&mut ctx)?; + Ok(input_mask.intersect_by_rank(&result_mask)) + })) + } +} + +fn expression_is_strict(expr: &BoundExpression) -> bool { + match expr.as_scalar() { + Some(scalar_fn) => { + scalar_fn.signature().is_strict() + && expr.children().iter().all(expression_is_strict) + } + None => true, + } +} + +fn retype_roots(expr: BoundExpression, dtype: DType) -> VortexResult { + Ok(expr + .transform_up(|node| { + if node.is_root() { + Ok(Transformed::yes(BoundExpression::new_root(dtype.clone()))) + } else { + Ok(Transformed::no(node)) + } + })? + .into_inner()) +} + +fn expanded_struct_root( + root_dtype: &DType, + fields: &StructFields, +) -> VortexResult { + let root = BoundExpression::new_root(root_dtype.clone()); + let children = fields + .names() + .iter() + .map(|name| get_item(name.clone(), root.clone())) + .collect::>(); + Ok(pack( + fields.names().iter().cloned().zip(children), + Nullability::NonNullable, + )) +} + +fn expand_struct_root( + expr: BoundExpression, + expanded_root: &BoundExpression, + fields: &StructFields, +) -> VortexResult { + Ok(expr + .transform_down(|node| { + if node.is_root() { + return Ok(Transformed { + value: expanded_root.clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + let Some(scalar_fn) = node.as_scalar() else { + return Ok(Transformed::no(node)); + }; + if !node + .children() + .first() + .is_some_and(BoundExpression::is_root) + { + return Ok(Transformed::no(node)); + } + + if let Some(field_name) = scalar_fn.as_opt::() { + let idx = fields.find(field_name).ok_or_else(|| { + vortex_err!("Field {field_name} not found while expanding struct root") + })?; + return Ok(Transformed { + value: expanded_root.children()[idx].clone(), + changed: true, + order: TraversalOrder::Skip, + }); + } + + if let Some(selection) = scalar_fn.as_opt::