diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index aee065a2a20..20b6156f8bd 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -8,6 +8,8 @@ use std::sync::Arc; use tracing::debug; use vortex::aggregate_fn::Accumulator; +use vortex::aggregate_fn::AggregateFnRef; +use vortex::aggregate_fn::AggregateFnVTableExt; use vortex::aggregate_fn::DynAccumulator; use vortex::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use vortex::aggregate_fn::NumericalAggregateOpts; @@ -538,6 +540,20 @@ impl PushedAggregate { Self::Count => Box::new(Accumulator::try_new(Count, opts, dtype)?), }) } + + /// If zone maps store information for this aggregate function, this + /// aggregate function, None otherwise. + /// + /// Example: Mean isn't stored in zone maps + pub fn zone_map_supply_fn(self) -> Option { + let opts = NumericalAggregateOpts::default(); + Some(match self { + Self::Min => Min.bind(opts), + Self::Max => Max.bind(opts), + Self::Sum => Sum.bind(opts), + Self::Mean | Self::First | Self::Count => return None, + }) + } } /// Check if this is an aggregate function we can handle in Vortex diff --git a/vortex-duckdb/src/projection.rs b/vortex-duckdb/src/projection.rs index df0481250bd..f28af916378 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -3,6 +3,7 @@ use std::ops::Range; use num_traits::AsPrimitive as _; +use vortex::array::stats::expr::stat; use vortex::dtype::DType; use vortex::error::VortexExpect; use vortex::error::VortexResult; @@ -17,7 +18,6 @@ use vortex::expr::root; use vortex::expr::select; use vortex::layout::layouts::row_idx::row_idx; use vortex::scan::selection::Selection; -use vortex_utils::aliases::hash_set::HashSet; use crate::convert::try_from_table_filter; use crate::convert::try_from_virtual_column_filter; @@ -25,6 +25,7 @@ use crate::duckdb::LogicalType; use crate::duckdb::TableFilterClass; use crate::duckdb::TableFilterSetRef; use crate::table_function::ColumnAggregate; +use crate::table_function::all_aggregates_read_zone_maps; // See MultiFileReader for constants @@ -192,21 +193,32 @@ impl Projection { // Create a projection for aggregate scan pub fn new_aggregate(aggregates: &[ColumnAggregate], fields: &[DuckdbField]) -> Self { - let mut names = Vec::with_capacity(aggregates.len()); - let mut seen: HashSet = HashSet::with_capacity(aggregates.len()); - for aggregate in aggregates { - let ColumnAggregate::Real { projection_id, .. } = aggregate else { - continue; - }; - if seen.contains(projection_id) { - continue; - } - seen.insert(*projection_id); - let projection_id: usize = projection_id.as_(); - names.push(fields[projection_id].name.as_str()); - } + let all_aggregates_read_zone_maps = all_aggregates_read_zone_maps(aggregates); + let fields: Vec<(String, Expression)> = aggregates + .iter() + .enumerate() + .filter_map(|(idx, aggregate)| { + let ColumnAggregate::Real { + projection_id, + aggregate, + } = aggregate + else { + return None; + }; + let projection_id: usize = (*projection_id).as_(); + let column = get_item(fields[projection_id].name.as_str(), root()); + let field_expr = match aggregate.zone_map_supply_fn() { + Some(aggregate_fn) if all_aggregates_read_zone_maps => { + stat(column, aggregate_fn) + } + _ => column, + }; + Some((idx.to_string(), field_expr)) + }) + .collect(); + Projection { - projection: select(names, root()), + projection: pack(fields, false.into()), file_index_column_pos: None, file_row_number_column_pos: None, } diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index 9457f00bc78..f1219fe433e 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -49,7 +49,6 @@ use vortex::scalar_fn::fns::operators::Operator; use vortex::scalar_fn::fns::pack::Pack; use vortex::scan::DataSource; use vortex::scan::ScanRequest; -use vortex_utils::aliases::hash_map::HashMap; use vortex_utils::parallelism::get_available_parallelism; use crate::RUNTIME; @@ -443,28 +442,20 @@ fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult bool { + aggregates.iter().all(|aggregate| match aggregate { + ColumnAggregate::Real { aggregate, .. } => aggregate.zone_map_supply_fn().is_some(), + ColumnAggregate::CountStar => true, + }) +} + fn scan_aggregate( local_state: &mut TableFunctionLocal, global_state: &TableFunctionGlobal, chunk: &mut DataChunkRef, ) -> VortexResult<()> { - let aggregates_len = global_state.aggregates.len(); - // seen[k] = output column for requested column k. - // If min(x), max(x), avg(y) are requested, seen = { 0: 0, 1: 1} - let mut seen: HashMap = HashMap::with_capacity(aggregates_len); - // positions[k] = column id for accumulator k - // If min(x), max(x), avg(y) are requested, positions = [0, 0, 1] - let mut positions: Vec = Vec::with_capacity(aggregates_len); - - for aggregate in &global_state.aggregates { - let ColumnAggregate::Real { projection_id, .. } = aggregate else { - continue; - }; - let len = seen.len(); - let pos = seen.entry_ref(projection_id).or_insert(len); - positions.push(*pos); - } - let has_count_star = local_state.partials.len() < aggregates_len; + let all_aggregates_read_zone_maps = all_aggregates_read_zone_maps(&global_state.aggregates); + let has_count_star = local_state.partials.len() < global_state.aggregates.len(); let mut ctx = SESSION.create_execution_ctx(); loop { @@ -497,8 +488,29 @@ fn scan_aggregate( }; let array = convert_result(result?.0, &mut ctx)?; - for (i, partial) in positions.iter().zip(local_state.partials.iter_mut()) { - partial.accumulate(array.unmasked_field(*i), &mut ctx)?; + for (i, partial) in local_state.partials.iter_mut().enumerate() { + let field = array.unmasked_field(i); + // If any aggregate can't read from zone map, all columns are + // decoded, so accumulating decoded fields is faster than + // combining partial scalars + if !all_aggregates_read_zone_maps { + partial.accumulate(field, &mut ctx)?; + continue; + } + if field.len() == 0 { + // filtered splits where all rows fail the filter + continue; + } + let field_scalar = field.execute_scalar(0, &mut ctx)?; + let target = partial.partial_scalar()?; + // field_scalar may be non-nullable, partial_scalar always returns + // nullable dtype. + let scalar = if field_scalar.dtype() == target.dtype() { + field_scalar + } else { + field_scalar.cast(target.dtype())? + }; + partial.combine_partials(scalar)?; } {