Skip to content

Commit 77982e0

Browse files
robert3005claude
robert3005
and
claude
committed
Pass one AggregateArgs to the aggregate vtable methods
Every `AggregateFnVTable` execution method took an options reference and an `AggregateDTypesRef` side by side, and every implementor and call site threaded both. They are now one `AggregateArgs<'a, O>` carrying the options alongside the input, return and partial dtypes. `AggregateDTypesRef` is gone; `AggregateDTypes::borrow` becomes `AggregateDTypes::args(&options)`, which binds the options an accumulator already stores to the dtypes it already resolved. `AggregateArgs` implements `Clone` and `Copy` by hand: the derives would bound them on `O`, which the borrowed fields do not need. Signed-off-by: Robert Kruszewski <github@robertk.io> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BAC54whRD3iHLBZfM4TCvd
1 parent a0f687c commit 77982e0

37 files changed

Lines changed: 532 additions & 802 deletions

File tree

‎vortex-array/src/aggregate_fn/accumulator.rs‎

Lines changed: 16 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,7 @@ impl<V: AggregateFnVTable> Accumulator<V> {
5757

5858
/// The state of a group with no accumulated values.
5959
pub(crate) fn empty_partial(&self) -> VortexResult<V::Partial> {
60-
self.vtable
61-
.empty_partial(&self.options, self.dtypes.borrow())
60+
self.vtable.empty_partial(self.dtypes.args(&self.options))
6261
}
6362

6463
fn ensure_partial(&mut self) -> VortexResult<()> {
@@ -74,7 +73,7 @@ impl<V: AggregateFnVTable> Accumulator<V> {
7473
None => other,
7574
Some(current) => {
7675
self.vtable
77-
.merge_partials(&self.options, self.dtypes.borrow(), current, other)?
76+
.merge_partials(self.dtypes.args(&self.options), current, other)?
7877
}
7978
});
8079
Ok(())
@@ -83,7 +82,7 @@ impl<V: AggregateFnVTable> Accumulator<V> {
8382
fn fold_partial_scalar(&mut self, scalar: Scalar) -> VortexResult<()> {
8483
let other = self
8584
.vtable
86-
.partial_from_scalar(&self.options, self.dtypes.borrow(), scalar)?;
85+
.partial_from_scalar(self.dtypes.args(&self.options), scalar)?;
8786
self.fold_partial(other)
8887
}
8988
}
@@ -208,7 +207,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
208207
let partial = self.partial.as_mut().vortex_expect("partial materialized");
209208
if self
210209
.vtable
211-
.try_accumulate(&self.options, self.dtypes.borrow(), partial, batch, ctx)?
210+
.try_accumulate(self.dtypes.args(&self.options), partial, batch, ctx)?
212211
{
213212
return Ok(());
214213
}
@@ -248,7 +247,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
248247
self.ensure_partial()?;
249248
let partial = self.partial.as_mut().vortex_expect("partial materialized");
250249
self.vtable
251-
.accumulate(&self.options, self.dtypes.borrow(), partial, &columnar, ctx)
250+
.accumulate(self.dtypes.args(&self.options), partial, &columnar, ctx)
252251
}
253252

254253
fn merge_from(&mut self, other: &mut dyn DynAccumulator) -> VortexResult<()> {
@@ -287,7 +286,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
287286
fn is_saturated(&self) -> bool {
288287
self.partial.as_ref().is_some_and(|partial| {
289288
self.vtable
290-
.is_saturated(&self.options, self.dtypes.borrow(), partial)
289+
.is_saturated(self.dtypes.args(&self.options), partial)
291290
})
292291
}
293292

@@ -296,20 +295,18 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
296295
}
297296

298297
fn partial_scalar(&self) -> VortexResult<Scalar> {
299-
let dtypes = self.dtypes.borrow();
298+
let args = self.dtypes.args(&self.options);
300299
let partial = match &self.partial {
301-
Some(partial) => self.vtable.to_scalar(&self.options, dtypes, partial)?,
302-
None => self
303-
.vtable
304-
.to_scalar(&self.options, dtypes, &self.empty_partial()?)?,
300+
Some(partial) => self.vtable.to_scalar(args, partial)?,
301+
None => self.vtable.to_scalar(args, &self.empty_partial()?)?,
305302
};
306303

307304
#[cfg(debug_assertions)]
308305
{
309306
vortex_ensure!(
310-
partial.dtype() == dtypes.partial_dtype,
307+
partial.dtype() == args.partial_dtype,
311308
"Aggregate returned incorrect DType on partial_scalar: expected {}, got {}",
312-
dtypes.partial_dtype,
309+
args.partial_dtype,
313310
partial.dtype(),
314311
);
315312
}
@@ -318,20 +315,16 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
318315
}
319316

