Fix BooleanBufferBuilder::resize

apache · Jan 16, 2022 · c7396ff · c7396ff
1 parent bc506a7
commit c7396ff
Showing 1 changed file with 22 additions and 1 deletion.
diff --git a/arrow/src/array/builder.rs b/arrow/src/array/builder.rs
@@ -372,7 +372,8 @@ impl BooleanBufferBuilder {
     #[inline]
     pub fn resize(&mut self, len: usize) {
         let len_bytes = bit_util::ceil(len, 8);
-        self.buffer.resize(len_bytes, 0)
+        self.buffer.resize(len_bytes, 0);
+        self.len = len;
     }
 
     #[inline]
@@ -2939,6 +2940,26 @@ mod tests {
         assert_eq!(arr1, arr2);
     }
 
+    #[test]
+    fn test_boolean_array_builder_resize() {
+        let mut builder = BooleanBufferBuilder::new(20);
+        builder.append_n(4, true);
+        builder.append_n(7, false);
+        builder.append_n(2, true);
+        builder.resize(20);
+
+        assert_eq!(builder.len, 20);
+        assert_eq!(builder.buffer.as_slice(), &[0b00001111, 0b00011000, 0b00000000]);
+
+        builder.resize(5);
+        assert_eq!(builder.len, 5);
+        assert_eq!(builder.buffer.as_slice(), &[0b00001111]);
+
+        builder.append_n(4, true);
+        assert_eq!(builder.len, 9);
+        assert_eq!(builder.buffer.as_slice(), &[0b11101111, 0b00000001]);
+    }
+
     #[test]
     fn test_boolean_builder_increases_buffer_len() {
         // 00000010 01001000