Skip to main content

vortex_array/arrays/constant/vtable/
operations.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5
6use crate::ExecutionCtx;
7use crate::array::ArrayView;
8use crate::array::OperationsVTable;
9use crate::arrays::Constant;
10use crate::scalar::Scalar;
11
12impl OperationsVTable<Constant> for Constant {
13    type ProbeState = ();
14
15    fn scalar_at(
16        array: ArrayView<'_, Constant>,
17        _index: usize,
18        _ctx: &mut ExecutionCtx,
19    ) -> VortexResult<Scalar> {
20        Ok(array.scalar.clone())
21    }
22}
23
24#[cfg(test)]
25mod tests {
26    use vortex_error::VortexResult;
27
28    use crate::IntoArray;
29    use crate::VortexSessionExecute;
30    use crate::arrays::ConstantArray;
31    use crate::dtype::DType;
32    use crate::dtype::Nullability;
33    use crate::dtype::PType;
34    use crate::dtype::UnionVariants;
35    use crate::scalar::Scalar;
36
37    #[test]
38    fn scalar_at_preserves_union_scalar() -> VortexResult<()> {
39        let variants = UnionVariants::try_new(
40            ["int", "string"].into(),
41            vec![
42                DType::Primitive(PType::I32, Nullability::Nullable),
43                DType::Utf8(Nullability::NonNullable),
44            ],
45            vec![5, 9],
46        )?;
47
48        let scalar = Scalar::union(
49            variants.clone(),
50            5,
51            Scalar::primitive(42_i32, Nullability::Nullable),
52            Nullability::Nullable,
53        )?;
54        let array = ConstantArray::new(scalar.clone(), 3).into_array();
55        let mut ctx = crate::array_session().create_execution_ctx();
56
57        assert_eq!(array.execute_scalar(1, &mut ctx)?, scalar);
58
59        let null = Scalar::null(DType::Union(variants, Nullability::Nullable));
60        let array = ConstantArray::new(null.clone(), 3).into_array();
61
62        assert_eq!(array.execute_scalar(1, &mut ctx)?, null);
63
64        Ok(())
65    }
66}