vortex_array/arrays/primitive/compute/
is_constant.rs

1use vortex_dtype::half::f16;
2use vortex_dtype::{NativePType, match_each_native_ptype};
3use vortex_error::VortexResult;
4
5use crate::arrays::{PrimitiveArray, PrimitiveEncoding};
6use crate::compute::{IsConstantFn, IsConstantOpts};
7use crate::variants::PrimitiveArrayTrait;
8
9cfg_if::cfg_if! {
10    if #[cfg(target_feature = "avx2")] {
11        pub const IS_CONST_LANE_WIDTH: usize = 32;
12    } else {
13        pub const IS_CONST_LANE_WIDTH: usize = 16;
14    }
15}
16
17impl IsConstantFn<&PrimitiveArray> for PrimitiveEncoding {
18    fn is_constant(
19        &self,
20        array: &PrimitiveArray,
21        _opts: &IsConstantOpts,
22    ) -> VortexResult<Option<bool>> {
23        let is_constant = match_each_native_ptype!(array.ptype(), integral: |$P| {
24            compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<$P>()}>(array.as_slice::<$P>())
25        } floating_point: |$P| {
26            compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<$P>()}>(unsafe { std::mem::transmute::<&[$P], &[<$P as EqFloat>::IntType]>(array.as_slice::<$P>()) })
27        });
28
29        Ok(Some(is_constant))
30    }
31}
32
33// Assumes any floating point has been cast into its bit representation for which != and !is_eq are the same
34// Assumes there's at least 1 value in the slice, which is an invariant of the entry level function.
35pub fn compute_is_constant<T: NativePType, const WIDTH: usize>(values: &[T]) -> bool {
36    let first_value = values[0];
37    let first_vec = &[first_value; WIDTH];
38
39    let mut chunks = values[1..].array_chunks::<WIDTH>();
40    for chunk in &mut chunks {
41        if first_vec != chunk {
42            return false;
43        }
44    }
45
46    for value in chunks.remainder() {
47        if !value.is_eq(first_value) {
48            return false;
49        }
50    }
51
52    true
53}
54
55trait EqFloat {
56    type IntType;
57}
58
59impl EqFloat for f16 {
60    type IntType = u16;
61}
62impl EqFloat for f32 {
63    type IntType = u32;
64}
65impl EqFloat for f64 {
66    type IntType = u64;
67}