vortex_vector/primitive/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Casting utilities for primitive vectors.
5
6use num_traits::NumCast;
7use vortex_buffer::Buffer;
8use vortex_buffer::BufferMut;
9use vortex_dtype::NativePType;
10use vortex_dtype::PType;
11use vortex_dtype::match_each_signed_integer_ptype;
12use vortex_dtype::match_each_unsigned_integer_ptype;
13use vortex_error::VortexExpect;
14use vortex_error::VortexResult;
15use vortex_error::vortex_err;
16use vortex_mask::AllOr;
17use vortex_mask::Mask;
18
19use super::PVector;
20use super::PrimitiveVector;
21use super::PrimitiveVectorMut;
22use crate::VectorMutOps;
23use crate::VectorOps;
24use crate::match_each_integer_pvector;
25use crate::match_each_integer_pvector_mut;
26
27/// Cast a [`PVector<Src>`] to a [`PVector<Dst>`] by converting each element.
28///
29/// # Errors
30///
31/// Returns an error if any valid element cannot be converted (e.g., overflow).
32pub fn cast_pvector<Src: NativePType, Dst: NativePType>(
33    src: &PVector<Src>,
34) -> VortexResult<PVector<Dst>> {
35    let elements: &[Src] = src.as_ref();
36    match src.validity().bit_buffer() {
37        AllOr::All => {
38            let mut buffer = BufferMut::with_capacity(elements.len());
39            for &item in elements {
40                let converted = <Dst as NumCast>::from(item).ok_or_else(
41                    || vortex_err!(ComputeError: "Failed to cast {} to {:?}", item, Dst::PTYPE),
42                )?;
43                // SAFETY: We pre-allocated the required capacity.
44                unsafe { buffer.push_unchecked(converted) }
45            }
46            Ok(PVector::from(buffer.freeze()))
47        }
48        AllOr::None => Ok(PVector::new(
49            Buffer::zeroed(elements.len()),
50            Mask::new_false(elements.len()),
51        )),
52        AllOr::Some(bit_buffer) => {
53            let mut buffer = BufferMut::with_capacity(elements.len());
54            for (&item, valid) in elements.iter().zip(bit_buffer.iter()) {
55                if valid {
56                    let converted = <Dst as NumCast>::from(item).ok_or_else(
57                        || vortex_err!(ComputeError: "Failed to cast {} to {:?}", item, Dst::PTYPE),
58                    )?;
59                    // SAFETY: We pre-allocated the required capacity.
60                    unsafe { buffer.push_unchecked(converted) }
61                } else {
62                    // SAFETY: We pre-allocated the required capacity.
63                    unsafe { buffer.push_unchecked(Dst::default()) }
64                }
65            }
66            Ok(PVector::new(buffer.freeze(), src.validity().clone()))
67        }
68    }
69}
70
71impl PrimitiveVectorMut {
72    /// Upcasts this integer vector to a wider integer type with matching signedness.
73    ///
74    /// Returns self unchanged if the target type is the same or smaller in byte width.
75    #[expect(
76        clippy::cognitive_complexity,
77        reason = "complexity from nested match_each_* macros"
78    )]
79    pub fn upcast(self, target: PType) -> Self {
80        debug_assert!(self.ptype().is_int());
81        debug_assert!(target.is_int());
82        debug_assert!(
83            (self.ptype().is_signed_int() && target.is_signed_int())
84                || (self.ptype().is_unsigned_int() && target.is_unsigned_int())
85        );
86
87        // No-op if already at target or target is same/smaller
88        if self.ptype() == target || target.byte_width() <= self.ptype().byte_width() {
89            return self;
90        }
91
92        // Freeze to immutable, cast, then convert back to mutable
93        let frozen = self.freeze();
94        match_each_integer_pvector!(&frozen, |src_vec| {
95            if target.is_unsigned_int() {
96                match_each_unsigned_integer_ptype!(target, |Dst| {
97                    let casted = cast_pvector::<_, Dst>(src_vec)
98                        .vortex_expect("upcast should never fail for widening casts");
99                    casted.into_mut().into()
100                })
101            } else {
102                match_each_signed_integer_ptype!(target, |Dst| {
103                    let casted = cast_pvector::<_, Dst>(src_vec)
104                        .vortex_expect("upcast should never fail for widening casts");
105                    casted.into_mut().into()
106                })
107            }
108        })
109    }
110
111    /// Extends this vector from another, automatically upcasting if needed.
112    ///
113    /// Unlike [`VectorMutOps::extend_from_vector`](VectorMutOps::extend_from_vector), this does
114    /// **NOT** panic on type mismatch. Instead, it upcasts `self` to the wider of the two types.
115    ///
116    /// Returns the new [`PType`] after any upcasting.
117    pub fn extend_from_vector_with_upcast(&mut self, other: &PrimitiveVector) -> PType {
118        debug_assert!(self.ptype().is_int());
119        debug_assert!(other.ptype().is_int());
120
121        let target = self
122            .ptype()
123            .to_unsigned()
124            .max_unsigned_ptype(other.ptype().to_unsigned());
125
126        if self.ptype() != target {
127            let old_self = std::mem::replace(self, Self::with_capacity(target, 0));
128            *self = old_self.upcast(target);
129        }
130
131        // Now extend with casting from other
132        self.reserve(other.len());
133        extend_with_cast(self, other);
134        self.ptype()
135    }
136}
137
138/// Extends `dst` with values from `src`, casting each element.
139#[expect(
140    clippy::cognitive_complexity,
141    reason = "complexity from nested match_each_* macros"
142)]
143fn extend_with_cast(dst: &mut PrimitiveVectorMut, src: &PrimitiveVector) {
144    match_each_integer_pvector_mut!(dst, |dst_vec| {
145        match_each_integer_pvector!(src, |src_vec| {
146            let src_slice = src_vec.as_ref();
147            let src_validity = src_vec.validity();
148            for i in 0..src_vec.len() {
149                if src_validity.value(i) {
150                    #[allow(clippy::unnecessary_cast)]
151                    let converted = <_ as NumCast>::from(src_slice[i])
152                        .vortex_expect("conversion should succeed after upcast");
153                    dst_vec.push_opt(Some(converted));
154                } else {
155                    dst_vec.push_opt(None);
156                }
157            }
158        });
159    });
160}
161
162#[cfg(test)]
163mod tests {
164    use vortex_dtype::PType;
165    use vortex_dtype::PTypeDowncast;
166
167    use super::*;
168    use crate::primitive::PVectorMut;
169
170    #[test]
171    fn test_upcast_unsigned() {
172        let mut vec: PrimitiveVectorMut =
173            PVectorMut::<u8>::from_iter([0u8, u8::MAX].map(Some)).into();
174        let other: PrimitiveVector = PVectorMut::<u32>::from_iter([u32::MAX].map(Some))
175            .freeze()
176            .into();
177
178        vec.extend_from_vector_with_upcast(&other);
179        assert_eq!(vec.ptype(), PType::U32);
180
181        let frozen = vec.freeze().into_u32();
182        assert_eq!(frozen.as_ref(), &[0, u8::MAX as u32, u32::MAX]);
183    }
184
185    #[test]
186    fn test_upcast_signed() {
187        let vec: PrimitiveVectorMut =
188            PVectorMut::<i8>::from_iter([i8::MIN, i8::MAX].map(Some)).into();
189
190        let mut vec = vec.upcast(PType::I32);
191        let other: PrimitiveVector = PVectorMut::<i32>::from_iter([i32::MIN, i32::MAX].map(Some))
192            .freeze()
193            .into();
194        extend_with_cast(&mut vec, &other);
195
196        let frozen = vec.freeze().into_i32();
197        assert_eq!(
198            frozen.as_ref(),
199            &[i8::MIN as i32, i8::MAX as i32, i32::MIN, i32::MAX]
200        );
201    }
202}
vortex_vector/primitive/cast.rs

vortex_vector/primitive/
cast.rs
