Skip to main content

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