diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 4e5a175b8b2..f85d432ede3 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -37,7 +37,6 @@ use vortex_array::scalar::Scalar; use vortex_array::stats::StatsSet; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_panic; use vortex_session::VortexSession; use crate::LayoutWriterContext; @@ -345,6 +344,11 @@ enum StatsNode { Leaf(StatsAccumulator), /// A dtype that does not support file stats (e.g. [`DType::Variant`]); contributes no entries. Skipped, + /// Stats accumulated only for the legacy top-level-only layout; contributes no nested entries. + /// + /// Used for the own stats of a non-nullable top-level struct field, which the nested layout + /// doesn't record but the legacy layout does. + LegacyOnly(StatsAccumulator), /// A subtree beneath a nullable struct, which contributes `entries` entries that record no /// aggregates. /// @@ -408,50 +412,31 @@ impl StatsNode { } } - fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { - match self { - Self::Skipped | Self::Unrecorded { .. } => Ok(()), - Self::Leaf(acc) => acc.push_chunk(array, ctx), - Self::Container { .. } => { - let struct_array = array.clone().execute::(ctx)?; - self.push_struct_chunk(array, &struct_array, ctx) - } - } - } - - /// Pushes a chunk into a `Container` node given an already-executed `StructArray` for `array`. - /// - /// Lets callers that already had to execute the chunk to a `StructArray` for another purpose - /// (e.g. [`FileStatsAccumulator::process`], which also feeds the legacy per-top-level-field - /// accumulators from the same execution) avoid doing so a second time. + /// Pushes a chunk of this node's dtype into its accumulators. /// /// `Container` is only ever built for `DType::Struct` today, so every child is addressed by /// [`Field::Name`] and extracted via [`StructArrayExt::iter_unmasked_fields`]. A future `List`/ /// `Map` container would need a different extraction here (e.g. flattened elements, or /// derived per-row shape arrays), dispatched per child's [`Field`] kind. - /// - /// # Panics - /// - /// Panics if `self` is not `Self::Container`. - fn push_struct_chunk( - &mut self, - array: &ArrayRef, - struct_array: &StructArray, - ctx: &mut ExecutionCtx, - ) -> VortexResult<()> { - let Self::Container { children, own } = self else { - vortex_panic!("push_struct_chunk is only called on Container nodes"); - }; - // The container's own `ArrayRef` already carries the validity needed to compute its own - // stats (e.g. null count), so we push it directly rather than building a synthetic array. - own.push_chunk(array, ctx)?; - for ((_, child), field) in children - .iter_mut() - .zip_eq(struct_array.iter_unmasked_fields()) - { - child.push_chunk(field, ctx)?; + fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + match self { + Self::Skipped | Self::Unrecorded { .. } => Ok(()), + Self::Leaf(acc) | Self::LegacyOnly(acc) => acc.push_chunk(array, ctx), + Self::Container { children, own } => { + // The container's own `ArrayRef` already carries the validity needed to compute + // its own stats (e.g. null count), so we push it directly rather than building a + // synthetic array. + own.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(()) + } } - Ok(()) } /// Appends this node's aggregates that `allowed` permits, in the same post-order as @@ -462,7 +447,7 @@ impl StatsNode { out: &mut Vec, ) -> VortexResult<()> { match self { - Self::Skipped => Ok(()), + Self::Skipped | Self::LegacyOnly(_) => Ok(()), Self::Unrecorded { entries } => { out.extend(iter::repeat_n(AggregateStats::default(), *entries)); Ok(()) @@ -479,20 +464,28 @@ impl StatsNode { } } } + + /// Returns this node's stats over its own, undecomposed dtype, as the legacy layout records + /// them for a top-level field. + fn legacy_aggregate_stats(&self) -> VortexResult { + match self { + Self::Skipped | Self::Unrecorded { .. } => Ok(AggregateStats::default()), + Self::Leaf(acc) | Self::LegacyOnly(acc) => acc.aggregate_stats(|_| true), + Self::Container { own, .. } => own.legacy_aggregate_stats(), + } + } } /// 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. /// -/// Alongside `root`'s post-order nested layout, this also maintains a `legacy` accumulation -/// covering only the top-level struct fields (or the whole array, for a non-struct root), -/// matching the layout file stats had before nested struct support: one entry per top-level -/// field, with no recursion into nested structs. When `write_legacy_stats` is `false`, `legacy` -/// keeps that shape but records nothing, and is never accumulated into. +/// The same accumulators also produce the legacy layout covering only the top-level struct fields +/// (or the whole array, for a non-struct root), matching the layout file stats had before nested +/// struct support: one entry per top-level field, with no recursion into nested structs. When +/// `write_legacy_stats` is `false`, the legacy layout keeps that shape but records nothing. #[derive(Clone)] pub struct FileStatsAccumulator { root: Arc>, - legacy: Arc>>, write_legacy_stats: bool, writer_ctx: LayoutWriterContext, ctx: Arc>, @@ -512,32 +505,29 @@ impl FileStatsAccumulator { } } - let root = Arc::new(Mutex::new(StatsNode::build(dtype, stats.as_deref()))); + let mut root = StatsNode::build(dtype, stats.as_deref()); - // Readers that predate nested stats reject a file whose legacy layout doesn't have one entry - // per top-level field, so entries that aren't recorded are still written, empty. - let legacy = Arc::new(Mutex::new(match dtype.as_struct_fields_opt() { - Some(struct_fields) => { - // The fields of a nullable root aren't recorded, like `StatsNode::Unrecorded`. - let record = write_legacy_stats && !dtype.is_nullable(); - struct_fields - .fields() - .map(|field_dtype| { - if record { - StatsAccumulator::new(&field_dtype, stats.as_deref()) - } else { - StatsAccumulator::empty() - } - }) - .collect() + // The legacy layout records each top-level field's stats over its own dtype. The nested + // layout already accumulates those for every top-level field except a non-nullable struct, + // whose own stats are only accumulated for the legacy layout. + if write_legacy_stats + && let (StatsNode::Container { children, .. }, Some(struct_fields)) = + (&mut root, dtype.as_struct_fields_opt()) + { + for ((_, child), field_dtype) in children.iter_mut().zip_eq(struct_fields.fields()) { + if let StatsNode::Container { own, .. } = child + && matches!(**own, StatsNode::Skipped) + { + **own = StatsNode::LegacyOnly(StatsAccumulator::new( + &field_dtype, + stats.as_deref(), + )); + } } - None if write_legacy_stats => vec![StatsAccumulator::new(dtype, stats.as_deref())], - None => vec![StatsAccumulator::empty()], - })); + } Ok(Self { - root, - legacy, + root: Arc::new(Mutex::new(root)), write_legacy_stats, writer_ctx: writer_ctx.clone(), ctx: Arc::new(Mutex::new(session.create_execution_ctx())), @@ -549,29 +539,7 @@ impl FileStatsAccumulator { chunk: VortexResult<(SequenceId, ArrayRef)>, ) -> 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)?; - // In the future we should avoid computing stats twice. - if self.write_legacy_stats { - for (acc, field) in self - .legacy - .lock() - .iter_mut() - .zip_eq(struct_chunk.iter_unmasked_fields()) - { - acc.push_chunk(field, &mut ctx)?; - } - } - self.root - .lock() - .push_struct_chunk(&chunk, &struct_chunk, &mut ctx)?; - } else { - if self.write_legacy_stats { - self.legacy.lock()[0].push_chunk(&chunk, &mut ctx)?; - } - self.root.lock().push_chunk(&chunk, &mut ctx)?; - } + self.root.lock().push_chunk(&chunk, &mut self.ctx.lock())?; Ok((sequence_id, chunk)) } @@ -594,10 +562,25 @@ impl FileStatsAccumulator { /// The legacy `ArrayStats` format has a fixed slot per [`Stat`] rather than open-ended /// aggregates, so it isn't restricted by the writer context. pub fn legacy_aggregate_stats(&self) -> VortexResult> { - self.legacy - .lock() - .iter() - .map(|acc| acc.aggregate_stats(|_| true)) + let root = self.root.lock(); + // Readers that predate nested stats reject a file whose legacy layout doesn't have one + // entry per top-level field, so entries that aren't recorded are still written, empty. The + // fields of a nullable root are `StatsNode::Unrecorded`, so they come out empty too. + let fields = match &*root { + StatsNode::Container { children, .. } => { + children.iter().map(|(_, child)| child).collect() + } + node => vec![node], + }; + fields + .into_iter() + .map(|node| { + if self.write_legacy_stats { + node.legacy_aggregate_stats() + } else { + Ok(AggregateStats::default()) + } + }) .collect() } } @@ -622,6 +605,7 @@ mod tests { use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_error::vortex_panic; use vortex_utils::aliases::hash_set::HashSet; use super::*; @@ -1036,6 +1020,42 @@ mod tests { Ok(()) } + #[test] + fn legacy_aggregate_stats_record_non_nullable_struct_field() -> VortexResult<()> { + // The nested layout has no entry for a non-nullable struct's own stats, but the legacy + // layout still records them for a top-level field. + let session = VortexSession::empty(); + let inner_dtype = DType::struct_([("b", i32_dtype())], Nullability::NonNullable); + let dtype = DType::struct_([("a", inner_dtype)], Nullability::NonNullable); + + let acc = FileStatsAccumulator::try_new( + &dtype, + Some(Arc::from([agg(Stat::NullCount)])), + &session, + true, + &LayoutWriterContext::new(ArrayContext::empty()), + )?; + + let b = buffer![1i32, 2, 3].into_array(); + let inner = + StructArray::new(FieldNames::from(["b"]), [b], 3, Validity::NonNullable).into_array(); + let outer = StructArray::new(FieldNames::from(["a"]), [inner], 3, Validity::NonNullable) + .into_array(); + + let (mut ptr, _eof) = SequenceId::root().split(); + acc.process(Ok((ptr.advance(), outer)))?; + + let legacy = acc.legacy_aggregate_stats()?; + assert_eq!(legacy.len(), 1); + assert_eq!( + legacy[0].get(&agg(Stat::NullCount)), + Precision::exact(Scalar::from(0u64)) + ); + // Only `a.b` has a nested entry. + assert_eq!(acc.aggregate_stats()?.len(), 1); + Ok(()) + } + #[test] fn legacy_aggregate_stats_of_nullable_root_are_unrecorded() -> VortexResult<()> { let session = VortexSession::empty();