Skip to content

Commit

Permalink
Move BooleanBufferBuilder and NullBufferBuilder to arrow_buffer
Browse files Browse the repository at this point in the history
  • Loading branch information
tustvold committed Jun 1, 2023
1 parent 30196d8 commit 4fb5b14
Show file tree
Hide file tree
Showing 17 changed files with 85 additions and 204 deletions.
11 changes: 5 additions & 6 deletions arrow-array/src/builder/boolean_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use arrow_buffer::NullBufferBuilder;
use crate::builder::{ArrayBuilder, BooleanBufferBuilder};
use crate::{ArrayRef, BooleanArray};
use arrow_buffer::Buffer;
Expand Down Expand Up @@ -150,7 +150,7 @@ impl BooleanBuilder {
let builder = ArrayData::builder(DataType::Boolean)
.len(len)
.add_buffer(self.values_builder.finish().into_inner())
.null_bit_buffer(null_bit_buffer);
.nulls(null_bit_buffer);

let array_data = unsafe { builder.build_unchecked() };
BooleanArray::from(array_data)
Expand All @@ -159,15 +159,14 @@ impl BooleanBuilder {
/// Builds the [BooleanArray] without resetting the builder.
pub fn finish_cloned(&self) -> BooleanArray {
let len = self.len();
let null_bit_buffer = self
let nulls = self
.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref);
.finish_cloned();
let value_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
let builder = ArrayData::builder(DataType::Boolean)
.len(len)
.add_buffer(value_buffer)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { builder.build_unchecked() };
BooleanArray::from(array_data)
Expand Down
10 changes: 3 additions & 7 deletions arrow-array/src/builder/fixed_size_binary_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::{ArrayBuilder, UInt8BufferBuilder};
use crate::{ArrayRef, FixedSizeBinaryArray};
use arrow_buffer::Buffer;
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType};
use std::any::Any;
Expand Down Expand Up @@ -98,7 +98,7 @@ impl FixedSizeBinaryBuilder {
let array_data_builder =
ArrayData::builder(DataType::FixedSizeBinary(self.value_length))
.add_buffer(self.values_builder.finish())
.null_bit_buffer(self.null_buffer_builder.finish())
.nulls(self.null_buffer_builder.finish())
.len(array_length);
let array_data = unsafe { array_data_builder.build_unchecked() };
FixedSizeBinaryArray::from(array_data)
Expand All @@ -111,11 +111,7 @@ impl FixedSizeBinaryBuilder {
let array_data_builder =
ArrayData::builder(DataType::FixedSizeBinary(self.value_length))
.add_buffer(values_buffer)
.null_bit_buffer(
self.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref),
)
.nulls(self.null_buffer_builder.finish_cloned())
.len(array_length);
let array_data = unsafe { array_data_builder.build_unchecked() };
FixedSizeBinaryArray::from(array_data)
Expand Down
14 changes: 5 additions & 9 deletions arrow-array/src/builder/fixed_size_list_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,9 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::ArrayBuilder;
use crate::{ArrayRef, FixedSizeListArray};
use arrow_buffer::Buffer;
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::{DataType, Field};
use std::any::Any;
Expand Down Expand Up @@ -167,14 +166,14 @@ where
len,
);

let null_bit_buffer = self.null_buffer_builder.finish();
let nulls = self.null_buffer_builder.finish();
let array_data = ArrayData::builder(DataType::FixedSizeList(
Arc::new(Field::new("item", values_data.data_type().clone(), true)),
self.list_len,
))
.len(len)
.add_child_data(values_data)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { array_data.build_unchecked() };

Expand All @@ -195,17 +194,14 @@ where
len,
);

let null_bit_buffer = self
.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref);
let nulls = self.null_buffer_builder.finish_cloned();
let array_data = ArrayData::builder(DataType::FixedSizeList(
Arc::new(Field::new("item", values_data.data_type().clone(), true)),
self.list_len,
))
.len(len)
.add_child_data(values_data)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { array_data.build_unchecked() };

