arrow_array/
trusted_len.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow_buffer::{bit_util, ArrowNativeType, Buffer, MutableBuffer};
19
20/// Creates two [`Buffer`]s from an iterator of `Option`.
21/// The first buffer corresponds to a bitmap buffer, the second one
22/// corresponds to a values buffer.
23/// # Safety
24/// The caller must ensure that `iterator` is `TrustedLen`.
25#[inline]
26pub(crate) unsafe fn trusted_len_unzip<I, P, T>(iterator: I) -> (Buffer, Buffer)
27where
28    T: ArrowNativeType,
29    P: std::borrow::Borrow<Option<T>>,
30    I: Iterator<Item = P>,
31{
32    let (_, upper) = iterator.size_hint();
33    let upper = upper.expect("trusted_len_unzip requires an upper limit");
34    let len = upper * std::mem::size_of::<T>();
35
36    let mut null = MutableBuffer::from_len_zeroed(upper.saturating_add(7) / 8);
37    let mut buffer = MutableBuffer::new(len);
38
39    let dst_null = null.as_mut_ptr();
40    let mut dst = buffer.as_mut_ptr() as *mut T;
41    for (i, item) in iterator.enumerate() {
42        let item = item.borrow();
43        if let Some(item) = item {
44            std::ptr::write(dst, *item);
45            bit_util::set_bit_raw(dst_null, i);
46        } else {
47            std::ptr::write(dst, T::default());
48        }
49        dst = dst.add(1);
50    }
51    assert_eq!(
52        dst.offset_from(buffer.as_ptr() as *mut T) as usize,
53        upper,
54        "Trusted iterator length was not accurately reported"
55    );
56    buffer.set_len(len);
57    (null.into(), buffer.into())
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn trusted_len_unzip_good() {
66        let vec = [Some(1u32), None];
67        let (null, buffer) = unsafe { trusted_len_unzip(vec.iter()) };
68        assert_eq!(null.as_slice(), &[0b00000001]);
69        assert_eq!(buffer.as_slice(), &[1u8, 0, 0, 0, 0, 0, 0, 0]);
70    }
71
72    #[test]
73    #[should_panic(expected = "trusted_len_unzip requires an upper limit")]
74    fn trusted_len_unzip_panic() {
75        let iter = std::iter::repeat(Some(4i32));
76        unsafe { trusted_len_unzip(iter) };
77    }
78}
arrow_array/trusted_len.rs

arrow_array/
trusted_len.rs
