(arrow-rs) branch master updated: fix: override `size_hint` for `BitIterator` to return the exact remaining size (#6495)

[jeffreyvo]

This is an automated email from the ASF dual-hosted git repository.

jeffreyvo pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 581c6470d fix: override `size_hint` for `BitIterator` to return the 
exact remaining size (#6495)
581c6470d is described below

commit 581c6470d36b12c5234555b5094ab6816044007a
Author: Beihao Zhou <68083940+beihao-z...@users.noreply.github.com>
AuthorDate: Wed Oct 2 08:05:33 2024 -0400

    fix: override `size_hint` for `BitIterator` to return the exact remaining 
size (#6495)
---
 arrow-buffer/src/util/bit_iterator.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/arrow-buffer/src/util/bit_iterator.rs 
b/arrow-buffer/src/util/bit_iterator.rs
index df40a8fba..2c4f246f2 100644
--- a/arrow-buffer/src/util/bit_iterator.rs
+++ b/arrow-buffer/src/util/bit_iterator.rs
@@ -66,6 +66,11 @@ impl<'a> Iterator for BitIterator<'a> {
         self.current_offset += 1;
         Some(v)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining_bits = self.end_offset - self.current_offset;
+        (remaining_bits, Some(remaining_bits))
+    }
 }
 
 impl<'a> ExactSizeIterator for BitIterator<'a> {}
@@ -263,6 +268,30 @@ pub fn try_for_each_valid_idx<E, F: FnMut(usize) -> 
Result<(), E>>(
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_bit_iterator_size_hint() {
+        let mut b = BitIterator::new(&[0b00000011], 0, 2);
+        assert_eq!(
+            b.size_hint(),
+            (2, Some(2)),
+            "Expected size_hint to be (2, Some(2))"
+        );
+
+        b.next();
+        assert_eq!(
+            b.size_hint(),
+            (1, Some(1)),
+            "Expected size_hint to be (1, Some(1)) after one bit consumed"
+        );
+
+        b.next();
+        assert_eq!(
+            b.size_hint(),
+            (0, Some(0)),
+            "Expected size_hint to be (0, Some(0)) after all bits consumed"
+        );
+    }
+
     #[test]
     fn test_bit_iterator() {
         let mask = &[0b00010010, 0b00100011, 0b00000101, 0b00010001, 
0b10010011];