16615125477

Committed 30 Jul 2025 06:26AM UTC coverage: 82.687% (-0.02%) from 82.703%

Build # 16615125477

Build Type

Pull #4050

github

Committed by

web-flow

Commit Message

Merge dd5f09ba6 into 5f86536fe

Pull Request Pull Request #4050: Multiple improvements to the C++ API

Run Details

45215 of 54682 relevant lines covered (82.69%)

184545.41 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.79

/vortex-array/src/compute/sum.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::sync::LazyLock;

use arcref::ArcRef;
use vortex_dtype::DType;
use vortex_error::{VortexResult, vortex_err, vortex_panic};
use vortex_scalar::Scalar;

use crate::Array;
use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, UnaryArgs};
use crate::stats::{Precision, Stat, StatsProvider};
use crate::vtable::VTable;

static SUM_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("sum".into(), ArcRef::new_ref(&Sum));
    for kernel in inventory::iter::<SumKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

/// Sum an array.
///
/// If the sum overflows, a null scalar will be returned.
/// If the sum is not supported for the array's dtype, an error will be raised.
/// If the array is all-invalid, the sum will be zero.
pub fn sum(array: &dyn Array) -> VortexResult<Scalar> {
    SUM_FN
        .invoke(&InvocationArgs {
            inputs: &[array.into()],
            options: &(),
        })?
        .unwrap_scalar()
}

struct Sum;

impl ComputeFnVTable for Sum {
    fn invoke(
        &self,
        args: &InvocationArgs,
        kernels: &[ArcRef<dyn Kernel>],
    ) -> VortexResult<Output> {
        let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;

        // Compute the expected dtype of the sum.
        let sum_dtype = self.return_dtype(args)?;

        // Short-circuit using array statistics.
        if let Some(Precision::Exact(sum)) = array.statistics().get(Stat::Sum) {
            return Ok(Scalar::new(sum_dtype, sum).into());
        }

        let sum_scalar = sum_impl(array, sum_dtype, kernels)?;

        // Update the statistics with the computed sum.
        array
            .statistics()
            .set(Stat::Sum, Precision::Exact(sum_scalar.value().clone()));

        Ok(sum_scalar.into())
    }

    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
        let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;
        Stat::Sum
            .dtype(array.dtype())
            .ok_or_else(|| vortex_err!("Sum not supported for dtype: {}", array.dtype()))
    }

    fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {
        // The sum function always returns a single scalar value.
        Ok(1)
    }

    fn is_elementwise(&self) -> bool {
        false
    }
}

pub struct SumKernelRef(ArcRef<dyn Kernel>);
inventory::collect!(SumKernelRef);

pub trait SumKernel: VTable {
    /// # Preconditions
    ///
    /// * The array's DType is summable
    /// * The array is not all-null
    fn sum(&self, array: &Self::Array) -> VortexResult<Scalar>;
}

#[derive(Debug)]
pub struct SumKernelAdapter<V: VTable>(pub V);

