jorgecarleitao commented on a change in pull request #8842: URL: https://github.com/apache/arrow/pull/8842#discussion_r540338798
########## File path: rust/arrow/src/array/array_boolean.rs ########## @@ -0,0 +1,331 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::borrow::Borrow; +use std::iter::{FromIterator, IntoIterator}; +use std::mem; +use std::{any::Any, fmt}; +use std::{convert::From, sync::Arc}; + +use super::*; +use super::{array::print_long_array, raw_pointer::RawPtrBox}; +use crate::buffer::{Buffer, MutableBuffer}; +use crate::memory; +use crate::util::bit_util; + +/// Array of bools +pub struct BooleanArray { + data: ArrayDataRef, + /// Pointer to the value array. The lifetime of this must be <= to the value buffer + /// stored in `data`, so it's safe to store. + raw_values: RawPtrBox<u8>, +} + +impl fmt::Debug for BooleanArray { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "BooleanArray\n[\n")?; + print_long_array(self, f, |array, index, f| { + fmt::Debug::fmt(&array.value(index), f) + })?; + write!(f, "]") + } +} + +impl BooleanArray { + /// Returns the length of this array. + pub fn len(&self) -> usize { + self.data.len() + } + + /// Returns whether this array is empty. + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + + /// Returns a raw pointer to the values of this array. + pub fn raw_values(&self) -> *const u8 { + unsafe { self.raw_values.get().add(self.data.offset()) } + } + + /// Returns a slice for the given offset and length + /// + /// Note this doesn't do any bound checking, for performance reason. + pub fn value_slice(&self, offset: usize, len: usize) -> &[u8] { + let raw = + unsafe { std::slice::from_raw_parts(self.raw_values().add(offset), len) }; + &raw[..] + } + + // Returns a new boolean array builder + pub fn builder(capacity: usize) -> BooleanBuilder { + BooleanBuilder::new(capacity) + } + + /// Returns a `Buffer` holding all the values of this array. + /// + /// Note this doesn't take the offset of this array into account. + pub fn values(&self) -> Buffer { + self.data.buffers()[0].clone() + } + + /// Returns the boolean value at index `i`. + /// + /// Note this doesn't do any bound checking, for performance reason. + pub fn value(&self, i: usize) -> bool { + let offset = i + self.offset(); + unsafe { bit_util::get_bit_raw(self.raw_values.get() as *const u8, offset) } + } +} + +impl Array for BooleanArray { + fn as_any(&self) -> &Any { + self + } + + fn data(&self) -> ArrayDataRef { + self.data.clone() + } + + fn data_ref(&self) -> &ArrayDataRef { + &self.data + } + + /// Returns the total number of bytes of memory occupied by the buffers owned by this [BooleanArray]. + fn get_buffer_memory_size(&self) -> usize { + self.data.get_buffer_memory_size() + } + + /// Returns the total number of bytes of memory occupied physically by this [BooleanArray]. + fn get_array_memory_size(&self) -> usize { + self.data.get_array_memory_size() + mem::size_of_val(self) + } +} + +impl From<Vec<bool>> for BooleanArray { + fn from(data: Vec<bool>) -> Self { + let mut mut_buf = MutableBuffer::new_null(data.len()); + { + let mut_slice = mut_buf.data_mut(); + for (i, b) in data.iter().enumerate() { + if *b { + bit_util::set_bit(mut_slice, i); + } + } + } + let array_data = ArrayData::builder(DataType::Boolean) + .len(data.len()) + .add_buffer(mut_buf.freeze()) + .build(); + BooleanArray::from(array_data) + } +} + +impl From<Vec<Option<bool>>> for BooleanArray { + fn from(data: Vec<Option<bool>>) -> Self { + let data_len = data.len(); + let num_byte = bit_util::ceil(data_len, 8); + let mut null_buf = MutableBuffer::new_null(data.len()); + let mut val_buf = MutableBuffer::new(num_byte).with_bitset(num_byte, false); + + { + let null_slice = null_buf.data_mut(); + let val_slice = val_buf.data_mut(); + + for (i, v) in data.iter().enumerate() { + if let Some(b) = v { + bit_util::set_bit(null_slice, i); + if *b { + bit_util::set_bit(val_slice, i); + } + } + } + } + + let array_data = ArrayData::builder(DataType::Boolean) + .len(data_len) + .add_buffer(val_buf.freeze()) + .null_bit_buffer(null_buf.freeze()) + .build(); + BooleanArray::from(array_data) + } +} + +impl From<ArrayDataRef> for BooleanArray { + fn from(data: ArrayDataRef) -> Self { + assert_eq!( + data.buffers().len(), + 1, + "BooleanArray data should contain a single buffer only (values buffer)" + ); + let raw_values = data.buffers()[0].raw_data(); + assert!( + memory::is_aligned::<u8>(raw_values, mem::align_of::<bool>()), + "memory is not aligned" + ); + Self { + data, + raw_values: RawPtrBox::new(raw_values as *const u8), + } + } +} + +impl<'a> IntoIterator for &'a BooleanArray { + type Item = Option<bool>; + type IntoIter = BooleanIter<'a>; + + fn into_iter(self) -> Self::IntoIter { + BooleanIter::<'a>::new(self) + } +} + +impl<'a> BooleanArray { + /// constructs a new iterator + pub fn iter(&'a self) -> BooleanIter<'a> { + BooleanIter::<'a>::new(&self) + } +} + +impl<Ptr: Borrow<Option<bool>>> FromIterator<Ptr> for BooleanArray { + fn from_iter<I: IntoIterator<Item = Ptr>>(iter: I) -> Self { Review comment: I agree. This was based on the implementation for the primitive types, that also does that. I made this decision on purpose at the time because rust does not offer a `From<FixedSizediter>` type thing afaik. However, in retrospect, we could remove this and instead allow the `Mutable` to grow (there is no reason not to since there will be a bound check anyways). ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
