Skip to content

Commit

Permalink
Further buffer constructors (#4402)
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold authored Jun 12, 2023
1 parent c1283f1 commit 2c71135
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 32 deletions.
36 changes: 7 additions & 29 deletions arrow-arith/src/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,9 @@
//! [here](https://doc.rust-lang.org/stable/core/arch/) for more information.
use arrow_array::*;
use arrow_buffer::bit_util::ceil;
use arrow_buffer::buffer::{bitwise_bin_op_helper, bitwise_quaternary_op_helper};
use arrow_buffer::{BooleanBuffer, MutableBuffer, NullBuffer};
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use arrow_buffer::{BooleanBuffer, NullBuffer};
use arrow_schema::ArrowError;

/// Logical 'and' boolean values with Kleene logic
///
Expand Down Expand Up @@ -314,7 +312,7 @@ pub fn not(left: &BooleanArray) -> Result<BooleanArray, ArrowError> {
/// ```
pub fn is_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
let values = match input.nulls() {
None => NullBuffer::new_null(input.len()).into_inner(),
None => BooleanBuffer::new_unset(input.len()),
Some(nulls) => !nulls.inner(),
};

Expand All @@ -333,31 +331,11 @@ pub fn is_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
/// assert_eq!(a_is_not_null, BooleanArray::from(vec![true, true, false]));
/// ```
pub fn is_not_null(input: &dyn Array) -> Result<BooleanArray, ArrowError> {
let len = input.len();

let output = match input.nulls() {
None => {
let len_bytes = ceil(len, 8);
MutableBuffer::new(len_bytes)
.with_bitset(len_bytes, true)
.into()
}
Some(nulls) => nulls.inner().sliced(),
};

let data = unsafe {
ArrayData::new_unchecked(
DataType::Boolean,
len,
None,
None,
0,
vec![output],
vec![],
)
let values = match input.nulls() {
None => BooleanBuffer::new_set(input.len()),
Some(n) => n.inner().clone(),
};

Ok(BooleanArray::from(data))
Ok(BooleanArray::new(values, None))
}

#[cfg(test)]
Expand Down
8 changes: 8 additions & 0 deletions arrow-array/src/array/boolean_array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,14 @@ impl BooleanArray {
Self { values, nulls }
}

/// Create a new [`BooleanArray`] with length `len` consisting only of nulls
pub fn new_null(len: usize) -> Self {
Self {
values: BooleanBuffer::new_unset(len),
nulls: Some(NullBuffer::new_null(len)),
}
}

/// Returns the length of this array.
pub fn len(&self) -> usize {
self.values.len()
Expand Down
17 changes: 17 additions & 0 deletions arrow-buffer/src/buffer/boolean.rs
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,23 @@ impl BooleanBuffer {
}
}

/// Create a new [`BooleanBuffer`] of `length` where all values are `true`
pub fn new_set(length: usize) -> Self {
let mut builder = BooleanBufferBuilder::new(length);
builder.append_n(length, true);
builder.finish()
}

/// Create a new [`BooleanBuffer`] of `length` where all values are `false`
pub fn new_unset(length: usize) -> Self {
let buffer = MutableBuffer::new_null(length).into_buffer();
Self {
buffer,
offset: 0,
len: length,
}
}

/// Invokes `f` with indexes `0..len` collecting the boolean results into a new `BooleanBuffer`
pub fn collect_bool<F: FnMut(usize) -> bool>(len: usize, f: F) -> Self {
let buffer = MutableBuffer::collect_bool(len, f);
Expand Down
14 changes: 11 additions & 3 deletions arrow-buffer/src/buffer/null.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,14 +41,22 @@ impl NullBuffer {

/// Create a new [`NullBuffer`] of length `len` where all values are null
pub fn new_null(len: usize) -> Self {
let buffer = MutableBuffer::new_null(len).into_buffer();
let buffer = BooleanBuffer::new(buffer, 0, len);
Self {
buffer,
buffer: BooleanBuffer::new_unset(len),
null_count: len,
}
}

/// Create a new [`NullBuffer`] of length `len` where all values are valid
///
/// Note: it is more efficient to not set the null buffer if it is known to be all valid
pub fn new_valid(len: usize) -> Self {
Self {
buffer: BooleanBuffer::new_set(len),
null_count: 0,
}
}

/// Create a new [`NullBuffer`] with the provided `buffer` and `null_count`
///
/// # Safety
Expand Down

0 comments on commit 2c71135

Please sign in to comment.