vortex_array/arrays/fixed_width/take/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
5mod avx2;
6mod records;
7mod scalar;
8mod slices;
9#[cfg(test)]
10mod tests;
11
12#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
13use std::sync::LazyLock;
14
15use vortex_buffer::Buffer;
16use vortex_error::VortexResult;
17use vortex_error::vortex_bail;
18use vortex_mask::Mask;
19
20use self::records::take_byte_records;
21use self::scalar::take_values_scalar;
22use self::slices::take_slices;
23use self::slices::take_slices_constant_length;
24use super::FixedWidthArray;
25use super::with_values;
26use crate::ArrayRef;
27use crate::Columnar;
28use crate::ExecutionCtx;
29use crate::IntoArray;
30use crate::array::ArrayView;
31use crate::arrays::ConstantArray;
32use crate::arrays::PiecewiseSequence;
33use crate::arrays::PrimitiveArray;
34use crate::arrays::dict::TakeExecute;
35use crate::arrays::piecewise_sequence::constant_unsigned_usize;
36use crate::arrays::piecewise_sequence::maybe_contiguous_slices;
37use crate::builtins::ArrayBuiltins;
38use crate::dtype::DType;
39use crate::dtype::UnsignedPType;
40use crate::dtype::half::f16;
41use crate::match_each_unsigned_integer_ptype;
42use crate::scalar::Scalar;
43
44#[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
45static HAS_AVX2: LazyLock<bool> = LazyLock::new(|| is_x86_feature_detected!("avx2"));
46
47impl<V: FixedWidthArray> TakeExecute for V {
48    fn take(
49        array: ArrayView<'_, Self>,
50        indices: &ArrayRef,
51        ctx: &mut ExecutionCtx,
52    ) -> VortexResult<Option<ArrayRef>> {
53        take(array, indices, ctx)
54    }
55}
56
57/// A fixed-width value whose initialized bytes may be moved through integer SIMD lanes.
58///
59/// # Safety
60///
61/// Implementors must have no uninitialized bytes. The shared AVX2 gather reads the complete
62/// representation through a same-width integer lane before writing those bytes back unchanged.
63pub(crate) unsafe trait FixedWidthTakeValue: Copy {}
64
65macro_rules! impl_fixed_width_take_value {
66    ($($ty:ty),+ $(,)?) => {
67        $(
68            // SAFETY: These scalar representations contain no padding or uninitialized bytes.
69            unsafe impl FixedWidthTakeValue for $ty {}
70        )+
71    };
72}
73
74impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,);
75
76// SAFETY: Byte arrays have no padding and every byte is initialized.
77unsafe impl<const N: usize> FixedWidthTakeValue for [u8; N] {}
78
79pub(crate) fn take_values<T: FixedWidthTakeValue, I: UnsignedPType>(
80    values: &[T],
81    indices: &[I],
82) -> Buffer<T> {
83    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
84    if *HAS_AVX2 {
85        // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte
86        // representation. The AVX2 dispatcher retains Primitive's existing scalar fallbacks and
87        // out-of-bounds behavior for every value width.
88        return unsafe { avx2::take_avx2(values, indices) };
89    }
90
91    take_values_scalar(values, indices)
92}
93
94pub(crate) fn take<V: FixedWidthArray>(
95    array: ArrayView<'_, V>,
96    indices: &ArrayRef,
97    ctx: &mut ExecutionCtx,
98) -> VortexResult<Option<ArrayRef>> {
99    if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequence>()
100        && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)?
101    {
102        return Ok(Some(taken));
103    }
104
105    let DType::Primitive(ptype, nullability) = indices.dtype() else {
106        vortex_bail!("Invalid indices dtype: {}", indices.dtype())
107    };
108    if !ptype.is_int() {
109        vortex_bail!("Invalid indices dtype: {}", indices.dtype())
110    }
111
112    let indices_validity = indices.validity()?;
113    let indices_nulls_zeroed = match indices_validity.execute_mask(indices.len(), ctx)? {
114        Mask::AllTrue(_) => indices.clone(),
115        Mask::AllFalse(_) => {
116            return Ok(Some(
117                ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len())
118                    .into_array(),
119            ));
120        }
121        Mask::Values(_) => indices
122            .clone()
123            .fill_null(Scalar::from(0).cast(indices.dtype())?)?,
124    };
125
126    let indices = if ptype.is_unsigned_int() {
127        indices_nulls_zeroed.execute::<PrimitiveArray>(ctx)?
128    } else {
129        indices_nulls_zeroed
130            .cast(DType::Primitive(ptype.to_unsigned(), *nullability))?
131            .execute::<PrimitiveArray>(ctx)?
132    };
133    let validity = array
134        .validity()?
135        .take(&indices.clone().into_array())?
136        .and(indices_validity)?;
137
138    let source = V::values(array);
139    let values = match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
140        take_byte_records(
141            &source,
142            V::byte_width(array),
143            array.len(),
144            indices.as_slice::<I>(),
145        )
146    })?;
147    Ok(Some(
148        with_values(array, values, indices.len(), validity)?.into_array(),
149    ))
150}
151
152fn take_contiguous_ranges<V: FixedWidthArray>(
153    array: ArrayView<'_, V>,
154    indices: ArrayView<'_, PiecewiseSequence>,
155    indices_ref: &ArrayRef,
156    ctx: &mut ExecutionCtx,
157) -> VortexResult<Option<ArrayRef>> {
158    let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else {
159        return Ok(None);
160    };
161
162    let values = V::values(array);
163    let byte_width = V::byte_width(array);
164    let output_len = indices_ref.len();
165    let taken = match lengths {
166        Columnar::Constant(lengths) => {
167            let length = constant_unsigned_usize(&lengths);
168            match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
169                take_slices_constant_length(
170                    &values,
171                    byte_width,
172                    array.len(),
173                    starts.as_slice::<S>(),
174                    length,
175                    output_len,
176                )
177            })
178        }
179        Columnar::Canonical(lengths) => {
180            let lengths = lengths.into_primitive();
181            match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
182                match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
183                    take_slices(
184                        &values,
185                        byte_width,
186                        array.len(),
187                        starts.as_slice::<S>(),
188                        lengths.as_slice::<L>(),
189                        output_len,
190                    )
191                })
192            })
193        }
194    }?;
195    let validity = array.validity()?.take(indices_ref)?;
196    Ok(Some(
197        with_values(array, taken, output_len, validity)?.into_array(),
198    ))
199}
vortex_array/arrays/fixed_width/take/mod.rs

vortex_array/arrays/fixed_width/take/
mod.rs