impl<V: VTable + SumKernel> SumKernelAdapter<V> {
    pub const fn lift(&'static self) -> SumKernelRef {
        SumKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + SumKernel> Kernel for SumKernelAdapter<V> {
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {
        let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;
        let Some(array) = array.as_opt::<V>() else {
            return Ok(None);
        };
        Ok(Some(V::sum(&self.0, array)?.into()))
    }
}

/// Sum an array.
///
/// If the sum overflows, a null scalar will be returned.
/// If the sum is not supported for the array's dtype, an error will be raised.
/// If the array is all-invalid, the sum will be zero.
pub fn sum_impl(
    array: &dyn Array,
    sum_dtype: DType,
    kernels: &[ArcRef<dyn Kernel>],
) -> VortexResult<Scalar> {
    if array.is_empty() {
        return if sum_dtype.is_float() {
            Ok(Scalar::new(sum_dtype, 0.0.into()))
        } else {
            Ok(Scalar::new(sum_dtype, 0.into()))
        };
    }

    // Sum of all null is null.
    if array.all_invalid()? {
        return Ok(Scalar::null(sum_dtype));
    }

    // Try to find a sum kernel
    let args = InvocationArgs {
        inputs: &[array.into()],
        options: &(),
    };
    for kernel in kernels {
        if let Some(output) = kernel.invoke(&args)? {
            return output.unwrap_scalar();
        }
    }
    if let Some(output) = array.invoke(&SUM_FN, &args)? {
        return output.unwrap_scalar();
    }

    // Otherwise, canonicalize and try again.
    log::debug!("No sum implementation found for {}", array.encoding_id());
    if array.is_canonical() {
        // Panic to avoid recursion, but it should never be hit.
        vortex_panic!(
            "No sum implementation found for canonical array: {}",
            array.encoding_id()
        );
    }
    sum(array.to_canonical()?.as_ref())
}

#[cfg(test)]
mod test {
    use vortex_dtype::{DType, Nullability, PType};
    use vortex_scalar::Scalar;

    use crate::arrays::{BoolArray, PrimitiveArray};
    use crate::compute::sum;

    #[test]
    fn sum_all_invalid() {
        let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]);
        let result = sum(array.as_ref()).unwrap();
        assert_eq!(
            result,
            Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))
        );
    }

    #[test]
    fn sum_all_invalid_float() {
        let array = PrimitiveArray::from_option_iter::<f32, _>([None, None, None]);
        let result = sum(array.as_ref()).unwrap();
        assert_eq!(
            result,
            Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))
        );
    }

    #[test]
    fn sum_constant() {
        let array = PrimitiveArray::from_iter([1, 1, 1, 1]);
        let result = sum(array.as_ref()).unwrap();
        assert_eq!(result.as_primitive().as_::<i32>().unwrap(), Some(4));
    }

    #[test]
    fn sum_constant_float() {
        let array = PrimitiveArray::from_iter([1., 1., 1., 1.]);
        let result = sum(array.as_ref()).unwrap();
        assert_eq!(result.as_primitive().as_::<f32>().unwrap(), Some(4.));
    }

    #[test]
    fn sum_boolean() {
        let array = BoolArray::from_iter([true, false, false, true]);
        let result = sum(array.as_ref()).unwrap();
        assert_eq!(result.as_primitive().as_::<i32>().unwrap(), Some(2));
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use std::sync::LazyLock;
5
6	use arcref::ArcRef;
7	use vortex_dtype::DType;
8	use vortex_error::{VortexResult, vortex_err, vortex_panic};
9	use vortex_scalar::Scalar;
10
11	use crate::Array;
12	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, UnaryArgs};
13	use crate::stats::{Precision, Stat, StatsProvider};
14	use crate::vtable::VTable;
15
16	static SUM_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	2,332✔
17	let compute = ComputeFn::new("sum".into(), ArcRef::new_ref(&Sum));	2,332✔
18	for kernel in inventory::iter::<SumKernelRef> {	16,324✔
19	compute.register_kernel(kernel.0.clone());	13,992✔
20	}	13,992✔
21	compute	2,332✔
22	});	2,332✔
23
24	/// Sum an array.
25	///
26	/// If the sum overflows, a null scalar will be returned.
27	/// If the sum is not supported for the array's dtype, an error will be raised.
28	/// If the array is all-invalid, the sum will be zero.
29	pub fn sum(array: &dyn Array) -> VortexResult<Scalar> {	13,784✔
30	SUM_FN	13,784✔
31	.invoke(&InvocationArgs {	13,784✔
32	inputs: &[array.into()],	13,784✔
33	options: &(),	13,784✔
34	})?	13,784✔
35	.unwrap_scalar()	13,784✔
36	}	13,784✔
37
38	struct Sum;
39
40	impl ComputeFnVTable for Sum {
41	fn invoke(	13,784✔
42	&self,	13,784✔
43	args: &InvocationArgs,	13,784✔
44	kernels: &[ArcRef<dyn Kernel>],	13,784✔
45	) -> VortexResult<Output> {	13,784✔
46	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	13,784✔
47
48	// Compute the expected dtype of the sum.
49	let sum_dtype = self.return_dtype(args)?;	13,784✔
50
51	// Short-circuit using array statistics.
52	if let Some(Precision::Exact(sum)) = array.statistics().get(Stat::Sum) {	13,784✔
53	return Ok(Scalar::new(sum_dtype, sum).into());	114✔
54	}	13,670✔
55
56	let sum_scalar = sum_impl(array, sum_dtype, kernels)?;	13,670✔
57
58	// Update the statistics with the computed sum.
59	array	13,670✔
60	.statistics()	13,670✔
61	.set(Stat::Sum, Precision::Exact(sum_scalar.value().clone()));	13,670✔
62
63	Ok(sum_scalar.into())	13,670✔
64	}	13,784✔
65
66	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	27,568✔
67	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	27,568✔
68	Stat::Sum	27,568✔
69	.dtype(array.dtype())	27,568✔
70	.ok_or_else(\|\| vortex_err!("Sum not supported for dtype: {}", array.dtype()))	27,568✔
71	}	27,568✔
72
73	fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {	13,784✔
74	// The sum function always returns a single scalar value.
75	Ok(1)	13,784✔
76	}	13,784✔
77
78	fn is_elementwise(&self) -> bool {	13,784✔
79	false	13,784✔
80	}	13,784✔
81	}
82
83	pub struct SumKernelRef(ArcRef<dyn Kernel>);
84	inventory::collect!(SumKernelRef);
85
86	pub trait SumKernel: VTable {
87	/// # Preconditions
88	///
89	/// * The array's DType is summable
90	/// * The array is not all-null
91	fn sum(&self, array: &Self::Array) -> VortexResult<Scalar>;
92	}
93
94	#[derive(Debug)]
95	pub struct SumKernelAdapter<V: VTable>(pub V);
96
97	impl<V: VTable + SumKernel> SumKernelAdapter<V> {
98	pub const fn lift(&'static self) -> SumKernelRef {	×
99	SumKernelRef(ArcRef::new_ref(self))	×
100	}	×
101	}
102
103	impl<V: VTable + SumKernel> Kernel for SumKernelAdapter<V> {
104	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	25,253✔
105	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	25,253✔
106	let Some(array) = array.as_opt::<V>() else {	25,253✔
107	return Ok(None);	11,663✔
108	};
109	Ok(Some(V::sum(&self.0, array)?.into()))	13,590✔
110	}	25,253✔
111	}
112
113	/// Sum an array.
114	///
115	/// If the sum overflows, a null scalar will be returned.
116	/// If the sum is not supported for the array's dtype, an error will be raised.
117	/// If the array is all-invalid, the sum will be zero.
118	pub fn sum_impl(	13,670✔
119	array: &dyn Array,	13,670✔
120	sum_dtype: DType,	13,670✔
121	kernels: &[ArcRef<dyn Kernel>],	13,670✔
122	) -> VortexResult<Scalar> {	13,670✔
123	if array.is_empty() {	13,670✔
124	return if sum_dtype.is_float() {	38✔
125	Ok(Scalar::new(sum_dtype, 0.0.into()))	×
126	} else {
127	Ok(Scalar::new(sum_dtype, 0.into()))	38✔
128	};
129	}	13,632✔
130
131	// Sum of all null is null.
132	if array.all_invalid()? {	13,632✔
133	return Ok(Scalar::null(sum_dtype));	4✔
134	}	13,628✔
135
136	// Try to find a sum kernel
137	let args = InvocationArgs {	13,628✔
138	inputs: &[array.into()],	13,628✔
139	options: &(),	13,628✔
140	};	13,628✔
141	for kernel in kernels {	25,291✔
142	if let Some(output) = kernel.invoke(&args)? {	25,253✔
143	return output.unwrap_scalar();	13,590✔
144	}	11,663✔
145	}
146	if let Some(output) = array.invoke(&SUM_FN, &args)? {	38✔
147	return output.unwrap_scalar();	×
148	}	38✔
149
150	// Otherwise, canonicalize and try again.
151	log::debug!("No sum implementation found for {}", array.encoding_id());	38✔
152	if array.is_canonical() {	38✔
153	// Panic to avoid recursion, but it should never be hit.
154	vortex_panic!(	×
155	"No sum implementation found for canonical array: {}",	×
156	array.encoding_id()	×
157	);
158	}	38✔
159	sum(array.to_canonical()?.as_ref())	38✔
160	}	13,670✔
161
162	#[cfg(test)]
163	mod test {
164	use vortex_dtype::{DType, Nullability, PType};
165	use vortex_scalar::Scalar;
166
167	use crate::arrays::{BoolArray, PrimitiveArray};
168	use crate::compute::sum;
169
170	#[test]
171	fn sum_all_invalid() {	1✔
172	let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]);	1✔
173	let result = sum(array.as_ref()).unwrap();	1✔
174	assert_eq!(	1✔
175	result,
176	Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))	1✔
177	);
178	}	1✔
179
180	#[test]
181	fn sum_all_invalid_float() {	1✔
182	let array = PrimitiveArray::from_option_iter::<f32, _>([None, None, None]);	1✔
183	let result = sum(array.as_ref()).unwrap();	1✔
184	assert_eq!(	1✔
185	result,
186	Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))	1✔
187	);
188	}	1✔
189
190	#[test]
191	fn sum_constant() {	1✔
192	let array = PrimitiveArray::from_iter([1, 1, 1, 1]);	1✔
193	let result = sum(array.as_ref()).unwrap();	1✔
194	assert_eq!(result.as_primitive().as_::<i32>().unwrap(), Some(4));	1✔
195	}	1✔
196
197	#[test]
198	fn sum_constant_float() {	1✔
199	let array = PrimitiveArray::from_iter([1., 1., 1., 1.]);	1✔
200	let result = sum(array.as_ref()).unwrap();	1✔
201	assert_eq!(result.as_primitive().as_::<f32>().unwrap(), Some(4.));	1✔
202	}	1✔
203
204	#[test]
205	fn sum_boolean() {	1✔
206	let array = BoolArray::from_iter([true, false, false, true]);	1✔
207	let result = sum(array.as_ref()).unwrap();	1✔
208	assert_eq!(result.as_primitive().as_::<i32>().unwrap(), Some(2));	1✔
209	}	1✔
210	}

vortex-data / vortex / 16615125477

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous