duckdb: use zone maps for min/max/sum aggregations

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<AggregateFnRef> {
+        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

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,14 +18,14 @@ 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;
 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<u64> = 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,
         }

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<Struc
     })
 }
 
+pub(crate) fn all_aggregates_read_zone_maps(aggregates: &[ColumnAggregate]) -> 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<u64, usize> = 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<usize> = 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)?;
         }
 
         {