17050456606

Committed 18 Aug 2025 07:32PM UTC coverage: 47.596%. First build

Build Type

github

Committed by GitHub

Commit Message

Merge afbf04d54 into 7eb8ac9fa

Pull Request Pull Request #4177:

Run Details

555 of 1372 new or added lines in 154 files covered. (40.45%)

18648 of 39180 relevant lines covered (47.6%)

234900.69 hits per line

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

0.0

/vortex-array/src/compute/nan_count.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::{VortexExpect, VortexResult, vortex_bail, vortex_err};
use vortex_scalar::{Scalar, ScalarValue};

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

static NAN_COUNT_FN: LazyLock<ComputeFn> = LazyLock::new(|| {
    let compute = ComputeFn::new("nan_count".into(), ArcRef::new_ref(&NaNCount));
    for kernel in inventory::iter::<NaNCountKernelRef> {
        compute.register_kernel(kernel.0.clone());
    }
    compute
});

/// Computes the number of NaN values in the array.
pub fn nan_count(array: &dyn Array) -> VortexResult<usize> {
    Ok(NAN_COUNT_FN
        .invoke(&InvocationArgs {
            inputs: &[array.into()],
            options: &(),
        })?
        .unwrap_scalar()?
        .as_primitive()
        .as_::<usize>()
        .vortex_expect("NaN count should not return null"))
}

struct NaNCount;

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

        let nan_count = nan_count_impl(array, kernels)?;

        // Update the stats set with the computed NaN count
        array.statistics().set(
            Stat::NaNCount,
            Precision::Exact(ScalarValue::from(nan_count as u64)),
        );

        Ok(Scalar::from(nan_count as u64).into())
    }

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

    fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {
        Ok(1)
    }

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

/// Computes the min and max of an array, returning the (min, max) values
pub trait NaNCountKernel: VTable {
    fn nan_count(&self, array: &Self::Array) -> VortexResult<usize>;
}

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

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

impl<V: VTable + NaNCountKernel> NaNCountKernelAdapter<V> {
    pub const fn lift(&'static self) -> NaNCountKernelRef {
        NaNCountKernelRef(ArcRef::new_ref(self))
    }
}

impl<V: VTable + NaNCountKernel> Kernel for NaNCountKernelAdapter<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);
        };
        let nan_count = V::nan_count(&self.0, array)?;
        Ok(Some(Scalar::from(nan_count as u64).into()))
    }
}

