Skip to content

Commit

Permalink
Add validation logic for StructBuilder::finish (#2413)
Browse files Browse the repository at this point in the history
* Validation for field builders lengths in StructBuilder::finish

* Code refactoring

* Adding null buffer validation in StructBuilder

* Removing unnecessary code
  • Loading branch information
psvri committed Aug 13, 2022
1 parent 98eeb01 commit f0d7d0b
Showing 1 changed file with 65 additions and 0 deletions.
65 changes: 65 additions & 0 deletions arrow/src/array/builder/struct_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -216,13 +216,16 @@ impl StructBuilder {

/// Builds the `StructArray` and reset this builder.
pub fn finish(&mut self) -> StructArray {
self.validate_content();

let mut child_data = Vec::with_capacity(self.field_builders.len());
for f in &mut self.field_builders {
let arr = f.finish();
child_data.push(arr.into_data());
}

let null_bit_buffer = self.null_buffer_builder.finish();

let builder = ArrayData::builder(DataType::Struct(self.fields.clone()))
.len(self.len)
.child_data(child_data)
Expand All @@ -233,6 +236,28 @@ impl StructBuilder {
let array_data = unsafe { builder.build_unchecked() };
StructArray::from(array_data)
}

/// Constructs and validates contents in the builder to ensure that
/// - fields and field_builders are of equal length
/// - the number of items in individual field_builders are equal to self.len
/// - the number of items in individual field_builders are equal to self.null_buffer_builder.len()
fn validate_content(&self) {
if self.fields.len() != self.field_builders.len() {
panic!("Number of fields is not equal to the number of field_builders.");
}
if !self.field_builders.iter().all(|x| x.len() == self.len) {
panic!("StructBuilder and field_builders are of unequal lengths.");
}
if !self
.field_builders
.iter()
.all(|x| x.len() == self.null_buffer_builder.len())
{
panic!(
"StructBuilder null buffer length and field_builders length are unequal."
);
}
}
}

#[cfg(test)]
Expand Down Expand Up @@ -411,4 +436,44 @@ mod tests {
let mut builder = StructBuilder::new(fields, field_builders);
assert!(builder.field_builder::<BinaryBuilder>(0).is_none());
}

#[test]
#[should_panic(expected = "StructBuilder and field_builders are of unequal lengths.")]
fn test_struct_array_builder_unequal_field_builders_lengths() {
let mut int_builder = Int32Builder::new(10);
let mut bool_builder = BooleanBuilder::new(10);

int_builder.append_value(1);
int_builder.append_value(2);
bool_builder.append_value(true);

let mut fields = Vec::new();
let mut field_builders = Vec::new();
fields.push(Field::new("f1", DataType::Int32, false));
field_builders.push(Box::new(int_builder) as Box<dyn ArrayBuilder>);
fields.push(Field::new("f2", DataType::Boolean, false));
field_builders.push(Box::new(bool_builder) as Box<dyn ArrayBuilder>);

let mut builder = StructBuilder::new(fields, field_builders);
builder.append(true);
builder.append(true);
builder.finish();
}

#[test]
#[should_panic(
expected = "Number of fields is not equal to the number of field_builders."
)]
fn test_struct_array_builder_unequal_field_field_builders() {
let int_builder = Int32Builder::new(10);

let mut fields = Vec::new();
let mut field_builders = Vec::new();
fields.push(Field::new("f1", DataType::Int32, false));
field_builders.push(Box::new(int_builder) as Box<dyn ArrayBuilder>);
fields.push(Field::new("f2", DataType::Boolean, false));

let mut builder = StructBuilder::new(fields, field_builders);
builder.finish();
}
}

0 comments on commit f0d7d0b

Please sign in to comment.