Skip to content

Commit b0ea625

Browse files
authored
Port Sparse sum to SumV2 (#9622)
## Summary - Tracking Issue: #9591 - Closes #9619 Port the Sparse aggregate pushdown to SumV2 while preserving SQL semantics for all-null input. ## Changes - register the Sparse aggregate kernel for SumV2 - construct the bound aggregate accumulator so Sum and SumV2 retain their respective partial-state, overflow, and NaN semantics - add direct registered-kernel tests for non-empty and all-null Sparse arrays Signed-off-by: Matt Katz <[email protected]>
1 parent 709ef16 commit b0ea625

2 files changed

Lines changed: 101 additions & 13 deletions

File tree

‎encodings/sparse/src/compute/sum.rs‎

Lines changed: 94 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,9 @@
44
use vortex_array::ArrayRef;
55
use vortex_array::ExecutionCtx;
66
use vortex_array::IntoArray;
7-
use vortex_array::aggregate_fn::Accumulator;
87
use vortex_array::aggregate_fn::AggregateFnRef;
9-
use vortex_array::aggregate_fn::DynAccumulator;
108
use vortex_array::aggregate_fn::fns::sum::Sum;
9+
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
1110
use vortex_array::aggregate_fn::kernels::DynAggregateKernel;
1211
use vortex_array::arrays::ConstantArray;
1312
use vortex_array::scalar::Scalar;
@@ -16,13 +15,13 @@ use vortex_error::VortexResult;
1615
use crate::Sparse;
1716
use crate::SparseExt as _;
1817

19-
/// Sparse-specific `sum` kernel.
18+
/// Sparse-specific `sum` and `sum_v2` kernel.
2019
///
2120
/// `sum(Sparse{ F, patches }) = sum(patches.values) + F * (N - patches.num_patches())`.
2221
///
23-
/// The constant contribution is computed via the existing `Sum` accumulator's constant
24-
/// short-circuit (`multiply_constant`), so overflow saturates to null exactly as in the
25-
/// baseline. The work is `O(P)` instead of `O(N)`.
22+
/// The constant contribution is computed via the aggregate accumulator's constant short-circuit
23+
/// (`multiply_constant`), so overflow and empty-input semantics match the baseline. The work is
24+
/// `O(P)` instead of `O(N)`.
2625
#[derive(Debug)]
2726
pub(crate) struct SparseSumKernel;
2827

@@ -33,9 +32,9 @@ impl DynAggregateKernel for SparseSumKernel {
3332
batch: &ArrayRef,
3433
ctx: &mut ExecutionCtx,
3534
) -> VortexResult<Option<Scalar>> {
36-
let Some(options) = aggregate_fn.as_opt::<Sum>() else {
35+
if !aggregate_fn.is::<Sum>() && !aggregate_fn.is::<SumV2>() {
3736
return Ok(None);
38-
};
37+
}
3938

4039
let Some(sparse) = batch.as_opt::<Sparse>() else {
4140
return Ok(None);
@@ -44,10 +43,9 @@ impl DynAggregateKernel for SparseSumKernel {
4443
let patches = sparse.patches();
4544
let n_fill = sparse.len() - patches.num_patches();
4645

47-
// Build a fresh Sum accumulator over the array dtype and fold in the fill and patch
48-
// contributions. The accumulator's existing semantics (checked overflow → null
49-
// partial, NaN handling per the options) are preserved.
50-
let mut acc = Accumulator::try_new(Sum, *options, batch.dtype().clone())?;
46+
// Build a fresh accumulator over the array dtype and fold in the fill and patch
47+
// contributions. This preserves each aggregate's partial state and semantics.
48+
let mut acc = aggregate_fn.accumulator(batch.dtype())?;
5149

5250
if n_fill > 0 {
5351
let fill_array = ConstantArray::new(sparse.fill_scalar().clone(), n_fill).into_array();
@@ -67,13 +65,22 @@ mod tests {
6765
use std::sync::LazyLock;
6866

6967
use rstest::rstest;
68+
use vortex_array::ArrayRef;
7069
use vortex_array::IntoArray;
7170
use vortex_array::VortexSessionExecute;
71+
use vortex_array::aggregate_fn::Accumulator;
72+
use vortex_array::aggregate_fn::AggregateFnVTableExt;
73+
use vortex_array::aggregate_fn::DynAccumulator;
74+
use vortex_array::aggregate_fn::NumericalAggregateOpts;
7275
use vortex_array::aggregate_fn::fns::sum::sum;
76+
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
77+
use vortex_array::aggregate_fn::fns::sum_v2::sum_v2;
78+
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
7379
use vortex_array::scalar::Scalar;
7480
use vortex_array::session::ArraySessionExt;
7581
use vortex_buffer::buffer;
7682
use vortex_error::VortexResult;
83+
use vortex_error::vortex_bail;
7784
use vortex_session::VortexSession;
7885

7986
use crate::Sparse;
@@ -103,6 +110,36 @@ mod tests {
103110
Ok(kernel_result)
104111
}
105112

113+
fn run_v2_kernel(arr: &ArrayRef, options: NumericalAggregateOpts) -> VortexResult<Scalar> {
114+
let aggregate = SumV2.bind(options);
115+
let Some(kernel) = SESSION
116+
.aggregate_fns()
117+
.find_aggregate_kernel(arr.encoding_id(), aggregate.id())
118+
else {
119+
vortex_bail!("Sparse SumV2 kernel is not registered");
120+
};
121+
let Some(partial) =
122+
kernel.aggregate(&aggregate, arr, &mut SESSION.create_execution_ctx())?
123+
else {
124+
vortex_bail!("Sparse SumV2 kernel declined the aggregate");
125+
};
126+
127+
let mut accumulator = Accumulator::try_new(SumV2, options, arr.dtype().clone())?;
128+
accumulator.combine_partials(partial)?;
129+
accumulator.finish()
130+
}
131+
132+
fn check_v2(array: SparseArray) -> VortexResult<Scalar> {
133+
let arr = array.into_array();
134+
let kernel_result = run_v2_kernel(&arr, NumericalAggregateOpts::default())?;
135+
let canonical_result = sum_v2(&arr, &mut CANONICAL_SESSION.create_execution_ctx())?;
136+
assert_eq!(
137+
kernel_result, canonical_result,
138+
"kernel and canonical sum_v2 paths disagree"
139+
);
140+
Ok(kernel_result)
141+
}
142+
106143
#[rstest]
107144
#[case::positive_fill(
108145
Sparse::try_new(
@@ -146,4 +183,49 @@ mod tests {
146183
let result = check(array).unwrap();
147184
assert_eq!(result.as_primitive().typed_value::<i64>(), Some(expected));
148185
}
186+
187+
#[test]
188+
fn sum_v2_kernel_nonempty() -> VortexResult<()> {
189+
let array = Sparse::try_new(
190+
buffer![0u64, 2].into_array(),
191+
buffer![10i32, 20].into_array(),
192+
5,
193+
Scalar::from(1i32),
194+
)?;
195+
let result = check_v2(array)?;
196+
assert_eq!(result.as_primitive().typed_value::<i64>(), Some(33));
197+
Ok(())
198+
}
199+
200+
#[test]
201+
fn sum_v2_kernel_all_null() -> VortexResult<()> {
202+
let array = Sparse::try_new(
203+
buffer![0u64, 3].into_array(),
204+
vortex_array::arrays::PrimitiveArray::from_option_iter([None::<i32>, None])
205+
.into_array(),
206+
6,
207+
Scalar::null_native::<i32>(),
208+
)?;
209+
assert!(check_v2(array)?.is_null());
210+
Ok(())
211+
}
212+
213+
#[test]
214+
fn sum_v2_kernel_preserves_options() -> VortexResult<()> {
215+
let array = Sparse::try_new(
216+
buffer![0u64].into_array(),
217+
buffer![1.0f64].into_array(),
218+
3,
219+
Scalar::from(f64::NAN),
220+
)?
221+
.into_array();
222+
let result = run_v2_kernel(&array, NumericalAggregateOpts::include_nans())?;
223+
assert!(
224+
result
225+
.as_primitive()
226+
.typed_value::<f64>()
227+
.is_some_and(f64::is_nan)
228+
);
229+
Ok(())
230+
}
149231
}

‎encodings/sparse/src/lib.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,14 @@ use vortex_array::aggregate_fn::fns::min_max::MinMax;
7777
use vortex_array::aggregate_fn::fns::nan_count::NanCount;
7878
use vortex_array::aggregate_fn::fns::null_count::NullCount;
7979
use vortex_array::aggregate_fn::fns::sum::Sum;
80+
use vortex_array::aggregate_fn::fns::sum_v2::SumV2;
8081
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
8182
use vortex_array::session::ArraySessionExt;
8283

8384
/// Initialize Sparse encoding in the given session.
8485
///
8586
/// Registers the Sparse array vtable, parent execution kernels, and aggregate kernels
86-
/// (`IsConstant`, `Sum`, `MinMax`, `NullCount`, `NanCount`).
87+
/// (`IsConstant`, `Sum`, `SumV2`, `MinMax`, `NullCount`, `NanCount`).
8788
pub fn initialize(session: &VortexSession) {
8889
session.arrays().register(Sparse);
8990
kernel::initialize(session);
@@ -99,6 +100,11 @@ pub fn initialize(session: &VortexSession) {
99100
Some(Sum.id()),
100101
&compute::sum::SparseSumKernel,
101102
);
103+
aggregate_fns.register_aggregate_kernel(
104+
Sparse.id(),
105+
Some(SumV2.id()),
106+
&compute::sum::SparseSumKernel,
107+
);
102108
aggregate_fns.register_aggregate_kernel(
103109
Sparse.id(),
104110
Some(MinMax.id()),

0 commit comments

Comments
 (0)