Skip to content

Compute legacy file stats from the aggregate stats accumulators #10846

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
robert3005 merged 1 commit into from
Oct 27, 2026
Merged

Compute legacy file stats from the aggregate stats accumulators #10846

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
210 changes: 115 additions & 95 deletions vortex-layout/src/layouts/file_stats.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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.
///
Expand Down Expand Up @@ -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::<StructArray>(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::<StructArray>(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
Expand All @@ -462,7 +447,7 @@ impl StatsNode {
out: &mut Vec<AggregateStats>,
) -> VortexResult<()> {
match self {
Self::Skipped => Ok(()),
Self::Skipped | Self::LegacyOnly(_) => Ok(()),
Self::Unrecorded { entries } => {
out.extend(iter::repeat_n(AggregateStats::default(), *entries));
Ok(())
Expand All @@ -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<AggregateStats> {
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<Mutex<StatsNode>>,
legacy: Arc<Mutex<Vec<StatsAccumulator>>>,
write_legacy_stats: bool,
writer_ctx: LayoutWriterContext,
ctx: Arc<Mutex<ExecutionCtx>>,
Expand All @@ -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())),
Expand All @@ -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::<StructArray>(&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))
}

Expand All @@ -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<Vec<AggregateStats>> {
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()
}
}
Expand All @@ -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::*;
Expand Down Expand Up @@ -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();
Expand Down
Loading