fn nan_count_impl(array: &dyn Array, kernels: &[ArcRef<dyn Kernel>]) -> VortexResult<usize> {
    if array.is_empty() || array.valid_count()? == 0 {
        return Ok(0);
    }

    if let Some(nan_count) = array
        .statistics()
        .get_as::<usize>(Stat::NaNCount)
        .and_then(Precision::as_exact)
    {
        // If the NaN count is already computed, return it
        return Ok(nan_count);
    }

    let args = InvocationArgs {
        inputs: &[array.into()],
        options: &(),
    };

    for kernel in kernels {
        if let Some(output) = kernel.invoke(&args)? {
            return output
                .unwrap_scalar()?
                .as_primitive()
                .as_::<usize>()
                .ok_or_else(|| vortex_err!("NaN count should not return null"));
        }
    }
    if let Some(output) = array.invoke(&NAN_COUNT_FN, &args)? {
        return output
            .unwrap_scalar()?
            .as_primitive()
            .as_::<usize>()
            .ok_or_else(|| vortex_err!("NaN count should not return null"));
    }

    if !array.is_canonical() {
        let canonical = array.to_canonical()?;
        return nan_count(canonical.as_ref());
    }

    vortex_bail!(
        "No NaN count kernel found for array type: {}",
        array.dtype()
    )
}

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::{VortexExpect, VortexResult, vortex_bail, vortex_err};
9	use vortex_scalar::{Scalar, ScalarValue};
10
11	use crate::Array;
12	use crate::compute::{ComputeFn, ComputeFnVTable, InvocationArgs, Kernel, Output, UnaryArgs};
13	use crate::stats::{Precision, Stat};
14	use crate::vtable::VTable;
15
16	static NAN_COUNT_FN: LazyLock<ComputeFn> = LazyLock::new(\|\| {	×
17	let compute = ComputeFn::new("nan_count".into(), ArcRef::new_ref(&NaNCount));	×
18	for kernel in inventory::iter::<NaNCountKernelRef> {	×
19	compute.register_kernel(kernel.0.clone());	×
20	}	×
21	compute	×
22	});	×
23
24	/// Computes the number of NaN values in the array.
25	pub fn nan_count(array: &dyn Array) -> VortexResult<usize> {	×
26	Ok(NAN_COUNT_FN	×
27	.invoke(&InvocationArgs {	×
28	inputs: &[array.into()],	×
29	options: &(),	×
30	})?	×
31	.unwrap_scalar()?	×
32	.as_primitive()	×
NEW 33	.as_::<usize>()	×
34	.vortex_expect("NaN count should not return null"))	×
35	}	×
36
37	struct NaNCount;
38
39	impl ComputeFnVTable for NaNCount {
40	fn invoke(	×
41	&self,	×
42	args: &InvocationArgs,	×
43	kernels: &[ArcRef<dyn Kernel>],	×
44	) -> VortexResult<Output> {	×
45	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	×
46
47	let nan_count = nan_count_impl(array, kernels)?;	×
48
49	// Update the stats set with the computed NaN count
50	array.statistics().set(	×
51	Stat::NaNCount,	×
52	Precision::Exact(ScalarValue::from(nan_count as u64)),	×
53	);
54
55	Ok(Scalar::from(nan_count as u64).into())	×
56	}	×
57
58	fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {	×
59	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	×
60	Stat::NaNCount	×
61	.dtype(array.dtype())	×
62	.ok_or_else(\|\| vortex_err!("Cannot compute NaN count for dtype {}", array.dtype()))	×
63	}	×
64
65	fn return_len(&self, _args: &InvocationArgs) -> VortexResult<usize> {	×
66	Ok(1)	×
67	}	×
68
69	fn is_elementwise(&self) -> bool {	×
70	false	×
71	}	×
72	}
73
74	/// Computes the min and max of an array, returning the (min, max) values
75	pub trait NaNCountKernel: VTable {
76	fn nan_count(&self, array: &Self::Array) -> VortexResult<usize>;
77	}
78
79	pub struct NaNCountKernelRef(ArcRef<dyn Kernel>);
80	inventory::collect!(NaNCountKernelRef);
81
82	#[derive(Debug)]
83	pub struct NaNCountKernelAdapter<V: VTable>(pub V);
84
85	impl<V: VTable + NaNCountKernel> NaNCountKernelAdapter<V> {
86	pub const fn lift(&'static self) -> NaNCountKernelRef {	×
87	NaNCountKernelRef(ArcRef::new_ref(self))	×
88	}	×
89	}
90
91	impl<V: VTable + NaNCountKernel> Kernel for NaNCountKernelAdapter<V> {
92	fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>> {	×
93	let UnaryArgs { array, .. } = UnaryArgs::<()>::try_from(args)?;	×
94	let Some(array) = array.as_opt::<V>() else {	×
95	return Ok(None);	×
96	};
97	let nan_count = V::nan_count(&self.0, array)?;	×
98	Ok(Some(Scalar::from(nan_count as u64).into()))	×
99	}	×
100	}
101
102	fn nan_count_impl(array: &dyn Array, kernels: &[ArcRef<dyn Kernel>]) -> VortexResult<usize> {	×
103	if array.is_empty() \|\| array.valid_count()? == 0 {	×
104	return Ok(0);	×
105	}	×
106
107	if let Some(nan_count) = array	×
108	.statistics()	×
109	.get_as::<usize>(Stat::NaNCount)	×
110	.and_then(Precision::as_exact)	×
111	{
112	// If the NaN count is already computed, return it
113	return Ok(nan_count);	×
114	}	×
115
116	let args = InvocationArgs {	×
117	inputs: &[array.into()],	×
118	options: &(),	×
119	};	×
120
121	for kernel in kernels {	×
122	if let Some(output) = kernel.invoke(&args)? {	×
123	return output	×
124	.unwrap_scalar()?	×
125	.as_primitive()	×
NEW 126	.as_::<usize>()	×
127	.ok_or_else(\|\| vortex_err!("NaN count should not return null"));	×
128	}	×
129	}
130	if let Some(output) = array.invoke(&NAN_COUNT_FN, &args)? {	×
131	return output	×
132	.unwrap_scalar()?	×
133	.as_primitive()	×
NEW 134	.as_::<usize>()	×
135	.ok_or_else(\|\| vortex_err!("NaN count should not return null"));	×
136	}	×
137
138	if !array.is_canonical() {	×
139	let canonical = array.to_canonical()?;	×
140	return nan_count(canonical.as_ref());	×
141	}	×
142
143	vortex_bail!(	×
144	"No NaN count kernel found for array type: {}",	×
145	array.dtype()	×
146	)
147	}	×

vortex-data / vortex / 17050456606

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