diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 9d4dbb90610..18e8f44a67d 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -62,6 +62,7 @@ pub struct WriteStrategyBuilder { allow_encodings: Option>, flat_strategy: Option>, probe_compressor: Option>, + include_sum_zone_stats: bool, /// Whether to write list fields using [`ListLayoutStrategy`]. /// /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy @@ -80,6 +81,7 @@ impl Default for WriteStrategyBuilder { allow_encodings: None, flat_strategy: None, probe_compressor: None, + include_sum_zone_stats: true, use_list_layout: use_experimental_list_layout(), } } @@ -173,6 +175,15 @@ impl WriteStrategyBuilder { self } + /// Configure whether zoned layouts include Sum in their default aggregate set. + /// + /// Disabling Sum can preserve write compatibility with readers whose Sum aggregate uses a + /// different partial-state representation. + pub fn with_sum_zone_stats(mut self, include: bool) -> Self { + self.include_sum_zone_stats = include; + self + } + /// Builds the canonical [`LayoutStrategy`] implementation, with the configured overrides /// applied. pub fn build(self) -> Arc { @@ -267,6 +278,7 @@ impl WriteStrategyBuilder { compress_then_flat.clone(), ZonedLayoutOptions { block_size: row_block_size, + include_sum: self.include_sum_zone_stats, ..Default::default() }, ); @@ -301,6 +313,7 @@ impl WriteStrategyBuilder { compress_then_flat.clone(), ZonedLayoutOptions { block_size: row_block_size, + include_sum: self.include_sum_zone_stats, ..Default::default() }, ); diff --git a/vortex-layout/src/layouts/zoned/writer.rs b/vortex-layout/src/layouts/zoned/writer.rs index 92c4b4c9870..086b9ddbfeb 100644 --- a/vortex-layout/src/layouts/zoned/writer.rs +++ b/vortex-layout/src/layouts/zoned/writer.rs @@ -59,6 +59,10 @@ pub struct ZonedLayoutOptions { /// /// If unset, the writer chooses pruning aggregates from the input dtype. pub aggregate_fns: Option>, + /// Whether the default aggregate set includes Sum. + /// + /// This has no effect when `aggregate_fns` is set explicitly. + pub include_sum: bool, /// Number of chunks to compute aggregate partials in parallel. pub concurrency: NonZeroUsize, } @@ -68,6 +72,7 @@ impl Default for ZonedLayoutOptions { Self { block_size: unsafe { NonZeroUsize::new_unchecked(8192) }, aggregate_fns: None, + include_sum: true, concurrency: unsafe { NonZeroUsize::new_unchecked(get_available_parallelism().unwrap_or(1)) }, @@ -106,11 +111,9 @@ impl LayoutStrategy for ZonedStrategy { mut eof: SequencePointer, session: &VortexSession, ) -> VortexResult { - let aggregate_fns = self - .options - .aggregate_fns - .clone() - .unwrap_or_else(|| default_zoned_aggregate_fns(stream.dtype(), session)); + let aggregate_fns = self.options.aggregate_fns.clone().unwrap_or_else(|| { + default_zoned_aggregate_fns(stream.dtype(), session, self.options.include_sum) + }); let compute_session = session.clone(); let stats_accumulator = Arc::new(Mutex::new(AggregateStatsAccumulator::new( @@ -192,7 +195,11 @@ impl LayoutStrategy for ZonedStrategy { } } -fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[AggregateFnRef]> { +fn default_zoned_aggregate_fns( + dtype: &DType, + session: &VortexSession, + include_sum: bool, +) -> Arc<[AggregateFnRef]> { let (max, min) = match dtype { DType::Utf8(_) | DType::Binary(_) => ( BoundedMax.bind(BoundedMaxOptions { @@ -209,9 +216,10 @@ fn default_zoned_aggregate_fns(dtype: &DType, session: &VortexSession) -> Arc<[A }; let mut aggregate_fns = vec![max, min]; - if Sum - .return_dtype(&SumAggregateOpts::skip_nans(), dtype) - .is_some() + if include_sum + && Sum + .return_dtype(&SumAggregateOpts::skip_nans(), dtype) + .is_some() { aggregate_fns.push(Sum.bind(SumAggregateOpts::skip_nans())); } @@ -243,6 +251,7 @@ mod tests { let aggregate_fns = default_zoned_aggregate_fns( &DType::Utf8(Nullability::NonNullable), &vortex_array::array_session(), + true, ); assert_eq!( @@ -258,7 +267,7 @@ mod tests { #[test] fn default_aggregates_keep_fixed_width_min_max_exact() { let aggregate_fns = - default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session()); + default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session(), true); assert!(aggregate_fns[0].is::()); assert!(aggregate_fns[1].is::()); @@ -270,7 +279,20 @@ mod tests { let dtype = DType::Extension( Timestamp::new(TimeUnit::Microseconds, Nullability::Nullable).erased(), ); - let aggregate_fns = default_zoned_aggregate_fns(&dtype, &vortex_array::array_session()); + let aggregate_fns = + default_zoned_aggregate_fns(&dtype, &vortex_array::array_session(), true); + + assert!( + aggregate_fns + .iter() + .all(|aggregate_fn| !aggregate_fn.is::()) + ); + } + + #[test] + fn default_aggregates_can_skip_sum() { + let aggregate_fns = + default_zoned_aggregate_fns(&PType::I32.into(), &vortex_array::array_session(), false); assert!( aggregate_fns