diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 7764378eb05..b6596bb6c75 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -34,7 +34,7 @@ use crate::VTable; use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::zone_map::ZoneMap; -type SharedZoneMap = Shared>>; +pub(super) type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = Shared>>>; type PredicateCache = Arc>>; @@ -135,6 +135,11 @@ impl PruningState { .clone() } + /// Shared future resolving to this layout's ZoneMap + pub(super) fn shared_zone_map(&self) -> SharedZoneMap { + self.zone_map() + } + fn zone_map(&self) -> SharedZoneMap { self.zone_map .get_or_init(move || { diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index dbcef1f7a68..9e248ad179d 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -6,13 +6,23 @@ use std::ops::Range; use std::sync::Arc; use futures::future::BoxFuture; +use futures::future::try_join_all; use itertools::Itertools; use tracing::trace; use vortex_array::ArrayRef; +use vortex_array::IntoArray; use vortex_array::MaskFuture; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::expr::Expression; +use vortex_array::expr::is_root; +use vortex_array::expr::root; +use vortex_array::scalar_fn::fns::pack::Pack; +use vortex_array::scalar_fn::fns::stat::StatFn; use vortex_buffer::BitBufferMut; use vortex_error::VortexError; use vortex_error::VortexResult; @@ -30,6 +40,18 @@ use crate::layouts::zoned::ZonedData; use crate::layouts::zoned::pruning::PruningState; use crate::segments::SegmentSource; +/// Slice of projection's "row_range" covered by a single zone +struct ZoneSegment { + zone_idx: u64, + /// Range within "row_range", i.e. local_range.start is offset from + /// row_range.start + offset: Range, + /// Absolute row range in file + absolute: Range, + /// True if zone is contained inside "row_range" + fully_covered: bool, +} + pub struct ZonedReader { dtype: DType, row_count: u64, @@ -37,6 +59,7 @@ pub struct ZonedReader { name: Arc, lazy_children: Arc, pruning: PruningState, + session: VortexSession, } impl ZonedReader { @@ -72,13 +95,14 @@ impl ZonedReader { zone_count, aggregate_fns, Arc::clone(&lazy_children), - session, + session.clone(), ), dtype, row_count, zone_len, name, lazy_children, + session, }) } @@ -221,13 +245,163 @@ impl LayoutReader for ZonedReader { expr: &Expression, mask: MaskFuture, ) -> VortexResult>> { - // TODO(ngates): there are some projection expressions that we may also be able to - // short-circuit with statistics. + if self.zone_len > 0 { + if is_zone_stat(expr) { + let aggregate_fn = expr.as_::().aggregate_fn().clone(); + let dtype = expr.return_dtype(&self.dtype)?; + return self.aggregate_projection(row_range, aggregate_fn, dtype, mask); + } + + // Multiple aggregations over a column are packed together per aggregate. + if expr.is::() + && (0..expr.as_::().names.len()).any(|i| is_zone_stat(expr.child(i))) + { + return self.aggregate_pack_projection(row_range, expr, mask); + } + } + self.data_child()? .projection_evaluation(row_range, expr, mask) } } +/// True if "expr" is an aggregate stat of form "stat(root(), agg)" +fn is_zone_stat(expr: &Expression) -> bool { + expr.is::() && is_root(expr.child(0)) +} + +impl ZonedReader { + fn aggregate_pack_projection( + &self, + row_range: &Range, + expr: &Expression, + mask: MaskFuture, + ) -> VortexResult>> { + let options = expr.as_::(); + let names = options.names.clone(); + let validity = options.nullability.into(); + + let futures = (0..names.len()) + .map(|i| { + let field = expr.child(i); + if is_zone_stat(field) { + let aggregate_fn = field.as_::().aggregate_fn().clone(); + let dtype = field.return_dtype(&self.dtype)?; + self.aggregate_projection(row_range, aggregate_fn, dtype, mask.clone()) + } else { + self.data_child()? + .projection_evaluation(row_range, field, mask.clone()) + } + }) + .try_collect::<_, Vec<_>, _>()?; + + Ok(Box::pin(async move { + let fields = try_join_all(futures).await?; + let len = fields.first().map(|f| f.len()).unwrap_or(0); + Ok(StructArray::new(names, fields, len, validity).into_array()) + })) + } + + fn aggregate_projection( + &self, + row_range: &Range, + aggregate_fn: AggregateFnRef, + out_dtype: DType, + mask: MaskFuture, + ) -> VortexResult>> { + let zone_len = self.zone_len as u64; + let row_count = self.row_count; + let zone_range = self.zone_range(row_range); + + // For every zone, slice of "row_range" this zone covers. + let segments: Vec = zone_range + .map(|zone_idx| { + let zone_start = zone_idx.saturating_mul(zone_len).min(row_count); + let zone_end = zone_idx + .saturating_add(1) + .saturating_mul(zone_len) + .min(row_count); + + let absolute_start = zone_start.max(row_range.start); + let absolute_end = zone_end.min(row_range.end); + let absolute = absolute_start..absolute_end; + + let offset_start: usize = usize::try_from(absolute_start - row_range.start)?; + let offset_end: usize = usize::try_from(absolute_end - row_range.start)?; + let offset = offset_start..offset_end; + + let fully_covered = zone_start >= row_range.start && zone_end <= row_range.end; + + Ok::<_, VortexError>(ZoneSegment { + zone_idx, + offset, + absolute, + fully_covered, + }) + }) + .try_collect()?; + + let data_child = Arc::clone(self.data_child()?); + let zone_map = self.pruning.shared_zone_map(); + let column_dtype = self.dtype.clone(); + let session = self.session.clone(); + + Ok(Box::pin(async move { + let mask = mask.await?; + let true_count = mask.true_count(); + let mut ctx = session.create_execution_ctx(); + let mut accumulator = aggregate_fn.accumulator(&column_dtype)?; + + let zone_map = match zone_map.await { + Ok(zone_map) if zone_map.supports_zone_partial(&aggregate_fn) => Some(zone_map), + Ok(_) | Err(_) => None, + }; + + let covered = |segment: &ZoneSegment| { + zone_map.is_some() + && segment.fully_covered + && mask.slice(segment.offset.clone()).all_true() + }; + let mut i = 0; + while i < segments.len() { + let start_covered = covered(&segments[i]); + let mut j = i + 1; + while j < segments.len() && covered(&segments[j]) == start_covered { + j += 1; + } + + if start_covered && let Some(zone_map) = zone_map.as_ref() { + let range = segments[i].zone_idx..segments[j - 1].zone_idx + 1; + let range = usize::try_from(range.start)?..usize::try_from(range.end)?; + zone_map.fold_zone_partials( + &aggregate_fn, + range, + &mut accumulator, + &session, + )?; + } else { + let sub_mask = mask.slice(segments[i].offset.start..segments[j - 1].offset.end); + if !sub_mask.all_false() { + let sub_range = segments[i].absolute.start..segments[j - 1].absolute.end; + let array = data_child + .projection_evaluation( + &sub_range, + &root(), + MaskFuture::ready(sub_mask), + )? + .await?; + accumulator.accumulate(&array, &mut ctx)?; + } + } + i = j; + } + + let partial = accumulator.partial_scalar()?.cast(&out_dtype)?; + Ok(ConstantArray::new(partial, true_count).into_array()) + })) + } +} + #[cfg(test)] mod test { use std::num::NonZeroUsize; @@ -247,6 +421,9 @@ mod test { use vortex_array::expr::is_not_null; use vortex_array::expr::lit; use vortex_array::expr::root; + use vortex_array::expr::stats::Stat; + use vortex_array::stats::expr::stat; + use vortex_array::stats::expr::sum; use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexExpect; @@ -342,6 +519,89 @@ mod test { }) } + /// Test aggregate projection answers from zone maps + #[rstest] + fn test_aggregate_projection( + #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let rows = layout.row_count(); + let range = &(0..rows); + let mask = MaskFuture::new_true(rows.try_into()?); + + block_on(|handle| async { + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; + + let sum = reader + .projection_evaluation(range, &sum(root()), mask.clone())? + .await?; + assert_eq!(sum.len(), usize::try_from(rows)?); + assert_eq!( + sum.execute_scalar(0, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(45) + ); + + let min_expr = &stat(root(), Stat::Min.aggregate_fn().unwrap()); + let min = reader + .projection_evaluation(range, min_expr, mask.clone())? + .await?; + assert_eq!( + min.execute_scalar(0, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(1) + ); + + let max_expr = &stat(root(), Stat::Max.aggregate_fn().unwrap()); + let max = reader.projection_evaluation(range, max_expr, mask)?.await?; + assert_eq!( + max.execute_scalar(0, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(9) + ); + Ok(()) + }) + } + + /// Test aggregate projection with filter decodes columns for zone maps + /// which are partially taken + #[rstest] + fn test_aggregate_projection_hybrid( + #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let rows = layout.row_count(); + + // zones: [1,2,3][4,5,6][7,8,9]. Cover zone 0 and 2 fully but select + // only row 3 (=4) of zone 1. We must count every selected row exactly + // once: 1+2+3 + 4 + 7+8+9 = 34. + let mask: Mask = [true, true, true, true, false, false, true, true, true] + .into_iter() + .collect(); + assert_eq!(mask.len(), usize::try_from(rows)?); + + block_on(|handle| async { + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; + let sum = reader + .projection_evaluation(&(0..rows), &sum(root()), MaskFuture::ready(mask))? + .await?; + + assert_eq!(sum.len(), 7); + assert_eq!( + sum.execute_scalar(0, &mut ctx)? + .as_primitive() + .typed_value::(), + Some(34) + ); + Ok(()) + }) + } + #[rstest] fn test_stats_pruning_mask( #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index e5f2af494de..792d3ece2bf 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -3,11 +3,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::AccumulatorRef; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnSatisfaction; use vortex_array::aggregate_fn::fns::all_nan::AllNan; @@ -16,7 +18,10 @@ use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; use vortex_array::aggregate_fn::fns::all_null::AllNull; use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; +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::sum::Sum; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -47,6 +52,7 @@ use vortex_mask::Mask; use vortex_runend::RunEnd; use vortex_session::VortexSession; +use crate::layouts::zoned::schema::aggregate_state_dtype; use crate::layouts::zoned::schema::aggregate_stats_table_dtype; use crate::layouts::zoned::schema::legacy_stats_table_dtype; @@ -194,6 +200,82 @@ impl ZoneMap { }) .map(Transformed::into_inner) } + + /// If this aggregate is self mergeable, zone map has a stat field for + /// "aggregate_fn", and column is a primitive with dtype equal to aggregate + /// function's partial dtype, return the per-zone partial column. + fn zone_partial_field(&self, aggregate_fn: &AggregateFnRef) -> Option<&ArrayRef> { + if !is_self_mergeable(aggregate_fn) { + return None; + } + let dtype = aggregate_state_dtype(&self.column_dtype, aggregate_fn)?; + + // We need exact boundaries in zone maps to supply this info as result + if self + .aggregate_fns + .iter() + .any(|stored| stored.can_satisfy(aggregate_fn).is_exact()) + && let Some(field) = self + .array + .unmasked_field_by_name_opt(aggregate_fn.to_string()) + { + return field.dtype().eq_ignore_nullability(&dtype).then_some(field); + } + + // Legacy stats schema, strings always carry inexact stats + let stat = Stat::from_aggregate_fn(aggregate_fn)?; + if matches!(stat, Stat::Min | Stat::Max) + && (self.column_dtype.is_utf8() || self.column_dtype.is_binary()) + { + return None; + } + let field = self.array.unmasked_field_by_name_opt(stat.name())?; + field.dtype().eq_ignore_nullability(&dtype).then_some(field) + } + + pub(super) fn supports_zone_partial(&self, aggregate_fn: &AggregateFnRef) -> bool { + self.zone_partial_field(aggregate_fn).is_some() + } + + /// Fold per-zone partials for "range" into "accumulator". + /// + /// If zone map can't answer aggregate_fn with Exact precision, returns + /// Ok(false) and doesn't change accumulator. + /// + /// Invariant: "accumulator" must be built for "aggregate_fn" and zone + /// map's column dtype. + pub(super) fn fold_zone_partials( + &self, + aggregate_fn: &AggregateFnRef, + range: Range, + accumulator: &mut AccumulatorRef, + session: &VortexSession, + ) -> VortexResult { + let Some(field) = self.zone_partial_field(aggregate_fn) else { + return Ok(false); + }; + + let start = range.start; + let end = range.end.min(self.array.len()); + if start >= end { + return Ok(true); + } + + let mut ctx = session.create_execution_ctx(); + let mut reducer = aggregate_fn.accumulator(field.dtype())?; + reducer.accumulate(&field.slice(start..end)?, &mut ctx)?; + accumulator.combine_partials(reducer.partial_scalar()?)?; + Ok(true) + } +} + +/// Let P_1 ... P_n be aggregate_fn's per-zone partials, and +/// C be aggregate_fn's per-zone partial column. +/// This predicate is true iff aggregate_fn([P_1....P_N]) == aggregate_fn(C) +/// +/// As an example, this predicate is true for min/max/sum but not for count(). +fn is_self_mergeable(aggregate_fn: &AggregateFnRef) -> bool { + aggregate_fn.is::() || aggregate_fn.is::() || aggregate_fn.is::() } struct ZoneMapStatsBinder<'a> { @@ -375,6 +457,7 @@ mod tests { 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::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; @@ -395,6 +478,7 @@ mod tests { use vortex_array::expr::not_eq; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; use vortex_array::stats::all_nan; use vortex_array::stats::all_non_nan; use vortex_array::stats::all_non_null; @@ -483,6 +567,61 @@ mod tests { ); } + #[test] + fn bounded_min_max() { + let max_bytes = default_bounded_stat_max_bytes(); + let max = BoundedMax.bind(BoundedMaxOptions { max_bytes }); + let min = BoundedMin.bind(BoundedMinOptions { max_bytes }); + let dtype = DType::Utf8(Nullability::NonNullable); + let array = + StructArray::from_fields(&[(max.to_string(), buffer![0i32].into_array())]).unwrap(); + let zone_map = unsafe { ZoneMap::new_unchecked(dtype, array, Arc::new([max, min]), 3, 3) }; + assert!(!zone_map.supports_zone_partial(&Max.bind(NumericalAggregateOpts::skip_nans()))); + assert!(!zone_map.supports_zone_partial(&Min.bind(NumericalAggregateOpts::skip_nans()))); + } + + #[test] + fn legacy_numeric_min_max_sum() { + let max = PrimitiveArray::new(buffer![3i32, 7, 5], Validity::AllValid).into_array(); + let max_is_truncated = BoolArray::from_iter([false, false, false]).into_array(); + let min = PrimitiveArray::new(buffer![1i32, 2, 4], Validity::AllValid).into_array(); + let min_is_truncated = BoolArray::from_iter([false, false, false]).into_array(); + let sum = PrimitiveArray::new(buffer![6i64, 15, 12], Validity::AllValid).into_array(); + let array = StructArray::from_fields(&[ + ("max", max), + ("max_is_truncated", max_is_truncated), + ("min", min), + ("min_is_truncated", min_is_truncated), + ("sum", sum), + ]) + .unwrap(); + let ptype = PType::I32.into(); + let stats = Arc::new([Stat::Max, Stat::Min, Stat::Sum]); + let zone_map = ZoneMap::try_new_legacy(ptype, array, stats, 4, 12).unwrap(); + + assert!(zone_map.supports_zone_partial(&Max.bind(NumericalAggregateOpts::skip_nans()))); + assert!(zone_map.supports_zone_partial(&Min.bind(NumericalAggregateOpts::skip_nans()))); + assert!(zone_map.supports_zone_partial(&Sum.bind(NumericalAggregateOpts::skip_nans()))); + + let sum_fn = Sum.bind(NumericalAggregateOpts::skip_nans()); + let mut acc = sum_fn.accumulator(&PType::I32.into()).unwrap(); + assert!( + zone_map + .fold_zone_partials(&sum_fn, 0..3, &mut acc, &SESSION) + .unwrap() + ); + assert_eq!(acc.partial_scalar().unwrap(), Scalar::from(33i64)); + } + + #[test] + fn legacy_string_min_max() { + let dtype = DType::Utf8(Nullability::NonNullable); + let array = StructArray::from_fields(&[("max", buffer![0i32].into_array())]).unwrap(); + let zone_map = unsafe { ZoneMap::new_unchecked(dtype, array, Arc::new([]), 3, 3) }; + assert!(!zone_map.supports_zone_partial(&Max.bind(NumericalAggregateOpts::skip_nans()))); + assert!(!zone_map.supports_zone_partial(&Min.bind(NumericalAggregateOpts::skip_nans()))); + } + #[test] fn bounded_display_names_satisfy_min_max_rewrites() { let bounded_max = BoundedMax.bind(BoundedMaxOptions {