320317
fn final_scalar(&self) -> VortexResult<Scalar> {
321-
let dtypes = self.dtypes.borrow();
318+
let args = self.dtypes.args(&self.options);
322319
let result = match &self.partial {
323-
Some(partial) => self
324-
.vtable
325-
.finalize_scalar(&self.options, dtypes, partial)?,
326-
None => self
327-
.vtable
328-
.finalize_scalar(&self.options, dtypes, &self.empty_partial()?)?,
320+
Some(partial) => self.vtable.finalize_scalar(args, partial)?,
321+
None => self.vtable.finalize_scalar(args, &self.empty_partial()?)?,
329322
};
330323

331324
vortex_ensure!(
332-
result.dtype() == dtypes.return_dtype,
325+
result.dtype() == args.return_dtype,
333326
"Aggregate returned incorrect DType on final_scalar: expected {}, got {}",
334-
dtypes.return_dtype,
327+
args.return_dtype,
335328
result.dtype(),
336329
);
337330

‎vortex-array/src/aggregate_fn/accumulator_grouped.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,7 +265,7 @@ impl<V: AggregateFnVTable> DynGroupedAccumulator for GroupedAccumulator<V> {
265265
let states = self.flush()?;
266266
let results = self
267267
.vtable
268-
.finalize(&self.options, self.dtypes.borrow(), states)?;
268+
.finalize(self.dtypes.args(&self.options), states)?;
269269

270270
vortex_ensure!(
271271
results.dtype() == &self.dtypes.return_dtype,

‎vortex-array/src/aggregate_fn/combined.rs‎

Lines changed: 28 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@ use crate::ArrayRef;
1919
use crate::Columnar;
2020
use crate::ExecutionCtx;
2121
use crate::aggregate_fn::Accumulator;
22+
use crate::aggregate_fn::AggregateArgs;
2223
use crate::aggregate_fn::AggregateDTypes;
23-
use crate::aggregate_fn::AggregateDTypesRef;
2424
use crate::aggregate_fn::AggregateFnId;
2525
use crate::aggregate_fn::AggregateFnVTable;
2626
use crate::aggregate_fn::DynAccumulator;
@@ -87,21 +87,19 @@ pub trait BinaryCombined: 'static + Send + Sync + Clone {
8787

8888
/// Combine the finalized left and right results into the final aggregate.
8989
///
90-
/// `dtypes` describes the combined aggregate; `left` and `right` already have their child's
91-
/// result dtype. The returned array must have dtype `dtypes.return_dtype`.
90+
/// `args` describes the combined aggregate; `left` and `right` already have their child's
91+
/// result dtype. The returned array must have dtype `args.return_dtype`.
9292
fn finalize(
9393
&self,
94-
options: &CombinedOptions<Self>,
95-
dtypes: AggregateDTypesRef<'_>,
94+
args: AggregateArgs<'_, CombinedOptions<Self>>,
9695
left: ArrayRef,
9796
right: ArrayRef,
9897
) -> VortexResult<ArrayRef>;
9998

100-
/// Combine the finalized child scalars into a result of dtype `dtypes.return_dtype`.
99+
/// Combine the finalized child scalars into a result of dtype `args.return_dtype`.
101100
fn finalize_scalar(
102101
&self,
103-
options: &CombinedOptions<Self>,
104-
dtypes: AggregateDTypesRef<'_>,
102+
args: AggregateArgs<'_, CombinedOptions<Self>>,
105103
left_scalar: Scalar,
106104
right_scalar: Scalar,
107105
) -> VortexResult<Scalar>;
@@ -153,11 +151,11 @@ impl<T: BinaryCombined> Combined<T> {
153151
/// Construct a pair of empty child accumulators.
154152
fn new_child_accumulators(
155153
&self,
156-
options: &CombinedOptions<T>,
157-
input_dtype: &DType,
154+
args: AggregateArgs<'_, CombinedOptions<T>>,
158155
) -> VortexResult<ChildAccumulators<T>> {
159-
let left = Accumulator::try_new(self.0.left(), options.0.clone(), input_dtype.clone())?;
160-
let right = Accumulator::try_new(self.0.right(), options.1.clone(), input_dtype.clone())?;
156+
let left = Accumulator::try_new(self.0.left(), args.options.0.clone(), args.dtype.clone())?;
157+
let right =
158+
Accumulator::try_new(self.0.right(), args.options.1.clone(), args.dtype.clone())?;
161159
Ok((left, right))
162160
}
163161
}
@@ -191,21 +189,16 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
191189
Some(self.0.partial_struct_dtype(l, r))
192190
}
193191

194-
fn empty_partial(
195-
&self,
196-
options: &Self::Options,
197-
dtypes: AggregateDTypesRef<'_>,
198-
) -> VortexResult<Self::Partial> {
199-
self.new_child_accumulators(options, dtypes.dtype)
192+
fn empty_partial(&self, args: AggregateArgs<'_, Self::Options>) -> VortexResult<Self::Partial> {
193+
self.new_child_accumulators(args)
200194
}
201195

202196
fn partial_from_scalar(
203197
&self,
204-
options: &Self::Options,
205-
dtypes: AggregateDTypesRef<'_>,
198+
args: AggregateArgs<'_, Self::Options>,
206199
scalar: Scalar,
207200
) -> VortexResult<Self::Partial> {
208-
let (mut left, mut right) = self.new_child_accumulators(options, dtypes.dtype)?;
201+
let (mut left, mut right) = self.new_child_accumulators(args)?;
209202
// A null partial represents an empty group and parses to empty child accumulators.
210203
if !scalar.is_null() {
211204
let s = scalar.as_struct();
@@ -225,8 +218,7 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
225218

226219
fn merge_partials(
227220
&self,
228-
_options: &Self::Options,
229-
_dtypes: AggregateDTypesRef<'_>,
221+
_args: AggregateArgs<'_, Self::Options>,
230222
(mut left, mut right): Self::Partial,
231223
(mut other_left, mut other_right): Self::Partial,
232224
) -> VortexResult<Self::Partial> {
@@ -238,22 +230,20 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
238230

239231
fn to_scalar(
240232
&self,
241-
_options: &Self::Options,
242-
dtypes: AggregateDTypesRef<'_>,
233+
args: AggregateArgs<'_, Self::Options>,
243234
partial: &Self::Partial,
244235
) -> VortexResult<Scalar> {
245236
let l_scalar = partial.0.partial_scalar()?;
246237
let r_scalar = partial.1.partial_scalar()?;
247238
Ok(Scalar::struct_(
248-
dtypes.partial_dtype.clone(),
239+
args.partial_dtype.clone(),
249240
vec![l_scalar, r_scalar],
250241
))
251242
}
252243

253244
fn is_saturated(
254245
&self,
255-
_options: &Self::Options,
256-
_dtypes: AggregateDTypesRef<'_>,
246+
_args: AggregateArgs<'_, Self::Options>,
257247
partial: &Self::Partial,
258248
) -> bool {
259249
partial.0.is_saturated() && partial.1.is_saturated()
@@ -266,8 +256,7 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
266256
/// `true` so [`Self::accumulate`] is unreachable.
267257
fn try_accumulate(
268258
&self,
269-
_options: &Self::Options,
270-
_dtypes: AggregateDTypesRef<'_>,
259+
_args: AggregateArgs<'_, Self::Options>,
271260
state: &mut Self::Partial,
272261
batch: &ArrayRef,
273262
ctx: &mut ExecutionCtx,
@@ -279,8 +268,7 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
279268

280269
fn accumulate(
281270
&self,
282-
_options: &Self::Options,
283-
_dtypes: AggregateDTypesRef<'_>,
271+
_args: AggregateArgs<'_, Self::Options>,
284272
_state: &mut Self::Partial,
285273
_batch: &Columnar,
286274
_ctx: &mut ExecutionCtx,
@@ -290,29 +278,27 @@ impl<T: BinaryCombined> AggregateFnVTable for Combined<T> {
290278

291279
fn finalize(
292280
&self,
293-
options: &Self::Options,
294-
dtypes: AggregateDTypesRef<'_>,
281+
args: AggregateArgs<'_, Self::Options>,
295282
states: ArrayRef,
296283
) -> VortexResult<ArrayRef> {
297284
let l_field = states.get_item(FieldName::from(self.0.left_name()))?;
298285
let r_field = states.get_item(FieldName::from(self.0.right_name()))?;
299286
let left = self.0.left();
300287
let right = self.0.right();
301-
let l_dtypes = AggregateDTypes::try_new(&left, &options.0, dtypes.dtype.clone())?;
302-
let r_dtypes = AggregateDTypes::try_new(&right, &options.1, dtypes.dtype.clone())?;
303-
let l_finalized = left.finalize(&options.0, l_dtypes.borrow(), l_field)?;
304-
let r_finalized = right.finalize(&options.1, r_dtypes.borrow(), r_field)?;
305-
BinaryCombined::finalize(&self.0, options, dtypes, l_finalized, r_finalized)
288+
let l_dtypes = AggregateDTypes::try_new(&left, &args.options.0, args.dtype.clone())?;
289+
let r_dtypes = AggregateDTypes::try_new(&right, &args.options.1, args.dtype.clone())?;
290+
let l_finalized = left.finalize(l_dtypes.args(&args.options.0), l_field)?;
291+
let r_finalized = right.finalize(r_dtypes.args(&args.options.1), r_field)?;
292+
BinaryCombined::finalize(&self.0, args, l_finalized, r_finalized)
306293
}
307294

308295
fn finalize_scalar(
309296
&self,
310-
options: &Self::Options,
311-
dtypes: AggregateDTypesRef<'_>,
297+
args: AggregateArgs<'_, Self::Options>,
312298
partial: &Self::Partial,
313299
) -> VortexResult<Scalar> {
314300
let l_scalar = partial.0.final_scalar()?;
315301
let r_scalar = partial.1.final_scalar()?;
316-
BinaryCombined::finalize_scalar(&self.0, options, dtypes, l_scalar, r_scalar)
302+
BinaryCombined::finalize_scalar(&self.0, args, l_scalar, r_scalar)
317303
}
318304
}

‎vortex-array/src/aggregate_fn/fns/all_nan/mod.rs‎

Lines changed: 11 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use crate::ArrayRef;
99
use crate::Columnar;
1010
use crate::ExecutionCtx;
1111
use crate::IntoArray;
12-
use crate::aggregate_fn::AggregateDTypesRef;
12+
use crate::aggregate_fn::AggregateArgs;
1313
use crate::aggregate_fn::AggregateFnId;
1414
use crate::aggregate_fn::AggregateFnVTable;
1515
use crate::aggregate_fn::EmptyOptions;
@@ -63,25 +63,22 @@ impl AggregateFnVTable for AllNan {
6363

6464
fn empty_partial(
6565
&self,
66-
_options: &Self::Options,
67-
_dtypes: AggregateDTypesRef<'_>,
66+
_args: AggregateArgs<'_, Self::Options>,
6867
) -> VortexResult<Self::Partial> {
6968
Ok(true)
7069
}
7170

7271
fn partial_from_scalar(
7372
&self,
74-
_options: &Self::Options,
75-
_dtypes: AggregateDTypesRef<'_>,
73+
_args: AggregateArgs<'_, Self::Options>,
7674
scalar: Scalar,
7775
) -> VortexResult<Self::Partial> {
7876
bool::try_from(&scalar)
7977
}
8078

8179
fn merge_partials(
8280
&self,
83-
_options: &Self::Options,
84-
_dtypes: AggregateDTypesRef<'_>,
81+
_args: AggregateArgs<'_, Self::Options>,
8582
first: Self::Partial,
8683
second: Self::Partial,
8784
) -> VortexResult<Self::Partial> {
@@ -90,26 +87,23 @@ impl AggregateFnVTable for AllNan {
9087

9188
fn to_scalar(
9289
&self,
93-
_options: &Self::Options,
94-
_dtypes: AggregateDTypesRef<'_>,
90+
_args: AggregateArgs<'_, Self::Options>,
9591
partial: &Self::Partial,
9692
) -> VortexResult<Scalar> {
9793
Ok(Scalar::bool(*partial, Nullability::Nullable))
9894
}
9995

10096
fn is_saturated(
10197
&self,
102-
_options: &Self::Options,
103-
_dtypes: AggregateDTypesRef<'_>,
98+
_args: AggregateArgs<'_, Self::Options>,
10499
partial: &Self::Partial,
105100
) -> bool {
106101
!*partial
107102
}
108103

109104
fn try_accumulate(
110105
&self,
111-
_options: &Self::Options,
112-
_dtypes: AggregateDTypesRef<'_>,
106+
_args: AggregateArgs<'_, Self::Options>,
113107
state: &mut Self::Partial,
114108
batch: &ArrayRef,
115109
ctx: &mut ExecutionCtx,
@@ -125,8 +119,7 @@ impl AggregateFnVTable for AllNan {
125119

126120
fn accumulate(
127121
&self,
128-
_options: &Self::Options,
129-
_dtypes: AggregateDTypesRef<'_>,
122+
_args: AggregateArgs<'_, Self::Options>,
130123
partial: &mut Self::Partial,
131124
batch: &Columnar,
132125
ctx: &mut ExecutionCtx,
@@ -148,20 +141,18 @@ impl AggregateFnVTable for AllNan {
148141

149142
fn finalize(
150143
&self,
151-
_options: &Self::Options,
152-
_dtypes: AggregateDTypesRef<'_>,
144+
_args: AggregateArgs<'_, Self::Options>,
153145
partials: ArrayRef,
154146
) -> VortexResult<ArrayRef> {
155147
Ok(partials)
156148
}
157149

158150
fn finalize_scalar(
159151
&self,
160-
options: &Self::Options,
161-
dtypes: AggregateDTypesRef<'_>,
152+
args: AggregateArgs<'_, Self::Options>,
162153
partial: &Self::Partial,
163154
) -> VortexResult<Scalar> {
164-
self.to_scalar(options, dtypes, partial)
155+
self.to_scalar(args, partial)
165156
}
166157
}
167158

0 commit comments

Comments
 (0)