16728097825

Committed 04 Aug 2025 04:00PM UTC coverage: 48.355% (-35.1%) from 83.429%

Build # 16728097825

Build Type

Pull #4108

github

Committed by

web-flow

Commit Message

Merge 1b2d27fd8 into 649ba9576

Pull Request Pull Request #4108: perf[vortex-array]: use all_valid instead of `invalid_count() == 0`

Coverage Stats

1 of 1 new or added line in 1 file covered. (100.0%)

11596 existing lines in 378 files now uncovered.

18635 of 38538 relevant lines covered (48.35%)

151786.4 hits per line

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

78.57

/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✔
17	let compute = ComputeFn::new("sum".into(), ArcRef::new_ref(&Sum));	2✔
18	for kernel in inventory::iter::<SumKernelRef> {	14✔
19	compute.register_kernel(kernel.0.clone());	12✔
20	}	12✔
21	compute	2✔
22	});	2✔
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> {	2,144✔
30	SUM_FN	2,144✔
31	.invoke(&InvocationArgs {	2,144✔
32	inputs: &[array.into()],	2,144✔
33	options: &(),	2,144✔
34	})?	2,144✔
35	.unwrap_scalar()	2,144✔
36	}	2,144✔
37
38	struct Sum;
39
40	impl ComputeFnVTable for Sum {
41	fn invoke(	2,144✔
42	&self,	2,144✔
43	args: &InvocationArgs,	2,144✔
44	kernels: &[ArcRef<dyn Kernel>],	2,144✔
45	) -> VortexResult<Output> {	2,144✔
46	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	2,144✔
47
48	// Compute the expected dtype of the sum.
49	let sum_dtype = self.return_dtype(args)?;	2,144✔
50
51	// Short-circuit using array statistics.
52	if let Some(Precision::Exact(sum)) = array.statistics().get(Stat::Sum) {	2,144✔
UNCOV 53	return Ok(Scalar::new(sum_dtype, sum).into());	×
54	}	2,144✔
55
56	let sum_scalar = sum_impl(array, sum_dtype, kernels)?;	2,144✔
57
58	// Update the statistics with the computed sum.
59	array	2,144✔
60	.statistics()	2,144✔
61	.set(Stat::Sum, Precision::Exact(sum_scalar.value().clone()));	2,144✔
62
63	Ok(sum_scalar.into())	2,144✔
64	}	2,144✔
65
66	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	4,288✔
67	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	4,288✔
68	Stat::Sum	4,288✔
69	.dtype(array.dtype())	4,288✔
70	.ok_or_else(\|\| vortex_err!("Sum not supported for dtype: {}", array.dtype()))	4,288✔
71	}	4,288✔
72
73	fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {	2,144✔
74	// The sum function always returns a single scalar value.
75	Ok(1)	2,144✔
76	}	2,144✔
77
78	fn is_elementwise(&self) -> bool {	2,144✔
79	false	2,144✔
80	}	2,144✔
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>> {	5,320✔
105	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	5,320✔
106	let Some(array) = array.as_opt::<V>() else {	5,320✔
107	return Ok(None);	3,176✔
108	};
109	Ok(Some(V::sum(&self.0, array)?.into()))	2,144✔
110	}	5,320✔
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(	2,144✔
119	array: &dyn Array,	2,144✔
120	sum_dtype: DType,	2,144✔
121	kernels: &[ArcRef<dyn Kernel>],	2,144✔
122	) -> VortexResult<Scalar> {	2,144✔
123	if array.is_empty() {	2,144✔
UNCOV 124	return if sum_dtype.is_float() {	×
125	Ok(Scalar::new(sum_dtype, 0.0.into()))	×
126	} else {
UNCOV 127	Ok(Scalar::new(sum_dtype, 0.into()))	×
128	};
129	}	2,144✔
130
131	// Sum of all null is null.
132	if array.all_invalid()? {	2,144✔
UNCOV 133	return Ok(Scalar::null(sum_dtype));	×
134	}	2,144✔
135
136	// Try to find a sum kernel
137	let args = InvocationArgs {	2,144✔
138	inputs: &[array.into()],	2,144✔
139	options: &(),	2,144✔
140	};	2,144✔
141	for kernel in kernels {	5,320✔
142	if let Some(output) = kernel.invoke(&args)? {	5,320✔
143	return output.unwrap_scalar();	2,144✔
144	}	3,176✔
145	}
UNCOV 146	if let Some(output) = array.invoke(&SUM_FN, &args)? {	×
147	return output.unwrap_scalar();	×
UNCOV 148	}	×
149
150	// Otherwise, canonicalize and try again.
UNCOV 151	log::debug!("No sum implementation found for {}", array.encoding_id());	×
UNCOV 152	if array.is_canonical() {	×
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	);
UNCOV 158	}	×
UNCOV 159	sum(array.to_canonical()?.as_ref())	×
160	}	2,144✔
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() {
172	let array = PrimitiveArray::from_option_iter::<i32, _>([None, None, None]);
173	let result = sum(array.as_ref()).unwrap();
174	assert_eq!(
175	result,
176	Scalar::null(DType::Primitive(PType::I64, Nullability::Nullable))
177	);
178	}
179
180	#[test]
181	fn sum_all_invalid_float() {
182	let array = PrimitiveArray::from_option_iter::<f32, _>([None, None, None]);
183	let result = sum(array.as_ref()).unwrap();
184	assert_eq!(
185	result,
186	Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable))
187	);
188	}
189
190	#[test]
191	fn sum_constant() {
192	let array = PrimitiveArray::from_iter([1, 1, 1, 1]);
193	let result = sum(array.as_ref()).unwrap();
194	assert_eq!(result.as_primitive().as_::<i32>().unwrap(), Some(4));
195	}
196
197	#[test]
198	fn sum_constant_float() {
199	let array = PrimitiveArray::from_iter([1., 1., 1., 1.]);
200	let result = sum(array.as_ref()).unwrap();
201	assert_eq!(result.as_primitive().as_::<f32>().unwrap(), Some(4.));
202	}
203
204	#[test]
205	fn sum_boolean() {
206	let array = BoolArray::from_iter([true, false, false, true]);
207	let result = sum(array.as_ref()).unwrap();
208	assert_eq!(result.as_primitive().as_::<i32>().unwrap(), Some(2));
209	}
210	}

vortex-data / vortex / 16728097825

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