Add is_valid and truncate methods to NullBufferBuilder

arrow-buffer/src/builder/null.rs

@@ -75,7 +75,6 @@ impl NullBufferBuilder {
     /// Creates a new builder from a `MutableBuffer`.
     pub fn new_from_buffer(buffer: MutableBuffer, len: usize) -> Self {
         let capacity = buffer.len() * 8;
-
         assert!(len <= capacity);
 
         let bitmap_builder = Some(BooleanBufferBuilder::new_from_buffer(buffer, len));
@@ -134,6 +133,42 @@ impl NullBufferBuilder {
         }
     }
 
+    /// Returns the capacity of the buffer
+    #[inline]
+    pub fn capacity(&self) -> usize {
+        if let Some(buf) = self.bitmap_builder.as_ref() {
+            buf.capacity()
+        } else {
+            0
+        }
+    }
+
+    /// Gets a bit in the buffer at `index`
+    #[inline]
+    pub fn is_valid(&mut self, index: usize) -> bool {
+        if let Some(buf) = self.bitmap_builder.as_mut() {
+            buf.get_bit(index)
+        } else {
+            true
+        }
+    }
+
+    /// Truncates the builder to the given length
+    ///
+    /// If `len` is greater than the buffer's current length, this has no effect
+    #[inline]
+    pub fn truncate(&mut self, len: usize) {
+        if len > self.len {
+            return;
+        }
+
+        if let Some(buf) = self.bitmap_builder.as_mut() {
+            buf.truncate(len);
+        } else {
+            self.len = len
+        }
+    }
+
     /// Appends a boolean slice into the builder
     /// to indicate the validations of these items.
     pub fn append_slice(&mut self, slice: &[bool]) {
@@ -218,6 +253,7 @@ mod tests {
         builder.append_n_nulls(2);
         builder.append_n_non_nulls(2);
         assert_eq!(6, builder.len());
+        assert_eq!(512, builder.capacity());
 
         let buf = builder.finish().unwrap();
         assert_eq!(&[0b110010_u8], buf.validity());
@@ -230,6 +266,7 @@ mod tests {
         builder.append_n_nulls(2);
         builder.append_slice(&[false, false, false]);
         assert_eq!(6, builder.len());
+        assert_eq!(512, builder.capacity());
 
         let buf = builder.finish().unwrap();
         assert_eq!(&[0b0_u8], buf.validity());
@@ -242,6 +279,7 @@ mod tests {
         builder.append_n_non_nulls(2);
         builder.append_slice(&[true, true, true]);
         assert_eq!(6, builder.len());
+        assert_eq!(0, builder.capacity());
 
         let buf = builder.finish();
         assert!(buf.is_none());
@@ -263,4 +301,35 @@ mod tests {
         let buf = builder.finish().unwrap();
         assert_eq!(&[0b1011_u8], buf.validity());
     }
+
+    #[test]
+    fn test_null_buffer_builder_is_valid() {
+        let mut builder = NullBufferBuilder::new(0);
+        builder.append_n_non_nulls(6);
+        assert!(builder.is_valid(0));
+
+        builder.append_null();
+        assert!(!builder.is_valid(6));
+
+        builder.append_non_null();
+        assert!(builder.is_valid(7));
+    }
+
+    #[test]
+    fn test_null_buffer_builder_truncate() {
+        let mut builder = NullBufferBuilder::new(10);
+        builder.append_n_non_nulls(16);
+        assert_eq!(builder.as_slice(), None);
+        builder.truncate(20);
+        assert_eq!(builder.as_slice(), None);
+        assert_eq!(builder.len(), 16);
+        assert_eq!(builder.capacity(), 0);
+        builder.truncate(14);
+        assert_eq!(builder.as_slice(), None);
+        assert_eq!(builder.len(), 14);
+        builder.append_null();
+        builder.append_non_null();
+        assert_eq!(builder.as_slice().unwrap(), &[0xFF, 0b10111111]);
+        assert_eq!(builder.capacity(), 512);
+    }
 }