Expand Down
10 changes: 3 additions & 7 deletions arrow-array/src/builder/generic_bytes_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::{ArrayBuilder, BufferBuilder, UInt8BufferBuilder};
use crate::types::{ByteArrayType, GenericBinaryType, GenericStringType};
use crate::{ArrayRef, GenericByteArray, OffsetSizeTrait};
use arrow_buffer::NullBufferBuilder;
use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer};
use arrow_data::ArrayDataBuilder;
use std::any::Any;
Expand Down Expand Up @@ -123,7 +123,7 @@ impl<T: ByteArrayType> GenericByteBuilder<T> {
.len(self.len())
.add_buffer(self.offsets_builder.finish())
.add_buffer(self.value_builder.finish())
.null_bit_buffer(self.null_buffer_builder.finish());
.nulls(self.null_buffer_builder.finish());

self.offsets_builder.append(self.next_offset());
let array_data = unsafe { array_builder.build_unchecked() };
Expand All @@ -139,11 +139,7 @@ impl<T: ByteArrayType> GenericByteBuilder<T> {
.len(self.len())
.add_buffer(offset_buffer)
.add_buffer(value_buffer)
.null_bit_buffer(
self.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref),
);
.nulls(self.null_buffer_builder.finish_cloned());

let array_data = unsafe { array_builder.build_unchecked() };
GenericByteArray::from(array_data)
Expand Down
11 changes: 4 additions & 7 deletions arrow-array/src/builder/generic_list_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::{ArrayBuilder, BufferBuilder};
use crate::{Array, ArrayRef, GenericListArray, OffsetSizeTrait};
use arrow_buffer::Buffer;
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::Field;
use std::any::Any;
Expand Down Expand Up @@ -243,7 +243,7 @@ where
.len(len)
.add_buffer(offset_buffer)
.add_child_data(values_data)
.null_bit_buffer(null_bit_buffer);
.nulls(null_bit_buffer);

let array_data = unsafe { array_data_builder.build_unchecked() };

Expand All @@ -257,10 +257,7 @@ where
let values_data = values_arr.to_data();

let offset_buffer = Buffer::from_slice_ref(self.offsets_builder.as_slice());
let null_bit_buffer = self
.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref);
let nulls = self.null_buffer_builder.finish_cloned();
let field = Arc::new(Field::new(
"item",
values_data.data_type().clone(),
Expand All @@ -271,7 +268,7 @@ where
.len(len)
.add_buffer(offset_buffer)
.add_child_data(values_data)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { array_data_builder.build_unchecked() };

Expand Down
14 changes: 5 additions & 9 deletions arrow-array/src/builder/map_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::{ArrayBuilder, BufferBuilder};
use crate::{Array, ArrayRef, MapArray, StructArray};
use arrow_buffer::Buffer;
use arrow_buffer::{NullBuffer, NullBufferBuilder};
use arrow_data::ArrayData;
use arrow_schema::{ArrowError, DataType, Field};
use std::any::Any;
Expand Down Expand Up @@ -160,20 +160,16 @@ impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
let keys_arr = self.key_builder.finish_cloned();
let values_arr = self.value_builder.finish_cloned();
let offset_buffer = Buffer::from_slice_ref(self.offsets_builder.as_slice());
let null_bit_buffer = self
.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref);

self.finish_helper(keys_arr, values_arr, offset_buffer, null_bit_buffer, len)
let nulls = self.null_buffer_builder.finish_cloned();
self.finish_helper(keys_arr, values_arr, offset_buffer, nulls, len)
}

