vortex_array/arrays/constant/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_mask::AllOr;
6
7use crate::ArrayRef;
8use crate::IntoArray;
9use crate::VortexSessionExecute;
10use crate::array::ArrayView;
11use crate::arrays::Constant;
12use crate::arrays::ConstantArray;
13use crate::arrays::MaskedArray;
14use crate::arrays::dict::TakeReduce;
15use crate::arrays::dict::TakeReduceAdaptor;
16use crate::legacy_session;
17use crate::optimizer::rules::ParentRuleSet;
18use crate::scalar::Scalar;
19use crate::validity::Validity;
20
21impl TakeReduce for Constant {
22    #[allow(clippy::disallowed_methods)]
23    fn take(array: ArrayView<'_, Constant>, indices: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
24        let mut ctx = legacy_session().create_execution_ctx();
25        let result = match indices
26            .validity()?
27            .execute_mask(indices.len(), &mut ctx)?
28            .bit_buffer()
29        {
30            AllOr::All => {
31                let scalar = Scalar::try_new(
32                    array
33                        .scalar()
34                        .dtype()
35                        .union_nullability(indices.dtype().nullability()),
36                    array.scalar().value().cloned(),
37                )?;
38                ConstantArray::new(scalar, indices.len()).into_array()
39            }
40            AllOr::None => ConstantArray::new(
41                Scalar::null(
42                    array
43                        .dtype()
44                        .union_nullability(indices.dtype().nullability()),
45                ),
46                indices.len(),
47            )
48            .into_array(),
49            AllOr::Some(v) => {
50                let arr = ConstantArray::new(array.scalar().clone(), indices.len()).into_array();
51
52                if array.scalar().is_null() {
53                    return Ok(Some(arr));
54                }
55
56                MaskedArray::try_new(arr, Validity::from(v.clone()))?.into_array()
57            }
58        };
59        Ok(Some(result))
60    }
61}
62
63impl Constant {
64    pub const TAKE_RULES: ParentRuleSet<Self> =
65        ParentRuleSet::new(&[ParentRuleSet::lift(&TakeReduceAdaptor::<Self>(Self))]);
66}
67
68#[cfg(test)]
69mod tests {
70    use std::f64;
71
72    use rstest::rstest;
73    use vortex_buffer::buffer;
74    use vortex_mask::AllOr;
75
76    use crate::IntoArray;
77    use crate::VortexSessionExecute;
78    use crate::array_session;
79    use crate::arrays::ConstantArray;
80    use crate::arrays::PrimitiveArray;
81    use crate::assert_arrays_eq;
82    use crate::compute::conformance::take::test_take_conformance;
83    use crate::dtype::Nullability;
84    use crate::scalar::Scalar;
85    use crate::validity::Validity;
86
87    #[test]
88    fn take_nullable_indices() {
89        let mut ctx = array_session().create_execution_ctx();
90        let array = ConstantArray::new(42, 10).into_array();
91        let taken = array
92            .take(
93                PrimitiveArray::new(
94                    buffer![0, 5, 7],
95                    Validity::from_iter(vec![false, true, false]),
96                )
97                .into_array(),
98            )
99            .unwrap();
100        let valid_indices: &[usize] = &[1usize];
101        assert_eq!(
102            &array.dtype().with_nullability(Nullability::Nullable),
103            taken.dtype()
104        );
105        assert_arrays_eq!(
106            taken.clone().execute::<PrimitiveArray>(&mut ctx).unwrap(),
107            PrimitiveArray::new(
108                buffer![42i32, 42, 42],
109                Validity::from_iter([false, true, false])
110            ),
111            &mut ctx
112        );
113        assert_eq!(
114            taken
115                .validity()
116                .unwrap()
117                .execute_mask(taken.len(), &mut ctx)
118                .unwrap()
119                .indices(),
120            AllOr::Some(valid_indices)
121        );
122    }
123
124    #[test]
125    fn take_all_valid_indices() {
126        let mut ctx = array_session().create_execution_ctx();
127        let array = ConstantArray::new(42, 10).into_array();
128        let taken = array
129            .take(PrimitiveArray::new(buffer![0, 5, 7], Validity::AllValid).into_array())
130            .unwrap();
131        assert_eq!(
132            &array.dtype().with_nullability(Nullability::Nullable),
133            taken.dtype()
134        );
135        assert_arrays_eq!(
136            taken.clone().execute::<PrimitiveArray>(&mut ctx).unwrap(),
137            PrimitiveArray::new(buffer![42i32, 42, 42], Validity::AllValid),
138            &mut ctx
139        );
140        assert_eq!(
141            taken
142                .validity()
143                .unwrap()
144                .execute_mask(taken.len(), &mut ctx)
145                .unwrap()
146                .indices(),
147            AllOr::All
148        );
149    }
150
151    #[rstest]
152    #[case(ConstantArray::new(42i32, 5))]
153    #[case(ConstantArray::new(f64::consts::PI, 10))]
154    #[case(ConstantArray::new(Scalar::from("hello"), 3))]
155    #[case(ConstantArray::new(Scalar::null_native::<i64>(), 5))]
156    #[case(ConstantArray::new(true, 1))]
157    fn test_take_constant_conformance(#[case] array: ConstantArray) {
158        test_take_conformance(
159            &array.into_array(),
160            &mut array_session().create_execution_ctx(),
161        );
162    }
163}
vortex_array/arrays/constant/compute/take.rs

vortex_array/arrays/constant/compute/
take.rs
