vortex_compute/lane_kernels/
map_into.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Out-of-place lane kernels: read from an [`IndexedSource`] and write into a
5//! caller-provided `&mut [MaybeUninit<R>]`.
6
7use std::mem::MaybeUninit;
8
9use vortex_buffer::BitBuffer;
10
11use crate::lane_kernels::CHUNK_LEN;
12use crate::lane_kernels::source::IndexedSource;
13
14/// Extension trait providing out-of-place lane-kernel methods on any [`IndexedSource`].
15///
16/// All methods have default implementations and are inherited via the blanket
17/// `impl<S: IndexedSource> IndexedSourceExt for S` below. Bring the trait into
18/// scope (`use vortex_compute::lane_kernels::IndexedSourceExt;`) to call
19/// them with method syntax: `values.try_map_masked_into(&mask, &mut out, f)`.
20pub trait IndexedSourceExt: IndexedSource + Sized {
21    /// Fallible map with mask-aware error attribution. `f` returns `Option<R>`;
22    /// `None` indicates a per-lane failure (e.g. range overflow on a narrowing cast).
23    ///
24    /// **Null-lane failures are filtered automatically.** The closure is called on
25    /// every lane regardless of validity; if a null lane's stored value causes `f(v)`
26    /// to return `None`, the kernel does *not* propagate that as `Err`. The per-lane
27    /// `is_none()` flags are bit-packed into a `u64` at the lane's position, then
28    /// AND-combined with the chunk's validity bitmap — null-lane bits vanish.
29    ///
30    /// The closure shape is the same as [`try_map_into`] (`FnMut(Item) -> Option<R>`);
31    /// the mask parameter is what makes this kernel mask-aware. Callers that need to
32    /// distinguish null lanes inside the closure (e.g. to short-circuit an expensive
33    /// computation) should construct their own per-lane validity check externally; for
34    /// the common case, the kernel's automatic filter is sufficient.
35    ///
36    /// On failure returns `Err(failing_lane_index)`. Lanes whose `f` returned `None`
37    /// write `R::default()` into `out`, but the contents of `out` must not be relied
38    /// upon when this function returns `Err`.
39    ///
40    /// [`try_map_into`]: IndexedSourceExt::try_map_into
41    ///
42    /// # Panics
43    ///
44    /// Panics if `self.len() != mask.len()` or `out.len() != self.len()`.
45    #[inline]
46    fn try_map_masked_into<R, F>(
47        self,
48        mask: &BitBuffer,
49        out: &mut [MaybeUninit<R>],
50        mut f: F,
51    ) -> Result<(), usize>
52    where
53        R: Copy + Default,
54        F: FnMut(Self::Item) -> Option<R>,
55    {
56        #[inline(always)]
57        fn chunk<S, R, F>(
58            values: &S,
59            out: &mut [MaybeUninit<R>],
60            f: &mut F,
61            src_chunk: u64,
62            base: usize,
63            count: usize,
64        ) -> Option<usize>
65        where
66            S: IndexedSource,
67            R: Copy + Default,
68            F: FnMut(S::Item) -> Option<R>,
69        {
70            let mut fail_bits: u64 = 0;
71            for bit_idx in 0..count {
72                let idx = base + bit_idx;
73                // SAFETY: caller guarantees base + count <= len.
74                let val = unsafe { values.get_unchecked(idx) };
75                let opt = f(val);
76                fail_bits |= (opt.is_none() as u64) << bit_idx;
77                let result = opt.unwrap_or_default();
78                unsafe { out.get_unchecked_mut(idx).write(result) };
79            }
80            let valid_failures = fail_bits & src_chunk;
81            (valid_failures != 0).then_some(base + valid_failures.trailing_zeros() as usize)
82        }
83
84        let values = self;
85        let len = values.len();
86        assert_eq!(len, mask.len(), "values and mask must have the same length");
87        assert_eq!(out.len(), len, "out must have the same length as values");
88
89        let chunks = mask.chunks();
90        let chunks_count = len / 64;
91        let remainder = len % 64;
92
93        for (chunk_idx, src_chunk) in chunks.iter().enumerate() {
94            if let Some(idx) = chunk(&values, out, &mut f, src_chunk, chunk_idx * 64, 64) {
95                return Err(idx);
96            }
97        }
98        if remainder != 0
99            && let Some(idx) = chunk(
100                &values,
101                out,
102                &mut f,
103                chunks.remainder_bits(),
104                chunks_count * 64,
105                remainder,
106            )
107        {
108            return Err(idx);
109        }
110        Ok(())
111    }
112
113    /// Apply `f(value)` lane-by-lane with **no validity awareness at all** — every
114    /// closure invocation is treated as "happened", regardless of whether the lane
115    /// is null. Use this only when the input is known non-nullable.
116    ///
117    /// # Panics
118    ///
119    /// Panics if `out.len() != self.len()`.
120    #[inline]
121    fn map_into<R, F>(self, out: &mut [MaybeUninit<R>], mut f: F)
122    where
123        F: FnMut(Self::Item) -> R,
124    {
125        #[inline(always)]
126        fn chunk<S, R, F>(
127            values: &S,
128            out: &mut [MaybeUninit<R>],
129            f: &mut F,
130            base: usize,
131            count: usize,
132        ) where
133            S: IndexedSource,
134            F: FnMut(S::Item) -> R,
135        {
136            for bit_idx in 0..count {
137                let idx = base + bit_idx;
138                // SAFETY: caller guarantees base + count <= len.
139                let val = unsafe { values.get_unchecked(idx) };
140                unsafe { out.get_unchecked_mut(idx).write(f(val)) };
141            }
142        }
143
144        let values = self;
145        let len = values.len();
146        assert_eq!(out.len(), len, "out must have the same length as values");
147
148        let chunks_count = len / CHUNK_LEN;
149        let remainder = len % CHUNK_LEN;
150
151        for chunk_idx in 0..chunks_count {
152            chunk(&values, out, &mut f, chunk_idx * CHUNK_LEN, CHUNK_LEN);
153        }
154        if remainder != 0 {
155            chunk(&values, out, &mut f, chunks_count * CHUNK_LEN, remainder);
156        }
157    }
158
159    /// Fallible map with **no validity awareness at all** — every `None` returned
160    /// by the closure is treated as a failure, even at null lanes.
161    ///
162    /// # Use this only for non-nullable inputs.
163    ///
164    /// For nullable inputs with a fallible closure, use [`try_map_masked_into`] —
165    /// it has the same value-only closure shape (and the same perf win) but
166    /// **correctly suppresses null-lane failures** via per-chunk
167    /// `fail_bits & mask_chunk`.
168    ///
169    /// Using this kernel on a nullable input where a null lane's stored value
170    /// would cause `f` to return `None` will produce a spurious `Err`. This is a
171    /// correctness footgun on purpose — the name and this doc are how the API
172    /// signals "you must know your input has no nulls."
173    ///
174    /// On failure returns `Err(failing_lane_index)`.
175    ///
176    /// [`try_map_masked_into`]: IndexedSourceExt::try_map_masked_into
177    ///
178    /// # Panics
179    ///
180    /// Panics if `out.len() != self.len()`.
181    #[inline]
182    fn try_map_into<R, F>(self, out: &mut [MaybeUninit<R>], mut f: F) -> Result<(), usize>
183    where
184        R: Copy + Default,
185        F: FnMut(Self::Item) -> Option<R>,
186    {
187        /// Returns `true` if any lane in `[base, base+count)` failed (OR-reduced);
188        /// the cold attribution path is called at the kernel level so it can be
189        /// inlined separately for full vs remainder.
190        #[inline(always)]
191        fn chunk<S, R, F>(
192            values: &S,
193            out: &mut [MaybeUninit<R>],
194            f: &mut F,
195            base: usize,
196            count: usize,
197        ) -> bool
198        where
199            S: IndexedSource,
200            R: Copy + Default,
201            F: FnMut(S::Item) -> Option<R>,
202        {
203            let mut fail_acc: u64 = 0;
204            for bit_idx in 0..count {
205                let idx = base + bit_idx;
206                // SAFETY: caller guarantees base + count <= len.
207                let val = unsafe { values.get_unchecked(idx) };
208                let opt = f(val);
209                fail_acc |= opt.is_none() as u64;
210                let result = opt.unwrap_or_default();
211                unsafe { out.get_unchecked_mut(idx).write(result) };
212            }
213            fail_acc != 0
214        }
215
216        let values = self;
217        let len = values.len();
218        assert_eq!(out.len(), len, "out must have the same length as values");
219
220        let chunks_count = len / CHUNK_LEN;
221        let remainder = len % CHUNK_LEN;
222
223        for chunk_idx in 0..chunks_count {
224            let base = chunk_idx * CHUNK_LEN;
225            if chunk(&values, out, &mut f, base, CHUNK_LEN) {
226                return Err(attribute_failure_no_mask(&values, base, CHUNK_LEN, &mut f));
227            }
228        }
229        if remainder != 0 {
230            let base = chunks_count * CHUNK_LEN;
231            if chunk(&values, out, &mut f, base, remainder) {
232                return Err(attribute_failure_no_mask(&values, base, remainder, &mut f));
233            }
234        }
235        Ok(())
236    }
237}
238
239impl<S: IndexedSource> IndexedSourceExt for S {}
240
241/// Shared cold scan: walks a chunk, returns the first lane index where
242/// `lane_fails(bit_idx, value)` returns `true`. Used by
243/// [`attribute_failure_no_mask`].
244///
245/// Caller guarantees `base + chunk_len <= values.len()`.
246#[cold]
247#[inline(never)]
248fn cold_scan<S>(
249    values: &S,
250    base: usize,
251    chunk_len: usize,
252    mut lane_fails: impl FnMut(usize /* bit_idx */, S::Item) -> bool,
253) -> usize
254where
255    S: IndexedSource,
256{
257    for bit_idx in 0..chunk_len {
258        let idx = base + bit_idx;
259        // SAFETY: caller guarantees idx < values.len().
260        let val = unsafe { values.get_unchecked(idx) };
261        if lane_fails(bit_idx, val) {
262            return idx;
263        }
264    }
265    unreachable!("cold_scan called without a failing lane")
266}
267
268/// Cold attribution for the no-mask variant. Replays `f` over the chunk to find
269/// the first lane that returns `None`.
270#[inline]
271fn attribute_failure_no_mask<S, R, F>(values: &S, base: usize, chunk_len: usize, f: &mut F) -> usize
272where
273    S: IndexedSource,
274    F: FnMut(S::Item) -> Option<R>,
275{
276    cold_scan(values, base, chunk_len, |_bit_idx, val| f(val).is_none())
277}
278
279#[cfg(test)]
280#[allow(clippy::cast_possible_truncation)]
281mod tests {
282    use vortex_buffer::BitBuffer;
283    use vortex_buffer::BitBufferMut;
284
285    use super::*;
286
287    fn write_t<T: Copy>(out: Vec<MaybeUninit<T>>) -> Vec<T> {
288        // SAFETY: tests always fully initialize the buffer.
289        unsafe { std::mem::transmute(out) }
290    }
291
292    #[test]
293    fn try_map_masked_into_all_ok() {
294        let values: Vec<u64> = (0..200).collect();
295        let mask = BitBuffer::new_set(200);
296        let mut out = vec![MaybeUninit::<u32>::uninit(); 200];
297        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
298            (v <= u32::MAX as u64).then_some(v as u32)
299        });
300        assert!(res.is_ok());
301        let got = write_t(out);
302        assert_eq!(got, (0..200u32).collect::<Vec<_>>());
303    }
304
305    #[test]
306    fn try_map_masked_into_overflow_fails() {
307        let mut values: Vec<u64> = (0..200).collect();
308        values[137] = (u32::MAX as u64) + 1;
309        let mask = BitBuffer::new_set(200);
310        let mut out = vec![MaybeUninit::<u32>::uninit(); 200];
311        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
312            (v <= u32::MAX as u64).then_some(v as u32)
313        });
314        assert_eq!(res, Err(137));
315    }
316
317    #[test]
318    fn try_map_masked_into_overflow_reports_first_failing_lane() {
319        let mut values: Vec<u64> = (0..200).collect();
320        values[50] = u64::MAX;
321        values[51] = u64::MAX;
322        values[137] = u64::MAX;
323        let mask = BitBuffer::new_set(200);
324        let mut out = vec![MaybeUninit::<u32>::uninit(); 200];
325        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
326            (v <= u32::MAX as u64).then_some(v as u32)
327        });
328        assert_eq!(res, Err(50));
329    }
330
331    #[test]
332    fn try_map_masked_into_value_only_closure_filters_null_overflow() {
333        let mut values: Vec<u64> = (0..200).collect();
334        values[5] = u64::MAX;
335        values[42] = u64::MAX;
336        let mask = {
337            let mut m = BitBufferMut::with_capacity(200);
338            for i in 0..200 {
339                m.append(i != 5 && i != 42);
340            }
341            m.freeze()
342        };
343        let mut out = vec![MaybeUninit::<u32>::uninit(); 200];
344        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
345            (v <= u32::MAX as u64).then_some(v as u32)
346        });
347        assert!(
348            res.is_ok(),
349            "null-lane overflow should be filtered by the cold path"
350        );
351    }
352
353    #[test]
354    fn try_map_masked_into_value_only_closure_reports_first_valid_failure() {
355        let mut values: Vec<u64> = (0..200).collect();
356        values[5] = u64::MAX;
357        values[42] = u64::MAX;
358        values[77] = u64::MAX;
359        values[100] = u64::MAX;
360        let mask = {
361            let mut m = BitBufferMut::with_capacity(200);
362            for i in 0..200 {
363                m.append(i != 5 && i != 42);
364            }
365            m.freeze()
366        };
367        let mut out = vec![MaybeUninit::<u32>::uninit(); 200];
368        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
369            (v <= u32::MAX as u64).then_some(v as u32)
370        });
371        assert_eq!(res, Err(77));
372    }
373
374    #[test]
375    fn try_map_masked_into_null_lane_bypasses_check() {
376        let mut values: Vec<u64> = (0..200).collect();
377        values[5] = u64::MAX;
378        let mask = {
379            let mut m = BitBufferMut::with_capacity(200);
380            for i in 0..200 {
381                m.append(i != 5);
382            }
383            m.freeze()
384        };
385        let mut out = vec![MaybeUninit::<u32>::uninit(); 200];
386        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
387            (v <= u32::MAX as u64).then_some(v as u32)
388        });
389        assert!(res.is_ok());
390        let got = write_t(out);
391        assert_eq!(got[5], 0);
392        assert_eq!(got[6], 6);
393    }
394
395    #[test]
396    fn try_map_masked_into_branchful_matches_branchless() {
397        let mut values: Vec<u64> = (0..130).map(|i| i as u64 * 7).collect();
398        values[2] = u64::MAX;
399        values[65] = u32::MAX as u64;
400        let mask = {
401            let mut m = BitBufferMut::with_capacity(130);
402            for i in 0..130 {
403                m.append(!matches!(i, 2 | 17 | 99));
404            }
405            m.freeze()
406        };
407
408        let mut branchless = vec![MaybeUninit::<u32>::uninit(); 130];
409        let mut branchful = vec![MaybeUninit::<u32>::uninit(); 130];
410        values
411            .as_slice()
412            .try_map_masked_into(&mask, &mut branchless, |v| {
413                (v <= u32::MAX as u64).then_some(v as u32)
414            })
415            .unwrap();
416        values
417            .as_slice()
418            .try_map_masked_into(&mask, &mut branchful, |v| u32::try_from(v).ok())
419            .unwrap();
420
421        assert_eq!(write_t(branchful), write_t(branchless));
422    }
423
424    #[test]
425    fn try_map_masked_into_partial_chunk() {
426        let values: Vec<u64> = (0..130).collect();
427        let mask = BitBuffer::new_set(130);
428        let mut out = vec![MaybeUninit::<u32>::uninit(); 130];
429        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
430            (v <= u32::MAX as u64).then_some(v as u32)
431        });
432        assert!(res.is_ok());
433        let got = write_t(out);
434        assert_eq!(got.len(), 130);
435        assert_eq!(got[129], 129);
436    }
437
438    #[test]
439    fn try_map_masked_into_sliced_mask_unaligned_offset() {
440        let big = BitBuffer::new_set(256);
441        let mask = big.slice(13..143);
442        assert_eq!(mask.len(), 130);
443
444        let values: Vec<u64> = (0..130).collect();
445        let mut out = vec![MaybeUninit::<u32>::uninit(); 130];
446        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
447            (v <= u32::MAX as u64).then_some(v as u32)
448        });
449        assert!(res.is_ok());
450        let got = write_t(out);
451        assert_eq!(got, (0..130u32).collect::<Vec<_>>());
452    }
453
454    #[test]
455    fn try_map_masked_into_sliced_mask_with_overflow() {
456        let big = BitBuffer::new_set(256);
457        let mask = big.slice(13..143);
458        assert_eq!(mask.len(), 130);
459
460        let mut values: Vec<u64> = (0..130).collect();
461        values[77] = u64::MAX;
462        let mut out = vec![MaybeUninit::<u32>::uninit(); 130];
463        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
464            (v <= u32::MAX as u64).then_some(v as u32)
465        });
466        assert_eq!(res, Err(77));
467    }
468
469    #[test]
470    fn try_map_masked_into_sliced_mask_null_lanes() {
471        let mut m = BitBufferMut::with_capacity(256);
472        for i in 0..256 {
473            m.append(i % 3 != 0);
474        }
475        let big = m.freeze();
476        let mask = big.slice(13..143);
477        assert_eq!(mask.len(), 130);
478
479        let mut values: Vec<u64> = (0..130).collect();
480        values[2] = u64::MAX;
481        let mut out = vec![MaybeUninit::<u32>::uninit(); 130];
482        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
483            (v <= u32::MAX as u64).then_some(v as u32)
484        });
485        assert!(res.is_ok(), "null lane should bypass the range check");
486    }
487
488    #[test]
489    fn try_map_masked_into_overflow_in_remainder() {
490        let mut values: Vec<u64> = (0..130).collect();
491        values[129] = (u32::MAX as u64) + 1;
492        let mask = BitBuffer::new_set(130);
493        let mut out = vec![MaybeUninit::<u32>::uninit(); 130];
494        let res = values.as_slice().try_map_masked_into(&mask, &mut out, |v| {
495            (v <= u32::MAX as u64).then_some(v as u32)
496        });
497        assert_eq!(res, Err(129));
498    }
499}
vortex_compute/lane_kernels/map_into.rs

vortex_compute/lane_kernels/
map_into.rs
