From 5815bdead86c07e2201491619c4e60c5db1fd150 Mon Sep 17 00:00:00 2001 From: Thor Date: Wed, 2 Sep 2026 14:49:03 -0500 Subject: [PATCH 1/4] 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" Signed-off-by: Thor --- vortex-array/src/aggregate_fn/accumulator.rs | 14 + .../src/aggregate_fn/fns/bounded_max/mod.rs | 177 +- .../src/aggregate_fn/fns/bounded_min/mod.rs | 175 +- vortex-array/src/aggregate_fn/vtable.rs | 15 + vortex-array/src/stats/mod.rs | 10 - vortex-datafusion/src/persistent/format.rs | 86 +- vortex-duckdb/cpp/copy_function.cpp | 23 +- vortex-duckdb/cpp/include/copy_function.hpp | 5 +- vortex-duckdb/cpp/include/table_function.h | 3 + vortex-duckdb/src/column_statistics.rs | 42 +- vortex-duckdb/src/copy.rs | 198 ++- vortex-duckdb/src/ffi.rs | 2 + vortex-duckdb/src/file_reader.rs | 21 +- .../flatbuffers/vortex-file/footer.fbs | 44 +- vortex-file/src/file.rs | 7 +- .../src/flatbuffers/generated/footer.rs | 386 ++++- vortex-file/src/footer/file_statistics.rs | 820 +++++++++- vortex-file/src/footer/serializer.rs | 2 +- vortex-file/src/pruning.rs | 49 +- vortex-file/src/tests.rs | 230 ++- vortex-file/src/v2/file_stats_reader.rs | 64 +- vortex-file/src/writer.rs | 51 +- vortex-jni/src/writer.rs | 60 +- vortex-layout/src/layouts/file_stats.rs | 1431 +++++++++++------ vortex-layout/src/layouts/zoned/mod.rs | 1 + .../slt/duckdb/copy_return_stats.slt | 6 +- 26 files changed, 3047 insertions(+), 875 deletions(-) diff --git a/vortex-array/src/aggregate_fn/accumulator.rs b/vortex-array/src/aggregate_fn/accumulator.rs index d2273a24cca..8e9ad228e1b 100644 --- a/vortex-array/src/aggregate_fn/accumulator.rs +++ b/vortex-array/src/aggregate_fn/accumulator.rs @@ -15,6 +15,7 @@ use crate::ExecutionCtx; use crate::aggregate_fn::AggregateDTypes; use crate::aggregate_fn::AggregateFn; use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnSatisfaction; use crate::aggregate_fn::AggregateFnVTable; use crate::aggregate_fn::session::AggregateFnSessionExt; use crate::columnar::AnyColumnar; @@ -192,6 +193,10 @@ pub trait DynAccumulator: 'static + Send { /// Whether the accumulator's result is fully determined. fn is_saturated(&self) -> bool; + /// Whether the current result can satisfy `requested`, taking the accumulated state into + /// account. See [`AggregateFnVTable::partial_can_satisfy`]. + fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction; + /// Reset the accumulator's state to the empty group. fn reset(&mut self); @@ -318,6 +323,15 @@ impl DynAccumulator for Accumulator { }) } + fn can_satisfy(&self, requested: &AggregateFnRef) -> AggregateFnSatisfaction { + let options = self.aggregate_fn.as_::(); + // Without a partial nothing was accumulated yet, so only the options can be consulted. + match &self.partial { + Some(partial) => self.vtable.partial_can_satisfy(options, partial, requested), + None => self.vtable.can_satisfy(options, requested), + } + } + fn reset(&mut self) { self.partial = None; } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 4fedcc48c74..40d1629a6fa 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cmp::Ordering; use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; @@ -33,7 +34,6 @@ use crate::dtype::DType; use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; -use crate::partial_ord::partial_max; use crate::scalar::Scalar; use crate::scalar::ScalarTruncation; use crate::scalar::upper_bound; @@ -65,27 +65,59 @@ pub struct BoundedMax; enum BoundedMaxState { Empty, - Value(Scalar), + /// The true maximum: it fit within the byte bound, so it was never truncated. + Exact(Scalar), + /// An upper bound of the true maximum, either because the maximum was truncated or because + /// it was combined from a serialized partial, which doesn't record exactness. + Inexact(Scalar), Unknown, } +impl BoundedMaxState { + fn bound(bound: Scalar, exact: bool) -> Self { + if exact { + Self::Exact(bound) + } else { + Self::Inexact(bound) + } + } +} + /// Partial accumulator state for the bounded maximum aggregate. pub struct BoundedMaxPartial { state: BoundedMaxState, } impl BoundedMaxPartial { - fn merge_bound(&mut self, max: Scalar) { - if max.is_null() { + fn merge_bound(&mut self, bound: Scalar, exact: bool) { + if bound.is_null() { return; } - self.state = match std::mem::replace(&mut self.state, BoundedMaxState::Empty) { - BoundedMaxState::Empty => BoundedMaxState::Value(max), - BoundedMaxState::Value(current) => BoundedMaxState::Value( - partial_max(max, current).vortex_expect("incomparable bounded max scalars"), - ), - BoundedMaxState::Unknown => BoundedMaxState::Unknown, + let (current, current_exact) = + match std::mem::replace(&mut self.state, BoundedMaxState::Empty) { + BoundedMaxState::Empty => { + self.state = BoundedMaxState::bound(bound, exact); + return; + } + BoundedMaxState::Unknown => { + self.state = BoundedMaxState::Unknown; + return; + } + BoundedMaxState::Exact(current) => (current, true), + BoundedMaxState::Inexact(current) => (current, false), + }; + + // The larger bound wins. It is the true maximum iff it was not truncated: every other + // value is at most its own bound, which is below the winner. On a tie, either side being + // exact means the true maximum reaches the bound. + self.state = match bound + .partial_cmp(¤t) + .vortex_expect("incomparable bounded max scalars") + { + Ordering::Greater => BoundedMaxState::bound(bound, exact), + Ordering::Less => BoundedMaxState::bound(current, current_exact), + Ordering::Equal => BoundedMaxState::bound(current, exact || current_exact), }; } @@ -96,7 +128,7 @@ impl BoundedMaxPartial { fn final_scalar(&self, args: AggregateArgs<'_, BoundedMaxOptions>) -> VortexResult { let dtype = args.return_dtype.clone(); match &self.state { - BoundedMaxState::Value(max) => max.cast(&dtype), + BoundedMaxState::Exact(bound) | BoundedMaxState::Inexact(bound) => bound.cast(&dtype), BoundedMaxState::Empty | BoundedMaxState::Unknown => Ok(Scalar::null(dtype)), } } @@ -183,6 +215,27 @@ impl AggregateFnVTable for BoundedMax { } } + fn partial_can_satisfy( + &self, + options: &Self::Options, + partial: &Self::Partial, + requested: &AggregateFnRef, + ) -> AggregateFnSatisfaction { + let satisfaction = self.can_satisfy(options, requested); + let untruncated = matches!( + partial.state, + BoundedMaxState::Empty | BoundedMaxState::Exact(_) + ); + if satisfaction == AggregateFnSatisfaction::Approximate + && untruncated + && requested.is::() + { + AggregateFnSatisfaction::Exact + } else { + satisfaction + } + } + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { supported_dtype(options, input_dtype).map(make_bounded_max_partial_dtype) } @@ -226,7 +279,8 @@ impl AggregateFnVTable for BoundedMax { } else if bound.is_null() { BoundedMaxState::Empty } else { - BoundedMaxState::Value(bound) + // Serialized partials don't record exactness. + BoundedMaxState::Inexact(bound) } }; Ok(BoundedMaxPartial { state }) @@ -240,7 +294,8 @@ impl AggregateFnVTable for BoundedMax { ) -> VortexResult { match second.state { BoundedMaxState::Empty => {} - BoundedMaxState::Value(max) => first.merge_bound(max), + BoundedMaxState::Exact(bound) => first.merge_bound(bound, true), + BoundedMaxState::Inexact(bound) => first.merge_bound(bound, false), BoundedMaxState::Unknown => first.unknown(), } Ok(first) @@ -255,10 +310,10 @@ impl AggregateFnVTable for BoundedMax { let bound_dtype = args.dtype.as_nullable(); match &partial.state { BoundedMaxState::Empty => Ok(Scalar::null(dtype)), - BoundedMaxState::Value(max) => Ok(Scalar::struct_( + BoundedMaxState::Exact(bound) | BoundedMaxState::Inexact(bound) => Ok(Scalar::struct_( dtype, vec![ - max.cast(&bound_dtype)?, + bound.cast(&bound_dtype)?, Scalar::bool(false, Nullability::NonNullable), ], )), @@ -293,7 +348,7 @@ impl AggregateFnVTable for BoundedMax { return Ok(()); }; match truncate_max(result.max, args.options.max_bytes.get())? { - Some(bound) => partial.merge_bound(bound), + Some((bound, truncated)) => partial.merge_bound(bound, !truncated), None => partial.unknown(), } Ok(()) @@ -322,22 +377,21 @@ fn supported_dtype<'a>(_options: &BoundedMaxOptions, input_dtype: &'a DType) -> .map(|_| input_dtype) } -fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult> { +/// Returns the upper bound and whether it was truncated, or `None` if no bound fits `max_bytes`. +fn truncate_max(value: Scalar, max_bytes: usize) -> VortexResult> { let nullability = value.dtype().nullability(); match value.dtype() { - DType::Utf8(_) => { - Ok( - upper_bound(BufferString::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - DType::Binary(_) => { - Ok( - upper_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - _ => Ok(Some(value)), + DType::Utf8(_) => Ok(upper_bound( + BufferString::from_scalar(value)?, + max_bytes, + nullability, + )), + DType::Binary(_) => Ok(upper_bound( + ByteBuffer::from_scalar(value)?, + max_bytes, + nullability, + )), + _ => Ok(Some((value, false))), } } @@ -584,4 +638,67 @@ mod tests { assert_eq!(roundtrip, options); Ok(()) } + + #[test] + fn bounded_max_exactly_satisfies_max_only_when_untruncated() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let requested = Max.bind(NumericalAggregateOpts::skip_nans()); + let new_acc = || { + Accumulator::try_new( + BoundedMax, + BoundedMaxOptions { + max_bytes: max_bytes(4), + }, + VarBinViewArray::from_iter_str(["x"]).dtype().clone(), + ) + }; + + // Every value fits the bound, so the bound is the true extremum. + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["b", "abc"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated value that isn't the extremum doesn't affect exactness. + acc.accumulate( + &VarBinViewArray::from_iter_str(["aaaaa"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated extremum makes the bound approximate. + acc.accumulate( + &VarBinViewArray::from_iter_str(["bbbbb"]).into_array(), + &mut ctx, + )?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + + // Exactness isn't serialized, so a winning partial combined from its scalar form is + // approximate even though it was exact where it was computed. + let mut source = new_acc()?; + source.accumulate( + &VarBinViewArray::from_iter_str(["b"]).into_array(), + &mut ctx, + )?; + assert_eq!( + source.can_satisfy(&requested), + AggregateFnSatisfaction::Exact + ); + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["a"]).into_array(), + &mut ctx, + )?; + acc.combine_partials(source.partial_scalar()?)?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index 8c4e91178b0..e3629e5cf99 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cmp::Ordering; use std::fmt::Display; use std::fmt::Formatter; use std::num::NonZeroUsize; @@ -27,7 +28,6 @@ use crate::aggregate_fn::fns::min::Min; use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::min_max::columnar_min_max; use crate::dtype::DType; -use crate::partial_ord::partial_min; use crate::scalar::Scalar; use crate::scalar::ScalarTruncation; use crate::scalar::lower_bound; @@ -51,7 +51,21 @@ pub struct BoundedMin; enum BoundedMinState { Empty, - Value(Scalar), + /// The true minimum: it fit within the byte bound, so it was never truncated. + Exact(Scalar), + /// A lower bound of the true minimum, either because the minimum was truncated or because + /// it was combined from a serialized partial, which doesn't record exactness. + Inexact(Scalar), +} + +impl BoundedMinState { + fn bound(bound: Scalar, exact: bool) -> Self { + if exact { + Self::Exact(bound) + } else { + Self::Inexact(bound) + } + } } /// Partial accumulator state for the bounded minimum aggregate. @@ -60,16 +74,31 @@ pub struct BoundedMinPartial { } impl BoundedMinPartial { - fn merge(&mut self, min: Scalar) { - if min.is_null() { + fn merge(&mut self, bound: Scalar, exact: bool) { + if bound.is_null() { return; } - self.state = match std::mem::replace(&mut self.state, BoundedMinState::Empty) { - BoundedMinState::Empty => BoundedMinState::Value(min), - BoundedMinState::Value(current) => BoundedMinState::Value( - partial_min(min, current).vortex_expect("incomparable bounded min scalars"), - ), + let (current, current_exact) = + match std::mem::replace(&mut self.state, BoundedMinState::Empty) { + BoundedMinState::Empty => { + self.state = BoundedMinState::bound(bound, exact); + return; + } + BoundedMinState::Exact(current) => (current, true), + BoundedMinState::Inexact(current) => (current, false), + }; + + // The smaller bound wins. It is the true minimum iff it was not truncated: every other + // value is at least its own bound, which is above the winner. On a tie, either side being + // exact means the true minimum reaches the bound. + self.state = match bound + .partial_cmp(¤t) + .vortex_expect("incomparable bounded min scalars") + { + Ordering::Less => BoundedMinState::bound(bound, exact), + Ordering::Greater => BoundedMinState::bound(current, current_exact), + Ordering::Equal => BoundedMinState::bound(current, exact || current_exact), }; } } @@ -137,6 +166,27 @@ impl AggregateFnVTable for BoundedMin { } } + fn partial_can_satisfy( + &self, + options: &Self::Options, + partial: &Self::Partial, + requested: &AggregateFnRef, + ) -> AggregateFnSatisfaction { + let satisfaction = self.can_satisfy(options, requested); + let untruncated = matches!( + partial.state, + BoundedMinState::Empty | BoundedMinState::Exact(_) + ); + if satisfaction == AggregateFnSatisfaction::Approximate + && untruncated + && requested.is::() + { + AggregateFnSatisfaction::Exact + } else { + satisfaction + } + } + fn partial_dtype(&self, options: &Self::Options, input_dtype: &DType) -> Option { self.return_dtype(options, input_dtype) } @@ -155,11 +205,12 @@ impl AggregateFnVTable for BoundedMin { _args: AggregateArgs<'_, Self::Options>, scalar: &Scalar, ) -> VortexResult { - // A null partial means the producing accumulator saw nothing valid. + // A null partial means the producing accumulator saw nothing valid. Serialized partials + // don't record exactness. let state = if scalar.is_null() { BoundedMinState::Empty } else { - BoundedMinState::Value(scalar.clone()) + BoundedMinState::Inexact(scalar.clone()) }; Ok(BoundedMinPartial { state }) } @@ -170,8 +221,10 @@ impl AggregateFnVTable for BoundedMin { mut first: Self::Partial, second: Self::Partial, ) -> VortexResult { - if let BoundedMinState::Value(min) = second.state { - first.merge(min); + match second.state { + BoundedMinState::Empty => {} + BoundedMinState::Exact(bound) => first.merge(bound, true), + BoundedMinState::Inexact(bound) => first.merge(bound, false), } Ok(first) } @@ -184,7 +237,7 @@ impl AggregateFnVTable for BoundedMin { let dtype = args.dtype.as_nullable(); match &partial.state { BoundedMinState::Empty => Ok(Scalar::null(dtype)), - BoundedMinState::Value(min) => min.cast(&dtype), + BoundedMinState::Exact(bound) | BoundedMinState::Inexact(bound) => bound.cast(&dtype), } } @@ -208,8 +261,8 @@ impl AggregateFnVTable for BoundedMin { let Some(result) = columnar_min_max(batch, NumericalAggregateOpts::default(), ctx)? else { return Ok(()); }; - if let Some(bound) = truncate_min(result.min, args.options.max_bytes.get())? { - partial.merge(bound); + if let Some((bound, truncated)) = truncate_min(result.min, args.options.max_bytes.get())? { + partial.merge(bound, !truncated); } Ok(()) } @@ -237,22 +290,21 @@ fn supported_dtype<'a>(_options: &BoundedMinOptions, input_dtype: &'a DType) -> .map(|_| input_dtype) } -fn truncate_min(value: Scalar, max_bytes: usize) -> VortexResult> { +/// Returns the lower bound and whether it was truncated. +fn truncate_min(value: Scalar, max_bytes: usize) -> VortexResult> { let nullability = value.dtype().nullability(); match value.dtype() { - DType::Utf8(_) => { - Ok( - lower_bound(BufferString::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - DType::Binary(_) => { - Ok( - lower_bound(ByteBuffer::from_scalar(value)?, max_bytes, nullability) - .map(|(bound, _)| bound), - ) - } - _ => Ok(Some(value)), + DType::Utf8(_) => Ok(lower_bound( + BufferString::from_scalar(value)?, + max_bytes, + nullability, + )), + DType::Binary(_) => Ok(lower_bound( + ByteBuffer::from_scalar(value)?, + max_bytes, + nullability, + )), + _ => Ok(Some((value, false))), } } #[cfg(test)] @@ -418,4 +470,67 @@ mod tests { assert_eq!(roundtrip, options); Ok(()) } + + #[test] + fn bounded_min_exactly_satisfies_min_only_when_untruncated() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let requested = Min.bind(NumericalAggregateOpts::skip_nans()); + let new_acc = || { + Accumulator::try_new( + BoundedMin, + BoundedMinOptions { + max_bytes: max_bytes(4), + }, + VarBinViewArray::from_iter_str(["x"]).dtype().clone(), + ) + }; + + // Every value fits the bound, so the bound is the true extremum. + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["c", "dde"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated value that isn't the extremum doesn't affect exactness. + acc.accumulate( + &VarBinViewArray::from_iter_str(["ddddd"]).into_array(), + &mut ctx, + )?; + assert_eq!(acc.can_satisfy(&requested), AggregateFnSatisfaction::Exact); + + // A truncated extremum makes the bound approximate. + acc.accumulate( + &VarBinViewArray::from_iter_str(["aaaaa"]).into_array(), + &mut ctx, + )?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + + // Exactness isn't serialized, so a winning partial combined from its scalar form is + // approximate even though it was exact where it was computed. + let mut source = new_acc()?; + source.accumulate( + &VarBinViewArray::from_iter_str(["a"]).into_array(), + &mut ctx, + )?; + assert_eq!( + source.can_satisfy(&requested), + AggregateFnSatisfaction::Exact + ); + let mut acc = new_acc()?; + acc.accumulate( + &VarBinViewArray::from_iter_str(["c"]).into_array(), + &mut ctx, + )?; + acc.combine_partials(source.partial_scalar()?)?; + assert_eq!( + acc.can_satisfy(&requested), + AggregateFnSatisfaction::Approximate + ); + Ok(()) + } } diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs index ba236ab1d5f..ade0cdd0889 100644 --- a/vortex-array/src/aggregate_fn/vtable.rs +++ b/vortex-array/src/aggregate_fn/vtable.rs @@ -164,6 +164,21 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync { } } + /// Return whether the result accumulated in `partial` can satisfy `requested`. + /// + /// Refines [`Self::can_satisfy`], which only sees the options, with what was actually + /// accumulated: e.g. a byte-bounded maximum whose inputs all fit within the bound exactly + /// satisfies the unbounded maximum. Must never report a stronger satisfaction than is true for + /// `partial`. The default implementation ignores `partial`. + fn partial_can_satisfy( + &self, + options: &Self::Options, + _partial: &Self::Partial, + requested: &AggregateFnRef, + ) -> AggregateFnSatisfaction { + self.can_satisfy(options, requested) + } + /// The return [`DType`] of the aggregate. /// /// Returns `None` if the aggregate function cannot be applied to the input dtype. diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs index f7610fc71e4..fd57031a96a 100644 --- a/vortex-array/src/stats/mod.rs +++ b/vortex-array/src/stats/mod.rs @@ -33,16 +33,6 @@ use vortex_error::VortexExpect; use crate::expr::stats::Stat; -/// Statistics that are used for pruning files (i.e., we want to ensure they are computed when compressing/writing). -/// Sum is included for boolean arrays. -pub const PRUNING_STATS: &[Stat] = &[ - Stat::Min, - Stat::Max, - Stat::Sum, - Stat::NullCount, - Stat::NaNCount, -]; - pub fn as_stat_bitset_bytes(stats: &[Stat]) -> Vec { let max_stat = u8::from(last::().vortex_expect("last stat")) as usize + 1; // TODO(ngates): use vortex-buffer::BitBuffer diff --git a/vortex-datafusion/src/persistent/format.rs b/vortex-datafusion/src/persistent/format.rs index 35b835ca763..4b621dea0af 100644 --- a/vortex-datafusion/src/persistent/format.rs +++ b/vortex-datafusion/src/persistent/format.rs @@ -45,15 +45,20 @@ use futures::stream; use object_store::ObjectMeta; use object_store::ObjectStore; use vortex::VortexSessionDefault; +use vortex::aggregate_fn::AggregateFnVTableExt; +use vortex::aggregate_fn::EmptyOptions; +use vortex::aggregate_fn::NumericalAggregateOpts; +use vortex::aggregate_fn::fns::is_constant::IsConstant; +use vortex::aggregate_fn::fns::max::Max; +use vortex::aggregate_fn::fns::min::Min; +use vortex::aggregate_fn::fns::null_count::NullCount; +use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex::array::memory::MemorySessionExt; -use vortex::dtype::DType; -use vortex::dtype::Nullability; -use vortex::dtype::PType; +use vortex::dtype::FieldPath; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_err; use vortex::expr::stats::Precision; -use vortex::expr::stats::Stat; use vortex::file::EOF_SIZE; use vortex::file::MAX_POSTSCRIPT_SIZE; use vortex::file::OpenOptionsSessionExt; @@ -61,7 +66,6 @@ use vortex::file::VORTEX_FILE_EXTENSION; use vortex::io::object_store::ObjectStoreReadAt; use vortex::io::session::RuntimeSessionExt; use vortex::scalar::Scalar; -use vortex::scalar::ScalarValue as VortexScalarValue; use vortex::session::VortexSession; use vortex_arrow::ArrowSessionExt; @@ -598,7 +602,7 @@ impl FileFormat for VortexFormat { }) }); - let (dtype, file_stats, row_count) = match cached_metadata { + let (_dtype, file_stats, row_count) = match cached_metadata { Some(metadata) => metadata, None => { // Not entry - open the file @@ -635,10 +639,6 @@ impl FileFormat for VortexFormat { } }; - let struct_dtype = dtype - .as_struct_fields_opt() - .vortex_expect("dtype is not a struct"); - // Evaluate the statistics for each column that we are able to return to DataFusion. let Some(file_stats) = file_stats else { // If the file has no column stats, the best we can do is return a row count. @@ -661,32 +661,32 @@ impl FileFormat for VortexFormat { for field in table_schema.fields().iter() { // If the column does not exist, continue. This can happen if the schema has evolved // but we have not yet updated the Vortex file. - let Some(col_idx) = struct_dtype.find(field.name()) else { + let Some((aggregates, _)) = + file_stats.get_by_path(&FieldPath::from_name(field.name().as_str())) + else { // The default sets all statistics to `Precision`. column_statistics.push(ColumnStatistics::default()); continue; }; - let (stats_set, stats_dtype) = file_stats.get(col_idx); // Update the total size in bytes. - let column_size = - stats_set.get_as::(Stat::UncompressedSizeInBytes, &PType::U64.into()); + let column_size = aggregates + .get(&UncompressedSizeInBytes.bind(EmptyOptions)) + .and_then(|size| usize::try_from(&size).ok()); let min = scalar_stat_to_df( - Stat::Min, - stats_set.get(Stat::Min), - stats_dtype, + aggregates.get(&Min.bind(NumericalAggregateOpts::skip_nans())), field.data_type(), ); let max = scalar_stat_to_df( - Stat::Max, - stats_set.get(Stat::Max), - stats_dtype, + aggregates.get(&Max.bind(NumericalAggregateOpts::skip_nans())), field.data_type(), ); - let null_count = stats_set.get_as::(Stat::NullCount, &PType::U64.into()); + let null_count = aggregates + .get(&NullCount.bind(EmptyOptions)) + .and_then(|count| usize::try_from(&count).ok()); column_statistics.push(ColumnStatistics { null_count: null_count.to_df(), @@ -694,10 +694,9 @@ impl FileFormat for VortexFormat { max_value: max, sum_value: DFPrecision::Absent, distinct_count: is_constant_to_distinct_count( - stats_set.get_as::( - Stat::IsConstant, - &DType::Bool(Nullability::NonNullable), - ), + aggregates + .get(&IsConstant.bind(EmptyOptions)) + .and_then(|is_constant| bool::try_from(&is_constant).ok()), ), byte_size: column_size.to_df(), }) @@ -770,23 +769,11 @@ impl FileFormat for VortexFormat { } fn scalar_stat_to_df( - stat: Stat, - value: Precision, - stats_dtype: &DType, + value: Precision, target_dtype: &DataType, ) -> DFPrecision { - let Some(stat_dtype) = stat.dtype(stats_dtype) else { - return DFPrecision::Absent; - }; - value - .and_then(|stat_value| { - let scalar = Scalar::try_new(stat_dtype, Some(stat_value)) - .ok()? - .try_to_df() - .ok()?; - scalar.cast_to(target_dtype).ok() - }) + .and_then(|scalar| scalar.try_to_df().ok()?.cast_to(target_dtype).ok()) .to_df() } @@ -851,26 +838,18 @@ mod tests { fn test_scalar_stat_to_df_cast( #[case] value: ScalarValue, #[case] expected: ScalarValue, - #[values(Stat::Min, Stat::Max)] stat: Stat, #[values(true, false)] exact: bool, ) -> VortexResult<()> { let session = VortexSession::default(); let scalar = scalar_from_df(&value, &session)?; - let value = scalar - .value() - .cloned() - .ok_or_else(|| vortex_err!("expected non-null scalar"))?; let target_dtype = expected.data_type(); let (value, expected) = if exact { - (Precision::Exact(value), DFPrecision::Exact(expected)) + (Precision::Exact(scalar), DFPrecision::Exact(expected)) } else { - (Precision::Inexact(value), DFPrecision::Inexact(expected)) + (Precision::Inexact(scalar), DFPrecision::Inexact(expected)) }; - assert_eq!( - scalar_stat_to_df(stat, value, scalar.dtype(), &target_dtype), - expected - ); + assert_eq!(scalar_stat_to_df(value, &target_dtype), expected); Ok(()) } @@ -888,16 +867,11 @@ mod tests { fn test_scalar_stat_to_df_failed_cast( #[case] value: ScalarValue, #[case] target_dtype: DataType, - #[values(Stat::Min, Stat::Max)] stat: Stat, ) -> VortexResult<()> { let session = VortexSession::default(); let scalar = scalar_from_df(&value, &session)?; - let value = scalar - .value() - .cloned() - .ok_or_else(|| vortex_err!("expected non-null scalar"))?; assert_eq!( - scalar_stat_to_df(stat, Precision::Exact(value), scalar.dtype(), &target_dtype), + scalar_stat_to_df(Precision::Exact(scalar), &target_dtype), DFPrecision::Absent ); Ok(()) diff --git a/vortex-duckdb/cpp/copy_function.cpp b/vortex-duckdb/cpp/copy_function.cpp index 840d159f266..9c61804b21f 100644 --- a/vortex-duckdb/cpp/copy_function.cpp +++ b/vortex-duckdb/cpp/copy_function.cpp @@ -11,7 +11,6 @@ #include "duckdb/main/capi/capi_internal.hpp" #include "duckdb/main/client_context.hpp" #include "duckdb/main/connection.hpp" -#include "duckdb/parser/keyword_helper.hpp" #include "duckdb/parser/parsed_data/create_copy_function_info.hpp" unique_ptr copy_to_bind(ClientContext &, @@ -40,7 +39,7 @@ unique_ptr copy_to_bind(ClientContext &, throw BinderException(IntoErrString(error_out)); } auto cdata = unique_ptr(reinterpret_cast(ffi_bind_data)); - return make_uniq(std::move(cdata), names); + return make_uniq(std::move(cdata)); } unique_ptr @@ -86,7 +85,7 @@ void copy_to_get_written_statistics(ClientContext &, gstate.Cast().written_stats = &statistics; } -void copy_to_finalize(ClientContext &, FunctionData &bind_data, GlobalFunctionData &gstate) { +void copy_to_finalize(ClientContext &, FunctionData &, GlobalFunctionData &gstate) { auto &global = gstate.Cast(); void *const ffi_global = global.ffi_global->DataPtr(); duckdb_vx_error error_out = nullptr; @@ -98,25 +97,18 @@ void copy_to_finalize(ClientContext &, FunctionData &bind_data, GlobalFunctionDa if (!global.written_stats) { return; } - auto &names = bind_data.Cast().column_names; duckdb_vx_written_file_statistics file_stats; if (!duckdb_copy_function_get_written_file_statistics(ffi_global, &file_stats)) { // Statistics were requested (written_stats is set) but the finished write produced none; // that is an internal inconsistency, not a silently empty result. throw InternalException("vortex COPY: written statistics were requested but not produced"); } - if (file_stats.num_columns != names.size()) { - throw InternalException("vortex COPY: %llu statistics columns for %llu written columns", - file_stats.num_columns, - names.size()); - } D_ASSERT(global.written_stats != nullptr); global.written_stats->row_count = file_stats.row_count; global.written_stats->file_size_bytes = file_stats.file_size_bytes; global.written_stats->footer_size_bytes = Value::UBIGINT(file_stats.footer_size_bytes); - // Keyed by top-level column name only. The vortex footer reports one statistics set per - // top-level field, so nested struct/list leaf columns get no statistics here (unlike parquet, - // which recurses to leaf paths). Flat tables are fully covered. + // One entry per leaf column: nested struct fields are reported at their full path and structs + // get no entry of their own, as in parquet. for (idx_t i = 0; i < file_stats.num_columns; i++) { duckdb_vx_written_column_statistics col_stats {}; duckdb_vx_error col_error = nullptr; @@ -145,10 +137,9 @@ void copy_to_finalize(ClientContext &, FunctionData &bind_data, GlobalFunctionDa if (col_stats.has_nan_stat) { column["has_nan"] = Value::BOOLEAN(col_stats.contains_nan); } - // DuckLake keys column statistics by a quoted, dot-separated path (see - // DuckLakeUtil::ParseQuotedList); match the parquet writer, which quotes each name. - global.written_stats->column_statistics.emplace(KeywordHelper::WriteQuoted(names[i], '"'), - std::move(column)); + auto column_key = reinterpret_cast(col_stats.column_key)->ToString(); + duckdb_destroy_value(&col_stats.column_key); + global.written_stats->column_statistics.emplace(std::move(column_key), std::move(column)); } } diff --git a/vortex-duckdb/cpp/include/copy_function.hpp b/vortex-duckdb/cpp/include/copy_function.hpp index ec8cccd7048..1fadd418611 100644 --- a/vortex-duckdb/cpp/include/copy_function.hpp +++ b/vortex-duckdb/cpp/include/copy_function.hpp @@ -8,12 +8,9 @@ using namespace duckdb; struct VortexCopyBindData final : TableFunctionData { - VortexCopyBindData(unique_ptr ffi_bind, vector column_names) - : ffi_bind(std::move(ffi_bind)), column_names(std::move(column_names)) { + VortexCopyBindData(unique_ptr ffi_bind) : ffi_bind(std::move(ffi_bind)) { } unique_ptr ffi_bind; - // Column names in write order, used to key WRITTEN_FILE_STATISTICS. - vector column_names; }; struct VortexCopyGlobalState final : GlobalFunctionData { diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index 5924a209fbc..30ab6b30bb1 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -83,6 +83,9 @@ typedef struct { // Per-column statistics of a written Vortex file. typedef struct { + // Owned VARCHAR value holding the column's quoted, dot-separated path (e.g. `"s"."b"`), the key + // format DuckLake expects; the caller must destroy it. + duckdb_value column_key; // Owned values, null if absent; the caller must destroy them. duckdb_value min; duckdb_value max; diff --git a/vortex-duckdb/src/column_statistics.rs b/vortex-duckdb/src/column_statistics.rs index 9be4ec2fefa..3e581c8db96 100644 --- a/vortex-duckdb/src/column_statistics.rs +++ b/vortex-duckdb/src/column_statistics.rs @@ -1,12 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex::array::stats::StatsSet; +use vortex::aggregate_fn::AggregateFnRef; +use vortex::aggregate_fn::AggregateFnVTableExt; +use vortex::aggregate_fn::EmptyOptions; +use vortex::aggregate_fn::NumericalAggregateOpts; +use vortex::aggregate_fn::fns::max::Max; +use vortex::aggregate_fn::fns::min::Min; +use vortex::aggregate_fn::fns::null_count::NullCount; +use vortex::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex::dtype::DType; use vortex::error::VortexExpect as _; use vortex::error::VortexResult; use vortex::expr::stats::Precision; -use vortex::expr::stats::Stat; +use vortex::layout::layouts::file_stats::AggregateStats; use vortex::scalar::Scalar; use vortex::scalar::ScalarValue; @@ -62,29 +69,22 @@ pub struct ColumnStatisticsAggregate { } impl ColumnStatisticsAggregate { - pub fn new(stats: &StatsSet) -> Self { - let min = match stats.get(Stat::Min) { - Precision::Exact(min) => Some(min), - _ => None, - }; - let max = match stats.get(Stat::Max) { - Precision::Exact(max) => Some(max), + pub fn new(aggregates: &AggregateStats) -> Self { + let exact = |aggregate_fn: AggregateFnRef| match aggregates.get(&aggregate_fn) { + Precision::Exact(value) => value.into_value(), _ => None, }; - let max_string_length = - if let Precision::Exact(value) = stats.get(Stat::UncompressedSizeInBytes) { - // DuckDB's string length is u32 - #[allow(clippy::cast_possible_truncation)] - Some(value.as_primitive().as_u64().vortex_expect("not a u64") as u32) - } else { - None - }; + let min = exact(Min.bind(NumericalAggregateOpts::skip_nans())); + let max = exact(Max.bind(NumericalAggregateOpts::skip_nans())); - let has_null = match stats.get(Stat::NullCount) { - Precision::Exact(cnt) => cnt.as_primitive().as_u64().vortex_expect("not a u64") > 0, - _ => true, - }; + // DuckDB's string length is u32 + #[allow(clippy::cast_possible_truncation)] + let max_string_length = exact(UncompressedSizeInBytes.bind(EmptyOptions)) + .map(|value| value.as_primitive().as_u64().vortex_expect("not a u64") as u32); + + let has_null = exact(NullCount.bind(EmptyOptions)) + .is_none_or(|count| count.as_primitive().as_u64().vortex_expect("not a u64") > 0); Self { min, diff --git a/vortex-duckdb/src/copy.rs b/vortex-duckdb/src/copy.rs index 6d3dbbbd952..4f179abbe9d 100644 --- a/vortex-duckdb/src/copy.rs +++ b/vortex-duckdb/src/copy.rs @@ -12,10 +12,19 @@ use object_store::ObjectStore; use object_store::registry::ObjectStoreRegistry; use parking_lot::Mutex; use static_assertions::assert_impl_all; +use vortex::aggregate_fn::AggregateFnRef; +use vortex::aggregate_fn::AggregateFnVTableExt; +use vortex::aggregate_fn::EmptyOptions; +use vortex::aggregate_fn::NumericalAggregateOpts; +use vortex::aggregate_fn::fns::max::Max; +use vortex::aggregate_fn::fns::min::Min; +use vortex::aggregate_fn::fns::nan_count::NanCount; +use vortex::aggregate_fn::fns::null_count::NullCount; use vortex::array::ArrayRef; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; use vortex::dtype::FieldName; +use vortex::dtype::FieldPath; use vortex::dtype::Nullability::NonNullable; use vortex::dtype::Nullability::Nullable; use vortex::dtype::StructFields; @@ -24,7 +33,8 @@ use vortex::error::VortexResult; use vortex::error::vortex_bail; use vortex::error::vortex_err; use vortex::expr::stats::Precision; -use vortex::expr::stats::Stat; +use vortex::file::CompressedFieldSizes; +use vortex::file::FileStatistics; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; use vortex::file::multi::parse_uri_or_path; @@ -54,11 +64,13 @@ pub struct CopyFunctionBind { } assert_impl_all!(CopyFunctionBind: Send, Clone); -/// The per-column compressed sizes are computed once here rather than per column, since DuckDB -/// queries statistics one column at a time. +/// The compressed field sizes and the statistics columns are computed once here rather than per +/// column, since DuckDB queries statistics one column at a time. struct FinishedWrite { summary: WriteSummary, - column_sizes: Vec, + column_sizes: Option, + /// Indices into the footer's [`FileStatistics`] of the entries reported as statistics columns. + stats_columns: Vec, } /// Write to a file has two phases, writing data chunks and then closing the file. @@ -159,12 +171,18 @@ pub fn copy_to_finalize(init_global: &mut CopyFunctionGlobal) -> VortexResult<() .take() .vortex_expect("no file to close"); // Keep the write summary (footer + size) so DuckLake can read per-file statistics back - // without re-opening the file. Compute the per-column compressed sizes once, up front. + // without re-opening the file. Compute the compressed field sizes once, up front. let summary = task.await?; - let column_sizes = summary.compressed_column_sizes().unwrap_or_default(); + let column_sizes = summary.footer().compressed_field_sizes().ok(); + let stats_columns = summary + .footer() + .statistics() + .map(leaf_stats_columns) + .unwrap_or_default(); *init_global.finished.lock() = Some(FinishedWrite { summary, column_sizes, + stats_columns, }); Ok(()) }) @@ -181,6 +199,8 @@ pub(crate) struct WrittenFileStats { /// Per-column statistics of the written Vortex file. `min`/`max` are DuckDB values converted from /// the Vortex scalar; every field is optional and omitted when the statistic is not available. pub(crate) struct WrittenColumnStats { + /// The column's path in DuckLake's key format (see [`ducklake_column_key`]). + pub column_key: String, pub min: Option, pub max: Option, pub null_count: Option, @@ -192,7 +212,11 @@ pub(crate) struct WrittenColumnStats { /// Read file-level statistics back from the finished write. `None` before finalize. pub(crate) fn written_file_stats(global: &CopyFunctionGlobal) -> Option { let guard = global.finished.lock(); - Some(file_stats_from_summary(&guard.as_ref()?.summary)) + let finished = guard.as_ref()?; + Some(file_stats_from_summary( + &finished.summary, + finished.stats_columns.len(), + )) } /// Read per-column statistics for `column_index` from the finished write. `Ok(None)` if the file is @@ -205,14 +229,45 @@ pub(crate) fn written_column_stats( let Some(finished) = guard.as_ref() else { return Ok(None); }; - column_stats_from_summary(&finished.summary, column_index, &finished.column_sizes).map(Some) + column_stats_from_summary( + &finished.summary, + &finished.stats_columns, + column_index, + finished.column_sizes.as_ref(), + ) + .map(Some) } -fn file_stats_from_summary(summary: &WriteSummary) -> WrittenFileStats { - let num_columns = summary - .footer() - .statistics() - .map_or(0, |s| s.stats_sets().len()); +/// The indices of the leaf entries of `file_stats`, which are the columns statistics are reported +/// for. Like parquet, which reports one entry per leaf column writer, struct columns get no entry of +/// their own; nested fields are reported at their full path. Lists and maps are opaque leaves of the +/// statistics layout, so they are reported at their own path rather than per element. +fn leaf_stats_columns(file_stats: &FileStatistics) -> Vec { + file_stats + .dtypes() + .iter() + .enumerate() + .filter(|(_, dtype)| !dtype.is_struct()) + .map(|(index, _)| index) + .collect() +} + +/// DuckLake keys column statistics by a quoted, dot-separated path (see +/// `DuckLakeUtil::ParseQuotedList`). This matches the parquet writer, which quotes each name with +/// `KeywordHelper::WriteQuoted(name, '"')`, escaping embedded quotes by doubling them. A list +/// element is named `element`, as in parquet's list encoding. +fn ducklake_column_key(path: &FieldPath) -> String { + path.parts() + .iter() + .map(|part| { + let name = part.as_name().unwrap_or("element"); + format!("\"{}\"", name.replace('"', "\"\"")) + }) + .collect::>() + .join(".") +} + +fn file_stats_from_summary(summary: &WriteSummary, num_columns: usize) -> WrittenFileStats { WrittenFileStats { row_count: summary.row_count(), file_size_bytes: summary.size(), @@ -222,40 +277,40 @@ fn file_stats_from_summary(summary: &WriteSummary) -> WrittenFileStats { } } -/// Per-column statistics from a finished write's summary and its precomputed compressed sizes -/// (`column_sizes`, indexed the same as the footer's stats sets). -/// -/// Only top-level columns are covered: the footer exposes one statistics set per top-level field, -/// so nested struct/list leaf columns are not reported (parquet, by contrast, recurses to leaf -/// paths). Flat tables - the common DuckLake case - are fully covered. +/// Statistics for the `column_index`-th statistics column of a finished write, where +/// `stats_columns` are the [leaf entries][leaf_stats_columns] of its footer statistics and +/// `column_sizes` its compressed field sizes. fn column_stats_from_summary( summary: &WriteSummary, + stats_columns: &[usize], column_index: usize, - column_sizes: &[u64], + column_sizes: Option<&CompressedFieldSizes>, ) -> VortexResult { let file_stats = summary .footer() .statistics() .ok_or_else(|| vortex_err!("written file has no statistics"))?; - let stats_sets = file_stats.stats_sets(); - if column_index >= stats_sets.len() { - vortex_bail!( - "column index {column_index} out of range for {} statistics sets", - stats_sets.len() - ); - } - let stats = &stats_sets[column_index]; - let dtype = &file_stats.dtypes()[column_index]; + let stats_index = *stats_columns.get(column_index).ok_or_else(|| { + vortex_err!( + "column index {column_index} out of range for {} statistics columns", + stats_columns.len() + ) + })?; + let (aggregates, dtype) = file_stats.get(stats_index); + let path = &file_stats.paths()[stats_index]; + let value = + |aggregate_fn: AggregateFnRef| aggregates.get(&aggregate_fn).and_then(Scalar::into_value); Ok(WrittenColumnStats { - min: exact_scalar_to_duckdb(stats.get(Stat::Min), dtype)?, - max: exact_scalar_to_duckdb(stats.get(Stat::Max), dtype)?, - null_count: exact_u64(stats.get(Stat::NullCount)), + column_key: ducklake_column_key(path), + min: exact_scalar_to_duckdb(value(Min.bind(NumericalAggregateOpts::skip_nans())), dtype)?, + max: exact_scalar_to_duckdb(value(Max.bind(NumericalAggregateOpts::skip_nans())), dtype)?, + null_count: exact_u64(value(NullCount.bind(EmptyOptions))), // NaNCount is exact only for float columns, so this is emitted just for them (as in parquet). - has_nan: exact_u64(stats.get(Stat::NaNCount)).map(|count| count > 0), + has_nan: exact_u64(value(NanCount.bind(EmptyOptions))).map(|count| count > 0), num_values: summary.row_count(), // On-disk compressed size; excludes bytes not attributable to a column (e.g. struct validity). - column_size_bytes: column_sizes.get(column_index).copied(), + column_size_bytes: column_sizes.and_then(|sizes| sizes.get(path)), }) } @@ -332,15 +387,23 @@ pub fn copy_to_initialize_global( mod tests { use vortex::array::IntoArray; use vortex::array::arrays::StructArray; - use vortex::array::stats::PRUNING_STATS; + use vortex::array::validity::Validity; use vortex::buffer::ByteBufferMut; use vortex::buffer::buffer; + use vortex::expr::stats::Stat; use super::*; + fn pruning_aggregate_fns() -> Vec { + [Stat::Min, Stat::Max, Stat::NullCount, Stat::NaNCount] + .into_iter() + .filter_map(|stat| stat.aggregate_fn()) + .collect() + } + /// Writes a one-column file and returns its summary, with `file_statistics` controlling which /// statistics the footer carries (empty means none at all). - fn write_summary(file_statistics: Vec) -> WriteSummary { + fn write_summary(file_statistics: Vec) -> WriteSummary { RUNTIME.block_on(async { let array = StructArray::from_fields(&[("i", buffer![1u32, 2, 3].into_array())]) .unwrap() @@ -355,16 +418,71 @@ mod tests { }) } + /// Statistics for the `column_index`-th statistics column of `summary`, resolved the same way + /// `copy_to_finalize` and [`written_column_stats`] do. + fn column_stats( + summary: &WriteSummary, + column_index: usize, + ) -> VortexResult { + let stats_columns = summary + .footer() + .statistics() + .map(leaf_stats_columns) + .unwrap_or_default(); + let column_sizes = summary.footer().compressed_field_sizes()?; + column_stats_from_summary(summary, &stats_columns, column_index, Some(&column_sizes)) + } + #[test] fn column_stats_out_of_range_is_an_error() { - let summary = write_summary(PRUNING_STATS.to_vec()); - assert!(column_stats_from_summary(&summary, 0, &[]).is_ok()); - assert!(column_stats_from_summary(&summary, 1, &[]).is_err()); + let summary = write_summary(pruning_aggregate_fns()); + assert!(column_stats(&summary, 0).is_ok()); + assert!(column_stats(&summary, 1).is_err()); } #[test] fn column_stats_without_file_statistics_is_an_error() { let summary = write_summary(vec![]); - assert!(column_stats_from_summary(&summary, 0, &[]).is_err()); + assert!(column_stats(&summary, 0).is_err()); + } + + #[test] + fn column_stats_reports_nested_leaves_by_quoted_path() -> VortexResult<()> { + // A nullable struct `s` contributes `s.b` and a trailing own entry to the post-order stats + // layout; only the leaves are statistics columns, keyed by their full quoted path. + let summary = RUNTIME.block_on(async { + let b = buffer![10i32, 20, 30].into_array(); + let inner = StructArray::new(["b"].into(), [b], 3, Validity::AllValid).into_array(); + let c = buffer![7u32, 8, 9].into_array(); + let outer = + StructArray::new(["s", "c\"q"].into(), [inner, c], 3, Validity::NonNullable) + .into_array(); + + let mut buf = ByteBufferMut::empty(); + let mut writer = SESSION + .write_options() + .with_file_statistics(pruning_aggregate_fns()) + .writer(&mut buf, outer.dtype().clone()); + writer.push(outer).await?; + writer.finish().await + })?; + + let file_stats = summary + .footer() + .statistics() + .ok_or_else(|| vortex_err!("no statistics"))?; + assert_eq!(leaf_stats_columns(file_stats).len(), 2); + let nested = column_stats(&summary, 0)?; + assert_eq!(nested.column_key, "\"s\".\"b\""); + assert!(nested.min.is_some()); + assert!(nested.column_size_bytes.is_some_and(|size| size > 0)); + + let top_level = column_stats(&summary, 1)?; + assert_eq!(top_level.column_key, "\"c\"\"q\""); + assert!(top_level.min.is_some()); + assert!(top_level.max.is_some()); + + assert!(column_stats(&summary, 2).is_err()); + Ok(()) } } diff --git a/vortex-duckdb/src/ffi.rs b/vortex-duckdb/src/ffi.rs index e774589746d..c926414e43e 100644 --- a/vortex-duckdb/src/ffi.rs +++ b/vortex-duckdb/src/ffi.rs @@ -33,6 +33,7 @@ use crate::duckdb::Expression; use crate::duckdb::LogicalType; use crate::duckdb::LogicalTypeRef; use crate::duckdb::TableInitInput; +use crate::duckdb::Value; use crate::duckdb::try_or; use crate::duckdb::try_or_null; use crate::file_reader::OpenFileReader; @@ -495,6 +496,7 @@ pub unsafe extern "C-unwind" fn duckdb_copy_function_get_written_column_statisti return Ok(false); }; let out = unsafe { &mut *out }; + out.column_key = Value::from(stats.column_key.as_str()).into_ptr(); out.min = stats.min.map_or(ptr::null_mut(), |v| v.into_ptr()); out.max = stats.max.map_or(ptr::null_mut(), |v| v.into_ptr()); out.has_null_count = stats.null_count.is_some(); diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index f4c9bb67995..2be61774161 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -11,6 +11,7 @@ use url::Url; use vortex::array::VortexSessionExecute as _; use vortex::array::arrays::struct_::StructArrayExt as _; use vortex::dtype::DType; +use vortex::dtype::FieldPath; use vortex::error::VortexExpect; use vortex::error::VortexResult; use vortex::error::vortex_panic; @@ -274,16 +275,12 @@ pub fn reader_get_statistics( .as_any() .downcast_ref::()?; - let DType::Struct(fields, _) = &file.reader.dtype() else { - return None; - }; - let index = fields.find(column)?; - let stats_sets = reader.file_stats().stats_sets(); + let (stats, dtype) = reader + .file_stats() + .get_by_path(&FieldPath::from_name(column))?; - let dtype = fields.field_by_index(index)?; - - let stats = ColumnStatisticsAggregate::new(stats_sets.get(index)?); - match ColumnStatistics::try_from(stats, dtype) { + let stats = ColumnStatisticsAggregate::new(stats); + match ColumnStatistics::try_from(stats, dtype.clone()) { Ok(stats) => Some(stats), Err(e) => vortex_panic!(e), } @@ -328,11 +325,11 @@ pub fn footer_get_statistics(footer: &Footer, index: usize) -> Option Some(stats), Err(e) => vortex_panic!(e), } diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index b123b0d5eeb..b5ff6a1fb26 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -54,11 +54,51 @@ table PostscriptSegment { // [postscript] // [footer] +/// An `AggregateSpec` describes an aggregate function, including its options, recorded in the +/// file statistics. +table AggregateSpec { + id: string (required); + /// The aggregate function's serialized options, interpreted by the function's plugin. + options: [ubyte]; +} + +/// The accumulated state of one aggregate function over one file statistics entry. +table AggregateState { + /// Index of the aggregate function in `FileStatistics.aggregate_specs`. + aggregate_spec: uint16; + /// Protobuf serialized `vortex.scalar.ScalarValue` of the aggregate's partial state, typed by + /// the aggregate's partial dtype for the entry's dtype. A null value means no value was + /// accumulated. + partial: [ubyte] (required); +} + +/// The aggregate states of one file statistics entry. +table FieldAggregates { + aggregates: [AggregateState]; +} + /// The `FileStatistics` object contains file-level statistics for the Vortex file. 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. + /// Statistics for each top-level field in the root schema, in schema order. If the root schema + /// is not a struct, there will be a single entry in this array. Does not recurse into nested + /// struct fields; see `nested_field_aggregates` for that. + /// + /// Kept in this legacy shape for backward compatibility with readers that predate + /// `nested_field_aggregates`, which strictly validate this array's length against the number + /// of top-level struct fields. field_stats: [ArrayStats]; + /// Dictionary-encoded aggregate functions referenced by `nested_field_aggregates`, up to + /// u16::MAX. Each distinct pair of ID and options is one entry. + aggregate_specs: [AggregateSpec]; + /// Aggregate states following a post-order walk of the root `DType` tree: each leaf field gets + /// an entry, and each nullable struct additionally gets a trailing entry for its own null + /// count, inserted after its children's entries. + /// + /// Readers that know about this field should prefer it over `field_stats` when present, + /// skipping aggregates they don't recognize. It is absent for files written before nested + /// file stats were supported, and readers that don't know about it fall back to + /// `field_stats`'s legacy top-level-fields-only layout. + nested_field_aggregates: [FieldAggregates]; } /// The `Registry` object stores dictionary-encoded configuration for segments, diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 72a5e8787e0..2c97d80611c 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/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index 86df41f0305..b1e67403280 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -798,6 +798,339 @@ impl ::core::fmt::Debug for PostscriptSegment<'_> { ds.finish() } } +pub enum AggregateSpecOffset {} +#[derive(Copy, Clone, PartialEq)] + +/// An `AggregateSpec` describes an aggregate function, including its options, recorded in the +/// file statistics. +pub struct AggregateSpec<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for AggregateSpec<'a> { + type Inner = AggregateSpec<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> AggregateSpec<'a> { + pub const VT_ID: ::flatbuffers::VOffsetT = 4; + pub const VT_OPTIONS: ::flatbuffers::VOffsetT = 6; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + AggregateSpec { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args AggregateSpecArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = AggregateSpecBuilder::new(_fbb); + if let Some(x) = args.options { builder.add_options(x); } + if let Some(x) = args.id { builder.add_id(x); } + builder.finish() + } + + + #[inline] + pub fn id(&self) -> &'a str { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<&str>>(AggregateSpec::VT_ID, None).unwrap()} + } + /// The aggregate function's serialized options, interpreted by the function's plugin. + #[inline] + pub fn options(&self) -> Option<::flatbuffers::Vector<'a, u8>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u8>>>(AggregateSpec::VT_OPTIONS, None)} + } +} + +impl ::flatbuffers::Verifiable for AggregateSpec<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::<::flatbuffers::ForwardsUOffset<&str>>("id", Self::VT_ID, true)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u8>>>("options", Self::VT_OPTIONS, false)? + .finish(); + Ok(()) + } +} +pub struct AggregateSpecArgs<'a> { + pub id: Option<::flatbuffers::WIPOffset<&'a str>>, + pub options: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u8>>>, +} +impl<'a> Default for AggregateSpecArgs<'a> { + #[inline] + fn default() -> Self { + AggregateSpecArgs { + id: None, // required field + options: None, + } + } +} + +pub struct AggregateSpecBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> AggregateSpecBuilder<'a, 'b, A> { + #[inline] + pub fn add_id(&mut self, id: ::flatbuffers::WIPOffset<&'b str>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(AggregateSpec::VT_ID, id); + } + #[inline] + pub fn add_options(&mut self, options: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u8>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(AggregateSpec::VT_OPTIONS, options); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> AggregateSpecBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + AggregateSpecBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + self.fbb_.required(o, AggregateSpec::VT_ID,"id"); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for AggregateSpec<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("AggregateSpec"); + ds.field("id", &self.id()); + ds.field("options", &self.options()); + ds.finish() + } +} +pub enum AggregateStateOffset {} +#[derive(Copy, Clone, PartialEq)] + +/// The accumulated state of one aggregate function over one file statistics entry. +pub struct AggregateState<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for AggregateState<'a> { + type Inner = AggregateState<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> AggregateState<'a> { + pub const VT_AGGREGATE_SPEC: ::flatbuffers::VOffsetT = 4; + pub const VT_PARTIAL: ::flatbuffers::VOffsetT = 6; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + AggregateState { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args AggregateStateArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = AggregateStateBuilder::new(_fbb); + if let Some(x) = args.partial { builder.add_partial(x); } + builder.add_aggregate_spec(args.aggregate_spec); + builder.finish() + } + + + /// Index of the aggregate function in `FileStatistics.aggregate_specs`. + #[inline] + pub fn aggregate_spec(&self) -> u16 { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::(AggregateState::VT_AGGREGATE_SPEC, Some(0)).unwrap()} + } + /// Protobuf serialized `vortex.scalar.ScalarValue` of the aggregate's partial state, typed by + /// the aggregate's partial dtype for the entry's dtype. A null value means no value was + /// accumulated. + #[inline] + pub fn partial(&self) -> ::flatbuffers::Vector<'a, u8> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u8>>>(AggregateState::VT_PARTIAL, None).unwrap()} + } +} + +impl ::flatbuffers::Verifiable for AggregateState<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::("aggregate_spec", Self::VT_AGGREGATE_SPEC, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u8>>>("partial", Self::VT_PARTIAL, true)? + .finish(); + Ok(()) + } +} +pub struct AggregateStateArgs<'a> { + pub aggregate_spec: u16, + pub partial: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u8>>>, +} +impl<'a> Default for AggregateStateArgs<'a> { + #[inline] + fn default() -> Self { + AggregateStateArgs { + aggregate_spec: 0, + partial: None, // required field + } + } +} + +pub struct AggregateStateBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> AggregateStateBuilder<'a, 'b, A> { + #[inline] + pub fn add_aggregate_spec(&mut self, aggregate_spec: u16) { + self.fbb_.push_slot::(AggregateState::VT_AGGREGATE_SPEC, aggregate_spec, 0); + } + #[inline] + pub fn add_partial(&mut self, partial: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u8>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(AggregateState::VT_PARTIAL, partial); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> AggregateStateBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + AggregateStateBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + self.fbb_.required(o, AggregateState::VT_PARTIAL,"partial"); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for AggregateState<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("AggregateState"); + ds.field("aggregate_spec", &self.aggregate_spec()); + ds.field("partial", &self.partial()); + ds.finish() + } +} +pub enum FieldAggregatesOffset {} +#[derive(Copy, Clone, PartialEq)] + +/// The aggregate states of one file statistics entry. +pub struct FieldAggregates<'a> { + pub _tab: ::flatbuffers::Table<'a>, +} + +impl<'a> ::flatbuffers::Follow<'a> for FieldAggregates<'a> { + type Inner = FieldAggregates<'a>; + #[inline] + unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { + Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } + } +} + +impl<'a> FieldAggregates<'a> { + pub const VT_AGGREGATES: ::flatbuffers::VOffsetT = 4; + + #[inline] + pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { + FieldAggregates { _tab: table } + } + #[allow(unused_mut)] + pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( + _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, + args: &'args FieldAggregatesArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = FieldAggregatesBuilder::new(_fbb); + if let Some(x) = args.aggregates { builder.add_aggregates(x); } + builder.finish() + } + + + #[inline] + pub fn aggregates(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FieldAggregates::VT_AGGREGATES, None)} + } +} + +impl ::flatbuffers::Verifiable for FieldAggregates<'_> { + #[inline] + fn run_verifier( + v: &mut ::flatbuffers::Verifier, pos: usize + ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { + v.visit_table(pos)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("aggregates", Self::VT_AGGREGATES, false)? + .finish(); + Ok(()) + } +} +pub struct FieldAggregatesArgs<'a> { + pub aggregates: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, +} +impl<'a> Default for FieldAggregatesArgs<'a> { + #[inline] + fn default() -> Self { + FieldAggregatesArgs { + aggregates: None, + } + } +} + +pub struct FieldAggregatesBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { + fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, + start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, +} +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FieldAggregatesBuilder<'a, 'b, A> { + #[inline] + pub fn add_aggregates(&mut self, aggregates: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FieldAggregates::VT_AGGREGATES, aggregates); + } + #[inline] + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FieldAggregatesBuilder<'a, 'b, A> { + let start = _fbb.start_table(); + FieldAggregatesBuilder { + fbb_: _fbb, + start_: start, + } + } + #[inline] + pub fn finish(self) -> ::flatbuffers::WIPOffset> { + let o = self.fbb_.end_table(self.start_); + ::flatbuffers::WIPOffset::new(o.value()) + } +} + +impl ::core::fmt::Debug for FieldAggregates<'_> { + fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { + let mut ds = f.debug_struct("FieldAggregates"); + ds.field("aggregates", &self.aggregates()); + ds.finish() + } +} pub enum FileStatisticsOffset {} #[derive(Copy, Clone, PartialEq)] @@ -816,6 +1149,8 @@ impl<'a> ::flatbuffers::Follow<'a> for FileStatistics<'a> { impl<'a> FileStatistics<'a> { pub const VT_FIELD_STATS: ::flatbuffers::VOffsetT = 4; + pub const VT_AGGREGATE_SPECS: ::flatbuffers::VOffsetT = 6; + pub const VT_NESTED_FIELD_AGGREGATES: ::flatbuffers::VOffsetT = 8; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -827,13 +1162,20 @@ impl<'a> FileStatistics<'a> { args: &'args FileStatisticsArgs<'args> ) -> ::flatbuffers::WIPOffset> { let mut builder = FileStatisticsBuilder::new(_fbb); + if let Some(x) = args.nested_field_aggregates { builder.add_nested_field_aggregates(x); } + if let Some(x) = args.aggregate_specs { builder.add_aggregate_specs(x); } if let Some(x) = args.field_stats { builder.add_field_stats(x); } 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. + /// Statistics for each top-level field in the root schema, in schema order. If the root schema + /// is not a struct, there will be a single entry in this array. Does not recurse into nested + /// struct fields; see `nested_field_aggregates` for that. + /// + /// Kept in this legacy shape for backward compatibility with readers that predate + /// `nested_field_aggregates`, which strictly validate this array's length against the number + /// of top-level struct fields. #[inline] pub fn field_stats(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { // Safety: @@ -841,6 +1183,30 @@ 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)} } + /// Dictionary-encoded aggregate functions referenced by `nested_field_aggregates`, up to + /// u16::MAX. Each distinct pair of ID and options is one entry. + #[inline] + pub fn aggregate_specs(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FileStatistics::VT_AGGREGATE_SPECS, None)} + } + /// Aggregate states following a post-order walk of the root `DType` tree: each leaf field gets + /// an entry, and each nullable struct additionally gets a trailing entry for its own null + /// count, inserted after its children's entries. + /// + /// Readers that know about this field should prefer it over `field_stats` when present, + /// skipping aggregates they don't recognize. It is absent for files written before nested + /// file stats were supported, and readers that don't know about it fall back to + /// `field_stats`'s legacy top-level-fields-only layout. + #[inline] + pub fn nested_field_aggregates(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FileStatistics::VT_NESTED_FIELD_AGGREGATES, None)} + } } impl ::flatbuffers::Verifiable for FileStatistics<'_> { @@ -850,18 +1216,24 @@ 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::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("aggregate_specs", Self::VT_AGGREGATE_SPECS, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("nested_field_aggregates", Self::VT_NESTED_FIELD_AGGREGATES, false)? .finish(); Ok(()) } } pub struct FileStatisticsArgs<'a> { pub field_stats: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub aggregate_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub nested_field_aggregates: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, } impl<'a> Default for FileStatisticsArgs<'a> { #[inline] fn default() -> Self { FileStatisticsArgs { field_stats: None, + aggregate_specs: None, + nested_field_aggregates: None, } } } @@ -876,6 +1248,14 @@ 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_aggregate_specs(&mut self, aggregate_specs: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_AGGREGATE_SPECS, aggregate_specs); + } + #[inline] + pub fn add_nested_field_aggregates(&mut self, nested_field_aggregates: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_NESTED_FIELD_AGGREGATES, nested_field_aggregates); + } + #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FileStatisticsBuilder<'a, 'b, A> { let start = _fbb.start_table(); FileStatisticsBuilder { @@ -894,6 +1274,8 @@ 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("aggregate_specs", &self.aggregate_specs()); + ds.field("nested_field_aggregates", &self.nested_field_aggregates()); ds.finish() } } diff --git a/vortex-file/src/footer/file_statistics.rs b/vortex-file/src/footer/file_statistics.rs index f78abf05de8..eaf8b954ab4 100644 --- a/vortex-file/src/footer/file_statistics.rs +++ b/vortex-file/src/footer/file_statistics.rs @@ -11,122 +11,224 @@ use std::sync::Arc; use flatbuffers::FlatBufferBuilder; use flatbuffers::WIPOffset; use itertools::Itertools; +use vortex_array::aggregate_fn::AggregateFnId; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::session::AggregateFnSessionExt; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldPath; use vortex_array::flatbuffers::FlatBufferRoot; use vortex_array::flatbuffers::WriteFlatBuffer; +use vortex_array::flatbuffers::array::ArrayStats; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; use vortex_array::stats::StatsSet; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_ensure_eq; +use vortex_error::vortex_err; +use vortex_layout::layouts::file_stats::AggregateStat; +use vortex_layout::layouts::file_stats::AggregateStats; +use vortex_layout::layouts::file_stats::postorder_stats_layout; use vortex_session::VortexSession; +use vortex_utils::aliases::hash_map::HashMap; use crate::flatbuffers::footer as fb; /// Contains statistical information about the data in a Vortex file. /// -/// This struct wraps an array of `StatsSet` objects, each containing statistics -/// for a field or column in the file. These statistics can be used for query -/// optimization and data exploration. +/// Each entry holds the aggregates recorded for a field of the file. These statistics can be used +/// for query optimization and data exploration. #[derive(Clone, Debug)] pub struct FileStatistics { - /// An array of statistics sets, one for each field or column in the file. - stats: Arc<[StatsSet]>, + /// The aggregates recorded for each field or column in the file, following the post-order + /// nested-struct layout. + aggregates: Arc<[AggregateStats]>, /// 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 `aggregates` + /// 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]>, + /// Maps each entry in `paths` to its index, so [`Self::get_by_path`] doesn't need to scan + /// `paths` linearly. + path_index: Arc>, + /// Legacy top-level-fields-only aggregates, one per top-level struct field (or a single entry + /// for a non-struct root dtype). Only populated by [`Self::new_with_dtype`], for serialization + /// into `field_stats`; empty for instances built from [`Self::from_flatbuffer`], which are + /// never re-serialized. + legacy_aggregates: Arc<[AggregateStats]>, } impl FileStatistics { - /// Creates a new [`FileStatistics`] from the given statistics and data types. + /// Builds `path_index` from `paths` and assembles the final struct. The single place that + /// constructs a [`FileStatistics`], so `path_index` can't drift out of sync with `paths`. + fn from_parts( + aggregates: Arc<[AggregateStats]>, + dtypes: Arc<[DType]>, + paths: Arc<[FieldPath]>, + legacy_aggregates: Arc<[AggregateStats]>, + ) -> Self { + let path_index = paths + .iter() + .cloned() + .enumerate() + .map(|(i, p)| (p, i)) + .collect(); + Self { + aggregates, + dtypes, + paths, + path_index: Arc::new(path_index), + legacy_aggregates, + } + } + + /// Creates a new [`FileStatistics`] from the given aggregates, 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 `aggregates`, `dtypes`, and `paths` have different lengths. + pub fn new( + aggregates: Arc<[AggregateStats]>, + dtypes: Arc<[DType]>, + paths: Arc<[FieldPath]>, + ) -> Self { assert_eq!( - stats.len(), + aggregates.len(), dtypes.len(), - "stats and dtypes must have the same length" + "aggregates and dtypes must have the same length" ); - - Self { stats, dtypes } + assert_eq!( + aggregates.len(), + paths.len(), + "aggregates and paths must have the same length" + ); + Self::from_parts(aggregates, dtypes, paths, Arc::new([])) } - /// Creates a new [`FileStatistics`] from the given statistics and file dtype. + /// Creates a new [`FileStatistics`] from the given nested aggregates, legacy top-level-only + /// aggregates, 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. + /// `aggregates` must follow the post-order nested-struct layout produced by + /// [`postorder_stats_layout`] for `file_dtype`. `legacy_aggregates` must have one entry per + /// top-level struct field (or a single entry for a non-struct root dtype); they are written + /// in the legacy `ArrayStats` format for readers that predate nested field statistics. /// /// # 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 dtypes = struct_fields.fields().collect(); + /// Panics if the number of entries doesn't match the expected number based on the dtype. + pub fn new_with_dtype( + aggregates: Arc<[AggregateStats]>, + legacy_aggregates: Arc<[AggregateStats]>, + file_dtype: &DType, + ) -> Self { + let layout = postorder_stats_layout(file_dtype); + assert_eq!( + aggregates.len(), + layout.len(), + "aggregates length must match the post-order stats layout for the file dtype" + ); - 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()]), - } - } + let (paths, dtypes): (Vec, Vec) = layout.into_iter().unzip(); + Self::from_parts(aggregates, dtypes.into(), paths.into(), legacy_aggregates) } /// 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. + /// Aggregates the session doesn't know are skipped: file statistics are optional, so a reader + /// without an aggregate loses only what that aggregate could tell it. pub fn from_flatbuffer<'a>( fb: &fb::FileStatistics<'a>, file_dtype: &DType, session: &VortexSession, ) -> VortexResult { + if let Some(nested_field_aggregates) = fb.nested_field_aggregates() { + let aggregate_fns: Vec> = fb + .aggregate_specs() + .iter() + .flat_map(|specs| specs.iter()) + .map(|spec| aggregate_fn_from_spec(&spec, session)) + .try_collect()?; + let layout = postorder_stats_layout(file_dtype); + vortex_ensure_eq!(nested_field_aggregates.len(), layout.len()); + + let mut aggregates = Vec::with_capacity(layout.len()); + let mut dtypes = Vec::with_capacity(layout.len()); + let mut paths = Vec::with_capacity(layout.len()); + for (field_aggregates, (path, dtype)) in nested_field_aggregates.iter().zip(layout) { + aggregates.push(aggregate_stats_from_flatbuffer( + &field_aggregates, + &aggregate_fns, + &dtype, + session, + )?); + dtypes.push(dtype); + paths.push(path); + } + + return Ok(Self::from_parts( + aggregates.into(), + dtypes.into(), + paths.into(), + Arc::new([]), + )); + } + + // Legacy (pre-nested-stats) layout: top-level struct fields only, or a single entry for a + // non-struct root dtype. let field_stats = fb.field_stats().unwrap_or_default(); + let mut array_stats: Vec = field_stats.iter().collect(); if let DType::Struct(struct_fields, _) = file_dtype { - vortex_ensure_eq!(field_stats.len(), struct_fields.nfields()); + vortex_ensure_eq!(array_stats.len(), struct_fields.nfields()); - let stats_sets: Arc<[StatsSet]> = field_stats + let aggregates: Arc<[AggregateStats]> = array_stats .into_iter() .zip(struct_fields.fields()) .map(|(array_stat, field_dtype)| { - StatsSet::from_flatbuffer(&array_stat, &field_dtype, session) + aggregates_from_array_stats(&array_stat, &field_dtype, session) }) .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, - }) + Ok(Self::new(aggregates, dtypes, paths)) } else { - vortex_ensure_eq!(field_stats.len(), 1); + vortex_ensure_eq!(array_stats.len(), 1); - let array_stat = field_stats.get(0); - let stats_set = StatsSet::from_flatbuffer(&array_stat, file_dtype, session)?; + let array_stat = array_stats + .pop() + .vortex_expect("we just checked that there was 1 field"); + let aggregates = aggregates_from_array_stats(&array_stat, file_dtype, session)?; - Ok(Self { - stats: Arc::new([stats_set]), - dtypes: Arc::new([file_dtype.clone()]), - }) + Ok(Self::new( + Arc::new([aggregates]), + Arc::new([file_dtype.clone()]), + Arc::new([FieldPath::root()]), + )) } } - /// Returns a reference to the statistics sets. - pub fn stats_sets(&self) -> &Arc<[StatsSet]> { - &self.stats + /// Returns a reference to the aggregates recorded for each entry. + pub fn aggregates(&self) -> &Arc<[AggregateStats]> { + &self.aggregates + } + + /// Returns `true` if there is no statistical information at all, in either the nested or the + /// legacy layout. + /// + /// These can disagree: a non-nullable struct field whose entire subtree is unsupported dtypes + /// (e.g. all-`Variant`) contributes no entries to the nested post-order layout (every leaf is + /// skipped, and there's no nullable struct along the way to emit an own entry), but still gets + /// a legacy top-level entry with a real `NullCount`, since that stat is dtype-agnostic. + pub fn is_empty(&self) -> bool { + self.aggregates.is_empty() && self.legacy_aggregates.is_empty() } /// Returns a reference to the data types. @@ -134,22 +236,120 @@ impl FileStatistics { &self.dtypes } - /// Returns the statistics and data type for a specific field. + /// Returns a reference to the field paths. + pub fn paths(&self) -> &Arc<[FieldPath]> { + &self.paths + } + + /// Returns the aggregates and data type for a specific field. /// /// # Panics /// /// Panics if `field_idx` is out of bounds. - pub fn get(&self, field_idx: usize) -> (&StatsSet, &DType) { - (&self.stats[field_idx], &self.dtypes[field_idx]) + pub fn get(&self, field_idx: usize) -> (&AggregateStats, &DType) { + (&self.aggregates[field_idx], &self.dtypes[field_idx]) + } + + /// Returns the aggregates and data type for the field at the given path, if present. + pub fn get_by_path(&self, path: &FieldPath) -> Option<(&AggregateStats, &DType)> { + self.path_index + .get(path) + .map(|&idx| (&self.aggregates[idx], &self.dtypes[idx])) } } +/// The aggregates a legacy [`ArrayStats`] over a field of `dtype` holds, for the stats that have an +/// aggregate function. +fn aggregates_from_array_stats( + array_stats: &ArrayStats<'_>, + dtype: &DType, + session: &VortexSession, +) -> VortexResult { + let stats_set = StatsSet::from_flatbuffer(array_stats, dtype, session)?; + Ok(AggregateStats::new( + stats_set + .iter() + .filter_map(|(stat, value)| { + let aggregate_fn = stat.aggregate_fn()?; + let stat_dtype = stat.dtype(dtype)?; + let value = value + .clone() + .map(|value| Scalar::try_new(stat_dtype.clone(), Some(value))) + .transpose() + .ok()?; + Some(AggregateStat::from_value(aggregate_fn, value)) + }) + .collect(), + )) +} + +/// Resolves an [`fb::AggregateSpec`]'s aggregate function, or `None` if the session doesn't +/// know it. +fn aggregate_fn_from_spec( + spec: &fb::AggregateSpec<'_>, + session: &VortexSession, +) -> VortexResult> { + #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")] + let id = AggregateFnId::new(spec.id()); + let Some(plugin) = session.aggregate_fns().find_plugin(&id) else { + return Ok(None); + }; + let options = spec + .options() + .map(|options| options.bytes()) + .unwrap_or_default(); + let aggregate_fn = plugin.deserialize(options, session)?; + if aggregate_fn.id() != id { + vortex_bail!( + "Aggregate function ID mismatch: expected {}, got {}", + id, + aggregate_fn.id() + ); + } + Ok(Some(aggregate_fn)) +} + +fn aggregate_stats_from_flatbuffer( + field_aggregates: &fb::FieldAggregates<'_>, + aggregate_fns: &[Option], + dtype: &DType, + session: &VortexSession, +) -> VortexResult { + let mut aggregates = Vec::new(); + for state in field_aggregates.aggregates().unwrap_or_default() { + let spec = usize::from(state.aggregate_spec()); + let Some(aggregate_fn) = aggregate_fns.get(spec).ok_or_else(|| { + vortex_err!( + "aggregate spec index {spec} out of range for {} specs", + aggregate_fns.len() + ) + })? + else { + continue; + }; + let aggregate_fn = aggregate_fn.clone(); + let partial_dtype = aggregate_fn.state_dtype(dtype).ok_or_else(|| { + vortex_err!("aggregate {aggregate_fn} does not support field dtype {dtype}") + })?; + let partial = + ScalarValue::from_proto_bytes(state.partial().bytes(), &partial_dtype, session)?; + let partial = Scalar::try_new(partial_dtype, partial)?; + aggregates.push(AggregateStat::try_from_partial( + aggregate_fn, + partial, + dtype, + )?); + } + Ok(AggregateStats::new(aggregates)) +} + impl<'a> IntoIterator for &'a FileStatistics { - type Item = (&'a StatsSet, &'a DType); - type IntoIter = std::iter::Zip, std::slice::Iter<'a, DType>>; + type Item = (&'a AggregateStats, &'a DType); + type IntoIter = + std::iter::Zip, std::slice::Iter<'a, DType>>; fn into_iter(self) -> Self::IntoIter { - self.stats.iter().zip(self.dtypes.iter()) + self.aggregates.iter().zip(self.dtypes.iter()) } } @@ -162,18 +362,508 @@ impl WriteFlatBuffer for FileStatistics { &self, fbb: &mut FlatBufferBuilder<'fb>, ) -> VortexResult>> { + // The legacy `ArrayStats` format has a fixed slot per stat rather than open-ended + // aggregates. let field_stats = self - .stats_sets() + .legacy_aggregates .iter() - .map(|s| s.write_flatbuffer(fbb)) + .map(|aggregates| aggregates.to_stats_set().write_flatbuffer(fbb)) .collect::>>()?; let field_stats = fbb.create_vector(field_stats.as_slice()); + // Aggregate functions are dictionary-encoded, like the footer's array and layout IDs, so + // each distinct function and its options are written once. + let mut specs: Vec<&AggregateFnRef> = Vec::new(); + let mut spec_indices: HashMap<&AggregateFnRef, u16> = HashMap::default(); + let mut nested_field_aggregates = Vec::with_capacity(self.aggregates.len()); + for field_aggregates in self.aggregates.iter() { + let mut states = Vec::new(); + // Values without a partial state (e.g. from legacy statistics) can't be written. + for stat in field_aggregates.iter() { + let Some(partial) = stat.partial() else { + continue; + }; + let aggregate_fn = stat.aggregate_fn(); + let aggregate_spec = match spec_indices.get(aggregate_fn) { + Some(&index) => index, + None => { + let index = u16::try_from(specs.len()).map_err(|_| { + vortex_err!( + "file statistics exceed {} distinct aggregate functions", + u16::MAX + ) + })?; + specs.push(aggregate_fn); + spec_indices.insert(aggregate_fn, index); + index + } + }; + let partial = + fbb.create_vector(&ScalarValue::to_proto_bytes::>(partial.value())); + states.push(fb::AggregateState::create( + fbb, + &fb::AggregateStateArgs { + aggregate_spec, + partial: Some(partial), + }, + )); + } + let aggregates = fbb.create_vector(states.as_slice()); + nested_field_aggregates.push(fb::FieldAggregates::create( + fbb, + &fb::FieldAggregatesArgs { + aggregates: Some(aggregates), + }, + )); + } + let nested_field_aggregates = fbb.create_vector(nested_field_aggregates.as_slice()); + + let aggregate_specs = specs + .into_iter() + .map(|aggregate_fn| { + let options = aggregate_fn.options().serialize()?.ok_or_else(|| { + vortex_err!( + "Aggregate function '{}' is not serializable", + aggregate_fn.id() + ) + })?; + let id = fbb.create_string(aggregate_fn.id().as_ref()); + let options = fbb.create_vector(options.as_slice()); + Ok(fb::AggregateSpec::create( + fbb, + &fb::AggregateSpecArgs { + id: Some(id), + options: Some(options), + }, + )) + }) + .collect::>>()?; + let aggregate_specs = fbb.create_vector(aggregate_specs.as_slice()); + Ok(fb::FileStatistics::create( fbb, &fb::FileStatisticsArgs { field_stats: Some(field_stats), + aggregate_specs: Some(aggregate_specs), + nested_field_aggregates: Some(nested_field_aggregates), }, )) } } + +#[cfg(test)] +mod tests { + use std::num::NonZeroUsize; + + use flatbuffers::FlatBufferBuilder; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; + use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; + use vortex_array::aggregate_fn::fns::max::Max; + use vortex_array::aggregate_fn::fns::min::Min; + use vortex_array::aggregate_fn::fns::null_count::NullCount; + use vortex_array::array_session; + use vortex_array::arrays::VarBinViewArray; + 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::flatbuffers::WriteFlatBufferExt; + use vortex_array::scalar::ScalarValue; + + use super::*; + + fn i32_dtype() -> DType { + DType::Primitive(PType::I32, Nullability::NonNullable) + } + + /// An aggregate over a field of `dtype` whose partial state is `value` itself, as for `Min` + /// and `NullCount`. + fn stat(aggregate_fn: AggregateFnRef, value: Scalar, dtype: &DType) -> AggregateStat { + let partial_dtype = aggregate_fn + .state_dtype(dtype) + .vortex_expect("aggregate supports dtype"); + let partial = value.cast(&partial_dtype).vortex_expect("castable partial"); + AggregateStat::try_from_partial(aggregate_fn, partial, dtype).vortex_expect("valid partial") + } + + fn min_fn() -> AggregateFnRef { + Min.bind(NumericalAggregateOpts::skip_nans()) + } + + /// The value `aggregates` resolves for `aggregate_fn`. + fn value(aggregates: &AggregateStats, aggregate_fn: &AggregateFnRef) -> Precision { + aggregates.get(aggregate_fn).and_then(Scalar::into_value) + } + + fn min(value: i32) -> AggregateStat { + stat(min_fn(), Scalar::from(value), &i32_dtype()) + } + + fn null_count(value: u64, dtype: &DType) -> AggregateStat { + stat(NullCount.bind(EmptyOptions), Scalar::from(value), dtype) + } + + fn read_back(file_stats: &FileStatistics, file_dtype: &DType) -> VortexResult { + let bytes = file_stats.write_flatbuffer_bytes()?; + let fb = flatbuffers::root::(bytes.as_ref()) + .vortex_expect("valid flatbuffer"); + FileStatistics::from_flatbuffer(&fb, file_dtype, &array_session()) + } + + #[test] + fn nested_round_trip_resolves_by_path() -> VortexResult<()> { + let inner = DType::struct_([("b", i32_dtype())], Nullability::Nullable); + let file_dtype = DType::struct_([("a", inner.clone())], Nullability::NonNullable); + + // Layout: [a.b, a] (a's own null-count entry trails its child). + let file_stats = FileStatistics::new_with_dtype( + Arc::from([ + AggregateStats::new(vec![min(1)]), + AggregateStats::new(vec![null_count(1, &inner)]), + ]), + Arc::from([AggregateStats::default()]), + &file_dtype, + ); + + let read_back = read_back(&file_stats, &file_dtype)?; + + let (b, _) = read_back + .get_by_path(&FieldPath::from_name("a").push("b")) + .expect("a.b stats"); + assert_eq!( + value(b, &min_fn()), + Precision::exact(ScalarValue::from(1i32)) + ); + + let (a, _) = read_back + .get_by_path(&FieldPath::from_name("a")) + .expect("a's own null-count stats"); + assert_eq!( + value(a, &NullCount.bind(EmptyOptions)), + Precision::exact(ScalarValue::from(1u64)) + ); + + assert!(read_back.get_by_path(&FieldPath::root()).is_none()); + + Ok(()) + } + + #[test] + fn aggregate_specs_are_written_once_per_function_and_options() -> VortexResult<()> { + let file_dtype = DType::struct_( + [("a", i32_dtype()), ("b", i32_dtype()), ("c", i32_dtype())], + Nullability::NonNullable, + ); + let min_including_nans = Min.bind(NumericalAggregateOpts::include_nans()); + let file_stats = FileStatistics::new_with_dtype( + Arc::from([ + AggregateStats::new(vec![min(1), null_count(0, &i32_dtype())]), + AggregateStats::new(vec![min(2), null_count(0, &i32_dtype())]), + // Same ID as `min`, different options: a spec of its own. + AggregateStats::new(vec![stat( + min_including_nans.clone(), + Scalar::from(3i32), + &i32_dtype(), + )]), + ]), + Arc::from([ + AggregateStats::default(), + AggregateStats::default(), + AggregateStats::default(), + ]), + &file_dtype, + ); + + let bytes = file_stats.write_flatbuffer_bytes()?; + let fb = flatbuffers::root::(bytes.as_ref()) + .vortex_expect("valid flatbuffer"); + let specs = fb + .aggregate_specs() + .vortex_expect("aggregate specs") + .iter() + .map(|spec| { + ( + spec.id().to_string(), + spec.options().map(|options| options.bytes().to_vec()), + ) + }) + .collect::>(); + let options = |aggregate_fn: &AggregateFnRef| aggregate_fn.options().serialize(); + assert_eq!( + specs, + [ + ("vortex.min".to_string(), options(&min_fn())?), + ( + "vortex.null_count".to_string(), + options(&NullCount.bind(EmptyOptions))? + ), + ("vortex.min".to_string(), options(&min_including_nans)?), + ] + ); + + let read_back = read_back(&file_stats, &file_dtype)?; + let (b, _) = read_back + .get_by_path(&FieldPath::from_name("b")) + .expect("b stats"); + assert_eq!( + value(b, &min_fn()), + Precision::exact(ScalarValue::from(2i32)) + ); + let (c, _) = read_back + .get_by_path(&FieldPath::from_name("c")) + .expect("c stats"); + assert_eq!( + c.iter().next().map(AggregateStat::aggregate_fn), + Some(&min_including_nans) + ); + Ok(()) + } + + #[test] + fn bounded_aggregate_round_trips_as_an_approximate_max() -> VortexResult<()> { + let file_dtype = DType::struct_( + [("s", DType::Utf8(Nullability::NonNullable))], + Nullability::NonNullable, + ); + let field_dtype = DType::Utf8(Nullability::NonNullable); + let bounded_max = BoundedMax.bind(BoundedMaxOptions { + max_bytes: NonZeroUsize::new(4).vortex_expect("non-zero"), + }); + let mut accumulator = bounded_max.accumulator(&field_dtype)?; + accumulator.accumulate( + &VarBinViewArray::from_iter_str(["a value past the bound"]).into_array(), + &mut array_session().create_execution_ctx(), + )?; + let partial = accumulator.partial_scalar()?; + let file_stats = FileStatistics::new_with_dtype( + Arc::from([AggregateStats::new(vec![AggregateStat::try_from_partial( + bounded_max.clone(), + partial, + &field_dtype, + )?])]), + Arc::from([AggregateStats::default()]), + &file_dtype, + ); + + let read_back = read_back(&file_stats, &file_dtype)?; + let path = FieldPath::from_name("s"); + let (aggregates, _) = read_back.get_by_path(&path).expect("s aggregates"); + assert!( + aggregates + .iter() + .all(|stat| stat.aggregate_fn() == &bounded_max) + ); + assert_eq!( + value(aggregates, &Max.bind(NumericalAggregateOpts::skip_nans())), + Precision::inexact(ScalarValue::from("a vb")) + ); + Ok(()) + } + + #[test] + fn unknown_aggregates_are_skipped() -> VortexResult<()> { + let file_dtype = DType::struct_([("col", i32_dtype())], Nullability::NonNullable); + let known = min(7); + let partial = ScalarValue::to_proto_bytes::>( + known.partial().vortex_expect("has a partial").value(), + ); + + let options = known + .aggregate_fn() + .options() + .serialize()? + .vortex_expect("serializable options"); + + // Spec 0 is an aggregate no session knows, spec 1 is the known `min`. + let mut fbb = FlatBufferBuilder::new(); + let aggregate_specs = ["test.unknown", known.aggregate_fn().id().as_ref()].map(|id| { + let id = fbb.create_string(id); + let options = fbb.create_vector(options.as_slice()); + fb::AggregateSpec::create( + &mut fbb, + &fb::AggregateSpecArgs { + id: Some(id), + options: Some(options), + }, + ) + }); + let aggregate_specs = fbb.create_vector(&aggregate_specs); + let states = [0u16, 1].map(|aggregate_spec| { + let partial = fbb.create_vector(partial.as_slice()); + fb::AggregateState::create( + &mut fbb, + &fb::AggregateStateArgs { + aggregate_spec, + partial: Some(partial), + }, + ) + }); + let aggregates = fbb.create_vector(&states); + let field_aggregates = fb::FieldAggregates::create( + &mut fbb, + &fb::FieldAggregatesArgs { + aggregates: Some(aggregates), + }, + ); + let nested_field_aggregates = fbb.create_vector(&[field_aggregates]); + let root = fb::FileStatistics::create( + &mut fbb, + &fb::FileStatisticsArgs { + field_stats: None, + aggregate_specs: Some(aggregate_specs), + nested_field_aggregates: Some(nested_field_aggregates), + }, + ); + fbb.finish_minimal(root); + let bytes = fbb.finished_data().to_vec(); + + let fb = flatbuffers::root::(&bytes).vortex_expect("valid flatbuffer"); + let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &array_session())?; + let (aggregates, _) = read_back + .get_by_path(&FieldPath::from_name("col")) + .expect("col aggregates"); + assert_eq!(aggregates.iter().count(), 1); + assert_eq!( + value(aggregates, &min_fn()), + Precision::exact(ScalarValue::from(7i32)) + ); + Ok(()) + } + + #[test] + fn get_by_path_resolves_a_later_sibling_past_a_nullable_struct_field() { + // Regression test: a nullable struct field ("s") contributes both an `s.b` entry and a + // trailing entry for its own null count to the nested (post-order) layout, so that layout + // has one more entry (3) than there are top-level fields (2: "s", "c"). Resolving the + // later top-level field "c" by path must not be thrown off by "s"'s extra own-entry. + let s_dtype = DType::struct_([("b", i32_dtype())], Nullability::Nullable); + let file_dtype = DType::struct_( + [("s", s_dtype.clone()), ("c", i32_dtype())], + Nullability::NonNullable, + ); + + // Post-order layout: [s.b, s, c]. + let file_stats = FileStatistics::new_with_dtype( + Arc::from([ + AggregateStats::default(), + AggregateStats::new(vec![null_count(0, &s_dtype)]), + AggregateStats::new(vec![min(42)]), + ]), + Arc::from([AggregateStats::default(), AggregateStats::default()]), + &file_dtype, + ); + + let (c, _) = file_stats + .get_by_path(&FieldPath::from_name("c")) + .expect("c stats"); + assert_eq!( + value(c, &min_fn()), + Precision::exact(ScalarValue::from(42i32)) + ); + } + + #[test] + fn legacy_non_nested_footer_still_parses() -> VortexResult<()> { + // Simulates a footer written before nested field stats existed: + // `nested_field_aggregates` is absent, 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), + aggregate_specs: None, + nested_field_aggregates: None, + }, + ); + fbb.finish_minimal(root); + let bytes = fbb.finished_data().to_vec(); + + let fb = flatbuffers::root::(&bytes).vortex_expect("valid flatbuffer"); + assert!(fb.nested_field_aggregates().is_none()); + + let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &session)?; + let path = FieldPath::from_name("col"); + let (col, _) = read_back.get_by_path(&path).expect("col aggregates"); + assert_eq!( + value(col, &min_fn()), + Precision::exact(ScalarValue::from(7i32)) + ); + + Ok(()) + } + + #[test] + fn writer_still_emits_legacy_field_stats_matching_top_level_field_count() -> VortexResult<()> { + // Regression test: an old reader (pre-nested-stats) only ever looks at `field_stats` and + // requires its length to match the number of top-level struct fields. A nested/nullable + // struct schema's post-order layout has a different length, so `field_stats` must keep + // carrying the legacy top-level-only shape, not the nested one. + let inner = DType::struct_([("b", i32_dtype())], Nullability::Nullable); + let file_dtype = + DType::struct_([("a", inner), ("c", i32_dtype())], Nullability::NonNullable); + let struct_fields = file_dtype + .as_struct_fields_opt() + .vortex_expect("file_dtype is a struct"); + + // Nested (post-order) layout: [a.b, a, c] - 3 entries, differs from the 2 top-level fields. + let nested: Arc<[AggregateStats]> = Arc::from([ + AggregateStats::default(), + AggregateStats::default(), + AggregateStats::default(), + ]); + // Legacy layout: one entry per top-level field ("a", "c"). + let legacy: Arc<[AggregateStats]> = + Arc::from([AggregateStats::default(), AggregateStats::default()]); + + let file_stats = FileStatistics::new_with_dtype(nested, legacy, &file_dtype); + let bytes = file_stats.write_flatbuffer_bytes()?; + let fb = flatbuffers::root::(bytes.as_ref()) + .vortex_expect("valid flatbuffer"); + + let field_stats_len = fb.field_stats().map_or(0, |field_stats| field_stats.len()); + assert_eq!(field_stats_len, struct_fields.nfields()); + + Ok(()) + } + + #[test] + fn is_empty_considers_legacy_stats_too() { + // Regression test: the nested (post-order) layout can be empty while the legacy layout + // still carries real content, e.g. a non-nullable struct field whose entire subtree is + // unsupported dtypes contributes no nested entries, but still gets a legacy NullCount + // (dtype-agnostic). `is_empty` must not report "nothing to write" in that case, or the + // caller (the footer serializer) would silently drop the legacy stats too. + // A non-nullable struct field whose only child is an unsupported dtype (`Variant`) + // contributes zero entries to the nested post-order layout: the leaf is skipped, and the + // struct itself is non-nullable so it gets no own entry either. + let file_dtype = DType::struct_( + [("a", DType::Variant(Nullability::NonNullable))], + Nullability::NonNullable, + ); + assert_eq!(postorder_stats_layout(&file_dtype), Vec::new()); + let legacy = AggregateStats::new(vec![null_count( + 0, + &DType::Variant(Nullability::NonNullable), + )]); + + let file_stats = + FileStatistics::new_with_dtype(Arc::from([]), Arc::from([legacy]), &file_dtype); + + assert!(!file_stats.is_empty()); + } +} diff --git a/vortex-file/src/footer/serializer.rs b/vortex-file/src/footer/serializer.rs index 6997c1d30b0..e72e58c0389 100644 --- a/vortex-file/src/footer/serializer.rs +++ b/vortex-file/src/footer/serializer.rs @@ -147,7 +147,7 @@ impl FooterSerializer { let statistics_segment = match self.footer.statistics() { None => None, - Some(stats) if stats.stats_sets().is_empty() => None, + Some(stats) if stats.is_empty() => None, Some(stats) => { let (buffer, stats_segment) = write_flatbuffer(&mut self.offset, stats)?; buffers.push(buffer); diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index df327638d00..8d1486cb7f9 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -9,11 +9,8 @@ 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; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::literal::Literal; @@ -29,17 +26,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 +55,6 @@ pub(crate) fn can_prune_file_stats( struct FileStatsBinder<'a> { file_stats: &'a FileStatistics, - struct_fields: &'a StructFields, } impl StatBinder for FileStatsBinder<'_> { @@ -70,35 +62,32 @@ impl StatBinder for FileStatsBinder<'_> { &self, input: &BoundExpression, aggregate_fn: &AggregateFnRef, - _stat_dtype: &DType, + stat_dtype: &DType, ) -> VortexResult> { - let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { - return Ok(None); - }; let Some(field_path) = direct_field_path(input) else { return Ok(None); }; - Ok(self.stat_ref(&field_path, stat)) + Ok(self.aggregate_ref(&field_path, aggregate_fn, stat_dtype)) } } 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 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_scalar = Scalar::try_new(stat_dtype, Some(stat_value)).ok()?; - - Some(lit(stat_scalar)) + /// Binds `aggregate_fn` to the value of the file aggregate that best satisfies it. Pruning + /// only needs a bound, so an approximate value (e.g. a byte-bounded max for `max`) works too. + fn aggregate_ref( + &self, + field_path: &FieldPath, + aggregate_fn: &AggregateFnRef, + stat_dtype: &DType, + ) -> Option { + let (aggregates, _) = self.file_stats.get_by_path(field_path)?; + let value = aggregates.get(aggregate_fn).into_inner()?; + let value = if value.dtype() == stat_dtype { + value + } else { + value.cast(stat_dtype).ok()? + }; + Some(lit(value)) } } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index feeaf67f3af..0061c68af01 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -20,6 +20,7 @@ use tempfile::tempdir; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; @@ -38,6 +39,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; @@ -59,16 +61,17 @@ 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::Precision; 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; -use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::validity::Validity; @@ -95,6 +98,7 @@ use vortex_layout::DynLayout; use vortex_layout::LayoutStrategy; use vortex_layout::layouts::buffered::BufferedStrategy; use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; +use vortex_layout::layouts::file_stats::AggregateStats; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::struct_::StructStrategy; use vortex_layout::layouts::table::TableStrategy; @@ -127,6 +131,21 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +fn pruning_aggregate_fns() -> Vec { + [Stat::Min, Stat::Max, Stat::NullCount, Stat::NaNCount] + .into_iter() + .filter_map(|stat| stat.aggregate_fn()) + .collect() +} + +/// The value the file statistics entry `aggregates` resolves for `stat`. +fn stat_value(aggregates: &AggregateStats, stat: Stat) -> Precision { + let aggregate_fn = stat + .aggregate_fn() + .vortex_expect("test only uses stats with an aggregate fn"); + aggregates.get(&aggregate_fn).and_then(Scalar::into_value) +} + fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { StrictSortedBuffer::try_new(indices).expect("test indices should be strictly increasing") } @@ -1449,27 +1468,89 @@ 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![]; + let mut buf = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_file_statistics(pruning_aggregate_fns()) + .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!( + stat_value(root_stats, Stat::NullCount), + Precision::exact(ScalarValue::from(2u64)) + ); + + Ok(()) +} + +#[tokio::test] +async fn exclude_legacy_statistics_omits_legacy_but_keeps_nested() -> VortexResult<()> { + let inner = StructArray::try_new( + ["b"].into(), + vec![PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array()], + 3, + Validity::NonNullable, + )? + .into_array(); + let c = PrimitiveArray::from_iter([4i32, 5, 6]).into_array(); + let array = StructArray::try_new(["a", "c"].into(), vec![inner, c], 3, Validity::NonNullable)? + .into_array(); + + let mut buf_with_legacy = ByteBufferMut::empty(); SESSION .write_options() - .write(&mut writer, array.to_array_stream()) - .await - .unwrap(); + .with_file_statistics(pruning_aggregate_fns()) + .write(&mut buf_with_legacy, array.to_array_stream()) + .await?; + + let mut buf_without_legacy = ByteBufferMut::empty(); + let summary = SESSION + .write_options() + .with_file_statistics(pruning_aggregate_fns()) + .exclude_legacy_statistics() + .write(&mut buf_without_legacy, array.to_array_stream()) + .await?; + + assert!( + buf_without_legacy.len() < buf_with_legacy.len(), + "excluding legacy statistics should shrink the footer" + ); + + // Nested stats should still be fully populated and resolvable by path. + let stats = summary + .footer() + .statistics() + .expect("file statistics should be present"); + let (b_stats, _) = stats + .get_by_path(&field_path!(a.b)) + .expect("nested field stats should still resolve by path"); + assert_eq!( + stat_value(b_stats, Stat::NullCount), + Precision::exact(ScalarValue::from(1u64)) + ); + + Ok(()) } async fn round_trip( @@ -2169,7 +2250,7 @@ async fn test_writer_with_statistics() -> VortexResult<()> { let mut buf = ByteBufferMut::empty(); let mut writer = SESSION .write_options() - .with_file_statistics(PRUNING_STATS.to_vec()) + .with_file_statistics(pruning_aggregate_fns()) .writer(&mut buf, array.dtype().clone()); writer.push(array).await?; @@ -2188,9 +2269,15 @@ async fn file_sum_is_absent_when_a_chunk_overflows() -> VortexResult<()> { Nullability::NonNullable, ); let mut buf = ByteBufferMut::empty(); + // `Sum` is not in the core editions. let mut writer = SESSION .write_options() - .with_file_statistics(vec![Stat::Sum]) + .disable_editions() + .with_file_statistics(vec![ + Stat::Sum + .aggregate_fn() + .vortex_expect("sum has an aggregate fn"), + ]) .writer(&mut buf, dtype); // The first chunk overflows, so the file sum must not be the second chunk's sum of 2. @@ -2205,13 +2292,64 @@ async fn file_sum_is_absent_when_a_chunk_overflows() -> VortexResult<()> { .footer() .statistics() .vortex_expect("file statistics were requested"); - assert!(footer_stats.stats_sets()[0].get(Stat::Sum).is_absent()); + let (numbers, _) = footer_stats.get(0); + assert!(stat_value(numbers, Stat::Sum).is_absent()); let file = SESSION.open_options().open_buffer(buf)?; let file_stats = file .file_stats() .vortex_expect("file statistics were written"); - assert!(file_stats.stats_sets()[0].get(Stat::Sum).is_absent()); + let (numbers, _) = file_stats.get(0); + assert!(stat_value(numbers, Stat::Sum).is_absent()); + + Ok(()) +} + +#[tokio::test] +async fn default_string_stats_are_exact_within_the_byte_bound_after_reopening() -> VortexResult<()> +{ + // The default string min/max are byte-bounded, and only exact when nothing was truncated. + // The footer stores partial states, which don't record that exactness, so it has to survive + // as the recorded aggregate itself. + let long = "x".repeat(100); + let array = StructArray::from_fields(&[ + ( + "short", + VarBinViewArray::from_iter_str(["apple", "banana"]).into_array(), + ), + ( + "long", + VarBinViewArray::from_iter_str(["a", long.as_str()]).into_array(), + ), + ])? + .into_array(); + + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, array.to_array_stream()) + .await?; + let file = SESSION.open_options().open_buffer(ByteBuffer::from(buf))?; + let stats = file + .file_stats() + .expect("file statistics should be present"); + + let (short, _) = stats.get_by_path(&field_path!(short)).expect("short stats"); + assert_eq!( + stat_value(short, Stat::Max), + Precision::exact(ScalarValue::from("banana")) + ); + assert_eq!( + stat_value(short, Stat::Min), + Precision::exact(ScalarValue::from("apple")) + ); + + let (long, _) = stats.get_by_path(&field_path!(long)).expect("long stats"); + assert!(stat_value(long, Stat::Max).as_inexact().is_some()); + assert_eq!( + stat_value(long, Stat::Min), + Precision::exact(ScalarValue::from("a")) + ); Ok(()) } @@ -2941,6 +3079,66 @@ 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 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-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 03ad2ab88d4..77a189d2408 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, ) } @@ -167,10 +152,17 @@ mod tests { use vortex_array::ArrayContext; use vortex_array::IntoArray as _; + use vortex_array::aggregate_fn::AggregateFnVTableExt; + use vortex_array::aggregate_fn::EmptyOptions; + use vortex_array::aggregate_fn::NumericalAggregateOpts; + use vortex_array::aggregate_fn::fns::max::Max; + use vortex_array::aggregate_fn::fns::min::Min; + use vortex_array::aggregate_fn::fns::null_count::NullCount; use vortex_array::arrays::PrimitiveArray; 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; @@ -181,10 +173,8 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::expr::stats::Precision; - use vortex_array::expr::stats::Stat; use vortex_array::extension::datetime::TimeUnit; - use vortex_array::scalar::ScalarValue; - use vortex_array::stats::StatsSet; + use vortex_array::scalar::Scalar; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; @@ -192,6 +182,8 @@ mod tests { use vortex_io::session::RuntimeSessionExt; use vortex_layout::LayoutReader; use vortex_layout::LayoutStrategy; + use vortex_layout::layouts::file_stats::AggregateStat; + use vortex_layout::layouts::file_stats::AggregateStats; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; use vortex_layout::layouts::table::TableStrategy; use vortex_layout::segments::SegmentSink; @@ -211,24 +203,32 @@ mod tests { }); fn test_file_stats(min: i32, max: i32) -> FileStatistics { - let mut stats = StatsSet::default(); - stats.set(Stat::Min, Precision::exact(ScalarValue::from(min))); - stats.set(Stat::Max, Precision::exact(ScalarValue::from(max))); + let stats = AggregateStats::new(vec![ + AggregateStat::from_value( + Min.bind(NumericalAggregateOpts::skip_nans()), + Precision::exact(Scalar::from(min)), + ), + AggregateStat::from_value( + Max.bind(NumericalAggregateOpts::skip_nans()), + Precision::exact(Scalar::from(max)), + ), + ]); FileStatistics::new( Arc::from([stats]), Arc::from([DType::Primitive(PType::I32, Nullability::NonNullable)]), + Arc::from([FieldPath::from_name("col")]), ) } fn test_file_null_count_stats(null_count: u64) -> FileStatistics { - let mut stats = StatsSet::default(); - stats.set( - Stat::NullCount, - Precision::exact(ScalarValue::from(null_count)), - ); + let stats = AggregateStats::new(vec![AggregateStat::from_value( + NullCount.bind(EmptyOptions), + Precision::exact(Scalar::from(null_count)), + )]); FileStatistics::new( Arc::from([stats]), Arc::from([DType::Primitive(PType::I32, Nullability::Nullable)]), + Arc::from([FieldPath::from_name("col")]), ) } @@ -387,9 +387,15 @@ mod tests { let child = layout.new_reader("".into(), segments, &SESSION, &Default::default())?; // 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 stats = AggregateStats::new(vec![AggregateStat::from_value( + NullCount.bind(EmptyOptions), + Precision::exact(Scalar::from(1u64)), + )]); + 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-file/src/writer.rs b/vortex-file/src/writer.rs index 5f356e0d07f..6c97dc458db 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -19,13 +19,12 @@ use itertools::Itertools; use vortex_array::ArrayContext; use vortex_array::ArrayId; use vortex_array::ArrayRef; +use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; -use vortex_array::expr::stats::Stat; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorExt; use vortex_array::session::ArraySessionExt; -use vortex_array::stats::PRUNING_STATS; use vortex_array::stream::ArrayStream; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; @@ -83,8 +82,11 @@ pub struct VortexWriteOptions { disable_editions: bool, buffered_bytes: BufferedBytesTracker, exclude_dtype: bool, - max_variable_length_statistics_size: usize, - file_statistics: Vec, + /// The aggregates to compute for file-level statistics. + /// + /// If unset, the writer chooses pruning aggregates from the file dtype. + file_statistics: Option>, + write_legacy_statistics: bool, metadata: HashMap, } @@ -106,8 +108,8 @@ impl VortexWriteOptions { buffered_bytes: BufferedBytesTracker::new(), session, exclude_dtype: false, - file_statistics: PRUNING_STATS.to_vec(), - max_variable_length_statistics_size: 64, + file_statistics: None, + write_legacy_statistics: true, metadata: HashMap::default(), } } @@ -156,9 +158,25 @@ impl VortexWriteOptions { /// Configure which statistics to compute at the file level. /// - /// Pass an empty vector to omit file-level statistics. - pub fn with_file_statistics(mut self, file_statistics: Vec) -> Self { - self.file_statistics = file_statistics; + /// Pass an empty vector to omit file-level statistics. If left unset, the writer chooses + /// pruning aggregates from the file dtype. + /// + /// Like zone maps, file statistics only record aggregates the enabled editions include: the + /// write fails if this asks for one they don't, while the defaults leave such aggregates out. + pub fn with_file_statistics(mut self, file_statistics: Vec) -> Self { + self.file_statistics = Some(file_statistics); + self + } + + /// Exclude legacy top-level-only statistics (`field_stats`) from the file, computing and + /// writing only the full nested post-order statistics (`nested_field_aggregates`). + /// + /// Readers built before nested file stats existed only look at `field_stats` and validate its + /// length against the number of top-level struct fields; omitting it means those readers find + /// no usable statistics in this file (though they can still read the data). Only use this if + /// you control every reader of the resulting files and don't need that compatibility. + pub fn exclude_legacy_statistics(mut self) -> Self { + self.write_legacy_statistics = false; self } @@ -273,12 +291,16 @@ impl VortexWriteOptions { .map(move |result| result.map(|chunk| (ptr.advance(), chunk))), ) .sendable(); + // When unset, `accumulate_stats` resolves each leaf's own default from its dtype, the way + // `default_zoned_aggregate_fns` does for zoned layouts (e.g. bounded min/max for + // variable-length columns), rather than a single dtype-blind default applied uniformly. let (file_stats, stream) = accumulate_stats( stream, - self.file_statistics.clone().into(), - self.max_variable_length_statistics_size, + self.file_statistics.clone().map(Arc::from), &self.session, - ); + self.write_legacy_statistics, + &ctx, + )?; // First, write the magic bytes. write.write_all(ByteBuffer::copy_from(MAGIC_BYTES)).await?; @@ -320,11 +342,12 @@ impl VortexWriteOptions { let (layout, segment_specs) = layout_fut.await?; // Assemble the Footer object now that we have all the segments. - let statistics = if self.file_statistics.is_empty() { + let statistics = if matches!(&self.file_statistics, Some(stats) if stats.is_empty()) { None } else { Some(FileStatistics::new_with_dtype( - file_stats.stats_sets().into(), + file_stats.aggregate_stats()?.into(), + file_stats.legacy_aggregate_stats()?.into(), &dtype, )) }; diff --git a/vortex-jni/src/writer.rs b/vortex-jni/src/writer.rs index 794aa1cc65c..fc7060195e3 100644 --- a/vortex-jni/src/writer.rs +++ b/vortex-jni/src/writer.rs @@ -33,18 +33,24 @@ use jni::sys::jlong; use jni::sys::jobject; use object_store::ObjectStore; use object_store::path::Path as ObjectStorePath; +use vortex::aggregate_fn::AggregateFnRef; +use vortex::aggregate_fn::AggregateFnVTableExt; +use vortex::aggregate_fn::EmptyOptions; +use vortex::aggregate_fn::NumericalAggregateOpts; +use vortex::aggregate_fn::fns::max::Max; +use vortex::aggregate_fn::fns::min::Min; +use vortex::aggregate_fn::fns::nan_count::NanCount; +use vortex::aggregate_fn::fns::null_count::NullCount; use vortex::array::ArrayRef; use vortex::array::scalar::PValue; use vortex::array::scalar::Scalar; use vortex::array::scalar::ScalarValue; -use vortex::array::stats::StatsSet; use vortex::array::stream::ArrayStreamAdapter; use vortex::dtype::DType; +use vortex::dtype::FieldPath; use vortex::error::VortexError; use vortex::error::VortexResult; use vortex::error::vortex_err; -use vortex::expr::stats::Stat; -use vortex::expr::stats::StatsProvider; use vortex::file::CountingVortexWrite; use vortex::file::WriteOptionsSessionExt; use vortex::file::WriteSummary; @@ -56,6 +62,7 @@ use vortex::io::runtime::BlockingRuntime; use vortex::io::runtime::Task; use vortex::io::session::RuntimeSessionExt; use vortex::layout::BufferedBytesTracker; +use vortex::layout::layouts::file_stats::AggregateStats; use vortex::session::VortexSession; use vortex::utils::aliases::hash_map::HashMap; use vortex_arrow::ArrowSessionExt; @@ -182,14 +189,13 @@ fn checked_jlong(value: u64, name: &str) -> VortexResult { } fn exact_count_jlong( - stats: Option<&StatsSet>, - dtype: Option<&DType>, - stat: Stat, + aggregates: Option<&AggregateStats>, + aggregate_fn: &AggregateFnRef, ) -> VortexResult { - stats - .zip(dtype.and_then(|dt| stat.dtype(dt))) - .and_then(|(stats, dt)| stats.get_as::(stat, &dt).as_exact()) - .map(|value| checked_jlong(value, stat.name())) + aggregates + .and_then(|aggregates| aggregates.get(aggregate_fn).as_exact()) + .and_then(|count| u64::try_from(&count).ok()) + .map(|value| checked_jlong(value, &aggregate_fn.to_string())) .transpose() .map(|value| value.unwrap_or(-1)) } @@ -285,6 +291,7 @@ fn write_summary_to_java<'local>( ) -> Result, JNIError> { let column_sizes = summary.compressed_column_sizes()?; let file_stats = summary.footer().statistics(); + let struct_fields = summary.footer().dtype().as_struct_fields_opt(); let columns = env.new_object_array( i32::try_from(column_sizes.len()) .map_err(|_| vortex_err!("column count exceeds Java array range"))?, @@ -293,22 +300,25 @@ fn write_summary_to_java<'local>( )?; for (column_index, compressed_size) in column_sizes.into_iter().enumerate() { - let (stats, dtype) = file_stats - .and_then(|all_stats| { - all_stats - .stats_sets() - .get(column_index) - .zip(all_stats.dtypes().get(column_index)) + let aggregates = struct_fields + .and_then(|fields| fields.names().get(column_index)) + .zip(file_stats) + .and_then(|(name, all_stats)| { + all_stats.get_by_path(&FieldPath::from_name(name.clone())) }) - .map_or((None, None), |(stats, dtype)| (Some(stats), Some(dtype))); - let null_count = exact_count_jlong(stats, dtype, Stat::NullCount)?; - let nan_count = exact_count_jlong(stats, dtype, Stat::NaNCount)?; - let lower_bound = stats - .zip(dtype) - .and_then(|(stats, dtype)| stats.as_typed_ref(dtype).get(Stat::Min).into_inner()); - let upper_bound = stats - .zip(dtype) - .and_then(|(stats, dtype)| stats.as_typed_ref(dtype).get(Stat::Max).into_inner()); + .map(|(aggregates, _)| aggregates); + let null_count = exact_count_jlong(aggregates, &NullCount.bind(EmptyOptions))?; + let nan_count = exact_count_jlong(aggregates, &NanCount.bind(EmptyOptions))?; + let lower_bound = aggregates.and_then(|aggregates| { + aggregates + .get(&Min.bind(NumericalAggregateOpts::skip_nans())) + .into_inner() + }); + let upper_bound = aggregates.and_then(|aggregates| { + aggregates + .get(&Max.bind(NumericalAggregateOpts::skip_nans())) + .into_inner() + }); let column = env.with_local_frame_returning_local::<_, JObject, JNIError>(16, |env| { let lower_bound = match lower_bound { Some(value) => scalar_to_java(env, value)?, diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 3d031a65c7a..9f5ca85b30d 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::future; -use std::marker::PhantomData; use std::sync::Arc; use futures::StreamExt; @@ -11,50 +10,62 @@ 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::aggregate_fn::AggregateFnSatisfaction; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; +use vortex_array::aggregate_fn::fns::bounded_max::BoundedMaxOptions; +use vortex_array::aggregate_fn::fns::bounded_min::BoundedMin; +use vortex_array::aggregate_fn::fns::bounded_min::BoundedMinOptions; +use vortex_array::aggregate_fn::fns::max::Max; +use vortex_array::aggregate_fn::fns::min::Min; +use vortex_array::aggregate_fn::fns::nan_count::NanCount; +use vortex_array::aggregate_fn::fns::null_count::NullCount; +use vortex_array::aggregate_fn::fns::sum::Sum; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::builders::ArrayBuilder; -use vortex_array::builders::BoolBuilder; -use vortex_array::builders::builder_with_capacity_in; use vortex_array::dtype::DType; -use vortex_array::dtype::FieldName; +use vortex_array::dtype::Field; +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; -use vortex_array::scalar::ScalarTruncation; -use vortex_array::scalar::lower_bound; -use vortex_array::scalar::upper_bound; use vortex_array::stats::StatsSet; -use vortex_array::validity::Validity; -use vortex_buffer::BufferString; -use vortex_buffer::ByteBuffer; -use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use vortex_error::vortex_panic; use vortex_session::VortexSession; -use crate::layouts::zoned::MAX_IS_TRUNCATED; -use crate::layouts::zoned::MIN_IS_TRUNCATED; +use crate::LayoutWriterContext; +use crate::layouts::zoned::default_bounded_stat_max_bytes; use crate::sequence::SendableSequentialStream; use crate::sequence::SequenceId; use crate::sequence::SequentialStreamAdapter; use crate::sequence::SequentialStreamExt; +/// Accumulates file statistics over `stream` as it passes through. +/// +/// `stats` are the aggregates to record for every field, or `None` for each field's default (see +/// `default_pruning_aggregate_fns`). Only aggregates `ctx` allows are recorded in the nested +/// statistics: an explicitly requested aggregate it forbids fails, like it does for zone maps, +/// while forbidden defaults are left out. +/// +/// # Errors +/// +/// Returns an error if `stats` contains an aggregate `ctx` forbids. pub fn accumulate_stats( stream: SendableSequentialStream, - stats: Arc<[Stat]>, - max_variable_length_statistics_size: usize, + stats: Option>, session: &VortexSession, -) -> (FileStatsAccumulator, SendableSequentialStream) { - let accumulator = FileStatsAccumulator::new( - stream.dtype(), - stats, - max_variable_length_statistics_size, - session, - ); + write_legacy_stats: bool, + ctx: &LayoutWriterContext, +) -> VortexResult<(FileStatsAccumulator, SendableSequentialStream)> { + let accumulator = + FileStatsAccumulator::try_new(stream.dtype(), stats, session, write_legacy_stats, ctx)?; let stream = SequentialStreamAdapter::new( stream.dtype().clone(), stream.scan(accumulator.clone(), |acc, item| { @@ -62,439 +73,459 @@ pub fn accumulate_stats( }), ) .sendable(); - (accumulator, stream) + Ok((accumulator, stream)) } -/// Accumulates write-time statistics for a single file column. -struct StatsAccumulator { - builders: Vec>, - length: usize, +/// One aggregate function's value over a file statistics entry. +#[derive(Clone, Debug)] +pub struct AggregateStat { + aggregate_fn: AggregateFnRef, + /// The aggregate's partial state, which is what the footer stores. `None` for values recovered + /// from the legacy `ArrayStats` footer format, which only stores results. + partial: Option, + value: Precision, } -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(), - length: 0, - }; - } +impl AggregateStat { + fn from_accumulator( + aggregate_fn: AggregateFnRef, + accumulator: &AccumulatorRef, + ) -> VortexResult { + Ok(Self { + aggregate_fn, + partial: Some(accumulator.partial_scalar()?), + value: Precision::exact(accumulator.final_scalar()?), + }) + } - let builders = stats - .iter() - .filter_map(|&stat| { - stat.dtype(dtype).map(|stat_dtype| { - stats_builder_with_capacity( - stat, - &stat_dtype.as_nullable(), - 1024, - max_variable_length_statistics_size, - ) - }) - }) - .collect::>(); + /// Rebuilds an aggregate's value from its `partial` state over a field of `dtype`. + pub fn try_from_partial( + aggregate_fn: AggregateFnRef, + partial: Scalar, + dtype: &DType, + ) -> VortexResult { + let mut accumulator = aggregate_fn.accumulator(dtype)?; + accumulator.combine_partials(partial.clone())?; + let value = Precision::exact(accumulator.final_scalar()?); + Ok(Self { + aggregate_fn, + partial: Some(partial), + value, + }) + } + /// An aggregate value without a partial state, such as one read from the legacy `ArrayStats` + /// footer format. + pub fn from_value(aggregate_fn: AggregateFnRef, value: Precision) -> Self { Self { - builders, - length: 0, + aggregate_fn, + partial: None, + value, } } - fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { - for builder in &mut self.builders { - if let Some(value) = array.statistics().compute_stat(builder.stat(), ctx)? { - builder.append_scalar(value.cast(&value.dtype().as_nullable())?)?; - } else { - builder.append_null(); - } - } - self.length += 1; - Ok(()) + /// The aggregate function this value was computed by. + pub fn aggregate_fn(&self) -> &AggregateFnRef { + &self.aggregate_fn } - fn as_array(&mut self, ctx: &mut ExecutionCtx) -> VortexResult> { - let mut names = Vec::new(); - let mut fields = Vec::new(); - - for builder in self - .builders - .iter_mut() - // We sort the stats so the DType is deterministic based on which stats are present. - .sorted_unstable_by_key(|builder| builder.stat()) - { - let values = builder.finish(); - - // We drop any all-null stats columns. - if values.all_invalid(ctx)? { - continue; - } + /// The aggregate's partial state, if known. + pub fn partial(&self) -> Option<&Scalar> { + self.partial.as_ref() + } - names.extend(values.names); - fields.extend(values.arrays); - } + /// The aggregate's value. A null value means the aggregate saw no valid input. + pub fn value(&self) -> &Precision { + &self.value + } +} - if names.is_empty() { - return Ok(None); - } +/// The aggregates recorded for one file statistics entry. +#[derive(Clone, Debug, Default)] +pub struct AggregateStats { + aggregates: Vec, +} - StructArray::try_new(names.into(), fields, self.length, Validity::NonNullable).map(Some) +impl AggregateStats { + pub fn new(aggregates: Vec) -> Self { + Self { aggregates } } - /// Returns an aggregated stats set for the table. - fn as_stats_set(&mut self, stats: &[Stat], ctx: &mut ExecutionCtx) -> VortexResult { - let mut stats_set = StatsSet::default(); - let Some(stats_table) = self.as_array(ctx)? else { - return Ok(stats_set); - }; + pub fn is_empty(&self) -> bool { + self.aggregates.is_empty() + } - for &stat in stats { - let Some(values) = stats_table.unmasked_field_by_name_opt(stat.name()) else { - continue; - }; + pub fn iter(&self) -> impl Iterator { + self.aggregates.iter() + } - 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::Sum if !values.all_valid(ctx)? => { - // `values` holds each chunk's `Stat::Sum` from `push_chunk`. The legacy `Sum` - // aggregate behind it returns zero for empty input and null only on overflow. - // Summing `values` below would skip those nulls and report a wrong exact - // total, so leave the file sum unset. - 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) - } + /// Returns the value of `requested`, from the recorded aggregate that best satisfies it, the + /// way zone maps resolve aggregates: one that satisfies it exactly over one that only + /// approximates it. The value is exact only if the recorded value is exact and satisfies + /// `requested` exactly. Absent if no recorded aggregate satisfies `requested`, or if the value + /// is null. + pub fn get(&self, requested: &AggregateFnRef) -> Precision { + let mut approximate = Precision::Absent; + for stat in &self.aggregates { + match stat.aggregate_fn.can_satisfy(requested) { + AggregateFnSatisfaction::Exact => return non_null(stat.value.clone()), + AggregateFnSatisfaction::Approximate => { + approximate = stat.value.clone().into_inexact(); } - 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 => {} + AggregateFnSatisfaction::No => {} } } - Ok(stats_set) + non_null(approximate) } -} -fn stat_was_truncated( - stats_table: &StructArray, - stat: Stat, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let field_name = match stat { - Stat::Min => MIN_IS_TRUNCATED, - Stat::Max => MAX_IS_TRUNCATED, - _ => return Ok(false), - }; - let Some(is_truncated) = stats_table.unmasked_field_by_name_opt(field_name) else { - return Ok(false); - }; - - Ok(is_truncated - .statistics() - .compute_stat(Stat::Max, ctx)? - .is_some_and(|max| max.as_bool().value() == Some(true))) -} - -fn supports_file_stats(dtype: &DType) -> bool { - !matches!(dtype, DType::Variant(_)) -} - -fn is_varlen_dtype(dtype: &DType) -> bool { - matches!(dtype, DType::Utf8(_) | DType::Binary(_)) -} - -fn stats_builder_with_capacity( - stat: Stat, - dtype: &DType, - capacity: usize, - max_length: usize, -) -> Box { - let values_builder = builder_with_capacity_in( - dtype, - capacity, - vortex_buffer::BufferAllocatorRef::static_ref(), - ); - match stat { - Stat::Max => match dtype { - DType::Utf8(_) => Box::new(TruncatedMaxBinaryStatsBuilder::::new( - values_builder, - BoolBuilder::with_capacity_in( - Nullability::NonNullable, - capacity, - vortex_buffer::BufferAllocatorRef::static_ref(), - ), - max_length, - )), - DType::Binary(_) => Box::new(TruncatedMaxBinaryStatsBuilder::::new( - values_builder, - BoolBuilder::with_capacity_in( - Nullability::NonNullable, - capacity, - vortex_buffer::BufferAllocatorRef::static_ref(), - ), - max_length, - )), - _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), - }, - Stat::Min => match dtype { - DType::Utf8(_) => Box::new(TruncatedMinBinaryStatsBuilder::::new( - values_builder, - BoolBuilder::with_capacity_in( - Nullability::NonNullable, - capacity, - vortex_buffer::BufferAllocatorRef::static_ref(), - ), - max_length, - )), - DType::Binary(_) => Box::new(TruncatedMinBinaryStatsBuilder::::new( - values_builder, - BoolBuilder::with_capacity_in( - Nullability::NonNullable, - capacity, - vortex_buffer::BufferAllocatorRef::static_ref(), - ), - max_length, - )), - _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), - }, - _ => Box::new(StatNameArrayBuilder::new(stat, values_builder)), + /// Returns the recorded aggregates as a [`StatsSet`], for consumers keyed by [`Stat`]. + pub fn to_stats_set(&self) -> StatsSet { + let mut stats_set = StatsSet::default(); + for stat in Stat::all() { + let Some(requested) = stat.aggregate_fn() else { + continue; + }; + let value = self.get(&requested).and_then(Scalar::into_value); + if !value.is_absent() { + stats_set.set(stat, value); + } + } + stats_set } } -/// Arrays with their associated names, reduced version of a `StructArray`. -struct NamedArrays { - names: Vec, - arrays: Vec, +fn non_null(value: Precision) -> Precision { + value.and_then(|value| (!value.is_null()).then_some(value)) } -impl NamedArrays { - fn all_invalid(&self, ctx: &mut ExecutionCtx) -> VortexResult { - self.arrays[0].all_invalid(ctx) - } +/// Accumulates write-time statistics for a single file column. +struct StatsAccumulator { + aggregates: Vec<(AggregateFnRef, AccumulatorRef)>, } -trait StatsArrayBuilder: Send { - fn stat(&self) -> Stat; - - fn append_scalar(&mut self, value: Scalar) -> VortexResult<()>; - - fn append_null(&mut self); +impl StatsAccumulator { + fn new(dtype: &DType, stats: Option<&[AggregateFnRef]>) -> Self { + if !supports_file_stats(dtype) { + return Self { + aggregates: Vec::new(), + }; + } - fn finish(&mut self) -> NamedArrays; -} + // When no explicit list is requested, each leaf picks its own default the way + // `default_zoned_aggregate_fns` does for zoned layouts: a byte-bounded min/max for + // variable-length columns, exact min/max otherwise. + let default_stats; + let stats = match stats { + Some(stats) => stats, + None => { + default_stats = default_pruning_aggregate_fns(dtype); + default_stats.as_slice() + } + }; -struct StatNameArrayBuilder { - stat: Stat, - builder: Box, -} + let mut aggregates = Vec::new(); + for aggregate_fn in stats { + // A dtype that doesn't support a given aggregate simply fails to build an + // accumulator, which is silently skipped, matching this stat's absence from the + // result. + if let Ok(accumulator) = aggregate_fn.accumulator(dtype) { + aggregates.push((aggregate_fn.clone(), accumulator)); + } + } -impl StatNameArrayBuilder { - fn new(stat: Stat, builder: Box) -> Self { - Self { stat, builder } + Self { aggregates } } -} -impl StatsArrayBuilder for StatNameArrayBuilder { - fn stat(&self) -> Stat { - self.stat + fn push_chunk(&mut self, array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<()> { + for (_, accumulator) in &mut self.aggregates { + accumulator.accumulate(array, ctx)?; + } + Ok(()) } - fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> { - self.builder.append_scalar(&value) + /// Returns the accumulated aggregates that `allowed` permits. + /// + /// A bounded aggregate whose accumulator ended up exact, because no value exceeded its byte + /// bound, is recorded as the exact aggregate it stands in for (e.g. `BoundedMax` as `Max`). + /// The footer stores partial states, which don't record that exactness, so recording the exact + /// aggregate is what keeps it. + fn aggregate_stats( + &self, + allowed: impl Fn(&AggregateFnRef) -> bool, + ) -> VortexResult { + let mut out = Vec::with_capacity(self.aggregates.len()); + for (aggregate_fn, accumulator) in &self.aggregates { + if let Some(exact) = exact_counterpart(aggregate_fn) + && allowed(&exact) + && accumulator.can_satisfy(&exact).is_exact() + { + // The exact counterparts' partial state is their result, which is the bounded + // aggregate's result here. + let value = accumulator.final_scalar()?; + out.push(AggregateStat { + aggregate_fn: exact, + partial: Some(value.clone()), + value: Precision::exact(value), + }); + } else if allowed(aggregate_fn) { + out.push(AggregateStat::from_accumulator( + aggregate_fn.clone(), + accumulator, + )?); + } + } + Ok(AggregateStats::new(out)) } +} - fn append_null(&mut self) { - self.builder.append_null() +/// The exact aggregate a bounded aggregate approximates, if any. +fn exact_counterpart(aggregate_fn: &AggregateFnRef) -> Option { + if aggregate_fn.is::() { + Some(Max.bind(NumericalAggregateOpts::skip_nans())) + } else if aggregate_fn.is::() { + Some(Min.bind(NumericalAggregateOpts::skip_nans())) + } else { + None } +} - fn finish(&mut self) -> NamedArrays { - NamedArrays { - names: vec![self.stat.name().into()], - arrays: vec![self.builder.finish()], - } - } +fn supports_file_stats(dtype: &DType) -> bool { + !matches!(dtype, DType::Variant(_)) } -struct TruncatedMaxBinaryStatsBuilder { - values: Box, - is_truncated: BoolBuilder, - max_value_length: usize, - _marker: PhantomData, +fn is_varlen_dtype(dtype: &DType) -> bool { + matches!(dtype, DType::Utf8(_) | DType::Binary(_)) } -impl TruncatedMaxBinaryStatsBuilder { - fn new( - values: Box, - is_truncated: BoolBuilder, - max_value_length: usize, - ) -> Self { - Self { - values, - is_truncated, - max_value_length, - _marker: PhantomData, - } - } +/// Default file-level pruning aggregates for `dtype`, chosen the way `default_zoned_aggregate_fns` +/// picks zoned aggregates: a byte-bounded min/max (capped at [`default_bounded_stat_max_bytes`]) +/// for variable-length columns, and exact min/max otherwise. Callers wanting a different bound pass +/// their own `BoundedMax`/`BoundedMin` instead. +fn default_pruning_aggregate_fns(dtype: &DType) -> Vec { + let (max, min) = if is_varlen_dtype(dtype) { + let max_bytes = default_bounded_stat_max_bytes(); + ( + BoundedMax.bind(BoundedMaxOptions { max_bytes }), + BoundedMin.bind(BoundedMinOptions { max_bytes }), + ) + } else { + ( + Max.bind(NumericalAggregateOpts::skip_nans()), + Min.bind(NumericalAggregateOpts::skip_nans()), + ) + }; + + vec![ + max, + min, + Sum.bind(NumericalAggregateOpts::skip_nans()), + NullCount.bind(EmptyOptions), + NanCount.bind(EmptyOptions), + ] } -struct TruncatedMinBinaryStatsBuilder { - values: Box, - is_truncated: BoolBuilder, - max_value_length: usize, - _marker: PhantomData, +/// Computes the post-order sequence of `(FieldPath, DType)` entries that file-level statistics +/// are stored against. +/// +/// Each leaf field (including opaque `List`/`FixedSizeList` columns, which are not recursed into) +/// gets one entry. Each **nullable** struct additionally gets a trailing entry, keyed by its own +/// path and dtype, inserted immediately after its children's entries — this carries the struct's +/// own null count. Dtypes that don't support file stats (see `supports_file_stats`), such as +/// [`DType::Variant`], are skipped entirely: no entry is emitted for them or anything beneath +/// them. +/// +/// This function is the single source of truth for the number and order of stats entries, used +/// both when accumulating stats at write time and when reconstructing them from the footer at +/// read time, so no paths or dtypes need to be persisted in the flatbuffer itself. +pub fn postorder_stats_layout(dtype: &DType) -> Vec<(FieldPath, DType)> { + let mut out = Vec::new(); + postorder_stats_layout_into(dtype, FieldPath::root(), &mut out); + out } -impl TruncatedMinBinaryStatsBuilder { - fn new( - values: Box, - is_truncated: BoolBuilder, - max_value_length: usize, - ) -> Self { - Self { - values, - is_truncated, - max_value_length, - _marker: PhantomData, +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())), } } -impl StatsArrayBuilder for TruncatedMaxBinaryStatsBuilder { - fn stat(&self) -> Stat { - Stat::Max - } +/// A node in the tree of accumulators mirroring [`postorder_stats_layout`]'s walk of a `DType`. +enum StatsNode { + /// An opaque leaf: a dtype with no addressable children. Currently includes List/FixedSizeList and Map. + Leaf(StatsAccumulator), + /// A dtype that does not support file stats (e.g. [`DType::Variant`]); contributes no entries. + Skipped, + /// A dtype with addressable children. Currently only built for `DType::Struct`. + Container { + /// One child per addressable sub-dtype, in declaration order (one per struct field, + /// today). + children: Vec<(Field, StatsNode)>, + /// Stats computed over the container's own, undecomposed dtype (e.g. a `Struct`'s null + /// count), reusing `Leaf`/`Skipped` rather than a separate accumulator-plus-flag pair: + /// `Leaf` when the container is nullable (matching the trailing entry + /// [`postorder_stats_layout_into`] emits for it), `Skipped` otherwise — which also means a non-nullable struct's null count is never even + /// accumulated, not just never emitted. + own: Box, + }, +} - fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> { - let nullability = value.dtype().nullability(); - if let Some((upper_bound, truncated)) = - upper_bound(T::from_scalar(value)?, self.max_value_length, nullability) - { - self.values.append_scalar(&upper_bound)?; - self.is_truncated.append_value(truncated); - } else { - self.append_null() +impl StatsNode { + fn build(dtype: &DType, stats: Option<&[AggregateFnRef]>) -> Self { + match dtype.as_struct_fields_opt() { + Some(struct_fields) => { + let children = struct_fields + .names() + .iter() + .zip(struct_fields.fields()) + .map(|(name, field_dtype)| { + (Field::Name(name.clone()), Self::build(&field_dtype, stats)) + }) + .collect(); + let own = if dtype.nullability() == Nullability::Nullable { + Self::Leaf(StatsAccumulator::new(dtype, stats)) + } else { + Self::Skipped + }; + Self::Container { + children, + own: Box::new(own), + } + } + None if !supports_file_stats(dtype) => Self::Skipped, + None => Self::Leaf(StatsAccumulator::new(dtype, stats)), } - Ok(()) - } - - fn append_null(&mut self) { - ArrayBuilder::append_null(self.values.as_mut()); - self.is_truncated.append_value(false); } - fn finish(&mut self) -> NamedArrays { - NamedArrays { - names: vec![Stat::Max.name().into(), MAX_IS_TRUNCATED.into()], - arrays: vec![ - ArrayBuilder::finish(self.values.as_mut()), - ArrayBuilder::finish(&mut self.is_truncated), - ], + 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::Container { .. } => { + let struct_array = array.clone().execute::(ctx)?; + self.push_struct_chunk(array, &struct_array, ctx) + } } } -} - -impl StatsArrayBuilder for TruncatedMinBinaryStatsBuilder { - fn stat(&self) -> Stat { - Stat::Min - } - fn append_scalar(&mut self, value: Scalar) -> VortexResult<()> { - let nullability = value.dtype().nullability(); - if let Some((lower_bound, truncated)) = - lower_bound(T::from_scalar(value)?, self.max_value_length, nullability) + /// 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. + /// + /// `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()) { - self.values.append_scalar(&lower_bound)?; - self.is_truncated.append_value(truncated); - } else { - self.append_null() + child.push_chunk(field, ctx)?; } Ok(()) } - fn append_null(&mut self) { - ArrayBuilder::append_null(self.values.as_mut()); - self.is_truncated.append_value(false); - } - - fn finish(&mut self) -> NamedArrays { - NamedArrays { - names: vec![Stat::Min.name().into(), MIN_IS_TRUNCATED.into()], - arrays: vec![ - ArrayBuilder::finish(self.values.as_mut()), - ArrayBuilder::finish(&mut self.is_truncated), - ], + /// Appends this node's aggregates that `allowed` permits, in the same post-order as + /// [`postorder_stats_layout`]. + fn collect_aggregate_stats( + &self, + allowed: &impl Fn(&AggregateFnRef) -> bool, + out: &mut Vec, + ) -> VortexResult<()> { + match self { + Self::Skipped => Ok(()), + Self::Leaf(acc) => { + out.push(acc.aggregate_stats(allowed)?); + Ok(()) + } + Self::Container { children, own } => { + for (_, child) in children { + child.collect_aggregate_stats(allowed, out)?; + } + own.collect_aggregate_stats(allowed, out) + } } } } -/// An array stream processor that computes aggregate statistics for all fields. +/// 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. /// -/// Note: for now this only collects top-level struct fields. +/// 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. `legacy` is empty (and never accumulated into) +/// when `write_legacy_stats` is `false`. #[derive(Clone)] pub struct FileStatsAccumulator { - stats: Arc<[Stat]>, - accumulators: Arc>>, + root: Arc>, + legacy: Arc>>, + write_legacy_stats: bool, + writer_ctx: LayoutWriterContext, ctx: Arc>, } impl FileStatsAccumulator { - fn new( + fn try_new( dtype: &DType, - stats: Arc<[Stat]>, - max_variable_length_statistics_size: usize, + stats: Option>, 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 - ); - } + write_legacy_stats: bool, + writer_ctx: &LayoutWriterContext, + ) -> VortexResult { + for aggregate_fn in stats.iter().flat_map(|stats| stats.iter()) { + if !writer_ctx.allows_aggregate(&aggregate_fn.id()) { + vortex_bail!("Aggregate {} not permitted by ctx", aggregate_fn.id()); + } + } + + let root = Arc::new(Mutex::new(StatsNode::build(dtype, stats.as_deref()))); - struct_dtype + let legacy = Arc::new(Mutex::new(if write_legacy_stats { + match dtype.as_struct_fields_opt() { + Some(struct_fields) => struct_fields .fields() - .map(|field_dtype| { - StatsAccumulator::new( - &field_dtype, - &stats, - max_variable_length_statistics_size, - ) - }) - .collect() + .map(|field_dtype| StatsAccumulator::new(&field_dtype, stats.as_deref())) + .collect(), + None => vec![StatsAccumulator::new(dtype, stats.as_deref())], } - None => [StatsAccumulator::new( - dtype, - &stats, - max_variable_length_statistics_size, - )] - .into(), + } else { + Vec::new() })); - Self { - stats, - accumulators, + Ok(Self { + root, + legacy, + write_legacy_stats, + writer_ctx: writer_ctx.clone(), ctx: Arc::new(Mutex::new(session.create_execution_ctx())), - } + }) } fn process( @@ -505,144 +536,616 @@ impl FileStatsAccumulator { 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)?; + 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 { - self.accumulators.lock()[0].push_chunk(&chunk, &mut ctx)?; + if self.write_legacy_stats { + self.legacy.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 + /// Returns the accumulated aggregates of every entry in the post-order nested layout (see + /// [`postorder_stats_layout`]), leaving out the aggregates the writer context forbids. + pub fn aggregate_stats(&self) -> VortexResult> { + let allowed = + |aggregate_fn: &AggregateFnRef| self.writer_ctx.allows_aggregate(&aggregate_fn.id()); + 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_aggregate_stats(&allowed, &mut out)?; + Ok(out) + } + + /// Returns the legacy top-level-fields-only aggregates (one per top-level struct field, or a + /// single entry for a non-struct root dtype). Empty if `write_legacy_stats` was `false`. + /// + /// 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)) .collect() } } #[cfg(test)] mod tests { + use std::num::NonZeroUsize; + use rstest::rstest; + use vortex_array::ArrayContext; use vortex_array::IntoArray; + use vortex_array::aggregate_fn::AggregateFnVTable; use vortex_array::array_session; use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::bool::BoolArrayExt; + use vortex_array::builders::ArrayBuilder; use vortex_array::builders::VarBinViewBuilder; - use vortex_buffer::BitBuffer; + use vortex_array::dtype::FieldNames; + use vortex_array::dtype::PType; + use vortex_array::scalar::PValue; + use vortex_array::scalar::ScalarValue; + use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexExpect; + use vortex_utils::aliases::hash_set::HashSet; use super::*; + fn stats_set(acc: &StatsAccumulator) -> VortexResult { + Ok(acc.aggregate_stats(|_| true)?.to_stats_set()) + } + + fn node_stats_sets(node: &StatsNode) -> VortexResult> { + let mut aggregates = Vec::new(); + node.collect_aggregate_stats(&|_| true, &mut aggregates)?; + Ok(aggregates + .iter() + .map(AggregateStats::to_stats_set) + .collect()) + } + + fn agg(stat: Stat) -> AggregateFnRef { + stat.aggregate_fn() + .vortex_expect("test only uses stats with an aggregate fn") + } + #[rstest] - #[case(DType::Utf8(Nullability::NonNullable))] - #[case(DType::Binary(Nullability::NonNullable))] - fn truncates_accumulated_stats(#[case] dtype: DType) { + #[case::all_within_bound(&["short", "shorter"], true, true)] + #[case::only_min_truncated(&["short", "a value longer than the bound"], true, false)] + #[case::both_truncated(&["zz value past the bound", "aa value past the bound"], false, false)] + fn bounded_varlen_min_max_exact_iff_extremum_fits_bound( + #[values( + DType::Utf8(Nullability::NonNullable), + DType::Binary(Nullability::NonNullable) + )] + dtype: DType, + #[case] values: &[&str], + #[case] max_exact: bool, + #[case] min_exact: bool, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let max_bytes = NonZeroUsize::new(8).vortex_expect("non-zero"); + let bounded = [ + BoundedMax.bind(BoundedMaxOptions { max_bytes }), + BoundedMin.bind(BoundedMinOptions { max_bytes }), + ]; + let mut acc = StatsAccumulator::new(&dtype, Some(&bounded)); + // One chunk per value, so exactness has to survive merging across chunks. + for value in values { + let mut builder = VarBinViewBuilder::with_capacity_in( + dtype.clone(), + 1, + vortex_buffer::BufferAllocatorRef::statically_allocated(), + ); + builder.append_value(value); + acc.push_chunk(&builder.finish(), &mut ctx)?; + } + + let stats = stats_set(&acc)?; + assert_eq!(stats.get(Stat::Max).is_exact(), max_exact); + assert_eq!(stats.get(Stat::Min).is_exact(), min_exact); + Ok(()) + } + + #[test] + fn default_pruning_aggregate_fns_uses_bounded_min_max_for_varlen() { + let dtype = DType::Utf8(Nullability::NonNullable); + let fns = default_pruning_aggregate_fns(&dtype); + assert!(fns.iter().any(|f| f.is::())); + assert!(fns.iter().any(|f| f.is::())); + assert!(!fns.iter().any(|f| f.is::())); + assert!(!fns.iter().any(|f| f.is::())); + } + + #[test] + fn default_pruning_aggregate_fns_uses_plain_min_max_for_fixed_width() { + let fns = default_pruning_aggregate_fns(&i32_dtype()); + assert!(fns.iter().any(|f| f.is::())); + assert!(fns.iter().any(|f| f.is::())); + assert!(!fns.iter().any(|f| f.is::())); + assert!(!fns.iter().any(|f| f.is::())); + } + + #[test] + fn explicit_min_max_on_varlen_is_exact_and_untruncated() -> VortexResult<()> { + // An explicit override applies literally: unlike the default, plain `Max`/`Min` on a + // varlen dtype now computes the exact value with no byte-bound truncation. let mut ctx = array_session().create_execution_ctx(); + let dtype = DType::Utf8(Nullability::NonNullable); let mut builder = VarBinViewBuilder::with_capacity_in( dtype.clone(), 2, vortex_buffer::BufferAllocatorRef::statically_allocated(), ); - builder.append_value("Value to be truncated"); - builder.append_value("untruncated"); - let mut builder2 = VarBinViewBuilder::with_capacity_in( - dtype, - 2, - vortex_buffer::BufferAllocatorRef::statically_allocated(), + builder.append_value("a long value that would have been truncated before"); + builder.append_value("short"); + + let mut acc = StatsAccumulator::new(&dtype, Some(&[agg(Stat::Max), agg(Stat::Min)])); + acc.push_chunk(&builder.finish(), &mut ctx)?; + + let stats = stats_set(&acc)?; + assert!(matches!(stats.get(Stat::Max), Precision::Exact(_))); + assert!(matches!(stats.get(Stat::Min), Precision::Exact(_))); + 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, + Some(&[ + agg(Stat::Max), + agg(Stat::Min), + agg(Stat::Sum), + agg(Stat::NullCount), + ]), + ); + + 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 = stats_set(&acc)?; + assert_eq!( + stats.get(Stat::Max).as_exact(), + Some(ScalarValue::from(9i32)) ); - builder2.append_value("Another"); - builder2.append_value("wait a minute"); - let mut acc = - StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min, Stat::Sum], 12); - acc.push_chunk(&builder.finish(), &mut ctx) - .vortex_expect("push_chunk should succeed for test data"); - acc.push_chunk(&builder2.finish(), &mut ctx) - .vortex_expect("push_chunk should succeed for test data"); - let stats_table = acc - .as_array(&mut ctx) - .unwrap() - .expect("Must have stats table"); assert_eq!( - stats_table.names().as_ref(), - &[ - Stat::Max.name(), - MAX_IS_TRUNCATED, - Stat::Min.name(), - MIN_IS_TRUNCATED, + 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 { + 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) + ), ] ); - let field1_bool = stats_table - .unmasked_field(1) - .clone() - .execute::(&mut ctx) - .unwrap(); + } + + #[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!( - field1_bool.to_bit_buffer(), - BitBuffer::from(vec![false, true]) + layout, + vec![ + (FieldPath::from_name("a").push("b"), i32_dtype()), + (FieldPath::from_name("a"), inner), + ] ); - let field3_bool = stats_table - .unmasked_field(3) - .clone() - .execute::(&mut ctx) - .unwrap(); + } + + #[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!( - field3_bool.to_bit_buffer(), - BitBuffer::from(vec![true, false]) + layout, + vec![(FieldPath::from_name("a").push("b"), i32_dtype())] ); } - #[rstest] - #[case(DType::Utf8(Nullability::NonNullable))] - #[case(DType::Binary(Nullability::NonNullable))] - fn truncated_accumulated_stats_are_inexact(#[case] dtype: 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 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 = [agg(Stat::NullCount), agg(Stat::Min), agg(Stat::Max)]; + let mut node = StatsNode::build(root.dtype(), Some(&requested)); + node.push_chunk(&root, &mut ctx)?; + + let stats_sets = node_stats_sets(&node)?; + + // `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(); + + 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 = [agg(Stat::NullCount), agg(Stat::Min), agg(Stat::Max)]; + let mut node = StatsNode::build(outer.dtype(), Some(&requested)); + node.push_chunk(&outer, &mut ctx)?; + + let stats_sets = node_stats_sets(&node)?; + + // `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(()) + } + + #[test] + fn legacy_aggregate_stats_cover_only_top_level_fields() -> VortexResult<()> { + // The legacy accumulation must have exactly one entry per top-level field, built against + // that field's own dtype without recursing into it, even though it's a nested struct. + let session = array_session(); + let inner_dtype = DType::struct_([("b", i32_dtype())], Nullability::Nullable); + let dtype = DType::struct_([("a", inner_dtype)], Nullability::NonNullable); + + let acc = FileStatsAccumulator::try_new( + &dtype, + Some(Arc::from([ + agg(Stat::NullCount), + agg(Stat::Min), + agg(Stat::Max), + ])), + &session, + true, + &LayoutWriterContext::new(ArrayContext::empty()), + )?; + + 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 (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(1u64)) + ); + assert!(legacy[0].get(&agg(Stat::Min)).is_absent()); + assert!(legacy[0].get(&agg(Stat::Max)).is_absent()); + Ok(()) + } + + #[rstest] + #[case(DType::struct_([("a", i32_dtype())], Nullability::NonNullable), StructArray::new(FieldNames::from(["a"]), [buffer![1i32, 2, 3].into_array()], 3, Validity::NonNullable).into_array())] + #[case(i32_dtype(), buffer![1i32, 2, 3].into_array())] + fn write_legacy_stats_false_skips_legacy_accumulation( + #[case] dtype: DType, + #[case] chunk: ArrayRef, + ) -> VortexResult<()> { + // With `write_legacy_stats: false`, `legacy` stays empty for both struct and non-struct + // roots. This must not panic — `process` indexes into `legacy[0]` for a non-struct root + // and `zip_eq`s it against the struct's fields otherwise, both of which would panic if + // `legacy` were left empty without also gating those code paths. + let session = array_session(); + let acc = FileStatsAccumulator::try_new( + &dtype, + Some(Arc::from([agg(Stat::Min), agg(Stat::Max)])), + &session, + false, + &LayoutWriterContext::new(ArrayContext::empty()), + )?; + + let (mut ptr, _eof) = SequenceId::root().split(); + acc.process(Ok((ptr.advance(), chunk)))?; + + assert!(acc.legacy_aggregate_stats()?.is_empty()); + Ok(()) + } + + fn utf8_chunk(values: &[&str]) -> ArrayRef { + let dtype = DType::Utf8(Nullability::NonNullable); let mut builder = VarBinViewBuilder::with_capacity_in( dtype, - 2, + values.len(), vortex_buffer::BufferAllocatorRef::statically_allocated(), ); - builder.append_value("Value to be truncated"); - builder.append_value("Another truncated value"); - let mut acc = StatsAccumulator::new(builder.dtype(), &[Stat::Max, Stat::Min], 12); - acc.push_chunk(&builder.finish(), &mut ctx) - .vortex_expect("push_chunk should succeed for test data"); + for value in values { + builder.append_value(value); + } + builder.finish() + } - let stats = acc - .as_stats_set(&[Stat::Max, Stat::Min], &mut ctx) - .vortex_expect("as_stats_set should succeed for test data"); + #[rstest] + #[case::untruncated(&["short", "shorter"], true)] + #[case::truncated(&["a value longer than the bound"], false)] + fn exact_bounded_max_is_recorded_as_max( + #[case] values: &[&str], + #[case] exact: bool, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let max_bytes = NonZeroUsize::new(8).vortex_expect("non-zero"); + let mut acc = StatsAccumulator::new( + &DType::Utf8(Nullability::NonNullable), + Some(&[BoundedMax.bind(BoundedMaxOptions { max_bytes })]), + ); + acc.push_chunk(&utf8_chunk(values), &mut ctx)?; - assert!(matches!(stats.get(Stat::Min), Precision::Inexact(_))); - assert!(matches!(stats.get(Stat::Max), Precision::Inexact(_))); + let aggregates = acc.aggregate_stats(|_| true)?; + let [recorded] = aggregates.iter().collect::>()[..] else { + vortex_panic!("expected one recorded aggregate"); + }; + assert_eq!(recorded.aggregate_fn().is::(), exact); + assert_eq!(recorded.aggregate_fn().is::(), !exact); + // The stored partial must rebuild the same value on read. + let partial = recorded + .partial() + .vortex_expect("accumulated aggregates have a partial"); + let rebuilt = AggregateStat::try_from_partial( + recorded.aggregate_fn().clone(), + partial.clone(), + &DType::Utf8(Nullability::NonNullable), + )?; + assert_eq!(rebuilt.value(), recorded.value()); + Ok(()) } #[test] - fn fixed_width_stats_omit_is_truncated_columns() { + fn exact_bounded_max_stays_bounded_when_max_is_not_allowed() -> 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"); + let max_bytes = NonZeroUsize::new(8).vortex_expect("non-zero"); + let mut acc = StatsAccumulator::new( + &DType::Utf8(Nullability::NonNullable), + Some(&[BoundedMax.bind(BoundedMaxOptions { max_bytes })]), + ); + acc.push_chunk(&utf8_chunk(&["short"]), &mut ctx)?; + + let aggregates = acc.aggregate_stats(|aggregate_fn| !aggregate_fn.is::())?; + assert!( + aggregates + .iter() + .all(|stat| stat.aggregate_fn().is::()) + ); + Ok(()) + } + + #[test] + fn get_prefers_an_exact_aggregate_and_skips_nulls() -> VortexResult<()> { + let max_bytes = NonZeroUsize::new(8).vortex_expect("non-zero"); + let requested = Max.bind(NumericalAggregateOpts::skip_nans()); + let bound = Scalar::utf8("b", Nullability::Nullable); + let max = Scalar::utf8("a", Nullability::Nullable); + let bounded = AggregateStat::from_value( + BoundedMax.bind(BoundedMaxOptions { max_bytes }), + Precision::exact(bound.clone()), + ); + + let only_bounded = AggregateStats::new(vec![bounded.clone()]); + assert_eq!(only_bounded.get(&requested), Precision::inexact(bound)); + + let both = AggregateStats::new(vec![ + bounded, + AggregateStat::from_value(requested.clone(), Precision::exact(max.clone())), + ]); + assert_eq!(both.get(&requested), Precision::exact(max)); + + let null = AggregateStats::new(vec![AggregateStat::from_value( + requested.clone(), + Precision::exact(Scalar::null(DType::Utf8(Nullability::Nullable))), + )]); + assert!(null.get(&requested).is_absent()); + Ok(()) + } + + #[test] + fn forbidden_aggregates_are_gated() -> VortexResult<()> { + let session = array_session(); + let dtype = DType::struct_([("a", i32_dtype())], Nullability::NonNullable); + let ctx = LayoutWriterContext::new(ArrayContext::empty()).with_allowed_aggregates( + HashSet::from_iter([Min.id(), Max.id(), NanCount.id(), NullCount.id()]), + ); + + // An explicitly requested aggregate the context forbids fails the write. + let error = FileStatsAccumulator::try_new( + &dtype, + Some(Arc::from([agg(Stat::Sum)])), + &session, + true, + &ctx, + ) + .err() + .vortex_expect("Sum is not permitted"); + assert!(error.to_string().contains("not permitted by ctx")); + + // The defaults leave it out of the nested aggregates, but the legacy stats keep it. + let acc = FileStatsAccumulator::try_new(&dtype, None, &session, true, &ctx)?; + let chunk = StructArray::new( + FieldNames::from(["a"]), + [buffer![1i32, 2, 3].into_array()], + 3, + Validity::NonNullable, + ) + .into_array(); + let (mut ptr, _eof) = SequenceId::root().split(); + acc.process(Ok((ptr.advance(), chunk)))?; + + let aggregates = acc.aggregate_stats()?; + assert!( + aggregates[0] + .iter() + .any(|stat| stat.aggregate_fn().is::()) + ); + assert!( + !aggregates[0] + .iter() + .any(|stat| stat.aggregate_fn().is::()) + ); assert_eq!( - stats_table.names().as_ref(), - &[Stat::Max.name(), Stat::Min.name(), Stat::Sum.name()] + acc.legacy_aggregate_stats()?[0] + .get(&agg(Stat::Sum)) + .and_then(Scalar::into_value), + Precision::exact(ScalarValue::from(6i64)) ); + Ok(()) } #[rstest] @@ -655,9 +1158,13 @@ mod tests { #[case] chunks: Vec>>, #[case] expected: Option, ) -> VortexResult<()> { + // An overflowing chunk saturates the sum's partial state, which stays saturated when + // merged with later chunks, so the file sum is absent rather than the sum of the chunks + // that didn't overflow. let mut ctx = array_session().create_execution_ctx(); let dtype = DType::Primitive(PType::I64, Nullability::Nullable); - let mut acc = StatsAccumulator::new(&dtype, &[Stat::Sum], 12); + let sum = agg(Stat::Sum); + let mut acc = StatsAccumulator::new(&dtype, Some(std::slice::from_ref(&sum))); for chunk in chunks { acc.push_chunk( &PrimitiveArray::from_option_iter(chunk).into_array(), @@ -665,12 +1172,10 @@ mod tests { )?; } - let stats = acc.as_stats_set(&[Stat::Sum], &mut ctx)?; assert_eq!( - stats.get(Stat::Sum), - expected.map_or(Precision::Absent, Precision::exact) + acc.aggregate_stats(|_| true)?.get(&sum), + expected.map_or(Precision::Absent, |sum| Precision::exact(Scalar::from(sum))) ); - Ok(()) } } diff --git a/vortex-layout/src/layouts/zoned/mod.rs b/vortex-layout/src/layouts/zoned/mod.rs index 07b25ddfe34..726aefdc197 100644 --- a/vortex-layout/src/layouts/zoned/mod.rs +++ b/vortex-layout/src/layouts/zoned/mod.rs @@ -27,6 +27,7 @@ pub(crate) use builder::aggregate_partials; use prost::Message; pub use schema::MAX_IS_TRUNCATED; pub use schema::MIN_IS_TRUNCATED; +pub(crate) use schema::default_bounded_stat_max_bytes; use vortex_array::DeserializeMetadata; use vortex_array::SerializeMetadata; use vortex_array::aggregate_fn::AggregateFnRef; diff --git a/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt b/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt index 11a8b7f6888..fb38928e086 100644 --- a/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt +++ b/vortex-sqllogictest/slt/duckdb/copy_return_stats.slt @@ -11,10 +11,10 @@ TO '${WORK_DIR}/copy_return_stats.vortex' (FORMAT vortex, RETURN_STATS); ---- :\$\{WORK_DIR\}/copy_return_stats\.vortex 2 \d+ 0 \{'"d"'=\{column_size_bytes=\d+, max=1\.5, min=1\.5, null_count=1, num_values=2\}, '"i"'=\{column_size_bytes=\d+, max=2, min=1, null_count=0, num_values=2\}, '"s"'=\{column_size_bytes=\d+, max=b, min=a, null_count=0, num_values=2\}\} NULL -# Vortex reports statistics per top-level field, so nested columns are covered by the file-level -# counts only; the COPY must still succeed. +# As in parquet, statistics are reported per leaf column: struct fields are keyed by their full +# quoted path and the struct itself gets no entry. Lists are reported as a single column. query TIII?? COPY (SELECT {'a': 1, 'b': 2} AS st, [1, 2, 3] AS lst, 7 AS i) TO '${WORK_DIR}/copy_return_stats_nested.vortex' (FORMAT vortex, RETURN_STATS); ---- -:\$\{WORK_DIR\}/copy_return_stats_nested\.vortex 1 \d+ 0 \{'"i"'=\{column_size_bytes=\d+, max=7, min=7, null_count=0, num_values=1\}, '"lst"'=\{column_size_bytes=\d+, null_count=0, num_values=1\}, '"st"'=\{column_size_bytes=\d+, null_count=0, num_values=1\}\} NULL +:\$\{WORK_DIR\}/copy_return_stats_nested\.vortex 1 \d+ 0 \{'"i"'=\{column_size_bytes=\d+, max=7, min=7, null_count=0, num_values=1\}, '"lst"'=\{column_size_bytes=\d+, null_count=0, num_values=1\}, '"st"\."a"'=\{column_size_bytes=\d+, max=1, min=1, null_count=0, num_values=1\}, '"st"\."b"'=\{column_size_bytes=\d+, max=2, min=2, null_count=0, num_values=1\}\} NULL From 9a9480791a85962cd98ecd836c703cf41bf9d82f Mon Sep 17 00:00:00 2001 From: Thor Date: Tue, 6 Oct 2026 10:45:00 -0500 Subject: [PATCH 2/4] Reduce footer size by deduplicating aggregate specs. Signed-off-by: Thor --- .../flatbuffers/vortex-file/footer.fbs | 72 +-- .../src/flatbuffers/generated/footer.rs | 288 +++++------ vortex-file/src/footer/file_statistics.rs | 455 +++++++++++++----- vortex-file/src/writer.rs | 2 +- 4 files changed, 491 insertions(+), 326 deletions(-) diff --git a/vortex-file/flatbuffers/vortex-file/footer.fbs b/vortex-file/flatbuffers/vortex-file/footer.fbs index b5ff6a1fb26..636e8d4b6b7 100644 --- a/vortex-file/flatbuffers/vortex-file/footer.fbs +++ b/vortex-file/flatbuffers/vortex-file/footer.fbs @@ -54,51 +54,63 @@ table PostscriptSegment { // [postscript] // [footer] -/// An `AggregateSpec` describes an aggregate function, including its options, recorded in the -/// file statistics. +/// An aggregate function and its options, as recorded in the file statistics. +/// +/// Specs are dictionary-encoded: each distinct (id, options) pair is stored once and referenced +/// by index from `AggregateSet.specs`. table AggregateSpec { + /// The aggregate function's globally unique ID, looked up in the session's aggregate registry + /// at read time. id: string (required); - /// The aggregate function's serialized options, interpreted by the function's plugin. + /// The aggregate function's serialized options. Absent when the options serialize to nothing. options: [ubyte]; } -/// The accumulated state of one aggregate function over one file statistics entry. -table AggregateState { - /// Index of the aggregate function in `FileStatistics.aggregate_specs`. - aggregate_spec: uint16; - /// Protobuf serialized `vortex.scalar.ScalarValue` of the aggregate's partial state, typed by - /// the aggregate's partial dtype for the entry's dtype. A null value means no value was - /// accumulated. - partial: [ubyte] (required); -} - -/// The aggregate states of one file statistics entry. -table FieldAggregates { - aggregates: [AggregateState]; +/// An ordered list of aggregates recorded together for one or more file statistics entries. +/// +/// Most columns of the same kind record the same aggregates, so a file typically has only a +/// handful of distinct sets regardless of how many columns it has. +table AggregateSet { + /// Indices into `FileStatistics.aggregate_specs`, in the order their partial states appear in + /// `FileStatistics.partials`. May be empty for an entry that records nothing. + specs: [uint16] (required); } /// The `FileStatistics` object contains file-level statistics for the Vortex file. table FileStatistics { /// Statistics for each top-level field in the root schema, in schema order. If the root schema /// is not a struct, there will be a single entry in this array. Does not recurse into nested - /// struct fields; see `nested_field_aggregates` for that. + /// struct fields. /// - /// Kept in this legacy shape for backward compatibility with readers that predate - /// `nested_field_aggregates`, which strictly validate this array's length against the number - /// of top-level struct fields. + /// Kept in this legacy shape for readers that predate the aggregate-based statistics below, + /// which strictly validate this array's length against the number of top-level struct fields. + /// Independent of the fields below: neither refers to the other. field_stats: [ArrayStats]; - /// Dictionary-encoded aggregate functions referenced by `nested_field_aggregates`, up to - /// u16::MAX. Each distinct pair of ID and options is one entry. + /// Dictionary of the (aggregate ID, options) pairs used by this file, up to u16::MAX. aggregate_specs: [AggregateSpec]; - /// Aggregate states following a post-order walk of the root `DType` tree: each leaf field gets - /// an entry, and each nullable struct additionally gets a trailing entry for its own null - /// count, inserted after its children's entries. + /// Dictionary of the distinct aggregate lists used by entries, up to u16::MAX. + aggregate_sets: [AggregateSet]; + /// One element per file statistics entry, in post-order of the root `DType`: each leaf field + /// gets an entry, and each nullable struct additionally gets a trailing entry for its own null + /// count, after its children's entries. Each element is an index into `aggregate_sets`. + /// + /// The entry paths and dtypes are not stored; readers derive them from the root `DType` with + /// the same post-order walk. + /// + /// Absent for files written before aggregate-based statistics, whose readers fall back to + /// `field_stats`. + entry_sets: [uint16]; + /// The partial states of every entry's aggregates, concatenated. For each entry in + /// `entry_sets` order, and for each spec in that entry's set in order: + /// + /// - an unsigned LEB128 varint byte length `n`, then + /// - `n` bytes of a protobuf-serialized `vortex.scalar.ScalarValue`, typed by the + /// aggregate's partial dtype for the entry's dtype. A null value means no value was + /// accumulated. /// - /// Readers that know about this field should prefer it over `field_stats` when present, - /// skipping aggregates they don't recognize. It is absent for files written before nested - /// file stats were supported, and readers that don't know about it fall back to - /// `field_stats`'s legacy top-level-fields-only layout. - nested_field_aggregates: [FieldAggregates]; + /// Readers skip the partial of an aggregate they don't recognize using its length prefix, so + /// unknown aggregates never desynchronize the stream. + partials: [ubyte]; } /// The `Registry` object stores dictionary-encoded configuration for segments, diff --git a/vortex-file/src/flatbuffers/generated/footer.rs b/vortex-file/src/flatbuffers/generated/footer.rs index b1e67403280..9e0238c47e7 100644 --- a/vortex-file/src/flatbuffers/generated/footer.rs +++ b/vortex-file/src/flatbuffers/generated/footer.rs @@ -801,8 +801,10 @@ impl ::core::fmt::Debug for PostscriptSegment<'_> { pub enum AggregateSpecOffset {} #[derive(Copy, Clone, PartialEq)] -/// An `AggregateSpec` describes an aggregate function, including its options, recorded in the -/// file statistics. +/// An aggregate function and its options, as recorded in the file statistics. +/// +/// Specs are dictionary-encoded: each distinct (id, options) pair is stored once and referenced +/// by index from `AggregateSet.specs`. pub struct AggregateSpec<'a> { pub _tab: ::flatbuffers::Table<'a>, } @@ -835,6 +837,8 @@ impl<'a> AggregateSpec<'a> { } + /// The aggregate function's globally unique ID, looked up in the session's aggregate registry + /// at read time. #[inline] pub fn id(&self) -> &'a str { // Safety: @@ -842,7 +846,7 @@ impl<'a> AggregateSpec<'a> { // which contains a valid value in this slot unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<&str>>(AggregateSpec::VT_ID, None).unwrap()} } - /// The aggregate function's serialized options, interpreted by the function's plugin. + /// The aggregate function's serialized options. Absent when the options serialize to nothing. #[inline] pub fn options(&self) -> Option<::flatbuffers::Vector<'a, u8>> { // Safety: @@ -915,219 +919,106 @@ impl ::core::fmt::Debug for AggregateSpec<'_> { ds.finish() } } -pub enum AggregateStateOffset {} -#[derive(Copy, Clone, PartialEq)] - -/// The accumulated state of one aggregate function over one file statistics entry. -pub struct AggregateState<'a> { - pub _tab: ::flatbuffers::Table<'a>, -} - -impl<'a> ::flatbuffers::Follow<'a> for AggregateState<'a> { - type Inner = AggregateState<'a>; - #[inline] - unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { - Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } - } -} - -impl<'a> AggregateState<'a> { - pub const VT_AGGREGATE_SPEC: ::flatbuffers::VOffsetT = 4; - pub const VT_PARTIAL: ::flatbuffers::VOffsetT = 6; - - #[inline] - pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { - AggregateState { _tab: table } - } - #[allow(unused_mut)] - pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( - _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, - args: &'args AggregateStateArgs<'args> - ) -> ::flatbuffers::WIPOffset> { - let mut builder = AggregateStateBuilder::new(_fbb); - if let Some(x) = args.partial { builder.add_partial(x); } - builder.add_aggregate_spec(args.aggregate_spec); - builder.finish() - } - - - /// Index of the aggregate function in `FileStatistics.aggregate_specs`. - #[inline] - pub fn aggregate_spec(&self) -> u16 { - // Safety: - // Created from valid Table for this object - // which contains a valid value in this slot - unsafe { self._tab.get::(AggregateState::VT_AGGREGATE_SPEC, Some(0)).unwrap()} - } - /// Protobuf serialized `vortex.scalar.ScalarValue` of the aggregate's partial state, typed by - /// the aggregate's partial dtype for the entry's dtype. A null value means no value was - /// accumulated. - #[inline] - pub fn partial(&self) -> ::flatbuffers::Vector<'a, u8> { - // Safety: - // Created from valid Table for this object - // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u8>>>(AggregateState::VT_PARTIAL, None).unwrap()} - } -} - -impl ::flatbuffers::Verifiable for AggregateState<'_> { - #[inline] - fn run_verifier( - v: &mut ::flatbuffers::Verifier, pos: usize - ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { - v.visit_table(pos)? - .visit_field::("aggregate_spec", Self::VT_AGGREGATE_SPEC, false)? - .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u8>>>("partial", Self::VT_PARTIAL, true)? - .finish(); - Ok(()) - } -} -pub struct AggregateStateArgs<'a> { - pub aggregate_spec: u16, - pub partial: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u8>>>, -} -impl<'a> Default for AggregateStateArgs<'a> { - #[inline] - fn default() -> Self { - AggregateStateArgs { - aggregate_spec: 0, - partial: None, // required field - } - } -} - -pub struct AggregateStateBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { - fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, - start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, -} -impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> AggregateStateBuilder<'a, 'b, A> { - #[inline] - pub fn add_aggregate_spec(&mut self, aggregate_spec: u16) { - self.fbb_.push_slot::(AggregateState::VT_AGGREGATE_SPEC, aggregate_spec, 0); - } - #[inline] - pub fn add_partial(&mut self, partial: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u8>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(AggregateState::VT_PARTIAL, partial); - } - #[inline] - pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> AggregateStateBuilder<'a, 'b, A> { - let start = _fbb.start_table(); - AggregateStateBuilder { - fbb_: _fbb, - start_: start, - } - } - #[inline] - pub fn finish(self) -> ::flatbuffers::WIPOffset> { - let o = self.fbb_.end_table(self.start_); - self.fbb_.required(o, AggregateState::VT_PARTIAL,"partial"); - ::flatbuffers::WIPOffset::new(o.value()) - } -} - -impl ::core::fmt::Debug for AggregateState<'_> { - fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { - let mut ds = f.debug_struct("AggregateState"); - ds.field("aggregate_spec", &self.aggregate_spec()); - ds.field("partial", &self.partial()); - ds.finish() - } -} -pub enum FieldAggregatesOffset {} +pub enum AggregateSetOffset {} #[derive(Copy, Clone, PartialEq)] -/// The aggregate states of one file statistics entry. -pub struct FieldAggregates<'a> { +/// An ordered list of aggregates recorded together for one or more file statistics entries. +/// +/// Most columns of the same kind record the same aggregates, so a file typically has only a +/// handful of distinct sets regardless of how many columns it has. +pub struct AggregateSet<'a> { pub _tab: ::flatbuffers::Table<'a>, } -impl<'a> ::flatbuffers::Follow<'a> for FieldAggregates<'a> { - type Inner = FieldAggregates<'a>; +impl<'a> ::flatbuffers::Follow<'a> for AggregateSet<'a> { + type Inner = AggregateSet<'a>; #[inline] unsafe fn follow(buf: &'a [u8], loc: usize) -> Self::Inner { Self { _tab: unsafe { ::flatbuffers::Table::new(buf, loc) } } } } -impl<'a> FieldAggregates<'a> { - pub const VT_AGGREGATES: ::flatbuffers::VOffsetT = 4; +impl<'a> AggregateSet<'a> { + pub const VT_SPECS: ::flatbuffers::VOffsetT = 4; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { - FieldAggregates { _tab: table } + AggregateSet { _tab: table } } #[allow(unused_mut)] pub fn create<'bldr: 'args, 'args: 'mut_bldr, 'mut_bldr, A: ::flatbuffers::Allocator + 'bldr>( _fbb: &'mut_bldr mut ::flatbuffers::FlatBufferBuilder<'bldr, A>, - args: &'args FieldAggregatesArgs<'args> - ) -> ::flatbuffers::WIPOffset> { - let mut builder = FieldAggregatesBuilder::new(_fbb); - if let Some(x) = args.aggregates { builder.add_aggregates(x); } + args: &'args AggregateSetArgs<'args> + ) -> ::flatbuffers::WIPOffset> { + let mut builder = AggregateSetBuilder::new(_fbb); + if let Some(x) = args.specs { builder.add_specs(x); } builder.finish() } + /// Indices into `FileStatistics.aggregate_specs`, in the order their partial states appear in + /// `FileStatistics.partials`. May be empty for an entry that records nothing. #[inline] - pub fn aggregates(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + pub fn specs(&self) -> ::flatbuffers::Vector<'a, u16> { // Safety: // Created from valid Table for this object // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FieldAggregates::VT_AGGREGATES, None)} + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u16>>>(AggregateSet::VT_SPECS, None).unwrap()} } } -impl ::flatbuffers::Verifiable for FieldAggregates<'_> { +impl ::flatbuffers::Verifiable for AggregateSet<'_> { #[inline] fn run_verifier( v: &mut ::flatbuffers::Verifier, pos: usize ) -> Result<(), ::flatbuffers::InvalidFlatbuffer> { v.visit_table(pos)? - .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("aggregates", Self::VT_AGGREGATES, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u16>>>("specs", Self::VT_SPECS, true)? .finish(); Ok(()) } } -pub struct FieldAggregatesArgs<'a> { - pub aggregates: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, +pub struct AggregateSetArgs<'a> { + pub specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u16>>>, } -impl<'a> Default for FieldAggregatesArgs<'a> { +impl<'a> Default for AggregateSetArgs<'a> { #[inline] fn default() -> Self { - FieldAggregatesArgs { - aggregates: None, + AggregateSetArgs { + specs: None, // required field } } } -pub struct FieldAggregatesBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { +pub struct AggregateSetBuilder<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> { fbb_: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>, start_: ::flatbuffers::WIPOffset<::flatbuffers::TableUnfinishedWIPOffset>, } -impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FieldAggregatesBuilder<'a, 'b, A> { +impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> AggregateSetBuilder<'a, 'b, A> { #[inline] - pub fn add_aggregates(&mut self, aggregates: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FieldAggregates::VT_AGGREGATES, aggregates); + pub fn add_specs(&mut self, specs: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u16>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(AggregateSet::VT_SPECS, specs); } #[inline] - pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FieldAggregatesBuilder<'a, 'b, A> { + pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> AggregateSetBuilder<'a, 'b, A> { let start = _fbb.start_table(); - FieldAggregatesBuilder { + AggregateSetBuilder { fbb_: _fbb, start_: start, } } #[inline] - pub fn finish(self) -> ::flatbuffers::WIPOffset> { + pub fn finish(self) -> ::flatbuffers::WIPOffset> { let o = self.fbb_.end_table(self.start_); + self.fbb_.required(o, AggregateSet::VT_SPECS,"specs"); ::flatbuffers::WIPOffset::new(o.value()) } } -impl ::core::fmt::Debug for FieldAggregates<'_> { +impl ::core::fmt::Debug for AggregateSet<'_> { fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result { - let mut ds = f.debug_struct("FieldAggregates"); - ds.field("aggregates", &self.aggregates()); + let mut ds = f.debug_struct("AggregateSet"); + ds.field("specs", &self.specs()); ds.finish() } } @@ -1150,7 +1041,9 @@ impl<'a> ::flatbuffers::Follow<'a> for FileStatistics<'a> { impl<'a> FileStatistics<'a> { pub const VT_FIELD_STATS: ::flatbuffers::VOffsetT = 4; pub const VT_AGGREGATE_SPECS: ::flatbuffers::VOffsetT = 6; - pub const VT_NESTED_FIELD_AGGREGATES: ::flatbuffers::VOffsetT = 8; + pub const VT_AGGREGATE_SETS: ::flatbuffers::VOffsetT = 8; + pub const VT_ENTRY_SETS: ::flatbuffers::VOffsetT = 10; + pub const VT_PARTIALS: ::flatbuffers::VOffsetT = 12; #[inline] pub unsafe fn init_from_table(table: ::flatbuffers::Table<'a>) -> Self { @@ -1162,7 +1055,9 @@ impl<'a> FileStatistics<'a> { args: &'args FileStatisticsArgs<'args> ) -> ::flatbuffers::WIPOffset> { let mut builder = FileStatisticsBuilder::new(_fbb); - if let Some(x) = args.nested_field_aggregates { builder.add_nested_field_aggregates(x); } + if let Some(x) = args.partials { builder.add_partials(x); } + if let Some(x) = args.entry_sets { builder.add_entry_sets(x); } + if let Some(x) = args.aggregate_sets { builder.add_aggregate_sets(x); } if let Some(x) = args.aggregate_specs { builder.add_aggregate_specs(x); } if let Some(x) = args.field_stats { builder.add_field_stats(x); } builder.finish() @@ -1171,11 +1066,11 @@ impl<'a> FileStatistics<'a> { /// Statistics for each top-level field in the root schema, in schema order. If the root schema /// is not a struct, there will be a single entry in this array. Does not recurse into nested - /// struct fields; see `nested_field_aggregates` for that. + /// struct fields. /// - /// Kept in this legacy shape for backward compatibility with readers that predate - /// `nested_field_aggregates`, which strictly validate this array's length against the number - /// of top-level struct fields. + /// Kept in this legacy shape for readers that predate the aggregate-based statistics below, + /// which strictly validate this array's length against the number of top-level struct fields. + /// Independent of the fields below: neither refers to the other. #[inline] pub fn field_stats(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { // Safety: @@ -1183,8 +1078,7 @@ 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)} } - /// Dictionary-encoded aggregate functions referenced by `nested_field_aggregates`, up to - /// u16::MAX. Each distinct pair of ID and options is one entry. + /// Dictionary of the (aggregate ID, options) pairs used by this file, up to u16::MAX. #[inline] pub fn aggregate_specs(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { // Safety: @@ -1192,20 +1086,46 @@ 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_AGGREGATE_SPECS, None)} } - /// Aggregate states following a post-order walk of the root `DType` tree: each leaf field gets - /// an entry, and each nullable struct additionally gets a trailing entry for its own null - /// count, inserted after its children's entries. + /// Dictionary of the distinct aggregate lists used by entries, up to u16::MAX. + #[inline] + pub fn aggregate_sets(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FileStatistics::VT_AGGREGATE_SETS, None)} + } + /// One element per file statistics entry, in post-order of the root `DType`: each leaf field + /// gets an entry, and each nullable struct additionally gets a trailing entry for its own null + /// count, after its children's entries. Each element is an index into `aggregate_sets`. + /// + /// The entry paths and dtypes are not stored; readers derive them from the root `DType` with + /// the same post-order walk. /// - /// Readers that know about this field should prefer it over `field_stats` when present, - /// skipping aggregates they don't recognize. It is absent for files written before nested - /// file stats were supported, and readers that don't know about it fall back to - /// `field_stats`'s legacy top-level-fields-only layout. + /// Absent for files written before aggregate-based statistics, whose readers fall back to + /// `field_stats`. #[inline] - pub fn nested_field_aggregates(&self) -> Option<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>> { + pub fn entry_sets(&self) -> Option<::flatbuffers::Vector<'a, u16>> { // Safety: // Created from valid Table for this object // which contains a valid value in this slot - unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>(FileStatistics::VT_NESTED_FIELD_AGGREGATES, None)} + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u16>>>(FileStatistics::VT_ENTRY_SETS, None)} + } + /// The partial states of every entry's aggregates, concatenated. For each entry in + /// `entry_sets` order, and for each spec in that entry's set in order: + /// + /// - an unsigned LEB128 varint byte length `n`, then + /// - `n` bytes of a protobuf-serialized `vortex.scalar.ScalarValue`, typed by the + /// aggregate's partial dtype for the entry's dtype. A null value means no value was + /// accumulated. + /// + /// Readers skip the partial of an aggregate they don't recognize using its length prefix, so + /// unknown aggregates never desynchronize the stream. + #[inline] + pub fn partials(&self) -> Option<::flatbuffers::Vector<'a, u8>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { self._tab.get::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'a, u8>>>(FileStatistics::VT_PARTIALS, None)} } } @@ -1217,7 +1137,9 @@ impl ::flatbuffers::Verifiable for FileStatistics<'_> { v.visit_table(pos)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("field_stats", Self::VT_FIELD_STATS, false)? .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("aggregate_specs", Self::VT_AGGREGATE_SPECS, false)? - .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("nested_field_aggregates", Self::VT_NESTED_FIELD_AGGREGATES, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, ::flatbuffers::ForwardsUOffset>>>("aggregate_sets", Self::VT_AGGREGATE_SETS, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u16>>>("entry_sets", Self::VT_ENTRY_SETS, false)? + .visit_field::<::flatbuffers::ForwardsUOffset<::flatbuffers::Vector<'_, u8>>>("partials", Self::VT_PARTIALS, false)? .finish(); Ok(()) } @@ -1225,7 +1147,9 @@ impl ::flatbuffers::Verifiable for FileStatistics<'_> { pub struct FileStatisticsArgs<'a> { pub field_stats: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, pub aggregate_specs: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, - pub nested_field_aggregates: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub aggregate_sets: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, ::flatbuffers::ForwardsUOffset>>>>, + pub entry_sets: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u16>>>, + pub partials: Option<::flatbuffers::WIPOffset<::flatbuffers::Vector<'a, u8>>>, } impl<'a> Default for FileStatisticsArgs<'a> { #[inline] @@ -1233,7 +1157,9 @@ impl<'a> Default for FileStatisticsArgs<'a> { FileStatisticsArgs { field_stats: None, aggregate_specs: None, - nested_field_aggregates: None, + aggregate_sets: None, + entry_sets: None, + partials: None, } } } @@ -1252,8 +1178,16 @@ impl<'a: 'b, 'b, A: ::flatbuffers::Allocator + 'a> FileStatisticsBuilder<'a, 'b, self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_AGGREGATE_SPECS, aggregate_specs); } #[inline] - pub fn add_nested_field_aggregates(&mut self, nested_field_aggregates: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { - self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_NESTED_FIELD_AGGREGATES, nested_field_aggregates); + pub fn add_aggregate_sets(&mut self, aggregate_sets: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , ::flatbuffers::ForwardsUOffset>>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_AGGREGATE_SETS, aggregate_sets); + } + #[inline] + pub fn add_entry_sets(&mut self, entry_sets: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u16>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_ENTRY_SETS, entry_sets); + } + #[inline] + pub fn add_partials(&mut self, partials: ::flatbuffers::WIPOffset<::flatbuffers::Vector<'b , u8>>) { + self.fbb_.push_slot_always::<::flatbuffers::WIPOffset<_>>(FileStatistics::VT_PARTIALS, partials); } #[inline] pub fn new(_fbb: &'b mut ::flatbuffers::FlatBufferBuilder<'a, A>) -> FileStatisticsBuilder<'a, 'b, A> { @@ -1275,7 +1209,9 @@ impl ::core::fmt::Debug for FileStatistics<'_> { let mut ds = f.debug_struct("FileStatistics"); ds.field("field_stats", &self.field_stats()); ds.field("aggregate_specs", &self.aggregate_specs()); - ds.field("nested_field_aggregates", &self.nested_field_aggregates()); + ds.field("aggregate_sets", &self.aggregate_sets()); + ds.field("entry_sets", &self.entry_sets()); + ds.field("partials", &self.partials()); ds.finish() } } diff --git a/vortex-file/src/footer/file_statistics.rs b/vortex-file/src/footer/file_statistics.rs index eaf8b954ab4..c7a6f9658a4 100644 --- a/vortex-file/src/footer/file_statistics.rs +++ b/vortex-file/src/footer/file_statistics.rs @@ -6,6 +6,7 @@ //! File statistics provide metadata about the data in the file, such as min/max values, //! null counts, and other statistical information that can be used for query optimization //! and data exploration. +use std::hash::Hash; use std::sync::Arc; use flatbuffers::FlatBufferBuilder; @@ -25,6 +26,7 @@ use vortex_array::stats::StatsSet; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_layout::layouts::file_stats::AggregateStat; @@ -143,29 +145,57 @@ impl FileStatistics { file_dtype: &DType, session: &VortexSession, ) -> VortexResult { - if let Some(nested_field_aggregates) = fb.nested_field_aggregates() { + if let Some(entry_sets) = fb.entry_sets() { let aggregate_fns: Vec> = fb .aggregate_specs() .iter() .flat_map(|specs| specs.iter()) .map(|spec| aggregate_fn_from_spec(&spec, session)) .try_collect()?; + let aggregate_sets: Vec> = fb + .aggregate_sets() + .iter() + .flat_map(|sets| sets.iter()) + .map(|set| { + set.specs() + .iter() + .map(|spec| { + vortex_ensure!( + usize::from(spec) < aggregate_fns.len(), + "aggregate spec index {spec} out of range for {} specs", + aggregate_fns.len() + ); + Ok(spec) + }) + .collect() + }) + .try_collect()?; let layout = postorder_stats_layout(file_dtype); - vortex_ensure_eq!(nested_field_aggregates.len(), layout.len()); + vortex_ensure_eq!(entry_sets.len(), layout.len()); + let mut partials = + PartialsReader::new(fb.partials().map(|p| p.bytes()).unwrap_or_default()); let mut aggregates = Vec::with_capacity(layout.len()); let mut dtypes = Vec::with_capacity(layout.len()); let mut paths = Vec::with_capacity(layout.len()); - for (field_aggregates, (path, dtype)) in nested_field_aggregates.iter().zip(layout) { - aggregates.push(aggregate_stats_from_flatbuffer( - &field_aggregates, + for (set, (path, dtype)) in entry_sets.iter().zip(layout) { + let set = aggregate_sets.get(usize::from(set)).ok_or_else(|| { + vortex_err!( + "aggregate set index {set} out of range for {} sets", + aggregate_sets.len() + ) + })?; + aggregates.push(aggregate_stats_from_partials( + set, &aggregate_fns, + &mut partials, &dtype, session, )?); dtypes.push(dtype); paths.push(path); } + partials.finish()?; return Ok(Self::from_parts( aggregates.into(), @@ -309,33 +339,30 @@ fn aggregate_fn_from_spec( Ok(Some(aggregate_fn)) } -fn aggregate_stats_from_flatbuffer( - field_aggregates: &fb::FieldAggregates<'_>, +/// Reads one entry's partial states, one per spec in `set`, from `partials`. +/// +/// `set`'s spec indices must already be checked against `aggregate_fns`. +fn aggregate_stats_from_partials( + set: &[u16], aggregate_fns: &[Option], + partials: &mut PartialsReader<'_>, dtype: &DType, session: &VortexSession, ) -> VortexResult { - let mut aggregates = Vec::new(); - for state in field_aggregates.aggregates().unwrap_or_default() { - let spec = usize::from(state.aggregate_spec()); - let Some(aggregate_fn) = aggregate_fns.get(spec).ok_or_else(|| { - vortex_err!( - "aggregate spec index {spec} out of range for {} specs", - aggregate_fns.len() - ) - })? - else { + let mut aggregates = Vec::with_capacity(set.len()); + for &spec in set { + // Read the partial even for an unknown aggregate, to stay in step with the stream. + let bytes = partials.next_partial()?; + let Some(aggregate_fn) = &aggregate_fns[usize::from(spec)] else { continue; }; - let aggregate_fn = aggregate_fn.clone(); let partial_dtype = aggregate_fn.state_dtype(dtype).ok_or_else(|| { vortex_err!("aggregate {aggregate_fn} does not support field dtype {dtype}") })?; - let partial = - ScalarValue::from_proto_bytes(state.partial().bytes(), &partial_dtype, session)?; + let partial = ScalarValue::from_proto_bytes(bytes, &partial_dtype, session)?; let partial = Scalar::try_new(partial_dtype, partial)?; aggregates.push(AggregateStat::try_from_partial( - aggregate_fn, + aggregate_fn.clone(), partial, dtype, )?); @@ -343,6 +370,103 @@ fn aggregate_stats_from_flatbuffer( Ok(AggregateStats::new(aggregates)) } +/// The maximum length of an unsigned LEB128 encoding of a `u64`. +const MAX_VARINT_LEN: usize = 10; + +/// Appends `value` to `out` as an unsigned LEB128 varint. +fn write_varint(out: &mut Vec, mut value: u64) { + while value >= 0x80 { + out.push((value as u8) | 0x80); + value >>= 7; + } + out.push(value as u8); +} + +/// Reads the length-prefixed partial states of `FileStatistics.partials` in order. +struct PartialsReader<'a> { + remaining: &'a [u8], +} + +impl<'a> PartialsReader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { remaining: bytes } + } + + /// Returns the bytes of the next partial state. + fn next_partial(&mut self) -> VortexResult<&'a [u8]> { + let len = self.next_varint()?; + let len = usize::try_from(len) + .ok() + .filter(|&len| len <= self.remaining.len()) + .ok_or_else(|| { + vortex_err!( + "file statistics partial of {len} bytes exceeds the {} remaining bytes", + self.remaining.len() + ) + })?; + let (partial, rest) = self.remaining.split_at(len); + self.remaining = rest; + Ok(partial) + } + + fn next_varint(&mut self) -> VortexResult { + let mut value = 0u64; + for (i, &byte) in self.remaining.iter().take(MAX_VARINT_LEN).enumerate() { + let bits = u64::from(byte & 0x7f); + // The tenth byte holds only the top bit of a `u64`. + vortex_ensure!( + i < MAX_VARINT_LEN - 1 || bits <= 1, + "file statistics partial length overflows u64" + ); + value |= bits << (7 * i); + if byte & 0x80 == 0 { + self.remaining = &self.remaining[i + 1..]; + return Ok(value); + } + } + vortex_bail!("truncated or overlong file statistics partial length") + } + + /// Fails unless every partial state has been read. + fn finish(self) -> VortexResult<()> { + vortex_ensure!( + self.remaining.is_empty(), + "{} unread bytes after the last file statistics partial", + self.remaining.len() + ); + Ok(()) + } +} + +/// Assigns each distinct value a `u16` index in order of first appearance. +struct Dictionary { + values: Vec, + indices: HashMap, + kind: &'static str, +} + +impl Dictionary { + fn new(kind: &'static str) -> Self { + Self { + values: Vec::new(), + indices: HashMap::default(), + kind, + } + } + + fn intern(&mut self, value: K) -> VortexResult { + if let Some(&index) = self.indices.get(&value) { + return Ok(index); + } + let index = u16::try_from(self.values.len()).map_err(|_| { + vortex_err!("file statistics exceed {} distinct {}", u16::MAX, self.kind) + })?; + self.values.push(value.clone()); + self.indices.insert(value, index); + Ok(index) + } +} + impl<'a> IntoIterator for &'a FileStatistics { type Item = (&'a AggregateStats, &'a DType); type IntoIter = @@ -371,54 +495,29 @@ impl WriteFlatBuffer for FileStatistics { .collect::>>()?; let field_stats = fbb.create_vector(field_stats.as_slice()); - // Aggregate functions are dictionary-encoded, like the footer's array and layout IDs, so - // each distinct function and its options are written once. - let mut specs: Vec<&AggregateFnRef> = Vec::new(); - let mut spec_indices: HashMap<&AggregateFnRef, u16> = HashMap::default(); - let mut nested_field_aggregates = Vec::with_capacity(self.aggregates.len()); + // Aggregate functions and the lists of them recorded per entry are dictionary-encoded, so + // each distinct function and list is written once however many entries use it. + let mut specs = Dictionary::<&AggregateFnRef>::new("aggregate functions"); + let mut sets = Dictionary::>::new("aggregate sets"); + let mut entry_sets = Vec::with_capacity(self.aggregates.len()); + let mut partials = Vec::new(); for field_aggregates in self.aggregates.iter() { - let mut states = Vec::new(); + let mut set = Vec::new(); // Values without a partial state (e.g. from legacy statistics) can't be written. for stat in field_aggregates.iter() { let Some(partial) = stat.partial() else { continue; }; - let aggregate_fn = stat.aggregate_fn(); - let aggregate_spec = match spec_indices.get(aggregate_fn) { - Some(&index) => index, - None => { - let index = u16::try_from(specs.len()).map_err(|_| { - vortex_err!( - "file statistics exceed {} distinct aggregate functions", - u16::MAX - ) - })?; - specs.push(aggregate_fn); - spec_indices.insert(aggregate_fn, index); - index - } - }; - let partial = - fbb.create_vector(&ScalarValue::to_proto_bytes::>(partial.value())); - states.push(fb::AggregateState::create( - fbb, - &fb::AggregateStateArgs { - aggregate_spec, - partial: Some(partial), - }, - )); + set.push(specs.intern(stat.aggregate_fn())?); + let partial = ScalarValue::to_proto_bytes::>(partial.value()); + write_varint(&mut partials, partial.len() as u64); + partials.extend_from_slice(&partial); } - let aggregates = fbb.create_vector(states.as_slice()); - nested_field_aggregates.push(fb::FieldAggregates::create( - fbb, - &fb::FieldAggregatesArgs { - aggregates: Some(aggregates), - }, - )); + entry_sets.push(sets.intern(set)?); } - let nested_field_aggregates = fbb.create_vector(nested_field_aggregates.as_slice()); let aggregate_specs = specs + .values .into_iter() .map(|aggregate_fn| { let options = aggregate_fn.options().serialize()?.ok_or_else(|| { @@ -428,24 +527,37 @@ impl WriteFlatBuffer for FileStatistics { ) })?; let id = fbb.create_string(aggregate_fn.id().as_ref()); - let options = fbb.create_vector(options.as_slice()); + let options = (!options.is_empty()).then(|| fbb.create_vector(options.as_slice())); Ok(fb::AggregateSpec::create( fbb, &fb::AggregateSpecArgs { id: Some(id), - options: Some(options), + options, }, )) }) .collect::>>()?; let aggregate_specs = fbb.create_vector(aggregate_specs.as_slice()); + let aggregate_sets = sets + .values + .iter() + .map(|set| { + let specs = fbb.create_vector(set.as_slice()); + fb::AggregateSet::create(fbb, &fb::AggregateSetArgs { specs: Some(specs) }) + }) + .collect::>(); + let aggregate_sets = fbb.create_vector(aggregate_sets.as_slice()); + let entry_sets = fbb.create_vector(entry_sets.as_slice()); + let partials = fbb.create_vector(partials.as_slice()); Ok(fb::FileStatistics::create( fbb, &fb::FileStatisticsArgs { field_stats: Some(field_stats), aggregate_specs: Some(aggregate_specs), - nested_field_aggregates: Some(nested_field_aggregates), + aggregate_sets: Some(aggregate_sets), + entry_sets: Some(entry_sets), + partials: Some(partials), }, )) } @@ -456,6 +568,7 @@ mod tests { use std::num::NonZeroUsize; use flatbuffers::FlatBufferBuilder; + use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; @@ -555,7 +668,7 @@ mod tests { } #[test] - fn aggregate_specs_are_written_once_per_function_and_options() -> VortexResult<()> { + fn specs_and_sets_are_written_once() -> VortexResult<()> { let file_dtype = DType::struct_( [("a", i32_dtype()), ("b", i32_dtype()), ("c", i32_dtype())], Nullability::NonNullable, @@ -590,11 +703,15 @@ mod tests { .map(|spec| { ( spec.id().to_string(), - spec.options().map(|options| options.bytes().to_vec()), + spec.options() + .map(|options| options.bytes().to_vec()) + .unwrap_or_default(), ) }) .collect::>(); - let options = |aggregate_fn: &AggregateFnRef| aggregate_fn.options().serialize(); + let options = |aggregate_fn: &AggregateFnRef| -> VortexResult> { + Ok(aggregate_fn.options().serialize()?.unwrap_or_default()) + }; assert_eq!( specs, [ @@ -607,6 +724,21 @@ mod tests { ] ); + // "a" and "b" record the same aggregates, so they share a set. + let sets = fb + .aggregate_sets() + .vortex_expect("aggregate sets") + .iter() + .map(|set| set.specs().iter().collect::>()) + .collect::>(); + assert_eq!(sets, [vec![0, 1], vec![2]]); + let entry_sets = fb + .entry_sets() + .vortex_expect("entry sets") + .iter() + .collect::>(); + assert_eq!(entry_sets, [0, 0, 1]); + let read_back = read_back(&file_stats, &file_dtype)?; let (b, _) = read_back .get_by_path(&FieldPath::from_name("b")) @@ -666,65 +798,93 @@ mod tests { Ok(()) } - #[test] - fn unknown_aggregates_are_skipped() -> VortexResult<()> { - let file_dtype = DType::struct_([("col", i32_dtype())], Nullability::NonNullable); - let known = min(7); - let partial = ScalarValue::to_proto_bytes::>( - known.partial().vortex_expect("has a partial").value(), + /// Hand-assembles the aggregate-based sections of a `FileStatistics` flatbuffer, so tests can + /// write footers this writer never would. + fn raw_file_stats( + specs: &[(&str, &[u8])], + sets: &[&[u16]], + entry_sets: &[u16], + partials: &[u8], + ) -> Vec { + let mut fbb = FlatBufferBuilder::new(); + let specs = specs + .iter() + .map(|(id, options)| { + let id = fbb.create_string(id); + let options = fbb.create_vector(options); + fb::AggregateSpec::create( + &mut fbb, + &fb::AggregateSpecArgs { + id: Some(id), + options: Some(options), + }, + ) + }) + .collect::>(); + let specs = fbb.create_vector(&specs); + let sets = sets + .iter() + .map(|set| { + let specs = fbb.create_vector(set); + fb::AggregateSet::create(&mut fbb, &fb::AggregateSetArgs { specs: Some(specs) }) + }) + .collect::>(); + let sets = fbb.create_vector(&sets); + let entry_sets = fbb.create_vector(entry_sets); + let partials = fbb.create_vector(partials); + let root = fb::FileStatistics::create( + &mut fbb, + &fb::FileStatisticsArgs { + field_stats: None, + aggregate_specs: Some(specs), + aggregate_sets: Some(sets), + entry_sets: Some(entry_sets), + partials: Some(partials), + }, ); + fbb.finish_minimal(root); + fbb.finished_data().to_vec() + } + fn read_raw(bytes: &[u8], file_dtype: &DType) -> VortexResult { + let fb = flatbuffers::root::(bytes).vortex_expect("valid flatbuffer"); + FileStatistics::from_flatbuffer(&fb, file_dtype, &array_session()) + } + + /// The options of `min` and its partial state for `value`, length-prefixed as in `partials`. + fn min_spec_and_partial(value: i32) -> VortexResult<(Vec, Vec)> { + let known = min(value); let options = known .aggregate_fn() .options() .serialize()? .vortex_expect("serializable options"); - - // Spec 0 is an aggregate no session knows, spec 1 is the known `min`. - let mut fbb = FlatBufferBuilder::new(); - let aggregate_specs = ["test.unknown", known.aggregate_fn().id().as_ref()].map(|id| { - let id = fbb.create_string(id); - let options = fbb.create_vector(options.as_slice()); - fb::AggregateSpec::create( - &mut fbb, - &fb::AggregateSpecArgs { - id: Some(id), - options: Some(options), - }, - ) - }); - let aggregate_specs = fbb.create_vector(&aggregate_specs); - let states = [0u16, 1].map(|aggregate_spec| { - let partial = fbb.create_vector(partial.as_slice()); - fb::AggregateState::create( - &mut fbb, - &fb::AggregateStateArgs { - aggregate_spec, - partial: Some(partial), - }, - ) - }); - let aggregates = fbb.create_vector(&states); - let field_aggregates = fb::FieldAggregates::create( - &mut fbb, - &fb::FieldAggregatesArgs { - aggregates: Some(aggregates), - }, + let partial = ScalarValue::to_proto_bytes::>( + known.partial().vortex_expect("has a partial").value(), ); - let nested_field_aggregates = fbb.create_vector(&[field_aggregates]); - let root = fb::FileStatistics::create( - &mut fbb, - &fb::FileStatisticsArgs { - field_stats: None, - aggregate_specs: Some(aggregate_specs), - nested_field_aggregates: Some(nested_field_aggregates), - }, + let mut prefixed = Vec::new(); + write_varint(&mut prefixed, partial.len() as u64); + prefixed.extend_from_slice(&partial); + Ok((options, prefixed)) + } + + #[test] + fn unknown_aggregates_are_skipped_by_length() -> VortexResult<()> { + let file_dtype = DType::struct_([("col", i32_dtype())], Nullability::NonNullable); + let (options, known_partial) = min_spec_and_partial(7)?; + + // Spec 0 is an aggregate no session knows, recorded before the known `min`. Its partial is + // bytes this reader can't interpret, which it has to skip to reach `min`'s. + let mut partials = vec![3, 0xde, 0xad, 0xbe]; + partials.extend_from_slice(&known_partial); + let bytes = raw_file_stats( + &[("test.unknown", &[]), ("vortex.min", &options)], + &[&[0, 1]], + &[0], + &partials, ); - fbb.finish_minimal(root); - let bytes = fbb.finished_data().to_vec(); - let fb = flatbuffers::root::(&bytes).vortex_expect("valid flatbuffer"); - let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &array_session())?; + let read_back = read_raw(&bytes, &file_dtype)?; let (aggregates, _) = read_back .get_by_path(&FieldPath::from_name("col")) .expect("col aggregates"); @@ -736,6 +896,64 @@ mod tests { Ok(()) } + #[rstest] + #[case::set_index_out_of_range(vec![vec![0]], vec![1], None, "aggregate set index 1 out of range")] + #[case::spec_index_out_of_range(vec![vec![1]], vec![0], None, "aggregate spec index 1 out of range")] + #[case::wrong_entry_count(vec![vec![0]], vec![0, 0], None, "")] + #[case::partial_missing(vec![vec![0]], vec![0], Some(vec![]), "truncated or overlong")] + #[case::truncated_varint(vec![vec![0]], vec![0], Some(vec![0x80]), "truncated or overlong")] + #[case::overlong_varint(vec![vec![0]], vec![0], Some(vec![0x80; 11]), "truncated or overlong")] + #[case::varint_overflows( + vec![vec![0]], + vec![0], + Some(vec![0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02]), + "overflows u64" + )] + #[case::length_past_end(vec![vec![0]], vec![0], Some(vec![5, 0]), "exceeds the 1 remaining bytes")] + #[case::trailing_bytes(vec![vec![]], vec![0], Some(vec![0]), "1 unread bytes")] + fn malformed_aggregate_statistics_are_rejected( + #[case] sets: Vec>, + #[case] entry_sets: Vec, + #[case] partials: Option>, + #[case] message: &str, + ) -> VortexResult<()> { + let file_dtype = DType::struct_([("col", i32_dtype())], Nullability::NonNullable); + let (options, valid_partial) = min_spec_and_partial(7)?; + let sets = sets.iter().map(Vec::as_slice).collect::>(); + let bytes = raw_file_stats( + &[("vortex.min", &options)], + &sets, + &entry_sets, + &partials.unwrap_or(valid_partial), + ); + + let error = read_raw(&bytes, &file_dtype) + .err() + .vortex_expect("malformed statistics must be rejected"); + assert!( + error.to_string().contains(message), + "unexpected error: {error}" + ); + Ok(()) + } + + #[rstest] + #[case(0)] + #[case(1)] + #[case(127)] + #[case(128)] + #[case(300)] + #[case(u64::from(u32::MAX))] + #[case(u64::MAX)] + fn varints_round_trip(#[case] value: u64) -> VortexResult<()> { + let mut bytes = Vec::new(); + write_varint(&mut bytes, value); + assert!(bytes.len() <= MAX_VARINT_LEN); + let mut reader = PartialsReader::new(&bytes); + assert_eq!(reader.next_varint()?, value); + reader.finish() + } + #[test] fn get_by_path_resolves_a_later_sibling_past_a_nullable_struct_field() { // Regression test: a nullable struct field ("s") contributes both an `s.b` entry and a @@ -770,8 +988,8 @@ mod tests { #[test] fn legacy_non_nested_footer_still_parses() -> VortexResult<()> { - // Simulates a footer written before nested field stats existed: - // `nested_field_aggregates` is absent, and `field_stats` holds one entry per top-level + // Simulates a footer written before aggregate-based stats existed: `entry_sets` is + // absent, 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); @@ -786,15 +1004,14 @@ mod tests { &mut fbb, &fb::FileStatisticsArgs { field_stats: Some(field_stats), - aggregate_specs: None, - nested_field_aggregates: None, + ..Default::default() }, ); fbb.finish_minimal(root); let bytes = fbb.finished_data().to_vec(); let fb = flatbuffers::root::(&bytes).vortex_expect("valid flatbuffer"); - assert!(fb.nested_field_aggregates().is_none()); + assert!(fb.entry_sets().is_none()); let read_back = FileStatistics::from_flatbuffer(&fb, &file_dtype, &session)?; let path = FieldPath::from_name("col"); diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 6c97dc458db..120a467835e 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -169,7 +169,7 @@ impl VortexWriteOptions { } /// Exclude legacy top-level-only statistics (`field_stats`) from the file, computing and - /// writing only the full nested post-order statistics (`nested_field_aggregates`). + /// writing only the full nested post-order statistics (`entry_sets` and `partials`). /// /// Readers built before nested file stats existed only look at `field_stats` and validate its /// length against the number of top-level struct fields; omitting it means those readers find From 2a9e85035e34632b85a7ba084a1f60112c6a421f Mon Sep 17 00:00:00 2001 From: Thor Date: Wed, 7 Oct 2026 10:59:40 -0500 Subject: [PATCH 3/4] rebase fixes; file size changes for new test Signed-off-by: Thor --- vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs | 2 +- vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs | 2 +- vortex-file/src/tests.rs | 4 ++-- vortex-layout/src/layouts/file_stats.rs | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs index 40d1629a6fa..4615931e995 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_max/mod.rs @@ -694,7 +694,7 @@ mod tests { &VarBinViewArray::from_iter_str(["a"]).into_array(), &mut ctx, )?; - acc.combine_partials(source.partial_scalar()?)?; + acc.combine_partials(&source.partial_scalar()?)?; assert_eq!( acc.can_satisfy(&requested), AggregateFnSatisfaction::Approximate diff --git a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs index e3629e5cf99..05dfabbc4c8 100644 --- a/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs +++ b/vortex-array/src/aggregate_fn/fns/bounded_min/mod.rs @@ -526,7 +526,7 @@ mod tests { &VarBinViewArray::from_iter_str(["c"]).into_array(), &mut ctx, )?; - acc.combine_partials(source.partial_scalar()?)?; + acc.combine_partials(&source.partial_scalar()?)?; assert_eq!( acc.can_satisfy(&requested), AggregateFnSatisfaction::Approximate diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 0061c68af01..e551d1e01cd 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -168,13 +168,13 @@ async fn test_eof_values() { #[rstest] #[case::default( BtrBlocksCompressorBuilder::from_session(&SESSION), - if cfg!(feature = "zstd") { 70_036 } else { 69_972 } + if cfg!(feature = "zstd") { 70_252 } else { 70_188 } )] #[cfg_attr( feature = "zstd", case::compact( BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(), - 55_112 + 55_328 ) )] #[tokio::test] diff --git a/vortex-layout/src/layouts/file_stats.rs b/vortex-layout/src/layouts/file_stats.rs index 9f5ca85b30d..7576d54b8d5 100644 --- a/vortex-layout/src/layouts/file_stats.rs +++ b/vortex-layout/src/layouts/file_stats.rs @@ -105,7 +105,7 @@ impl AggregateStat { dtype: &DType, ) -> VortexResult { let mut accumulator = aggregate_fn.accumulator(dtype)?; - accumulator.combine_partials(partial.clone())?; + accumulator.combine_partials(&partial)?; let value = Precision::exact(accumulator.final_scalar()?); Ok(Self { aggregate_fn, From cf42ed608c6520f29f0482499a0a3731c45b7b1e Mon Sep 17 00:00:00 2001 From: Thor Date: Thu, 8 Oct 2026 11:46:14 -0500 Subject: [PATCH 4/4] Review comment fixups Signed-off-by: Thor --- vortex-duckdb/cpp/copy_function.cpp | 2 -- vortex-duckdb/cpp/include/table_function.h | 3 +-- vortex-duckdb/src/copy.rs | 12 ++++-------- 3 files changed, 5 insertions(+), 12 deletions(-) diff --git a/vortex-duckdb/cpp/copy_function.cpp b/vortex-duckdb/cpp/copy_function.cpp index 9c61804b21f..efde5045d98 100644 --- a/vortex-duckdb/cpp/copy_function.cpp +++ b/vortex-duckdb/cpp/copy_function.cpp @@ -107,8 +107,6 @@ void copy_to_finalize(ClientContext &, FunctionData &, GlobalFunctionData &gstat global.written_stats->row_count = file_stats.row_count; global.written_stats->file_size_bytes = file_stats.file_size_bytes; global.written_stats->footer_size_bytes = Value::UBIGINT(file_stats.footer_size_bytes); - // One entry per leaf column: nested struct fields are reported at their full path and structs - // get no entry of their own, as in parquet. for (idx_t i = 0; i < file_stats.num_columns; i++) { duckdb_vx_written_column_statistics col_stats {}; duckdb_vx_error col_error = nullptr; diff --git a/vortex-duckdb/cpp/include/table_function.h b/vortex-duckdb/cpp/include/table_function.h index 30ab6b30bb1..f0c6850a632 100644 --- a/vortex-duckdb/cpp/include/table_function.h +++ b/vortex-duckdb/cpp/include/table_function.h @@ -83,8 +83,7 @@ typedef struct { // Per-column statistics of a written Vortex file. typedef struct { - // Owned VARCHAR value holding the column's quoted, dot-separated path (e.g. `"s"."b"`), the key - // format DuckLake expects; the caller must destroy it. + // Owned VARCHAR holding column's dot-separated path. duckdb_value column_key; // Owned values, null if absent; the caller must destroy them. duckdb_value min; diff --git a/vortex-duckdb/src/copy.rs b/vortex-duckdb/src/copy.rs index 4f179abbe9d..388aaf3fd4d 100644 --- a/vortex-duckdb/src/copy.rs +++ b/vortex-duckdb/src/copy.rs @@ -238,10 +238,8 @@ pub(crate) fn written_column_stats( .map(Some) } -/// The indices of the leaf entries of `file_stats`, which are the columns statistics are reported -/// for. Like parquet, which reports one entry per leaf column writer, struct columns get no entry of -/// their own; nested fields are reported at their full path. Lists and maps are opaque leaves of the -/// statistics layout, so they are reported at their own path rather than per element. +/// The indices of the leaf entries of `file_stats`. +/// One entry per leaf column writer. fn leaf_stats_columns(file_stats: &FileStatistics) -> Vec { file_stats .dtypes() @@ -252,14 +250,12 @@ fn leaf_stats_columns(file_stats: &FileStatistics) -> Vec { .collect() } -/// DuckLake keys column statistics by a quoted, dot-separated path (see -/// `DuckLakeUtil::ParseQuotedList`). This matches the parquet writer, which quotes each name with -/// `KeywordHelper::WriteQuoted(name, '"')`, escaping embedded quotes by doubling them. A list -/// element is named `element`, as in parquet's list encoding. +/// DuckLake keys column statistics by a quoted, dot-separated path. fn ducklake_column_key(path: &FieldPath) -> String { path.parts() .iter() .map(|part| { + // List element is the only case where a part will be empty. let name = part.as_name().unwrap_or("element"); format!("\"{}\"", name.replace('"', "\"\"")) })