vortex_fastlanes/for/compute/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4mod cast;
5mod compare;
6mod is_constant;
7mod is_sorted;
8
9use vortex_array::Array;
10use vortex_array::ArrayRef;
11use vortex_array::IntoArray;
12use vortex_array::compute::FilterKernel;
13use vortex_array::compute::FilterKernelAdapter;
14use vortex_array::compute::TakeKernel;
15use vortex_array::compute::TakeKernelAdapter;
16use vortex_array::compute::filter;
17use vortex_array::compute::take;
18use vortex_array::register_kernel;
19use vortex_error::VortexResult;
20use vortex_mask::Mask;
21
22use crate::FoRArray;
23use crate::FoRVTable;
24
25impl TakeKernel for FoRVTable {
26    fn take(&self, array: &FoRArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
27        FoRArray::try_new(
28            take(array.encoded(), indices)?,
29            array.reference_scalar().clone(),
30        )
31        .map(|a| a.into_array())
32    }
33}
34
35register_kernel!(TakeKernelAdapter(FoRVTable).lift());
36
37impl FilterKernel for FoRVTable {
38    fn filter(&self, array: &FoRArray, mask: &Mask) -> VortexResult<ArrayRef> {
39        FoRArray::try_new(
40            filter(array.encoded(), mask)?,
41            array.reference_scalar().clone(),
42        )
43        .map(|a| a.into_array())
44    }
45}
46
47register_kernel!(FilterKernelAdapter(FoRVTable).lift());
48
49#[cfg(test)]
50mod test {
51    use rstest::rstest;
52    use vortex_array::IntoArray;
53    use vortex_array::arrays::PrimitiveArray;
54    use vortex_array::compute::conformance::filter::test_filter_conformance;
55    use vortex_buffer::buffer;
56    use vortex_scalar::Scalar;
57
58    use crate::FoRArray;
59
60    #[test]
61    fn test_filter_for_array() {
62        // Test with i32 values
63        let values = buffer![100i32, 101, 102, 103, 104].into_array();
64        let reference = Scalar::from(100i32);
65        let for_array = FoRArray::try_new(values, reference).unwrap();
66        test_filter_conformance(for_array.as_ref());
67
68        // Test with u64 values
69        let values = buffer![1000u64, 1001, 1002, 1003, 1004].into_array();
70        let reference = Scalar::from(1000u64);
71        let for_array = FoRArray::try_new(values, reference).unwrap();
72        test_filter_conformance(for_array.as_ref());
73
74        // Test with nullable values
75        let values =
76            PrimitiveArray::from_option_iter([Some(50i16), None, Some(52), Some(53), None]);
77        let reference = Scalar::from(50i16);
78        let for_array = FoRArray::try_new(values.into_array(), reference).unwrap();
79        test_filter_conformance(for_array.as_ref());
80    }
81
82    #[rstest]
83    #[case(FoRArray::try_new(buffer![100i32, 101, 102, 103, 104].into_array(), Scalar::from(100i32)).unwrap())]
84    #[case(FoRArray::try_new(buffer![1000u64, 1001, 1002, 1003, 1004].into_array(), Scalar::from(1000u64)).unwrap())]
85    #[case(FoRArray::try_new(
86        PrimitiveArray::from_option_iter([Some(50i16), None, Some(52), Some(53), None]).into_array(),
87        Scalar::from(50i16)
88    ).unwrap())]
89    #[case(FoRArray::try_new(buffer![-100i32, -99, -98, -97, -96].into_array(), Scalar::from(-100i32)).unwrap())]
90    #[case(FoRArray::try_new(buffer![42i64].into_array(), Scalar::from(40i64)).unwrap())]
91    fn test_take_for_conformance(#[case] for_array: FoRArray) {
92        use vortex_array::compute::conformance::take::test_take_conformance;
93        test_take_conformance(for_array.as_ref());
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use rstest::rstest;
100    use vortex_array::IntoArray;
101    use vortex_array::arrays::PrimitiveArray;
102    use vortex_array::compute::conformance::binary_numeric::test_binary_numeric_array;
103    use vortex_array::compute::conformance::consistency::test_array_consistency;
104    use vortex_buffer::buffer;
105    use vortex_scalar::Scalar;
106
107    use crate::FoRArray;
108
109    #[rstest]
110    // Basic FoR arrays
111    #[case::for_i32(FoRArray::try_new(
112        buffer![100i32, 101, 102, 103, 104].into_array(),
113        Scalar::from(100i32)
114    ).unwrap())]
115    #[case::for_u64(FoRArray::try_new(
116        buffer![1000u64, 1001, 1002, 1003, 1004].into_array(),
117        Scalar::from(1000u64)
118    ).unwrap())]
119    // Nullable arrays
120    #[case::for_nullable_i16(FoRArray::try_new(
121        PrimitiveArray::from_option_iter([Some(50i16), None, Some(52), Some(53), None]).into_array(),
122        Scalar::from(50i16)
123    ).unwrap())]
124    #[case::for_nullable_i32(FoRArray::try_new(
125        PrimitiveArray::from_option_iter([Some(200i32), None, Some(202), Some(203), None]).into_array(),
126        Scalar::from(200i32)
127    ).unwrap())]
128    // Negative values
129    #[case::for_negative(FoRArray::try_new(
130        buffer![-100i32, -99, -98, -97, -96].into_array(),
131        Scalar::from(-100i32)
132    ).unwrap())]
133    // Edge cases
134    #[case::for_single(FoRArray::try_new(
135        buffer![42i64].into_array(),
136        Scalar::from(40i64)
137    ).unwrap())]
138    #[case::for_zero_ref(FoRArray::try_new(
139        buffer![0u32, 1, 2, 3, 4].into_array(),
140        Scalar::from(0u32)
141    ).unwrap())]
142    // Large arrays (> 1024 elements for fastlanes chunking)
143    #[case::for_large(FoRArray::try_new(
144        PrimitiveArray::from_iter((0..1500).map(|i| 5000 + i)).into_array(),
145        Scalar::from(5000i32)
146    ).unwrap())]
147    #[case::for_very_large(FoRArray::try_new(
148        PrimitiveArray::from_iter((0..3072).map(|i| 10000 + i as i64)).into_array(),
149        Scalar::from(10000i64)
150    ).unwrap())]
151    #[case::for_large_nullable(FoRArray::try_new(
152        PrimitiveArray::from_option_iter((0..2048).map(|i| (i % 15 == 0).then_some(1000 + i))).into_array(),
153        Scalar::from(1000i32)
154    ).unwrap())]
155    // Arrays with large deltas from reference
156    #[case::for_large_deltas(FoRArray::try_new(
157        buffer![100i64, 200, 300, 400, 500].into_array(),
158        Scalar::from(100i64)
159    ).unwrap())]
160
161    fn test_for_consistency(#[case] array: FoRArray) {
162        test_array_consistency(array.as_ref());
163    }
164
165    #[rstest]
166    #[case::for_i32_basic(FoRArray::try_new(
167        buffer![100i32, 101, 102, 103, 104].into_array(),
168        Scalar::from(100i32)
169    ).unwrap())]
170    #[case::for_u32_basic(FoRArray::try_new(
171        buffer![1000u32, 1001, 1002, 1003, 1004].into_array(),
172        Scalar::from(1000u32)
173    ).unwrap())]
174    #[case::for_i64_basic(FoRArray::try_new(
175        buffer![5000i64, 5001, 5002, 5003, 5004].into_array(),
176        Scalar::from(5000i64)
177    ).unwrap())]
178    #[case::for_u64_basic(FoRArray::try_new(
179        buffer![10000u64, 10001, 10002, 10003, 10004].into_array(),
180        Scalar::from(10000u64)
181    ).unwrap())]
182    #[case::for_i32_large(FoRArray::try_new(
183        PrimitiveArray::from_iter((0..100).map(|i| 2000 + i)).into_array(),
184        Scalar::from(2000i32)
185    ).unwrap())]
186    fn test_for_binary_numeric(#[case] array: FoRArray) {
187        test_binary_numeric_array(array.into_array());
188    }
189}
vortex_fastlanes/for/compute/mod.rs

vortex_fastlanes/for/compute/
mod.rs
