vortex_sequence/compute/
cast.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::IntoArray;
7use vortex_array::dtype::DType;
8use vortex_array::dtype::Nullability;
9use vortex_array::scalar::Scalar;
10use vortex_array::scalar::ScalarValue;
11use vortex_array::scalar_fn::fns::cast::CastReduce;
12use vortex_error::VortexResult;
13use vortex_error::vortex_err;
14
15use crate::Sequence;
16impl CastReduce for Sequence {
17    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
18        // SequenceArray represents arithmetic sequences (base + i * multiplier) which
19        // only makes sense for integer types. Floating-point sequences would accumulate
20        // rounding errors, and other types don't support arithmetic operations.
21        let DType::Primitive(target_ptype, target_nullability) = dtype else {
22            return Ok(None);
23        };
24
25        if !target_ptype.is_int() {
26            return Ok(None);
27        }
28
29        // Check if this is just a nullability change
30        if array.ptype() == *target_ptype && array.dtype().nullability() != *target_nullability {
31            // For SequenceArray, we can just create a new one with the same parameters
32            // but different nullability
33            return Ok(Some(
34                Sequence::try_new(
35                    array.base(),
36                    array.multiplier(),
37                    *target_ptype,
38                    *target_nullability,
39                    array.len(),
40                )?
41                .into_array(),
42            ));
43        }
44
45        // For type changes, we need to cast the base and multiplier
46        if array.ptype() != *target_ptype {
47            // Create scalars from PValues and cast them
48            let base_scalar = Scalar::try_new(
49                DType::Primitive(array.ptype(), Nullability::NonNullable),
50                Some(ScalarValue::Primitive(array.base())),
51            )?;
52            let multiplier_scalar = Scalar::try_new(
53                DType::Primitive(array.ptype(), Nullability::NonNullable),
54                Some(ScalarValue::Primitive(array.multiplier())),
55            )?;
56
57            let new_base_scalar =
58                base_scalar.cast(&DType::Primitive(*target_ptype, Nullability::NonNullable))?;
59            let new_multiplier_scalar = multiplier_scalar
60                .cast(&DType::Primitive(*target_ptype, Nullability::NonNullable))?;
61
62            // Extract PValues from the casted scalars
63            let new_base = new_base_scalar
64                .as_primitive()
65                .pvalue()
66                .ok_or_else(|| vortex_err!("Cast resulted in null base value"))?;
67            let new_multiplier = new_multiplier_scalar
68                .as_primitive()
69                .pvalue()
70                .ok_or_else(|| vortex_err!("Cast resulted in null multiplier value"))?;
71
72            return Ok(Some(
73                Sequence::try_new(
74                    new_base,
75                    new_multiplier,
76                    *target_ptype,
77                    *target_nullability,
78                    array.len(),
79                )?
80                .into_array(),
81            ));
82        }
83
84        Ok(None)
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use std::sync::LazyLock;
91
92    use rstest::rstest;
93    use vortex_array::IntoArray;
94    use vortex_array::VortexSessionExecute;
95    use vortex_array::arrays::PrimitiveArray;
96    use vortex_array::assert_arrays_eq;
97    use vortex_array::builtins::ArrayBuiltins;
98    use vortex_array::compute::conformance::cast::test_cast_conformance;
99    use vortex_array::dtype::DType;
100    use vortex_array::dtype::Nullability;
101    use vortex_array::dtype::PType;
102    use vortex_session::VortexSession;
103
104    use crate::Sequence;
105    use crate::SequenceArray;
106
107    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
108        let session = vortex_array::array_session();
109        crate::initialize(&session);
110        session
111    });
112
113    #[test]
114    fn test_cast_sequence_nullability() {
115        let sequence = Sequence::try_new_typed(0u32, 1u32, Nullability::NonNullable, 4).unwrap();
116
117        // Cast to nullable
118        let casted = sequence
119            .into_array()
120            .cast(DType::Primitive(PType::U32, Nullability::Nullable))
121            .unwrap();
122        assert_eq!(
123            casted.dtype(),
124            &DType::Primitive(PType::U32, Nullability::Nullable)
125        );
126    }
127
128    #[test]
129    fn test_cast_sequence_u32_to_i64() {
130        let mut ctx = SESSION.create_execution_ctx();
131        let sequence = Sequence::try_new_typed(100u32, 10u32, Nullability::NonNullable, 4).unwrap();
132
133        let casted = sequence
134            .into_array()
135            .cast(DType::Primitive(PType::I64, Nullability::NonNullable))
136            .unwrap();
137        assert_eq!(
138            casted.dtype(),
139            &DType::Primitive(PType::I64, Nullability::NonNullable)
140        );
141
142        // Verify the values
143        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
144        assert_arrays_eq!(
145            decoded,
146            PrimitiveArray::from_iter([100i64, 110, 120, 130]),
147            &mut ctx
148        );
149    }
150
151    #[test]
152    fn test_cast_sequence_i16_to_i32_nullable() {
153        let mut ctx = SESSION.create_execution_ctx();
154        // Test ptype change AND nullability change in one cast
155        let sequence = Sequence::try_new_typed(5i16, 3i16, Nullability::NonNullable, 3).unwrap();
156
157        let casted = sequence
158            .into_array()
159            .cast(DType::Primitive(PType::I32, Nullability::Nullable))
160            .unwrap();
161        assert_eq!(
162            casted.dtype(),
163            &DType::Primitive(PType::I32, Nullability::Nullable)
164        );
165
166        // Verify the values
167        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
168        assert_arrays_eq!(
169            decoded,
170            PrimitiveArray::from_option_iter([Some(5i32), Some(8), Some(11)]),
171            &mut ctx
172        );
173    }
174
175    #[test]
176    fn test_cast_sequence_to_float_delegates_to_canonical() {
177        let mut ctx = SESSION.create_execution_ctx();
178        let sequence = Sequence::try_new_typed(0i32, 1i32, Nullability::NonNullable, 5).unwrap();
179
180        // Cast to float should delegate to canonical (SequenceArray doesn't support float)
181        let casted = sequence
182            .into_array()
183            .cast(DType::Primitive(PType::F32, Nullability::NonNullable))
184            .unwrap();
185        // Should still succeed by decoding to canonical first
186        assert_eq!(
187            casted.dtype(),
188            &DType::Primitive(PType::F32, Nullability::NonNullable)
189        );
190
191        // Verify the values were correctly converted
192        let decoded = casted.execute::<PrimitiveArray>(&mut ctx).unwrap();
193        assert_arrays_eq!(
194            decoded,
195            PrimitiveArray::from_iter([0.0f32, 1.0, 2.0, 3.0, 4.0]),
196            &mut ctx
197        );
198    }
199
200    #[rstest]
201    #[case::i32(Sequence::try_new_typed(0i32, 1i32, Nullability::NonNullable, 5).unwrap())]
202    #[case::u64(Sequence::try_new_typed(1000u64, 100u64, Nullability::NonNullable, 4).unwrap())]
203    // TODO(DK): SequenceArray does not actually conform. You cannot cast this array to u8 even
204    // though all its values are representable therein.
205    //
206    // #[case::negative_step(Sequence::try_new_typed(100i32, -10i32, Nullability::NonNullable,
207    // 5).unwrap())]
208    #[case::single(Sequence::try_new_typed(42i64, 0i64, Nullability::NonNullable, 1).unwrap())]
209    #[case::constant(Sequence::try_new_typed(
210        100i32,
211        0i32, // multiplier of 0 means constant array
212        Nullability::NonNullable,
213        5,
214    ).unwrap())]
215    fn test_cast_sequence_conformance(#[case] sequence: SequenceArray) {
216        test_cast_conformance(&sequence.into_array());
217    }
218}
vortex_sequence/compute/cast.rs

vortex_sequence/compute/
cast.rs
