| 24 | /// The caller must ensure that `iterator` is `TrustedLen`. |
| 25 | #[inline] |
| 26 | pub(crate) unsafe fn trusted_len_unzip<I, P, T>(iterator: I) -> (Buffer, Buffer) |
| 27 | where |
| 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 | unsafe { std::ptr::write(dst, *item) }; |
| 45 | unsafe { bit_util::set_bit_raw(dst_null, i) }; |
| 46 | } else { |
| 47 | unsafe { std::ptr::write(dst, T::default()) }; |
| 48 | } |
| 49 | dst = unsafe { dst.add(1) }; |
| 50 | } |
| 51 | assert_eq!( |
| 52 | unsafe { dst.offset_from(buffer.as_ptr() as *mut T) as usize }, |
| 53 | upper, |
| 54 | "Trusted iterator length was not accurately reported" |
| 55 | ); |
| 56 | unsafe { buffer.set_len(len) }; |
| 57 | (null.into(), buffer.into()) |
| 58 | } |
| 59 | |
| 60 | #[cfg(test)] |
| 61 | mod tests { |