Skip index write interface

vortex-array/src/stats/rewrite.rs

@@ -11,6 +11,7 @@ use vortex_error::vortex_ensure;
 use vortex_session::VortexSession;
 use vortex_utils::iter::ReduceBalancedIterExt;
 
+use crate::aggregate_fn::AggregateFnRef;
 use crate::dtype::DType;
 use crate::expr::BoundExpression;
 use crate::scalar_fn::ScalarFnId;
@@ -88,12 +89,16 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static {
 /// Context passed to stats rewrite rules.
 pub struct StatsRewriteCtx<'a> {
     session: &'a VortexSession,
+    aggregate_fns: &'a [AggregateFnRef],
 }
 
 impl<'a> StatsRewriteCtx<'a> {
     /// Create a rewrite context for `session`.
     pub fn new(session: &'a VortexSession) -> Self {
-        Self { session }
+        Self {
+            session,
+            aggregate_fns: &[],
+        }
     }
 
     /// Returns the session that owns the rewrite registry.
@@ -126,6 +131,17 @@ impl<'a> StatsRewriteCtx<'a> {
         );
         Ok(())
     }
+
+    /// Sets the aggregate functions available to stats rewrite rules.
+    pub fn with_aggregate_fns(mut self, aggregate_fns: &'a [AggregateFnRef]) -> Self {
+        self.aggregate_fns = aggregate_fns;
+        self
+    }
+
+    /// Returns the aggregate functions available to stats rewrite rules.
+    pub fn aggregate_fns(&self) -> &'a [AggregateFnRef] {
+        self.aggregate_fns
+    }
 }
 
 fn rewrite(

vortex-file/src/strategy.rs

@@ -6,6 +6,7 @@
 use std::num::NonZeroUsize;
 use std::sync::Arc;
 
+use vortex_array::aggregate_fn::AggregateFnRef;
 use vortex_array::dtype::FieldPath;
 use vortex_btrblocks::BtrBlocksCompressorBuilder;
 use vortex_btrblocks::SchemeExt;
@@ -56,6 +57,7 @@ pub struct WriteStrategyBuilder {
     row_block_size: usize,
     data_block_target_bytes: Option<u64>,
     field_writers: HashMap<FieldPath, Arc<dyn LayoutStrategy>>,
+    field_aggregates: HashMap<FieldPath, Arc<[AggregateFnRef]>>,
     flat_strategy: Option<Arc<dyn LayoutStrategy>>,
     probe_compressor: Option<Arc<dyn CompressorPlugin>>,
     /// Whether to write list fields using [`ListLayoutStrategy`].
@@ -73,6 +75,7 @@ impl Default for WriteStrategyBuilder {
             row_block_size: 8192,
             data_block_target_bytes: Some(ONE_MEG),
             field_writers: HashMap::new(),
+            field_aggregates: HashMap::new(),
             flat_strategy: None,
             probe_compressor: None,
             use_list_layout: use_experimental_list_layout(),
@@ -114,6 +117,9 @@ impl WriteStrategyBuilder {
     ///
     /// The field path is matched after the root struct is split into columns. This is useful when a
     /// column needs a custom compression/layout policy while the rest of the file uses defaults.
+    ///
+    /// **Note** if a [`WriteStrategyBuilder::with_field_aggregates`] was called over the same field,
+    /// the build process will panic.
     pub fn with_field_writer(
         mut self,
         field: impl Into<FieldPath>,
@@ -123,6 +129,20 @@ impl WriteStrategyBuilder {
         self
     }
 
+    /// Override default aggregates with custom zoned aggregates for a field.
+    ///
+    /// **Note** Calling [`Self::build`] panics if this path conflicts with another field
+    /// override.
+    pub fn with_field_aggregates(
+        mut self,
+        field: impl Into<FieldPath>,
+        aggregates: impl IntoIterator<Item = AggregateFnRef>,
+    ) -> Self {
+        self.field_aggregates
+            .insert(field.into(), aggregates.into_iter().collect());
+        self
+    }
+
     /// Override the flat layout strategy used for leaf chunks.
     ///
     /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom
@@ -227,36 +247,56 @@ impl WriteStrategyBuilder {
 
         let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0");
 
-        // 2. calculate stats for each row group
-        let stats = ZonedStrategy::new(
-            dict,
-            compress_then_flat.clone(),
-            ZonedLayoutOptions {
-                block_size: row_block_size,
-                ..Default::default()
-            },
-        );
+        // Default and per-field aggregates share the same zoning and repartitioning pipeline.
+        let build_repartition =
+            |zone_layout_options: ZonedLayoutOptions| -> Arc<dyn LayoutStrategy> {
+                // 2. calculate stats for each row group
+                let stats = ZonedStrategy::new(
+                    dict.clone(),
+                    compress_then_flat.clone(),
+                    zone_layout_options,
+                );
 
-        // 1. repartition each column to fixed row counts
-        let repartition = RepartitionStrategy::new(
-            stats,
-            RepartitionWriterOptions {
-                // No minimum block size in bytes
-                block_size_minimum: 0,
-                // Always repartition into 8K row blocks
-                block_len_multiple: self.row_block_size,
-                block_size_target: None,
-                canonicalize: false,
-            },
-        );
+                // 1. repartition each column to fixed row counts
+                Arc::new(RepartitionStrategy::new(
+                    stats,
+                    RepartitionWriterOptions {
+                        // No minimum block size in bytes
+                        block_size_minimum: 0,
+                        block_len_multiple: row_block_size.get(),
+                        block_size_target: None,
+                        canonicalize: false,
+                    },
+                ))
+            };
+
+        let repartition_strategy = build_repartition(ZonedLayoutOptions {
+            block_size: row_block_size,
+            ..Default::default()
+        });
+
+        // Field aggregates and field writers use the same field-path namespace.
+        // Conflicting overrides are unexpected and are rejected by TableStrategy.
+        let field_writers = self
+            .field_aggregates
+            .into_iter()
+            .map(|(field, aggregates)| {
+                let options = ZonedLayoutOptions {
+                    block_size: row_block_size,
+                    aggregate_fns: Some(aggregates),
+                    ..Default::default()
+                };
+                (field, build_repartition(options))
+            })
+            .chain(self.field_writers);
 
         // 0. start with splitting columns
         let validity_strategy = CollectStrategy::new(compress_then_flat.clone());
 
         // Take any field overrides from the builder and apply them to the final strategy.
         let mut table_strategy =
-            TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition))
-                .with_field_writers(self.field_writers);
+            TableStrategy::new(Arc::new(validity_strategy), repartition_strategy)
+                .with_field_writers(field_writers);
 
         if self.use_list_layout {
             // We need a closure here to enable recursive application of list layout.

vortex-file/src/tests.rs

@@ -2800,3 +2800,23 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> {
     assert_eq!(result.len(), 1);
     Ok(())
 }
+
+#[rstest]
+#[case(true)]
+#[case(false)]
+#[should_panic(expected = "Override for field_path $a conflicts with existing override for $a")]
+fn test_field_writer_conflicts_with_aggregates(#[case] writer_first: bool) {
+    let builder = crate::strategy::WriteStrategyBuilder::default();
+    let writer = Arc::new(FlatLayoutStrategy::default());
+    let builder = if writer_first {
+        builder
+            .with_field_writer(field_path!(a), writer)
+            .with_field_aggregates(field_path!(a), [])
+    } else {
+        builder
+            .with_field_aggregates(field_path!(a), [])
+            .with_field_writer(field_path!(a), writer)
+    };
+
+    builder.build();
+}