Dandandan commented on a change in pull request #8842: URL: https://github.com/apache/arrow/pull/8842#discussion_r536898750
########## 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 { Review comment: Would it be possible/better to have one for slices instead of Vec? ---------------------------------------------------------------- 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]