fn finish_helper(
&self,
keys_arr: Arc<dyn Array>,
values_arr: Arc<dyn Array>,
offset_buffer: Buffer,
null_bit_buffer: Option<Buffer>,
nulls: Option<NullBuffer>,
len: usize,
) -> MapArray {
assert!(
Expand Down Expand Up @@ -205,7 +201,7 @@ impl<K: ArrayBuilder, V: ArrayBuilder> MapBuilder<K, V> {
.len(len)
.add_buffer(offset_buffer)
.add_child_data(struct_array.into_data())
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { array_data.build_unchecked() };

Expand Down
4 changes: 1 addition & 3 deletions arrow-array/src/builder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -148,8 +148,7 @@
//! }
//! ```
mod boolean_buffer_builder;
pub use boolean_buffer_builder::*;
pub use arrow_buffer::BooleanBufferBuilder;

mod boolean_builder;
pub use boolean_builder::*;
Expand All @@ -165,7 +164,6 @@ mod generic_list_builder;
pub use generic_list_builder::*;
mod map_builder;
pub use map_builder::*;
mod null_buffer_builder;
mod primitive_builder;
pub use primitive_builder::*;
mod primitive_dictionary_builder;
Expand Down
13 changes: 5 additions & 8 deletions arrow-array/src/builder/primitive_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::{ArrayBuilder, BufferBuilder};
use crate::types::*;
use crate::{ArrayRef, ArrowPrimitiveType, PrimitiveArray};
use arrow_buffer::NullBufferBuilder;
use arrow_buffer::{Buffer, MutableBuffer};
use arrow_data::ArrayData;
use arrow_schema::DataType;
Expand Down Expand Up @@ -278,11 +278,11 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Builds the [`PrimitiveArray`] and reset this builder.
pub fn finish(&mut self) -> PrimitiveArray<T> {
let len = self.len();
let null_bit_buffer = self.null_buffer_builder.finish();
let nulls = self.null_buffer_builder.finish();
let builder = ArrayData::builder(self.data_type.clone())
.len(len)
.add_buffer(self.values_builder.finish())
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { builder.build_unchecked() };
PrimitiveArray::<T>::from(array_data)
Expand All @@ -291,15 +291,12 @@ impl<T: ArrowPrimitiveType> PrimitiveBuilder<T> {
/// Builds the [`PrimitiveArray`] without resetting the builder.
pub fn finish_cloned(&self) -> PrimitiveArray<T> {
let len = self.len();
let null_bit_buffer = self
.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref);
let nulls = self.null_buffer_builder.finish_cloned();
let values_buffer = Buffer::from_slice_ref(self.values_builder.as_slice());
let builder = ArrayData::builder(self.data_type.clone())
.len(len)
.add_buffer(values_buffer)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { builder.build_unchecked() };
PrimitiveArray::<T>::from(array_data)
Expand Down
14 changes: 5 additions & 9 deletions arrow-array/src/builder/struct_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,9 @@
// specific language governing permissions and limitations
// under the License.

use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::*;
use crate::{Array, ArrayRef, StructArray};
use arrow_buffer::Buffer;
use arrow_buffer::NullBufferBuilder;
use arrow_data::ArrayData;
use arrow_schema::{DataType, Fields, IntervalUnit, TimeUnit};
use std::any::Any;
Expand Down Expand Up @@ -247,12 +246,12 @@ impl StructBuilder {
child_data.push(arr.to_data());
}
let length = self.len();
let null_bit_buffer = self.null_buffer_builder.finish();
let nulls = self.null_buffer_builder.finish();

let builder = ArrayData::builder(DataType::Struct(self.fields.clone()))
.len(length)
.child_data(child_data)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { builder.build_unchecked() };
StructArray::from(array_data)
Expand All @@ -268,15 +267,12 @@ impl StructBuilder {
child_data.push(arr.to_data());
}
let length = self.len();
let null_bit_buffer = self
.null_buffer_builder
.as_slice()
.map(Buffer::from_slice_ref);
let nulls = self.null_buffer_builder.finish_cloned();

let builder = ArrayData::builder(DataType::Struct(self.fields.clone()))
.len(length)
.child_data(child_data)
.null_bit_buffer(null_bit_buffer);
.nulls(nulls);

let array_data = unsafe { builder.build_unchecked() };
StructArray::from(array_data)
Expand Down
4 changes: 2 additions & 2 deletions arrow-array/src/builder/union_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@
// under the License.

use crate::builder::buffer_builder::{Int32BufferBuilder, Int8BufferBuilder};
use crate::builder::null_buffer_builder::NullBufferBuilder;
use crate::builder::BufferBuilder;
use crate::{make_array, ArrowPrimitiveType, UnionArray};
use arrow_buffer::NullBufferBuilder;
use arrow_buffer::{ArrowNativeType, Buffer};
use arrow_data::ArrayDataBuilder;
use arrow_schema::{ArrowError, DataType, Field};
Expand Down Expand Up @@ -292,7 +292,7 @@ impl UnionBuilder {
let arr_data_builder = ArrayDataBuilder::new(data_type.clone())
.add_buffer(buffer)
.len(slots)
.null_bit_buffer(bitmap_builder.finish());
.nulls(bitmap_builder.finish());

let arr_data_ref = unsafe { arr_data_builder.build_unchecked() };
let array_ref = make_array(arr_data_ref);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,7 @@
// specific language governing permissions and limitations
// under the License.

use arrow_buffer::{bit_util, BooleanBuffer, Buffer, MutableBuffer};
use arrow_data::bit_mask;
use crate::{bit_mask, bit_util, BooleanBuffer, Buffer, MutableBuffer};
use std::ops::Range;

/// Builder for [`BooleanBuffer`]
Expand Down Expand Up @@ -221,6 +220,11 @@ impl BooleanBufferBuilder {
let len = std::mem::replace(&mut self.len, 0);
BooleanBuffer::new(buf.into(), 0, len)
}

/// Builds the [BooleanBuffer] without resetting the builder.
pub fn finish_cloned(&self) -> BooleanBuffer {
BooleanBuffer::new(Buffer::from_slice_ref(self.as_slice()), 0, self.len)
}
}

impl From<BooleanBufferBuilder> for Buffer {
Expand Down
Loading

0 comments on commit 4fb5b14

Please sign in to comment.