Clarify dense grouped aggregate ids

gatesn · gatesn · commit 9bd157fe5670 · 2026-06-11T17:42:21.000-04:00
Signed-off-by: "Nicholas Gates" &lt;nick@nickgates.com&gt;
diff --git a/vortex-array/src/aggregate_fn/accumulator_grouped.rs b/vortex-array/src/aggregate_fn/accumulator_grouped.rs
@@ -28,8 +28,10 @@ pub type GroupedAccumulatorRef = Box<dyn DynGroupedAccumulator>;
 
 /// An accumulator used for computing aggregates over dense group ids.
 ///
-/// Group ids are dense `u32` slots in the range `0..num_groups`. The accumulator keeps one partial
-/// state per slot, so ordered and unordered grouping only differ in how the caller assigns ids.
+/// Group ids are caller-assigned `u32` ordinals in the dense range `0..num_groups`. Input batches
+/// may repeat, omit, and reorder those ids, but every id must identify a state slot rather than a
+/// raw group key. The accumulator keeps one partial state per slot, so ordered and unordered
+/// grouping only differ in how the caller assigns ids.
 pub struct GroupedAccumulator<V: AggregateFnVTable> {
     /// The vtable of the aggregate function.
     vtable: V,
@@ -167,6 +169,9 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
 /// aggregate function is not known at compile time.
 pub trait DynGroupedAccumulator: 'static + Send {
     /// Accumulate a values batch into dense group state.
+    ///
+    /// `group_ids` is parallel to `batch`. Each id must be a caller-assigned group ordinal in
+    /// `0..num_groups`; ids may repeat, appear out of order, or be absent from a given batch.
     fn accumulate(
         &mut self,
         batch: &ArrayRef,
@@ -176,6 +181,9 @@ pub trait DynGroupedAccumulator: 'static + Send {
     ) -> VortexResult<()>;
 
     /// Fold columnar partial states into dense group state.
+    ///
+    /// `group_ids` is parallel to `partials` and follows the same dense ordinal contract as
+    /// [`Self::accumulate`].
     fn accumulate_partials(
         &mut self,
         partials: &ArrayRef,
diff --git a/vortex-array/src/aggregate_fn/fns/count/mod.rs b/vortex-array/src/aggregate_fn/fns/count/mod.rs
@@ -263,6 +263,16 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn grouped_count_rejects_out_of_range_group_id() -> VortexResult<()> {
+        let values = PrimitiveArray::new(buffer![1i32, 2], Validity::NonNullable).into_array();
+        let mut acc = GroupedAccumulator::try_new(Count, EmptyOptions, values.dtype().clone())?;
+        let mut ctx = LEGACY_SESSION.create_execution_ctx();
+
+        assert!(acc.accumulate(&values, &[0, 2], 2, &mut ctx).is_err());
+        Ok(())
+    }
+
     #[test]
     fn grouped_count_accumulate_partials_and_merge_group() -> VortexResult<()> {
         let dtype = DType::Primitive(PType::I32, Nullability::Nullable);
diff --git a/vortex-array/src/aggregate_fn/kernels.rs b/vortex-array/src/aggregate_fn/kernels.rs
@@ -30,8 +30,9 @@ pub trait DynAggregateKernel: 'static + Send + Sync + Debug {
 /// Partial grouped aggregate output produced by an encoding-specific grouped kernel.
 ///
 /// `group_ids` is parallel to `partials`: each row in `partials` is a partial state for the
-/// corresponding dense group id. The grouped accumulator merges this batch through
-/// `accumulate_partials`.
+/// corresponding dense group ordinal. The ids may repeat, omit, and reorder groups, but must be
+/// valid slots in the accumulator's `0..num_groups` range. The grouped accumulator merges this
+/// batch through `accumulate_partials`.
 #[derive(Clone, Debug)]
 pub struct GroupedAggregateKernelResult {
     group_ids: Buffer<u32>,
@@ -57,13 +58,13 @@ impl GroupedAggregateKernelResult {
 
 /// A pluggable kernel for batch aggregation of many groups.
 ///
-/// The kernel is matched on the encoding of the values array. It receives the same dense group ids
-/// that the caller passed to the grouped accumulator and may aggregate directly in the encoded
-/// domain.
+/// The kernel is matched on the encoding of the values array. It receives the same dense group
+/// ordinals that the caller passed to the grouped accumulator and may aggregate directly in the
+/// encoded domain.
 ///
 /// Return `Ok(None)` if the kernel cannot be applied to the given aggregate function.
 pub trait DynGroupedAggregateKernel: 'static + Send + Sync + Debug {
-    /// Aggregate values into a partial-state batch keyed by dense group id.
+    /// Aggregate values into a partial-state batch keyed by dense group ordinal.
     fn grouped_aggregate(
         &self,
         aggregate_fn: &AggregateFnRef,
diff --git a/vortex-array/src/aggregate_fn/vtable.rs b/vortex-array/src/aggregate_fn/vtable.rs
@@ -148,8 +148,9 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
 
     /// Try to accumulate a raw values batch into dense per-group states before decompression.
     ///
-    /// `group_ids` is parallel to `batch` and contains dense ids in `0..states.len()`. Returns
-    /// `true` when the batch was fully handled.
+    /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in
+    /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch.
+    /// Returns `true` when the batch was fully handled.
     fn try_accumulate_grouped(
         &self,
         _states: &mut [Self::Partial],
@@ -162,8 +163,9 @@ pub trait AggregateFnVTable: 'static + Sized + Clone + Send + Sync {
 
     /// Accumulate a canonical values batch into dense per-group states.
     ///
-    /// `group_ids` is parallel to `batch` and contains dense ids in `0..states.len()`. Returns
-    /// `true` when the batch was fully handled. The provided default preserves universal
+    /// `group_ids` is parallel to `batch` and contains caller-assigned dense ordinals in
+    /// `0..states.len()`. Ids may repeat, appear out of order, or be absent from the batch.
+    /// Returns `true` when the batch was fully handled. The provided default preserves universal
     /// correctness through [`GroupedAccumulator`]'s fallback.
     fn accumulate_grouped(
         &self,
