From f0d7d0b8d41e5a13005f15b41c8f09f01e09ebed Mon Sep 17 00:00:00 2001
From: Palladium <psvrishabh@gmail.com>
Date: Sat, 13 Aug 2022 18:41:25 +0530
Subject: [PATCH] Add validation logic for StructBuilder::finish (#2413)

* Validation for field builders lengths in StructBuilder::finish

* Code refactoring

* Adding null buffer validation in StructBuilder

* Removing unnecessary code
---
 arrow/src/array/builder/struct_builder.rs | 65 +++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/arrow/src/array/builder/struct_builder.rs b/arrow/src/array/builder/struct_builder.rs
index 554e3c553db5..59bd78398bfe 100644
--- a/arrow/src/array/builder/struct_builder.rs
+++ b/arrow/src/array/builder/struct_builder.rs
@@ -216,6 +216,8 @@ 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();
@@ -223,6 +225,7 @@ impl StructBuilder {
         }
 
         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)
@@ -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)]
@@ -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();
+    }
 }