From eb46880c175a350bb1798d1dc9a77e50e7d4ccd2 Mon Sep 17 00:00:00 2001 From: Thor Date: Wed, 2 Sep 2026 14:49:03 -0500 Subject: [PATCH 1/3] Support file stats on nested struct fields (#6389) Store file-level statistics via a post-order walk of the DType tree, so nested struct fields get whole-file pruning instead of only top-level columns. Each leaf gets a stats entry, and each nullable struct gets a trailing null-count entry. A new `is_nested` footer flag keeps old files readable via the legacy top-level-only layout. Signed-off-by: "Thor" --- vortex-file/src/file.rs | 7 +- vortex-file/src/footer/file_statistics.rs | 197 +++++++++-- vortex-file/src/pruning.rs | 20 +- vortex-file/src/tests.rs | 67 +++- vortex-file/src/v2/file_stats_reader.rs | 24 +- .../flatbuffers/vortex-file/footer.fbs | 8 + vortex-flatbuffers/src/generated/footer.rs | 25 ++ vortex-layout/src/layouts/file_stats.rs | 322 +++++++++++++++--- 8 files changed, 538 insertions(+), 132 deletions(-) diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index a235011b9c0..adf081f0f23 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -227,11 +227,7 @@ impl VortexFile { /// Row-count-aware pruning predicates are evaluated with the file's total /// row count as their scope. pub fn can_prune(&self, filter: &Expression) -> VortexResult { - let Some((stats, fields)) = self - .footer - .statistics() - .zip(self.footer.dtype().as_struct_fields_opt()) - else { + let Some(stats) = self.footer.statistics() else { return Ok(false); }; @@ -239,7 +235,6 @@ impl VortexFile { &filter.bind(self.footer.dtype())?, self.footer.row_count(), stats, - fields, &self.session, ) } diff --git a/vortex-file/src/footer/file_statistics.rs b/vortex-file/src/footer/file_statistics.rs index 4fac3ad8482..b23c21d05a1 100644 --- a/vortex-file/src/footer/file_statistics.rs +++ b/vortex-file/src/footer/file_statistics.rs @@ -12,6 +12,7 @@ use flatbuffers::FlatBufferBuilder; use flatbuffers::WIPOffset; use itertools::Itertools; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldPath; use vortex_array::stats::StatsSet; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -20,6 +21,7 @@ use vortex_flatbuffers::FlatBufferRoot; use vortex_flatbuffers::WriteFlatBuffer; use vortex_flatbuffers::array::ArrayStats; use vortex_flatbuffers::footer as fb; +use vortex_layout::layouts::file_stats::postorder_stats_layout; use vortex_session::VortexSession; /// Contains statistical information about the data in a Vortex file. @@ -33,61 +35,63 @@ pub struct FileStatistics { stats: Arc<[StatsSet]>, /// An array of `DType`s, one for each field or column in the file. dtypes: Arc<[DType]>, + /// An array of field paths, one for each field or column in the file. Parallel to `stats` and + /// `dtypes`. For files written before nested field stats, every path has depth 1 (or is the + /// root path, for a non-struct file dtype). + paths: Arc<[FieldPath]>, } impl FileStatistics { - /// Creates a new [`FileStatistics`] from the given statistics and data types. + /// Creates a new [`FileStatistics`] from the given statistics, data types, and field paths. /// /// # Panics /// - /// Panics if `stats` and `dtypes` have different lengths. - pub fn new(stats: Arc<[StatsSet]>, dtypes: Arc<[DType]>) -> Self { + /// Panics if `stats`, `dtypes`, and `paths` have different lengths. + pub fn new(stats: Arc<[StatsSet]>, dtypes: Arc<[DType]>, paths: Arc<[FieldPath]>) -> Self { assert_eq!( stats.len(), dtypes.len(), "stats and dtypes must have the same length" ); + assert_eq!( + stats.len(), + paths.len(), + "stats and paths must have the same length" + ); - Self { stats, dtypes } + Self { + stats, + dtypes, + paths, + } } /// Creates a new [`FileStatistics`] from the given statistics and file dtype. /// - /// If the [`DType`] of the file is a [`DType::Struct`], then there must be the same number of - /// stats as struct fields. Otherwise, there must be only 1 statistic. + /// `stats` must follow the post-order nested-struct layout produced by + /// [`postorder_stats_layout`] for `file_dtype`. /// /// # Panics /// /// Panics if the number of stats doesn't match the expected number based on the dtype. pub fn new_with_dtype(stats: Arc<[StatsSet]>, file_dtype: &DType) -> Self { - if let DType::Struct(struct_fields, _) = file_dtype { - assert_eq!( - stats.len(), - struct_fields.nfields(), - "stats length must match number of struct fields" - ); + let layout = postorder_stats_layout(file_dtype); + assert_eq!( + stats.len(), + layout.len(), + "stats length must match the post-order stats layout for the file dtype" + ); - let dtypes = struct_fields.fields().collect(); + let (paths, dtypes): (Vec, Vec) = layout.into_iter().unzip(); - Self { stats, dtypes } - } else { - assert_eq!( - stats.len(), - 1, - "non-struct dtype must have exactly 1 statistic" - ); - - Self { - stats, - dtypes: Arc::new([file_dtype.clone()]), - } + Self { + stats, + dtypes: dtypes.into(), + paths: paths.into(), } } /// Creates [`FileStatistics`] from a flatbuffers [`fb::FileStatistics<'a>`]. - /// - /// If the [`DType`] of the file is a [`DType::Struct`], then there must be the same number of - /// file stats in the flatbuffer. Otherwise, there must be only 1 statistic. pub fn from_flatbuffer<'a>( fb: &fb::FileStatistics<'a>, file_dtype: &DType, @@ -96,6 +100,28 @@ impl FileStatistics { let field_stats = fb.field_stats().unwrap_or_default(); let mut array_stats: Vec = field_stats.iter().collect(); + if fb.is_nested() { + let layout = postorder_stats_layout(file_dtype); + vortex_ensure_eq!(array_stats.len(), layout.len()); + + let mut stats_sets = Vec::with_capacity(array_stats.len()); + let mut dtypes = Vec::with_capacity(layout.len()); + let mut paths = Vec::with_capacity(layout.len()); + for (array_stat, (path, dtype)) in array_stats.into_iter().zip(layout) { + stats_sets.push(StatsSet::from_flatbuffer(&array_stat, &dtype, session)?); + dtypes.push(dtype); + paths.push(path); + } + + return Ok(Self { + stats: stats_sets.into(), + dtypes: dtypes.into(), + paths: paths.into(), + }); + } + + // Legacy (pre-nested-stats) layout: top-level struct fields only, or a single entry for a + // non-struct root dtype. if let DType::Struct(struct_fields, _) = file_dtype { vortex_ensure_eq!(array_stats.len(), struct_fields.nfields()); @@ -108,10 +134,16 @@ impl FileStatistics { .try_collect()?; let dtypes = struct_fields.fields().collect(); + let paths = struct_fields + .names() + .iter() + .map(|name| FieldPath::from_name(name.clone())) + .collect(); Ok(Self { stats: stats_sets, dtypes, + paths, }) } else { vortex_ensure_eq!(array_stats.len(), 1); @@ -124,6 +156,7 @@ impl FileStatistics { Ok(Self { stats: Arc::new([stats_set]), dtypes: Arc::new([file_dtype.clone()]), + paths: Arc::new([FieldPath::root()]), }) } } @@ -138,6 +171,11 @@ impl FileStatistics { &self.dtypes } + /// Returns a reference to the field paths. + pub fn paths(&self) -> &Arc<[FieldPath]> { + &self.paths + } + /// Returns the statistics and data type for a specific field. /// /// # Panics @@ -146,6 +184,14 @@ impl FileStatistics { pub fn get(&self, field_idx: usize) -> (&StatsSet, &DType) { (&self.stats[field_idx], &self.dtypes[field_idx]) } + + /// Returns the statistics and data type for the field at the given path, if present. + pub fn get_by_path(&self, path: &FieldPath) -> Option<(&StatsSet, &DType)> { + self.paths + .iter() + .position(|p| p == path) + .map(|idx| (&self.stats[idx], &self.dtypes[idx])) + } } impl<'a> IntoIterator for &'a FileStatistics { @@ -177,7 +223,102 @@ impl WriteFlatBuffer for FileStatistics { fbb, &fb::FileStatisticsArgs { field_stats: Some(field_stats), + is_nested: true, }, )) } } + +#[cfg(test)] +mod tests { + use flatbuffers::FlatBufferBuilder; + use vortex_array::array_session; + use vortex_array::dtype::FieldPath; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::stats::Precision; + use vortex_array::expr::stats::Stat; + use vortex_array::scalar::ScalarValue; + use vortex_flatbuffers::WriteFlatBuffer; + use vortex_flatbuffers::WriteFlatBufferExt; + + use super::*; + + fn i32_dtype() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) + } + + #[test] + fn nested_round_trip_resolves_by_path() -> VortexResult<()> { + let session = array_session(); + let inner = DType::struct_([("b", i32_dtype())], Nullability::Nullable); + let file_dtype = DType::struct_([("a", inner)], Nullability::NonNullable); + + // Layout: [a.b, a] (a's own null-count entry trails its child). + let mut b_stats = StatsSet::default(); + b_stats.set(Stat::Min, Precision::exact(ScalarValue::from(1i32))); + let mut a_stats = StatsSet::default(); + a_stats.set(Stat::NullCount, Precision::exact(ScalarValue::from(1u64))); + + let file_stats = FileStatistics::new_with_dtype(Arc::from([b_stats, a_stats]), &file_dtype); + + let bytes = file_stats.write_flatbuffer_bytes()?; + let fb = flatbuffers::root::(bytes.as_ref()) + .vortex_expect("valid flatbuffer"); + assert!(fb.is_nested()); + + let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &session)?; + + let (b, _) = read_back + .get_by_path(&FieldPath::from_name("a").push("b")) + .expect("a.b stats"); + assert_eq!(b.get(Stat::Min).as_exact(), Some(ScalarValue::from(1i32))); + + let (a, _) = read_back + .get_by_path(&FieldPath::from_name("a")) + .expect("a's own null-count stats"); + assert_eq!( + a.get(Stat::NullCount).as_exact(), + Some(ScalarValue::from(1u64)) + ); + + assert!(read_back.get_by_path(&FieldPath::root()).is_none()); + + Ok(()) + } + + #[test] + fn legacy_non_nested_footer_still_parses() -> VortexResult<()> { + // Simulates a footer written before nested field stats existed: `is_nested` is absent + // (defaults to false), and `field_stats` holds one entry per top-level struct field. + let session = array_session(); + let file_dtype = DType::struct_([("col", i32_dtype())], Nullability::NonNullable); + + let mut stats = StatsSet::default(); + stats.set(Stat::Min, Precision::exact(ScalarValue::from(7i32))); + + let mut fbb = FlatBufferBuilder::new(); + let array_stats = stats.write_flatbuffer(&mut fbb)?; + let field_stats = fbb.create_vector(&[array_stats]); + let root = fb::FileStatistics::create( + &mut fbb, + &fb::FileStatisticsArgs { + field_stats: Some(field_stats), + is_nested: false, + }, + ); + fbb.finish_minimal(root); + let bytes = fbb.finished_data().to_vec(); + + let fb = flatbuffers::root::(&bytes).vortex_expect("valid flatbuffer"); + assert!(!fb.is_nested()); + + let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &session)?; + let (col, _) = read_back + .get_by_path(&FieldPath::from_name("col")) + .expect("col stats"); + assert_eq!(col.get(Stat::Min).as_exact(), Some(ScalarValue::from(7i32))); + + Ok(()) + } +} diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index df327638d00..5f0d1b1570b 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -9,7 +9,6 @@ use vortex_array::arrays::ConstantArray; use vortex_array::arrays::NullArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; -use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::bound::lit; use vortex_array::expr::stats::Stat; @@ -29,17 +28,13 @@ pub(crate) fn can_prune_file_stats( expr: &BoundExpression, row_count: u64, file_stats: &FileStatistics, - struct_fields: &StructFields, session: &VortexSession, ) -> VortexResult { let Some(pruning_expr) = expr.falsify(session)? else { return Ok(false); }; - let binder = FileStatsBinder { - file_stats, - struct_fields, - }; + let binder = FileStatsBinder { file_stats }; let pruning_expr = bind_stats(pruning_expr, &binder)?; if let Some(result) = pruning_expr.as_opt::() { @@ -62,7 +57,6 @@ pub(crate) fn can_prune_file_stats( struct FileStatsBinder<'a> { file_stats: &'a FileStatistics, - struct_fields: &'a StructFields, } impl StatBinder for FileStatsBinder<'_> { @@ -84,18 +78,10 @@ impl StatBinder for FileStatsBinder<'_> { impl FileStatsBinder<'_> { fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { - // FileStats currently only holds top-level field statistics. - if field_path.parts().len() != 1 { - return None; - } - - let field_name = field_path.parts()[0].as_name()?; - let field_idx = self.struct_fields.find(field_name)?; - let field_stats = self.file_stats.stats_sets().get(field_idx)?; + let (field_stats, field_dtype) = self.file_stats.get_by_path(field_path)?; let stat_value = field_stats.get(stat).as_exact()?; - let field_dtype = self.struct_fields.field_by_index(field_idx)?; - let stat_dtype = stat.dtype(&field_dtype)?; + let stat_dtype = stat.dtype(field_dtype)?; let stat_scalar = Scalar::try_new(stat_dtype, Some(stat_value)).ok()?; Some(lit(stat_scalar)) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..161b33efb83 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -33,6 +33,7 @@ use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::FieldPath; use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; @@ -53,11 +54,13 @@ use vortex_array::expr::lt_eq; use vortex_array::expr::or; use vortex_array::expr::root; use vortex_array::expr::select; +use vortex_array::expr::stats::Stat; use vortex_array::extension::datetime::TimeUnit; use vortex_array::extension::datetime::Timestamp; use vortex_array::extension::datetime::TimestampOptions; use vortex_array::field_path; use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::pack::Pack; use vortex_array::scalar_fn::fns::pack::PackOptions; @@ -1303,27 +1306,40 @@ async fn file_take() -> VortexResult<()> { } #[tokio::test] -#[should_panic( - expected = "FileStatsAccumulator temporarily does not support nullable top-level structs" -)] -async fn write_nullable_top_level_struct() { +async fn write_nullable_top_level_struct() -> VortexResult<()> { let ages = PrimitiveArray::from_option_iter([Some(25), Some(31), None, Some(57), None]); + let row_validity = BoolArray::from_iter([true, true, false, true, false]).into_array(); let array = StructArray::try_new( ["age"].into(), vec![ages.into_array()], 5, - Validity::AllValid, - ) - .unwrap() + Validity::Array(row_validity), + )? .into_array(); - let mut writer = vec![]; - SESSION + let mut buf = ByteBufferMut::empty(); + let summary = SESSION .write_options() - .write(&mut writer, array.to_array_stream()) - .await - .unwrap(); + .with_file_statistics(PRUNING_STATS.to_vec()) + .write(&mut buf, array.to_array_stream()) + .await?; + + // The root struct is nullable and has 2 null rows, so its own null-count entry (keyed by the + // root field path) should reflect that, in addition to the leaf `age` field's stats. + let stats = summary + .footer() + .statistics() + .expect("file statistics should be present"); + let (root_stats, _) = stats + .get_by_path(&FieldPath::root()) + .expect("root struct should have its own null-count stats entry"); + assert_eq!( + root_stats.get(Stat::NullCount).as_exact(), + Some(ScalarValue::from(2u64)) + ); + + Ok(()) } async fn round_trip( @@ -2721,6 +2737,33 @@ async fn test_can_prune_composite_predicates() -> VortexResult<()> { Ok(()) } +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_can_prune_nested_struct_field() -> VortexResult<()> { + // Regression test for vortex-data/vortex#6389: whole-file stats now cover nested struct + // fields, not just top-level ones, so `can_prune` should resolve `person.age`. + let person = StructArray::from_fields(&[("age", buffer![15i32, 18, 22, 25].into_array())])?; + let st = StructArray::try_new( + ["person"].into(), + vec![person.into_array()], + 4, + Validity::NonNullable, + )?; + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, st.into_array().to_array_stream()) + .await?; + let file = SESSION.open_options().open_buffer(buf)?; + + let age = get_item("age", col("person")); + assert!(file.can_prune(>(age.clone(), lit(30)))?); + assert!(!file.can_prune(>(age, lit(20)))?); + + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 03ad2ab88d4..6a25e39d354 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -13,7 +13,6 @@ use std::sync::Arc; use vortex_array::MaskFuture; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; -use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; use vortex_error::VortexResult; @@ -40,29 +39,16 @@ use crate::pruning::can_prune_file_stats; pub struct FileStatsLayoutReader { child: LayoutReaderRef, file_stats: FileStatistics, - struct_fields: StructFields, session: VortexSession, prune_cache: DashMap, } impl FileStatsLayoutReader { /// Creates a new `FileStatsLayoutReader` wrapping the given child reader. - /// - /// The `struct_fields` are derived from the child reader's dtype. If the dtype is not a - /// struct, the available stats will be empty and no pruning will occur. - /// - /// Pre-computes the set of available stat field paths from the struct fields and file stats. pub fn new(child: LayoutReaderRef, file_stats: FileStatistics, session: VortexSession) -> Self { - let struct_fields = child - .dtype() - .as_struct_fields_opt() - .cloned() - .unwrap_or_default(); - Self { child, file_stats, - struct_fields, session, prune_cache: Default::default(), } @@ -77,7 +63,6 @@ impl FileStatsLayoutReader { expr, self.child.row_count(), &self.file_stats, - &self.struct_fields, &self.session, ) } @@ -171,6 +156,7 @@ mod tests { use vortex_array::arrays::StructArray; use vortex_array::arrays::datetime::TemporalData; use vortex_array::dtype::DType; + use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::checked_add; @@ -217,6 +203,7 @@ mod tests { FileStatistics::new( Arc::from([stats]), Arc::from([DType::Primitive(PType::I32, Nullability::NonNullable)]), + Arc::from([FieldPath::from_name("col")]), ) } @@ -229,6 +216,7 @@ mod tests { FileStatistics::new( Arc::from([stats]), Arc::from([DType::Primitive(PType::I32, Nullability::Nullable)]), + Arc::from([FieldPath::from_name("col")]), ) } @@ -389,7 +377,11 @@ mod tests { // File-level stats: 1 null in deleted_at. let mut stats = StatsSet::default(); stats.set(Stat::NullCount, Precision::exact(ScalarValue::from(1u64))); - let file_stats = FileStatistics::new(Arc::from([stats]), Arc::from([ts_dtype])); + let file_stats = FileStatistics::new( + Arc::from([stats]), + Arc::from([ts_dtype]), + Arc::from([FieldPath::from_name("deleted_at")]), + ); let reader = FileStatsLayoutReader::new(child, file_stats, SESSION.clone()); diff --git a/vortex-flatbuffers/flatbuffers/vortex-file/footer.fbs b/vortex-flatbuffers/flatbuffers/vortex-file/footer.fbs index b123b0d5eeb..e23f820faa6 100644 --- a/vortex-flatbuffers/flatbuffers/vortex-file/footer.fbs +++ b/vortex-flatbuffers/flatbuffers/vortex-file/footer.fbs @@ -58,7 +58,15 @@ table PostscriptSegment { table FileStatistics { /// Statistics for each field in the root schema. If the root schema is not a struct, there will /// be a single entry in this array. + /// + /// When `is_nested` is true, entries follow a post-order walk of the root `DType` tree: each + /// leaf field gets an entry, and each nullable struct additionally gets a trailing null-count + /// entry inserted after its children's entries. field_stats: [ArrayStats]; + /// Whether `field_stats` follows the post-order nested-struct layout (true) or the legacy + /// top-level-fields-only layout (false). Defaults to false for backward compatibility with + /// files written before nested field stats were supported. + is_nested: bool = false; } /// The `Registry` object stores dictionary-encoded configuration for segments, diff --git a/vortex-flatbuffers/src/generated/footer.rs b/vortex-flatbuffers/src/generated/footer.rs index 62ad85542e1..d55ba59f15a 100644 --- a/vortex-flatbuffers/src/generated/footer.rs +++ b/vortex-flatbuffers/src/generated/footer.rs @@ -1,5 +1,6 @@ // automatically generated by the FlatBuffers compiler, do not modify // @generated + extern crate alloc; use crate::array::*; @@ -811,6 +812,7 @@ impl<'a> ::flatbuffers::Follow<'a> for FileStatistics<'a> { impl<'a> FileStatistics<'a> { pub const VT_FIELD_STATS: ::flatbuffers::VOffsetT = 4; + pub const VT_IS_NESTED: ::flatbuffers::VOffsetT = 6; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -823,12 +825,17 @@ impl<'a> FileStatistics<'a> { ) -> ::flatbuffers::WIPOffset> { let mut builder = FileStatisticsBuilder::new(_fbb); if let Some(x) = args.field_stats { builder.add_field_stats(x); } + builder.add_is_nested(args.is_nested); builder.finish() } /// Statistics for each field in the root schema. If the root schema is not a struct, there will /// be a single entry in this array. + /// + /// When `is_nested` is true, entries follow a post-order walk of the root `DType` tree: each + /// leaf field gets an entry, and each nullable struct additionally gets a trailing null-count + /// entry inserted after its children's entries. #[inline] pub fn field_stats(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { // Safety: @@ -836,6 +843,16 @@ impl<'a> FileStatistics<'a> { // which contains a valid value in this slot unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FileStatistics::VT_FIELD_STATS, None)} } + /// Whether `field_stats` follows the post-order nested-struct layout (true) or the legacy + /// top-level-fields-only layout (false). Defaults to false for backward compatibility with + /// files written before nested field stats were supported. + #[inline] + pub fn is_nested(&self) -> bool { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(FileStatistics::VT_IS_NESTED, Some(false)).unwrap()} + } } impl ::flatbuffers::Verifiable for FileStatistics<'_> { @@ -845,18 +862,21 @@ impl ::flatbuffers::Verifiable for FileStatistics<'_> { ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { v.visit_table(pos)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("field_stats", Self::VT_FIELD_STATS, false)? + .visit_field::("is_nested", Self::VT_IS_NESTED, false)? .finish(); Ok(()) } } pub struct FileStatisticsArgs<'a> { pub field_stats: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub is_nested: bool, } impl<'a> Default for FileStatisticsArgs<'a> { #[inline] fn default() -> Self { FileStatisticsArgs { field_stats: None, + is_nested: false, } } } @@ -871,6 +891,10 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FileStatisticsBuilder<'a, 'b, self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_FIELD_STATS, field_stats); } #[inline] + pub fn add_is_nested(&mut self, is_nested: bool) { + self.fbb_.push_slot::(FileStatistics::VT_IS_NESTED, is_nested, false); + } + #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FileStatisticsBuilder<'a, 'b, A> { let start = _fbb.start_table(); FileStatisticsBuilder { @@ -889,6 +913,7 @@ impl ::core::fmt::Debug for FileStatistics<'_> { fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { let mut ds = f.debug_struct("FileStatistics"); ds.field("field_stats", &self.field_stats()); + ds.field("is_nested", &self.is_nested()); ds.finish() } } diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index dd0c4d6c13c..83226cbe856 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -19,6 +19,7 @@ use vortex_array::builders::BoolBuilder; use vortex_array::builders::builder_with_capacity; use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; +use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::stats::Precision; @@ -33,7 +34,6 @@ use vortex_buffer::BufferString; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; -use vortex_error::vortex_panic; use vortex_session::VortexSession; use crate::layouts::zoned::MAX_IS_TRUNCATED; @@ -417,13 +417,146 @@ impl StatsArrayBuilder for TruncatedMinBinaryStatsBuilder Vec<(FieldPath, DType)> { + let mut out = Vec::new(); + postorder_stats_layout_into(dtype, FieldPath::root(), &mut out); + out +} + +fn postorder_stats_layout_into(dtype: &DType, path: FieldPath, out: &mut Vec<(FieldPath, DType)>) { + match dtype.as_struct_fields_opt() { + Some(struct_fields) => { + for (name, field_dtype) in struct_fields.names().iter().zip(struct_fields.fields()) { + postorder_stats_layout_into(&field_dtype, path.clone().push(name.clone()), out); + } + if dtype.nullability() == Nullability::Nullable { + out.push((path, dtype.clone())); + } + } + None if !supports_file_stats(dtype) => {} + None => out.push((path, dtype.clone())), + } +} + +/// A node in the tree of accumulators mirroring [`postorder_stats_layout`]'s walk of a `DType`. +enum StatsNode { + /// An opaque leaf: a non-struct dtype, including `List`/`FixedSizeList` (not recursed into). + Leaf(StatsAccumulator), + /// A dtype that does not support file stats (e.g. [`DType::Variant`]); contributes no entries. + Skipped, + Struct { + /// One child per struct field, in declaration order. + children: Vec<(FieldName, StatsNode)>, + /// Accumulates the struct's own null count. `Some` iff the struct itself is nullable. + null_count: Option, + }, +} + +impl StatsNode { + fn build(dtype: &DType, stats: &[Stat], max_variable_length_statistics_size: usize) -> Self { + match dtype.as_struct_fields_opt() { + Some(struct_fields) => { + let children = struct_fields + .names() + .iter() + .zip(struct_fields.fields()) + .map(|(name, field_dtype)| { + ( + name.clone(), + Self::build(&field_dtype, stats, max_variable_length_statistics_size), + ) + }) + .collect(); + let null_count = (dtype.nullability() == Nullability::Nullable).then(|| { + StatsAccumulator::new(dtype, stats, max_variable_length_statistics_size) + }); + Self::Struct { + children, + null_count, + } + } + None if !supports_file_stats(dtype) => Self::Skipped, + None => Self::Leaf(StatsAccumulator::new( + dtype, + stats, + max_variable_length_statistics_size, + )), + } + } + + fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + match self { + Self::Skipped => Ok(()), + Self::Leaf(acc) => acc.push_chunk(array, ctx), + Self::Struct { + children, + null_count, + } => { + // The struct's own `ArrayRef` already carries the validity needed to compute its + // null count, so we push it directly rather than building a synthetic array. + if let Some(null_count) = null_count { + null_count.push_chunk(array, ctx)?; + } + let struct_array = array.clone().execute::(ctx)?; + for ((_, child), field) in children + .iter_mut() + .zip_eq(struct_array.iter_unmasked_fields()) + { + child.push_chunk(field, ctx)?; + } + Ok(()) + } + } + } + + /// Appends this node's `StatsSet`s, in the same post-order as [`postorder_stats_layout`]. + fn collect_stats_sets( + &mut self, + stats: &[Stat], + ctx: &mut ExecutionCtx, + out: &mut Vec, + ) -> VortexResult<()> { + match self { + Self::Skipped => Ok(()), + Self::Leaf(acc) => { + out.push(acc.as_stats_set(stats, ctx)?); + Ok(()) + } + Self::Struct { + children, + null_count, + } => { + for (_, child) in children.iter_mut() { + child.collect_stats_sets(stats, ctx, out)?; + } + if let Some(null_count) = null_count { + out.push(null_count.as_stats_set(stats, ctx)?); + } + Ok(()) + } + } + } +} + +/// An array stream processor that computes aggregate statistics for every field, recursing into +/// nested (possibly nullable) structs. See [`postorder_stats_layout`] for the entry ordering. #[derive(Clone)] pub struct FileStatsAccumulator { stats: Arc<[Stat]>, - accumulators: Arc>>, + root: Arc>, ctx: Arc>, } @@ -434,38 +567,15 @@ impl FileStatsAccumulator { max_variable_length_statistics_size: usize, session: &VortexSession, ) -> Self { - let accumulators = Arc::new(Mutex::new(match dtype.as_struct_fields_opt() { - Some(struct_dtype) => { - if dtype.nullability() == Nullability::Nullable { - // top level dtype could be nullable, but we don't support it yet - vortex_panic!( - "FileStatsAccumulator temporarily does not support nullable top-level structs, got: {}. Use Validity::NonNullable", - dtype - ); - } - - struct_dtype - .fields() - .map(|field_dtype| { - StatsAccumulator::new( - &field_dtype, - &stats, - max_variable_length_statistics_size, - ) - }) - .collect() - } - None => [StatsAccumulator::new( - dtype, - &stats, - max_variable_length_statistics_size, - )] - .into(), - })); + let root = Arc::new(Mutex::new(StatsNode::build( + dtype, + &stats, + max_variable_length_statistics_size, + ))); Self { stats, - accumulators, + root, ctx: Arc::new(Mutex::new(session.create_execution_ctx())), } } @@ -476,32 +586,18 @@ impl FileStatsAccumulator { ) -> VortexResult<(SequenceId, ArrayRef)> { let (sequence_id, chunk) = chunk?; let mut ctx = self.ctx.lock(); - if chunk.dtype().is_struct() { - let struct_chunk = chunk.clone().execute::(&mut ctx)?; - for (acc, field) in self - .accumulators - .lock() - .iter_mut() - .zip_eq(struct_chunk.iter_unmasked_fields()) - { - acc.push_chunk(field, &mut ctx)?; - } - } else { - self.accumulators.lock()[0].push_chunk(&chunk, &mut ctx)?; - } + self.root.lock().push_chunk(&chunk, &mut ctx)?; Ok((sequence_id, chunk)) } pub fn stats_sets(&self) -> Vec { let mut ctx = self.ctx.lock(); - self.accumulators + let mut out = Vec::new(); + self.root .lock() - .iter_mut() - .map(|acc| { - acc.as_stats_set(&self.stats, &mut ctx) - .vortex_expect("as_stats_table should not fail") - }) - .collect() + .collect_stats_sets(&self.stats, &mut ctx, &mut out) + .vortex_expect("collect_stats_sets should not fail"); + out } } @@ -513,6 +609,9 @@ mod tests { use vortex_array::arrays::BoolArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::builders::VarBinViewBuilder; + use vortex_array::dtype::FieldNames; + use vortex_array::scalar::PValue; + use vortex_array::scalar::ScalarValue; use vortex_buffer::BitBuffer; use vortex_buffer::buffer; @@ -604,4 +703,121 @@ mod tests { &[Stat::Max.name(), Stat::Min.name(), Stat::Sum.name()] ); } + + fn i32_dtype() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) + } + + #[test] + fn postorder_layout_flat_struct() { + let dtype = DType::struct_( + [ + ("a", i32_dtype()), + ("b", DType::Bool(Nullability::Nullable)), + ], + Nullability::NonNullable, + ); + let layout = postorder_stats_layout(&dtype); + assert_eq!( + layout, + vec![ + (FieldPath::from_name("a"), i32_dtype()), + ( + FieldPath::from_name("b"), + DType::Bool(Nullability::Nullable) + ), + ] + ); + } + + #[test] + fn postorder_layout_nested_nullable_struct_trails_children() { + let inner = DType::struct_([("b", i32_dtype())], Nullability::Nullable); + let dtype = DType::struct_([("a", inner.clone())], Nullability::NonNullable); + + let layout = postorder_stats_layout(&dtype); + assert_eq!( + layout, + vec![ + (FieldPath::from_name("a").push("b"), i32_dtype()), + (FieldPath::from_name("a"), inner), + ] + ); + } + + #[test] + fn postorder_layout_non_nullable_nested_struct_has_no_own_entry() { + let inner = DType::struct_([("b", i32_dtype())], Nullability::NonNullable); + let dtype = DType::struct_([("a", inner)], Nullability::NonNullable); + + let layout = postorder_stats_layout(&dtype); + assert_eq!( + layout, + vec![(FieldPath::from_name("a").push("b"), i32_dtype())] + ); + } + + #[test] + fn postorder_layout_nullable_root_struct_gets_trailing_root_entry() { + let dtype = DType::struct_([("a", i32_dtype())], Nullability::Nullable); + + let layout = postorder_stats_layout(&dtype); + assert_eq!( + layout, + vec![ + (FieldPath::from_name("a"), i32_dtype()), + (FieldPath::root(), dtype), + ] + ); + } + + #[test] + fn postorder_layout_list_field_is_opaque_leaf() { + let list_dtype = DType::list(i32_dtype(), Nullability::NonNullable); + let dtype = DType::struct_([("a", list_dtype.clone())], Nullability::NonNullable); + + let layout = postorder_stats_layout(&dtype); + assert_eq!(layout, vec![(FieldPath::from_name("a"), list_dtype)]); + } + + #[test] + fn postorder_layout_variant_field_is_skipped() { + let dtype = DType::struct_( + [ + ("a", i32_dtype()), + ("v", DType::Variant(Nullability::NonNullable)), + ], + Nullability::NonNullable, + ); + + let layout = postorder_stats_layout(&dtype); + assert_eq!(layout, vec![(FieldPath::from_name("a"), i32_dtype())]); + } + + #[test] + fn nested_nullable_struct_accumulates_its_own_null_count() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + + let b = buffer![1i32, 2, 3].into_array(); + let inner_validity = + Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let inner = StructArray::new(FieldNames::from(["b"]), [b], 3, inner_validity).into_array(); + let outer = StructArray::new(FieldNames::from(["a"]), [inner], 3, Validity::NonNullable) + .into_array(); + + let requested = [Stat::NullCount, Stat::Min, Stat::Max]; + let mut node = StatsNode::build(outer.dtype(), &requested, 1024); + node.push_chunk(&outer, &mut ctx)?; + + let mut stats_sets = Vec::new(); + node.collect_stats_sets(&requested, &mut ctx, &mut stats_sets)?; + + // `a.b`'s stats come first (post-order), then `a`'s own null-count entry. + assert_eq!(stats_sets.len(), 2); + assert_eq!( + stats_sets[1].get(Stat::NullCount).as_exact(), + Some(ScalarValue::Primitive(PValue::U64(1))) + ); + Ok(()) + } } From 10a8a45045ede406e64e6fbf7b7c8941d74614d8 Mon Sep 17 00:00:00 2001 From: Thor Date: Wed, 2 Sep 2026 15:31:11 -0500 Subject: [PATCH 2/3] Unify file_stats write-time accumulation onto AggregateFnRef Migrate Min/Max/Sum/NullCount/NaNCount/UncompressedSizeInBytes onto the same persistent AggregateFnRef/Accumulator framework the zoned layout's zone-map builder uses, removing the duplicate Stat-enum-keyed accumulation and its now-dead StatNameArrayBuilder. Utf8/Binary min/max truncation stays on the bespoke Precision-tracked path, since it needs per-value exactness that BoundedMax/BoundedMin can't express. Signed-off-by: "Thor" --- vortex-layout/src/layouts/file_stats.rs | 268 ++++++++++++++---------- 1 file changed, 156 insertions(+), 112 deletions(-) diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 83226cbe856..9ac96107b4f 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -11,7 +11,8 @@ use parking_lot::Mutex; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::VortexSessionExecute; -use vortex_array::aggregate_fn::fns::sum::sum; +use vortex_array::aggregate_fn::AccumulatorRef; +use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::builders::ArrayBuilder; @@ -21,7 +22,6 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldName; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; use vortex_array::expr::stats::Precision; use vortex_array::expr::stats::Stat; use vortex_array::scalar::Scalar; @@ -67,7 +67,12 @@ pub fn accumulate_stats( /// Accumulates write-time statistics for a single file column. struct StatsAccumulator { - builders: Vec>, + aggregates: Vec<(AggregateFnRef, AccumulatorRef)>, + /// Bespoke truncating builders for Utf8/Binary Min/Max, which track per-value exactness + /// (`Precision::Exact` vs `Inexact`) that the generic `BoundedMax`/`BoundedMin` aggregate fns + /// can't express (they treat truncatability as a static, config-level property, not a + /// per-value fact). + truncated: Vec>, length: usize, } @@ -75,33 +80,51 @@ impl StatsAccumulator { fn new(dtype: &DType, stats: &[Stat], max_variable_length_statistics_size: usize) -> Self { if !supports_file_stats(dtype) { return Self { - builders: Vec::new(), + aggregates: Vec::new(), + truncated: Vec::new(), length: 0, }; } - let builders = stats - .iter() - .filter_map(|&stat| { - stat.dtype(dtype).map(|stat_dtype| { - stats_builder_with_capacity( + let is_varlen = is_varlen_dtype(dtype); + let mut aggregates = Vec::new(); + let mut truncated: Vec> = Vec::new(); + + for &stat in stats { + if is_varlen && matches!(stat, Stat::Min | Stat::Max) { + if let Some(stat_dtype) = stat.dtype(dtype) { + truncated.push(stats_builder_with_capacity( stat, &stat_dtype.as_nullable(), 1024, max_variable_length_statistics_size, - ) - }) - }) - .collect::>(); + )); + } + continue; + } + + // `IsConstant`/`IsSorted`/`IsStrictSorted` have no aggregate-fn equivalent, and a + // dtype that doesn't support a given aggregate simply fails to build an accumulator — + // both cases are silently skipped, matching this stat's absence from the result. + if let Some(aggregate_fn) = stat.aggregate_fn() + && let Ok(accumulator) = aggregate_fn.accumulator(dtype) + { + aggregates.push((aggregate_fn, accumulator)); + } + } Self { - builders, + aggregates, + truncated, length: 0, } } fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { - for builder in &mut self.builders { + for (_, accumulator) in &mut self.aggregates { + accumulator.accumulate(array, ctx)?; + } + for builder in &mut self.truncated { if let Some(value) = array.statistics().compute_stat(builder.stat(), ctx)? { builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?; } else { @@ -112,12 +135,15 @@ impl StatsAccumulator { Ok(()) } - fn as_array(&mut self, ctx: &mut ExecutionCtx) -> VortexResult> { + /// Builds the intermediate per-chunk table backing the `truncated` (Utf8/Binary Min/Max) + /// builders, from which the file-wide truncated aggregate is re-derived. Returns `None` if + /// there are no such builders, or all their columns ended up all-null. + fn truncated_array(&mut self, ctx: &mut ExecutionCtx) -> VortexResult> { let mut names = Vec::new(); let mut fields = Vec::new(); for builder in self - .builders + .truncated .iter_mut() // We sort the stats so the DType is deterministic based on which stats are present. .sorted_unstable_by_key(|builder| builder.stat()) @@ -141,46 +167,45 @@ impl StatsAccumulator { } /// Returns an aggregated stats set for the table. - fn as_stats_set(&mut self, stats: &[Stat], ctx: &mut ExecutionCtx) -> VortexResult { + fn as_stats_set(&mut self, ctx: &mut ExecutionCtx) -> VortexResult { let mut stats_set = StatsSet::default(); - let Some(stats_table) = self.as_array(ctx)? else { + + for (aggregate_fn, accumulator) in &self.aggregates { + if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) + && let Some(v) = accumulator.final_scalar()?.into_value() + { + stats_set.set(stat, Precision::exact(v)); + } + } + + let Some(stats_table) = self.truncated_array(ctx)? else { return Ok(stats_set); }; - for &stat in stats { + for builder in &self.truncated { + let stat = builder.stat(); let Some(values) = stats_table.unmasked_field_by_name_opt(stat.name()) else { continue; }; - match stat { - Stat::Max if is_varlen_dtype(values.dtype()) && !values.all_valid(ctx)? => { - // A null truncated varlen max can mean either an empty chunk or no finite - // upper bound, so aggregating by skipping nulls would be unsound. - continue; - } - Stat::Min | Stat::Max | Stat::Sum => { - if let Some(s) = values.statistics().compute_stat(stat, ctx)? - && let Some(v) = s.into_value() - { - let precision = if stat_was_truncated(&stats_table, stat, ctx)? { - Precision::inexact(v) - } else { - Precision::exact(v) - }; - stats_set.set(stat, precision) - } - } - Stat::NullCount | Stat::NaNCount | Stat::UncompressedSizeInBytes => { - if let Some(sum_value) = sum(values, ctx)? - .cast(&DType::Primitive(PType::U64, Nullability::Nullable))? - .into_value() - { - stats_set.set(stat, Precision::exact(sum_value)); - } - } - Stat::IsConstant | Stat::IsSorted | Stat::IsStrictSorted => {} + if stat == Stat::Max && !values.all_valid(ctx)? { + // A null truncated varlen max can mean either an empty chunk or no finite + // upper bound, so aggregating by skipping nulls would be unsound. + continue; + } + + if let Some(s) = values.statistics().compute_stat(stat, ctx)? + && let Some(v) = s.into_value() + { + let precision = if stat_was_truncated(&stats_table, stat, ctx)? { + Precision::inexact(v) + } else { + Precision::exact(v) + }; + stats_set.set(stat, precision); } } + Ok(stats_set) } } @@ -213,6 +238,12 @@ fn is_varlen_dtype(dtype: &DType) -> bool { matches!(dtype, DType::Utf8(_) | DType::Binary(_)) } +/// Builds a bespoke truncating builder for Utf8/Binary Min/Max. +/// +/// # Panics +/// +/// Panics if `stat` is not `Min` or `Max`, or `dtype` is not `Utf8`/`Binary` — callers only +/// reach this for varlen Min/Max (see [`StatsAccumulator::new`]). fn stats_builder_with_capacity( stat: Stat, dtype: &DType, @@ -220,34 +251,36 @@ fn stats_builder_with_capacity( max_length: usize, ) -> Box { let values_builder = builder_with_capacity(dtype, capacity); - match stat { - Stat::Max => match dtype { - DType::Utf8(_) => Box::new(TruncatedMaxBinaryStatsBuilder::::new( + match (stat, dtype) { + (Stat::Max, DType::Utf8(_)) => { + Box::new(TruncatedMaxBinaryStatsBuilder::::new( values_builder, BoolBuilder::with_capacity(Nullability::NonNullable, capacity), max_length, - )), - DType::Binary(_) => Box::new(TruncatedMaxBinaryStatsBuilder::::new( + )) + } + (Stat::Max, DType::Binary(_)) => { + Box::new(TruncatedMaxBinaryStatsBuilder::::new( values_builder, BoolBuilder::with_capacity(Nullability::NonNullable, capacity), max_length, - )), - _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), - }, - Stat::Min => match dtype { - DType::Utf8(_) => Box::new(TruncatedMinBinaryStatsBuilder::::new( + )) + } + (Stat::Min, DType::Utf8(_)) => { + Box::new(TruncatedMinBinaryStatsBuilder::::new( values_builder, BoolBuilder::with_capacity(Nullability::NonNullable, capacity), max_length, - )), - DType::Binary(_) => Box::new(TruncatedMinBinaryStatsBuilder::::new( + )) + } + (Stat::Min, DType::Binary(_)) => { + Box::new(TruncatedMinBinaryStatsBuilder::::new( values_builder, BoolBuilder::with_capacity(Nullability::NonNullable, capacity), max_length, - )), - _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), - }, - _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), + )) + } + _ => unreachable!("stats_builder_with_capacity is only called for varlen Min/Max"), } } @@ -273,38 +306,6 @@ trait StatsArrayBuilder: Send { fn finish(&mut self) -> NamedArrays; } -struct StatNameArrayBuilder { - stat: Stat, - builder: Box, -} - -impl StatNameArrayBuilder { - fn new(stat: Stat, builder: Box) -> Self { - Self { stat, builder } - } -} - -impl StatsArrayBuilder for StatNameArrayBuilder { - fn stat(&self) -> Stat { - self.stat - } - - fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> { - self.builder.append_scalar(&value) - } - - fn append_null(&mut self) { - self.builder.append_null() - } - - fn finish(&mut self) -> NamedArrays { - NamedArrays { - names: vec![self.stat.name().into()], - arrays: vec![self.builder.finish()], - } - } -} - struct TruncatedMaxBinaryStatsBuilder { values: Box, is_truncated: BoolBuilder, @@ -525,14 +526,13 @@ impl StatsNode { /// Appends this node's `StatsSet`s, in the same post-order as [`postorder_stats_layout`]. fn collect_stats_sets( &mut self, - stats: &[Stat], ctx: &mut ExecutionCtx, out: &mut Vec, ) -> VortexResult<()> { match self { Self::Skipped => Ok(()), Self::Leaf(acc) => { - out.push(acc.as_stats_set(stats, ctx)?); + out.push(acc.as_stats_set(ctx)?); Ok(()) } Self::Struct { @@ -540,10 +540,10 @@ impl StatsNode { null_count, } => { for (_, child) in children.iter_mut() { - child.collect_stats_sets(stats, ctx, out)?; + child.collect_stats_sets(ctx, out)?; } if let Some(null_count) = null_count { - out.push(null_count.as_stats_set(stats, ctx)?); + out.push(null_count.as_stats_set(ctx)?); } Ok(()) } @@ -555,7 +555,6 @@ impl StatsNode { /// nested (possibly nullable) structs. See [`postorder_stats_layout`] for the entry ordering. #[derive(Clone)] pub struct FileStatsAccumulator { - stats: Arc<[Stat]>, root: Arc>, ctx: Arc>, } @@ -574,7 +573,6 @@ impl FileStatsAccumulator { ))); Self { - stats, root, ctx: Arc::new(Mutex::new(session.create_execution_ctx())), } @@ -595,7 +593,7 @@ impl FileStatsAccumulator { let mut out = Vec::new(); self.root .lock() - .collect_stats_sets(&self.stats, &mut ctx, &mut out) + .collect_stats_sets(&mut ctx, &mut out) .vortex_expect("collect_stats_sets should not fail"); out } @@ -610,6 +608,7 @@ mod tests { use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::builders::VarBinViewBuilder; use vortex_array::dtype::FieldNames; + use vortex_array::dtype::PType; use vortex_array::scalar::PValue; use vortex_array::scalar::ScalarValue; use vortex_buffer::BitBuffer; @@ -635,7 +634,7 @@ mod tests { acc.push_chunk(&builder2.finish(), &mut ctx) .vortex_expect("push_chunk should succeed for test data"); let stats_table = acc - .as_array(&mut ctx) + .truncated_array(&mut ctx) .unwrap() .expect("Must have stats table"); assert_eq!( @@ -680,7 +679,7 @@ mod tests { .vortex_expect("push_chunk should succeed for test data"); let stats = acc - .as_stats_set(&[Stat::Max, Stat::Min], &mut ctx) + .as_stats_set(&mut ctx) .vortex_expect("as_stats_set should succeed for test data"); assert!(matches!(stats.get(Stat::Min), Precision::Inexact(_))); @@ -688,20 +687,65 @@ mod tests { } #[test] - fn fixed_width_stats_omit_is_truncated_columns() { + fn fixed_width_stats_omit_is_truncated_columns() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = buffer![0, 1, 2].into_array(); let mut acc = StatsAccumulator::new(array.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); acc.push_chunk(&array, &mut ctx) .vortex_expect("push_chunk should succeed for test array"); - let stats_table = acc - .as_array(&mut ctx) - .unwrap() - .expect("Must have stats table"); + + // Fixed-width Min/Max/Sum are AggRef-backed now, so there's no intermediate + // truncated-stats table at all — they resolve directly through `as_stats_set`. + assert!(acc.truncated_array(&mut ctx)?.is_none()); + + let stats = acc.as_stats_set(&mut ctx)?; assert_eq!( - stats_table.names().as_ref(), - &[Stat::Max.name(), Stat::Min.name(), Stat::Sum.name()] + stats.get(Stat::Max).as_exact(), + Some(ScalarValue::from(2i32)) + ); + assert_eq!( + stats.get(Stat::Min).as_exact(), + Some(ScalarValue::from(0i32)) + ); + assert_eq!( + stats.get(Stat::Sum).as_exact(), + Some(ScalarValue::from(3i64)) + ); + Ok(()) + } + + #[test] + fn aggregates_persist_across_multiple_chunks() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let dtype = i32_dtype(); + let mut acc = StatsAccumulator::new( + &dtype, + &[Stat::Max, Stat::Min, Stat::Sum, Stat::NullCount], + 12, + ); + + acc.push_chunk(&buffer![0, 5, 2].into_array(), &mut ctx)?; + acc.push_chunk(&buffer![7, 1, 3].into_array(), &mut ctx)?; + acc.push_chunk(&buffer![-4, 9].into_array(), &mut ctx)?; + + let stats = acc.as_stats_set(&mut ctx)?; + assert_eq!( + stats.get(Stat::Max).as_exact(), + Some(ScalarValue::from(9i32)) + ); + assert_eq!( + stats.get(Stat::Min).as_exact(), + Some(ScalarValue::from(-4i32)) ); + assert_eq!( + stats.get(Stat::Sum).as_exact(), + Some(ScalarValue::from(23i64)) + ); + assert_eq!( + stats.get(Stat::NullCount).as_exact(), + Some(ScalarValue::from(0u64)) + ); + Ok(()) } fn i32_dtype() -> DType { @@ -810,7 +854,7 @@ mod tests { node.push_chunk(&outer, &mut ctx)?; let mut stats_sets = Vec::new(); - node.collect_stats_sets(&requested, &mut ctx, &mut stats_sets)?; + node.collect_stats_sets(&mut ctx, &mut stats_sets)?; // `a.b`'s stats come first (post-order), then `a`'s own null-count entry. assert_eq!(stats_sets.len(), 2); From d29bb2a8fc2aa11cbdd2b933f9b71cb961757cb7 Mon Sep 17 00:00:00 2001 From: Thor Date: Fri, 4 Sep 2026 12:51:58 -0500 Subject: [PATCH 3/3] Add tests for three-level nested struct file stats Cover arbitrary-depth nesting, not just the single level exercised by the existing tests: postorder layout with mixed nullability across three levels, StatsNode accumulation contributing a null-count entry at each nullable level, and whole-file pruning resolving a three-level field path (x.y.z). Signed-off-by: "Thor" --- vortex-file/src/tests.rs | 33 +++++++++++ vortex-layout/src/layouts/file_stats.rs | 75 +++++++++++++++++++++++++ 2 files changed, 108 insertions(+) diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 161b33efb83..7dfd526fa80 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2764,6 +2764,39 @@ async fn test_can_prune_nested_struct_field() -> VortexResult<()> { Ok(()) } +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_can_prune_three_level_nested_struct_field() -> VortexResult<()> { + // Regression test for vortex-data/vortex#6389: whole-file stats resolve field paths at + // arbitrary nesting depth, not just one level. + let struct_z = StructArray::from_fields(&[("z", buffer![15i32, 18, 22, 25].into_array())])?; + let struct_y = StructArray::try_new( + ["y"].into(), + vec![struct_z.into_array()], + 4, + Validity::NonNullable, + )?; + let struct_x = StructArray::try_new( + ["x"].into(), + vec![struct_y.into_array()], + 4, + Validity::NonNullable, + )?; + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, struct_x.into_array().to_array_stream()) + .await?; + let file = SESSION.open_options().open_buffer(buf)?; + + let z_field = get_item("z", get_item("y", col("x"))); + assert!(file.can_prune(>(z_field.clone(), lit(30)))?); + assert!(!file.can_prune(>(z_field, lit(20)))?); + + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 9ac96107b4f..574c1dfd6fe 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -838,6 +838,81 @@ mod tests { assert_eq!(layout, vec![(FieldPath::from_name("a"), i32_dtype())]); } + #[test] + fn postorder_layout_three_level_nesting_with_mixed_nullability() { + // root (nullable) -> a (non-nullable) -> b (nullable) -> c (leaf). Exercises composition + // across more than one level, including a non-nullable struct sandwiched between two + // nullable ones, which should get no entry of its own. + let b_dtype = DType::struct_([("c", i32_dtype())], Nullability::Nullable); + let a_dtype = DType::struct_([("b", b_dtype.clone())], Nullability::NonNullable); + let root_dtype = DType::struct_([("a", a_dtype)], Nullability::Nullable); + + let layout = postorder_stats_layout(&root_dtype); + assert_eq!( + layout, + vec![ + (FieldPath::from_name("a").push("b").push("c"), i32_dtype()), + (FieldPath::from_name("a").push("b"), b_dtype), + (FieldPath::root(), root_dtype), + ] + ); + } + + #[test] + fn three_level_nested_struct_accumulates_stats_at_each_level() -> VortexResult<()> { + // Same shape as `postorder_layout_three_level_nesting_with_mixed_nullability`, but + // exercises actual accumulation: each nullable level along the path should independently + // contribute its own null-count entry, in post-order. + let mut ctx = array_session().create_execution_ctx(); + + let leaf_c = buffer![10i32, 20, 30].into_array(); + let b_validity = Validity::Array(BoolArray::from_iter([true, false, true]).into_array()); + let struct_b = + StructArray::new(FieldNames::from(["c"]), [leaf_c], 3, b_validity).into_array(); + let struct_a = StructArray::new( + FieldNames::from(["b"]), + [struct_b], + 3, + Validity::NonNullable, + ) + .into_array(); + let root_validity = Validity::Array(BoolArray::from_iter([true, true, false]).into_array()); + let root = + StructArray::new(FieldNames::from(["a"]), [struct_a], 3, root_validity).into_array(); + + let requested = [Stat::NullCount, Stat::Min, Stat::Max]; + let mut node = StatsNode::build(root.dtype(), &requested, 1024); + node.push_chunk(&root, &mut ctx)?; + + let mut stats_sets = Vec::new(); + node.collect_stats_sets(&mut ctx, &mut stats_sets)?; + + // `a.b.c`'s own stats come first (post-order), then `a.b`'s null-count entry, then the + // root's own null-count entry. + assert_eq!(stats_sets.len(), 3); + assert_eq!( + stats_sets[0].get(Stat::Min).as_exact(), + Some(ScalarValue::from(10i32)) + ); + assert_eq!( + stats_sets[0].get(Stat::Max).as_exact(), + Some(ScalarValue::from(30i32)) + ); + assert_eq!( + stats_sets[0].get(Stat::NullCount).as_exact(), + Some(ScalarValue::from(0u64)) + ); + assert_eq!( + stats_sets[1].get(Stat::NullCount).as_exact(), + Some(ScalarValue::from(1u64)) + ); + assert_eq!( + stats_sets[2].get(Stat::NullCount).as_exact(), + Some(ScalarValue::from(1u64)) + ); + Ok(()) + } + #[test] fn nested_nullable_struct_accumulates_its_own_null_count() -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx();