Skip to content

Commit

Permalink
Add LargeListBuilder to make_builder (#5171)
Browse files Browse the repository at this point in the history
  • Loading branch information
@viirya
viirya authored Dec 6, 2023
1 parent ea3aca1 commit 32b68f7
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 37 deletions.
91 changes: 54 additions & 37 deletions arrow-array/src/builder/generic_list_builder.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -584,14 +584,31 @@ mod tests {
&DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
10,
);
let mut builder = ListBuilder::new(values_builder);
test_boxed_generic_list_generic_list_array_builder::<i32>(values_builder);
}

#[test]
fn test_boxed_large_list_large_list_array_builder() {
// This test is same as `test_list_list_array_builder` but uses boxed builders.
let values_builder = make_builder(
&DataType::LargeList(Arc::new(Field::new("item", DataType::Int32, true))),
10,
);
test_boxed_generic_list_generic_list_array_builder::<i64>(values_builder);
}

fn test_boxed_generic_list_generic_list_array_builder<O: OffsetSizeTrait + PartialEq>(
values_builder: Box<dyn ArrayBuilder>,
) {
let mut builder: GenericListBuilder<O, Box<dyn ArrayBuilder>> =
GenericListBuilder::<O, Box<dyn ArrayBuilder>>::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")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -600,8 +617,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -610,14 +627,14 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -626,8 +643,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -636,16 +653,16 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder.append(true);

builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -654,8 +671,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -664,30 +681,30 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
.expect("should be an Int32Builder")
.expect("should be an (Large)ListBuilder")
.append_value(7);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(false);
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -696,8 +713,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder.append(true);

Expand All @@ -706,8 +723,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -716,8 +733,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.values()
.as_any_mut()
.downcast_mut::<Int32Builder>()
Expand All @@ -726,8 +743,8 @@ mod tests {
builder
.values()
.as_any_mut()
.downcast_mut::<ListBuilder<Box<dyn ArrayBuilder>>>()
.expect("should be an ListBuilder")
.downcast_mut::<GenericListBuilder<O, Box<dyn ArrayBuilder>>>()
.expect("should be an (Large)ListBuilder")
.append(true);
builder.append(true);

Expand All @@ -736,12 +753,12 @@ mod tests {
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!(l1.value_offsets(), &[0, 2, 5, 5, 6].map(O::usize_as));
let l2 = l1.values().as_list::<O>();

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

let i1 = l2.values().as_primitive::<Int32Type>();
assert_eq!(10, i1.len());
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 @@ -173,6 +173,10 @@ pub fn make_builder(datatype: &DataType, capacity: usize) -> Box<dyn ArrayBuilde
let builder = make_builder(field.data_type(), capacity);
Box::new(ListBuilder::with_capacity(builder, capacity))
}
DataType::LargeList(field) => {
let builder = make_builder(field.data_type(), capacity);
Box::new(LargeListBuilder::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 32b68f7

Please sign in to comment.