vortex_sparse/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::ArrayRef;
5use vortex_array::ArrayView;
6use vortex_array::ExecutionCtx;
7use vortex_array::IntoArray;
8use vortex_array::arrays::dict::TakeExecute;
9use vortex_error::VortexResult;
10
11use crate::ConstantArray;
12use crate::Sparse;
13use crate::SparseExt as _;
14impl TakeExecute for Sparse {
15    fn take(
16        array: ArrayView<'_, Self>,
17        indices: &ArrayRef,
18        ctx: &mut ExecutionCtx,
19    ) -> VortexResult<Option<ArrayRef>> {
20        let patches_take = if array.fill_scalar().is_null() {
21            array.patches().take(indices, ctx)?
22        } else {
23            array.patches().take_with_nulls(indices, ctx)?
24        };
25
26        let Some(new_patches) = patches_take else {
27            let result_fill_scalar = array.fill_scalar().cast(
28                &array
29                    .dtype()
30                    .union_nullability(indices.dtype().nullability()),
31            )?;
32            return Ok(Some(
33                ConstantArray::new(result_fill_scalar, indices.len()).into_array(),
34            ));
35        };
36
37        // See `SparseEncoding::slice`.
38        if new_patches.array_len() == new_patches.values().len() {
39            return Ok(Some(new_patches.into_values()));
40        }
41
42        Ok(Some(
43            Sparse::try_new_from_patches(
44                new_patches,
45                array.fill_scalar().cast(
46                    &array
47                        .dtype()
48                        .union_nullability(indices.dtype().nullability()),
49                )?,
50            )?
51            .into_array(),
52        ))
53    }
54}
55
56#[cfg(test)]
57mod test {
58    use std::sync::LazyLock;
59
60    use rstest::rstest;
61    use vortex_array::ArrayRef;
62    use vortex_array::IntoArray;
63    use vortex_array::VortexSessionExecute;
64    use vortex_array::arrays::ConstantArray;
65    use vortex_array::arrays::PrimitiveArray;
66    use vortex_array::assert_arrays_eq;
67    use vortex_array::compute::conformance::take::test_take_conformance;
68    use vortex_array::dtype::Nullability;
69    use vortex_array::scalar::Scalar;
70    use vortex_array::validity::Validity;
71    use vortex_buffer::buffer;
72    use vortex_session::VortexSession;
73
74    use crate::Sparse;
75    use crate::SparseArray;
76
77    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
78        let session = vortex_array::array_session();
79        crate::initialize(&session);
80        session
81    });
82
83    fn test_array_fill_value() -> Scalar {
84        // making this const is annoying
85        Scalar::null_native::<f64>()
86    }
87
88    fn sparse_array() -> ArrayRef {
89        Sparse::try_new(
90            buffer![0u64, 37, 47, 99].into_array(),
91            PrimitiveArray::new(buffer![1.23f64, 0.47, 9.99, 3.5], Validity::AllValid).into_array(),
92            100,
93            test_array_fill_value(),
94        )
95        .unwrap()
96        .into_array()
97    }
98
99    #[test]
100    fn take_with_non_zero_offset() {
101        let sparse = sparse_array();
102        let sparse = sparse.slice(30..40).unwrap();
103        let taken = sparse.take(buffer![6, 7, 8].into_array()).unwrap();
104        let expected = PrimitiveArray::from_option_iter([Option::<f64>::None, Some(0.47), None]);
105        assert_arrays_eq!(
106            taken,
107            expected.into_array(),
108            &mut SESSION.create_execution_ctx()
109        );
110    }
111
112    #[test]
113    fn sparse_take() {
114        let sparse = sparse_array();
115        let taken = sparse.take(buffer![0, 47, 47, 0, 99].into_array()).unwrap();
116        let expected = PrimitiveArray::from_option_iter([
117            Some(1.23f64),
118            Some(9.99),
119            Some(9.99),
120            Some(1.23),
121            Some(3.5),
122        ]);
123        assert_arrays_eq!(
124            taken,
125            expected.into_array(),
126            &mut SESSION.create_execution_ctx()
127        );
128    }
129
130    #[test]
131    fn nonexistent_take() {
132        let sparse = sparse_array();
133        let taken = sparse.take(buffer![69].into_array()).unwrap();
134        let expected = ConstantArray::new(test_array_fill_value(), 1).into_array();
135        assert_arrays_eq!(taken, expected, &mut SESSION.create_execution_ctx());
136    }
137
138    #[test]
139    fn ordered_take() {
140        let sparse = sparse_array();
141        // Note: take returns a canonical array, not SparseArray
142        let taken = sparse.take(buffer![69, 37].into_array()).unwrap();
143        // Index 69 is not in sparse array (fill value is null), index 37 has value 0.47
144        let expected = PrimitiveArray::from_option_iter([Option::<f64>::None, Some(0.47f64)]);
145        assert_arrays_eq!(
146            taken,
147            expected.into_array(),
148            &mut SESSION.create_execution_ctx()
149        );
150    }
151
152    #[test]
153    fn nullable_take() {
154        let arr = Sparse::try_new(
155            buffer![1u32].into_array(),
156            buffer![10].into_array(),
157            10,
158            Scalar::primitive(1, Nullability::NonNullable),
159        )
160        .unwrap();
161
162        let taken = arr
163            .take(
164                PrimitiveArray::from_option_iter([Some(2u32), Some(1u32), Option::<u32>::None])
165                    .into_array(),
166            )
167            .unwrap();
168
169        let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::<i32>::None]);
170        assert_arrays_eq!(
171            taken,
172            expected.into_array(),
173            &mut SESSION.create_execution_ctx()
174        );
175    }
176
177    #[test]
178    fn nullable_take_with_many_patches() {
179        let arr = Sparse::try_new(
180            buffer![1u32, 3, 7, 8, 9].into_array(),
181            buffer![10, 8, 3, 2, 1].into_array(),
182            10,
183            Scalar::primitive(1, Nullability::NonNullable),
184        )
185        .unwrap();
186
187        let taken = arr
188            .take(
189                PrimitiveArray::from_option_iter([Some(2u32), Some(1u32), Option::<u32>::None])
190                    .into_array(),
191            )
192            .unwrap();
193
194        let expected = PrimitiveArray::from_option_iter([Some(1), Some(10), Option::<i32>::None]);
195        assert_arrays_eq!(
196            taken,
197            expected.into_array(),
198            &mut SESSION.create_execution_ctx()
199        );
200    }
201
202    #[rstest]
203    #[case(Sparse::try_new(
204        buffer![0u64, 37, 47, 99].into_array(),
205        PrimitiveArray::new(buffer![1.23f64, 0.47, 9.99, 3.5], Validity::AllValid).into_array(),
206        100,
207        Scalar::null_native::<f64>(),
208    ).unwrap())]
209    #[case(Sparse::try_new(
210        buffer![1u32, 3, 7, 8, 9].into_array(),
211        buffer![10, 8, 3, 2, 1].into_array(),
212        10,
213        Scalar::from(0i32),
214    ).unwrap())]
215    #[case({
216        let nullable_values = PrimitiveArray::from_option_iter([Some(100i64), None, Some(300)]);
217        Sparse::try_new(
218            buffer![2u64, 4, 6].into_array(),
219            nullable_values.into_array(),
220            10,
221            Scalar::null_native::<i64>(),
222        ).unwrap()
223    })]
224    #[case(Sparse::try_new(
225        buffer![5u64].into_array(),
226        buffer![999i32].into_array(),
227        20,
228        Scalar::from(-1i32),
229    ).unwrap())]
230    fn test_take_sparse_conformance(#[case] sparse: SparseArray) {
231        test_take_conformance(&sparse.into_array(), &mut SESSION.create_execution_ctx());
232    }
233}
vortex_sparse/compute/take.rs

vortex_sparse/compute/
take.rs
