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

vortex_array/arrays/fixed_width/take/
mod.rs
