Skip to content

Commit

Permalink
Implementing ArrayBuilder for Box<dyn ArrayBuilder>
Browse files Browse the repository at this point in the history
  • Loading branch information
@viirya
viirya committed Nov 22, 2023
1 parent 61da64a commit 9c8aff3
Show file tree
Hide file tree
Showing 3 changed files with 235 additions and 1 deletion.
202 changes: 201 additions & 1 deletion arrow-array/src/builder/generic_list_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -353,7 +353,7 @@ where
#[cfg(test)]
mod tests {
use super::*;
use crate::builder::{Int32Builder, ListBuilder};
use crate::builder::{make_builder, Int32Builder, ListBuilder};
use crate::cast::AsArray;
use crate::types::Int32Type;
use crate::{Array, Int32Array};
Expand Down Expand Up @@ -548,4 +548,204 @@ mod tests {
assert_eq!(elements.null_count(), 1);
assert!(elements.is_null(3));
}

#[test]
fn test_boxed_primitive_aray_builder() {
let values_builder = make_builder(&DataType::Int32, 5);
let mut builder = ListBuilder::new(values_builder);

builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_slice(&[1, 2, 3]);
builder.append(true);

builder
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_slice(&[4, 5, 6]);
builder.append(true);

let arr = builder.finish();
assert_eq!(2, arr.len());

let elements = arr.values().as_primitive::<Int32Type>();
assert_eq!(elements.values(), &[1, 2, 3, 4, 5, 6]);
}

#[test]
fn test_boxed_list_list_array_builder() {
// This test is same as `test_list_list_array_builder` but uses boxed builders.
let values_builder = make_builder(
&DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
10,
);
let mut builder = ListBuilder::new(values_builder);

// [[[1, 2], [3, 4]], [[5, 6, 7], null, [8]], null, [[9, 10]]]
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(1);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(2);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(3);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(4);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.append(true);
builder.append(true);

builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(5);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(6);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(7);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.append(false);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(8);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.append(true);
builder.append(true);

builder.append(false);

builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(9);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.append_value(10);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.append(true);
builder.append(true);

let l1 = builder.finish();

assert_eq!(4, l1.len());
assert_eq!(1, l1.null_count());

assert_eq!(l1.value_offsets(), &[0, 2, 5, 5, 6]);
let l2 = l1.values().as_list::<i32>();

assert_eq!(6, l2.len());
assert_eq!(1, l2.null_count());
assert_eq!(l2.value_offsets(), &[0, 2, 4, 7, 7, 8, 10]);

let i1 = l2.values().as_primitive::<Int32Type>();
assert_eq!(10, i1.len());
assert_eq!(0, i1.null_count());
assert_eq!(i1.values(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
}
}
30 changes: 30 additions & 0 deletions arrow-array/src/builder/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -265,6 +265,36 @@ pub trait ArrayBuilder: Any + Send {
fn into_box_any(self: Box<Self>) -> Box<dyn Any>;
}

impl ArrayBuilder for Box<dyn ArrayBuilder> {
fn len(&self) -> usize {
(**self).len()
}

fn is_empty(&self) -> bool {
(**self).is_empty()
}

fn finish(&mut self) -> ArrayRef {
(**self).finish()
}

fn finish_cloned(&self) -> ArrayRef {
(**self).finish_cloned()
}

fn as_any(&self) -> &dyn Any {
(**self).as_any()
}

fn as_any_mut(&mut self) -> &mut dyn Any {
(**self).as_any_mut()
}

fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
self
}
}

/// Builder for [`ListArray`](crate::array::ListArray)
pub type ListBuilder<T> = GenericListBuilder<i32, T>;

Expand Down
4 changes: 4 additions & 0 deletions arrow-array/src/builder/struct_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,10 @@ pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuilde
DataType::Duration(TimeUnit::Nanosecond) => {
Box::new(DurationNanosecondBuilder::with_capacity(capacity))
}
DataType::List(field) => {
let builder = make_builder(field.data_type(), capacity);
Box::new(ListBuilder::with_capacity(builder, capacity))
}
DataType::Struct(fields) => Box::new(StructBuilder::from_fields(fields.clone(), capacity)),
t => panic!("Data type {t:?} is not currently supported"),
}
Expand Down

0 comments on commit 9c8aff3

Please sign in to comment.