44use vortex_array:: ArrayRef ;
55use vortex_array:: ExecutionCtx ;
66use vortex_array:: IntoArray ;
7- use vortex_array:: aggregate_fn:: Accumulator ;
87use vortex_array:: aggregate_fn:: AggregateFnRef ;
9- use vortex_array:: aggregate_fn:: DynAccumulator ;
108use vortex_array:: aggregate_fn:: fns:: sum:: Sum ;
9+ use vortex_array:: aggregate_fn:: fns:: sum_v2:: SumV2 ;
1110use vortex_array:: aggregate_fn:: kernels:: DynAggregateKernel ;
1211use vortex_array:: arrays:: ConstantArray ;
1312use vortex_array:: scalar:: Scalar ;
@@ -16,13 +15,13 @@ use vortex_error::VortexResult;
1615use crate :: Sparse ;
1716use 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 ) ]
2726pub ( 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}
0 commit comments