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

vortex_array/arrays/constant/compute/
take.rs
