vortex_fastlanes/bitpacking/compute/
is_constant.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use itertools::Itertools;
7use lending_iterator::LendingIterator;
8use vortex_array::ToCanonical;
9use vortex_array::arrays::IS_CONST_LANE_WIDTH;
10use vortex_array::arrays::PrimitiveArray;
11use vortex_array::arrays::compute_is_constant;
12use vortex_array::compute::IsConstantKernel;
13use vortex_array::compute::IsConstantKernelAdapter;
14use vortex_array::compute::IsConstantOpts;
15use vortex_array::register_kernel;
16use vortex_dtype::IntegerPType;
17use vortex_dtype::match_each_integer_ptype;
18use vortex_dtype::match_each_unsigned_integer_ptype;
19use vortex_error::VortexResult;
20
21use crate::BitPackedArray;
22use crate::BitPackedVTable;
23use crate::unpack_iter::BitPacked;
24
25impl IsConstantKernel for BitPackedVTable {
26    fn is_constant(
27        &self,
28        array: &BitPackedArray,
29        opts: &IsConstantOpts,
30    ) -> VortexResult<Option<bool>> {
31        if opts.is_negligible_cost() {
32            return Ok(None);
33        }
34        match_each_integer_ptype!(array.ptype(), |P| {
35            bitpacked_is_constant::<P, { IS_CONST_LANE_WIDTH / size_of::<P>() }>(array)
36        })
37        .map(Some)
38    }
39}
40
41register_kernel!(IsConstantKernelAdapter(BitPackedVTable).lift());
42
43fn bitpacked_is_constant<T: BitPacked, const WIDTH: usize>(
44    array: &BitPackedArray,
45) -> VortexResult<bool> {
46    let mut bit_unpack_iterator = array.unpacked_chunks::<T>();
47    let patches = array.patches().map(|p| {
48        let values = p.values().to_primitive();
49        let indices = p.indices().to_primitive();
50        let offset = p.offset();
51        (indices, values, offset)
52    });
53
54    let mut header_constant_value = None;
55    let mut current_idx = 0;
56    if let Some(header) = bit_unpack_iterator.initial() {
57        if let Some((indices, patches, offset)) = &patches {
58            apply_patches(
59                header,
60                current_idx..header.len(),
61                indices,
62                patches.as_slice::<T>(),
63                *offset,
64            )
65        }
66
67        if !compute_is_constant::<_, WIDTH>(header) {
68            return Ok(false);
69        }
70        header_constant_value = Some(header[0]);
71        current_idx = header.len();
72    }
73
74    let mut first_chunk_value = None;
75    let mut chunks_iter = bit_unpack_iterator.full_chunks();
76    while let Some(chunk) = chunks_iter.next() {
77        if let Some((indices, patches, offset)) = &patches {
78            let chunk_len = chunk.len();
79            apply_patches(
80                chunk,
81                current_idx..current_idx + chunk_len,
82                indices,
83                patches.as_slice::<T>(),
84                *offset,
85            )
86        }
87
88        if !compute_is_constant::<_, WIDTH>(chunk) {
89            return Ok(false);
90        }
91
92        if let Some(chunk_value) = first_chunk_value {
93            if chunk_value != chunk[0] {
94                return Ok(false);
95            }
96        } else {
97            if let Some(header_value) = header_constant_value
98                && header_value != chunk[0]
99            {
100                return Ok(false);
101            }
102            first_chunk_value = Some(chunk[0]);
103        }
104
105        current_idx += chunk.len();
106    }
107
108    if let Some(trailer) = bit_unpack_iterator.trailer() {
109        if let Some((indices, patches, offset)) = &patches {
110            let chunk_len = trailer.len();
111            apply_patches(
112                trailer,
113                current_idx..current_idx + chunk_len,
114                indices,
115                patches.as_slice::<T>(),
116                *offset,
117            )
118        }
119
120        if !compute_is_constant::<_, WIDTH>(trailer) {
121            return Ok(false);
122        }
123
124        if let Some(previous_const_value) = header_constant_value.or(first_chunk_value)
125            && previous_const_value != trailer[0]
126        {
127            return Ok(false);
128        }
129    }
130
131    Ok(true)
132}
133
134fn apply_patches<T: BitPacked>(
135    values: &mut [T],
136    values_range: Range<usize>,
137    patch_indices: &PrimitiveArray,
138    patch_values: &[T],
139    indices_offset: usize,
140) {
141    match_each_unsigned_integer_ptype!(patch_indices.ptype(), |I| {
142        apply_patches_idx_typed(
143            values,
144            values_range,
145            patch_indices.as_slice::<I>(),
146            patch_values,
147            indices_offset,
148        )
149    });
150}
151
152fn apply_patches_idx_typed<T: BitPacked, I: IntegerPType>(
153    values: &mut [T],
154    values_range: Range<usize>,
155    patch_indices: &[I],
156    patch_values: &[T],
157    indices_offset: usize,
158) {
159    for (i, &v) in patch_indices
160        .iter()
161        .map(|i| i.as_() - indices_offset)
162        .zip_eq(patch_values)
163        .skip_while(|(i, _)| i < &values_range.start)
164        .take_while(|(i, _)| i < &values_range.end)
165    {
166        values[i - values_range.start] = v
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use vortex_array::IntoArray;
173    use vortex_array::compute::is_constant;
174    use vortex_buffer::buffer;
175
176    use crate::BitPackedArray;
177
178    #[test]
179    fn is_constant_with_patches() {
180        let array = BitPackedArray::encode(&buffer![4; 1025].into_array(), 2).unwrap();
181        assert!(is_constant(array.as_ref()).unwrap().unwrap());
182    }
183}
vortex_fastlanes/bitpacking/compute/is_constant.rs

vortex_fastlanes/bitpacking/compute/
is_constant.rs
