Repository: jorgecarleitao/arrow2
Pull request #1022 (2022-05-28 16:41:44)

Improved docs for `bitmap`
No backward incompatible changes, just docs easier to follow

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1022](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (c079528) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/e976c1e8874082838460314dc304369cd4b29ee9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (e976c1e) will **decrease** coverage by `0.00%`.
> The diff coverage is `76.92%`.

```diff
@@            Coverage Diff             @@
##             main    #1022      +/-   ##
==========================================
- Coverage   71.63%   71.62%   -0.01%     
==========================================
  Files         359      359              
  Lines       19904    19915      +11     
==========================================
+ Hits        14259    14265       +6     
- Misses       5645     5650       +5     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/types/bit\_chunk.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL3R5cGVzL2JpdF9jaHVuay5ycw==) | `100.00% <ø> (ø)` | |
| [src/bitmap/immutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC9pbW11dGFibGUucnM=) | `82.79% <73.52%> (-5.01%)` | :arrow_down: |
| [src/bitmap/mutable.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC9tdXRhYmxlLnJz) | `90.25% <100.00%> (ø)` | |
| [src/bitmap/utils/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9tb2QucnM=) | `100.00% <100.00%> (ø)` | |
| [src/io/ipc/read/read\_basic.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3JlYWRfYmFzaWMucnM=) | `67.64% <100.00%> (ø)` | |
| [src/bitmap/utils/slice\_iterator.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9zbGljZV9pdGVyYXRvci5ycw==) | `87.93% <0.00%> (+1.72%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [e976c1e...c079528](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1022?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).


Patch:
diff --git a/src/bitmap/immutable.rs b/src/bitmap/immutable.rs
index c51a7335e5..bee9e62c63 100644
--- a/src/bitmap/immutable.rs
+++ b/src/bitmap/immutable.rs
@@ -2,26 +2,51 @@ use either::Either;
 use std::iter::FromIterator;
 use std::sync::Arc;
 
-use crate::{buffer::bytes::Bytes, trusted_len::TrustedLen};
+use crate::{buffer::bytes::Bytes, error::Error, trusted_len::TrustedLen};
 
 use super::{
     utils::{count_zeros, fmt, get_bit, get_bit_unchecked, BitChunk, BitChunks, BitmapIter},
     MutableBitmap,
 };
 
-/// An immutable container whose API is optimized to handle bitmaps. All quantities on this
-/// container's API are measured in bits.
-/// # Implementation
-/// * memory on this container is sharable across thread boundaries
-/// * Cloning [`Bitmap`] is `O(1)`
-/// * Slicing [`Bitmap`] is `O(1)`
+/// An immutable container semantically equivalent to `Arc<Vec<bool>>` but represented as `Arc<Vec<u8>>` where
+/// each boolean is represented as a single bit.
+///
+/// # Examples
+/// ```
+/// use arrow2::bitmap::{Bitmap, MutableBitmap};
+///
+/// let bitmap = Bitmap::from([true, false, true]);
+/// assert_eq!(bitmap.iter().collect::<Vec<_>>(), vec![true, false, true]);
+///
+/// // creation directly from bytes
+/// let bitmap = Bitmap::try_new(vec![0b00001101], 5).unwrap();
+/// // note: the first bit is the left-most of the first byte
+/// assert_eq!(bitmap.iter().collect::<Vec<_>>(), vec![true, false, true, true, false]);
+/// // we can also get the slice:
+/// assert_eq!(bitmap.as_slice(), ([0b00001101u8].as_ref(), 0, 5));
+/// // debug helps :)
+/// assert_eq!(format!("{:?}", bitmap), "[0b___01101]".to_string());
+///
+/// // it supports copy-on-write semantics (to a `MutableBitmap`)
+/// let bitmap: MutableBitmap = bitmap.into_mut().right().unwrap();
+/// assert_eq!(bitmap, MutableBitmap::from([true, false, true, true, false]));
+///
+/// // slicing is 'O(1)' (data is shared)
+/// let bitmap = Bitmap::try_new(vec![0b00001101], 5).unwrap();
+/// let sliced = bitmap.slice(1, 4);
+/// assert_eq!(sliced.as_slice(), ([0b00001101u8].as_ref(), 1, 4)); // 1 here is the offset:
+/// assert_eq!(format!("{:?}", sliced), "[0b___0110_]".to_string());
+/// // when sliced (or cloned), it is no longer possible to `into_mut`.
+/// let same: Bitmap = sliced.into_mut().left().unwrap();
+/// ```
 #[derive(Clone)]
 pub struct Bitmap {
     bytes: Arc<Bytes<u8>>,
     // both are measured in bits. They are used to bound the bitmap to a region of Bytes.
     offset: usize,
     length: usize,
-    // this is a cache: it must be computed on initialization
+    // this is a cache: it is computed on initialization
     null_count: usize,
 }
 
@@ -45,10 +70,25 @@ impl Bitmap {
         Self::default()
     }
 
-    /// Initializes an new [`Bitmap`] filled with unset values.
+    /// Initializes a new [`Bitmap`] from vector of bytes and a length.
+    /// # Errors
+    /// This function errors iff `length > bytes.len() * 8`
     #[inline]
-    pub fn new_zeroed(length: usize) -> Self {
-        MutableBitmap::from_len_zeroed(length).into()
+    pub fn try_new(bytes: Vec<u8>, length: usize) -> Result<Self, Error> {
+        if length > bytes.len().saturating_mul(8) {
+            return Err(Error::InvalidArgumentError(format!(
+                "The length of the bitmap ({}) must be `<=` to the number of bytes times 8 ({})",
+                length,
+                bytes.len().saturating_mul(8)
+            )));
+        }
+        let null_count = count_zeros(&bytes, 0, length);
+        Ok(Self {
+            length,
+            offset: 0,
+            bytes: Arc::new(bytes.into()),
+            null_count,
+        })
     }
 
     /// Returns the length of the [`Bitmap`].
@@ -63,6 +103,18 @@ impl Bitmap {
         self.len() == 0
     }
 
+    /// Returns a new iterator of `bool` over this bitmap
+    pub fn iter(&self) -> BitmapIter {
+        BitmapIter::new(&self.bytes, self.offset, self.length)
+    }
+
+    /// Returns an iterator over bits in bit chunks [`BitChunk`].
+    ///
+    /// This iterator is useful to operate over multiple bits via e.g. bitwise.
+    pub fn chunks<T: BitChunk>(&self) -> BitChunks<T> {
+        BitChunks::new(&self.bytes, self.offset, self.length)
+    }
+
     /// Creates a new [`Bitmap`] from [`Bytes`] and a length.
     /// # Panic
     /// Panics iff `length <= bytes.len() * 8`
@@ -78,28 +130,22 @@ impl Bitmap {
         }
     }
 
-    /// Creates a new [`Bitmap`] from [`Vec`] and a length.
-    /// This function is `O(1)`
-    /// # Panic
-    /// Panics iff `length <= buffer.len() * 8`
-    #[inline]
-    pub fn from_u8_vec(vec: Vec<u8>, length: usize) -> Self {
-        Bitmap::from_bytes(vec.into(), length)
-    }
-
-    /// Creates a new [`Bitmap`] from a slice and length.
-    /// # Panic
-    /// Panics iff `length <= bytes.len() * 8`
-    #[inline]
-    pub fn from_u8_slice<T: AsRef<[u8]>>(buffer: T, length: usize) -> Self {
-        let buffer = Vec::<u8>::from(buffer.as_ref());
-        Bitmap::from_u8_vec(buffer, length)
-    }
-
-    /// Counts the nulls (unset bits) starting from `offset` bits and for `length` bits.
+    /// Returns the byte slice of this [`Bitmap`].
+    ///
+    /// The returned tuple contains:
+    /// * `.1`: The byte slice, truncated to the start of the first bit. So the start of the slice
+    ///       is within the first 8 bits.
+    /// * `.2`: The start offset in bits on a range `0 <= offsets < 8`.
+    /// * `.3`: The length in number of bits.
     #[inline]
-    pub fn null_count_range(&self, offset: usize, length: usize) -> usize {
-        count_zeros(&self.bytes, self.offset + offset, length)
+    pub fn as_slice(&self) -> (&[u8], usize, usize) {
+        let start = self.offset / 8;
+        let len = (self.offset % 8 + self.length).saturating_add(7) / 8;
+        (
+            &self.bytes[start..start + len],
+            self.offset % 8,
+            self.length,
+        )
     }
 
     /// Returns the number of unset bits on this [`Bitmap`].
@@ -110,9 +156,10 @@ impl Bitmap {
 
     /// Slices `self`, offsetting by `offset` and truncating up to `length` bits.
     /// # Panic
-    /// Panics iff `self.offset + offset + length >= self.bytes.len() * 8`, i.e. if the offset and `length`
+    /// Panics iff `offset + length > self.length`, i.e. if the offset and `length`
     /// exceeds the allocated capacity of `self`.
     #[inline]
+    #[must_use]
     pub fn slice(self, offset: usize, length: usize) -> Self {
         assert!(offset + length <= self.length);
         unsafe { self.slice_unchecked(offset, length) }
@@ -147,16 +194,6 @@ impl Bitmap {
         get_bit(&self.bytes, self.offset + i)
     }
 
-    /// Returns whether the bit at position `i` is set.
-    #[inline]
-    pub fn get(&self, i: usize) -> Option<bool> {
-        if i < self.len() {
-            Some(unsafe { self.get_bit_unchecked(i) })
-        } else {
-            None
-        }
-    }
-
     /// Unsafely returns whether the bit at position `i` is set.
     /// # Safety
     /// Unsound iff `i >= self.len()`.
@@ -196,6 +233,45 @@ impl Bitmap {
             _ => Either::Left(self),
         }
     }
+
+    /// Initializes an new [`Bitmap`] filled with unset values.
+    #[inline]
+    pub fn new_zeroed(length: usize) -> Self {
+        MutableBitmap::from_len_zeroed(length).into()
+    }
+
+    /// Counts the nulls (unset bits) starting from `offset` bits and for `length` bits.
+    #[inline]
+    pub fn null_count_range(&self, offset: usize, length: usize) -> usize {
+        count_zeros(&self.bytes, self.offset + offset, length)
+    }
+
+    /// Creates a new [`Bitmap`] from a slice and length.
+    /// # Panic
+    /// Panics iff `length <= bytes.len() * 8`
+    #[inline]
+    pub fn from_u8_slice<T: AsRef<[u8]>>(slice: T, length: usize) -> Self {
+        Bitmap::try_new(slice.as_ref().to_vec(), length).unwrap()
+    }
+
+    /// Alias for `Bitmap::try_new().unwrap()`
+    /// This function is `O(1)`
+    /// # Panic
+    /// This function panics iff `length <= bytes.len() * 8`
+    #[inline]
+    pub fn from_u8_vec(vec: Vec<u8>, length: usize) -> Self {
+        Bitmap::try_new(vec, length).unwrap()
+    }
+
+    /// Returns whether the bit at position `i` is set.
+    #[inline]
+    pub fn get(&self, i: usize) -> Option<bool> {
+        if i < self.len() {
+            Some(unsafe { self.get_bit_unchecked(i) })
+        } else {
+            None
+        }
+    }
 }
 
 impl<P: AsRef<[bool]>> From<P> for Bitmap {
@@ -213,14 +289,6 @@ impl FromIterator<bool> for Bitmap {
     }
 }
 
-impl Bitmap {
-    /// Returns an iterator over bits in chunks of `T`, which is useful for
-    /// bit operations.
-    pub fn chunks<T: BitChunk>(&self) -> BitChunks<T> {
-        BitChunks::new(&self.bytes, self.offset, self.length)
-    }
-}
-
 impl Bitmap {
     /// Creates a new [`Bitmap`] from an iterator of booleans.
     /// # Safety
@@ -258,26 +326,6 @@ impl Bitmap {
     }
 }
 
-impl Bitmap {
-    /// Returns the byte slice of this Bitmap.
-    ///
-    /// The returned tuple contains:
-    /// .1 -> The byte slice, truncated to the start of the first bit. So the start of the slice
-    ///       is within the first 8 bits.
-    /// .2 -> The start offset in bits, given what described above `0 <= offsets < 8`.
-    /// .3 -> The length in bits.
-    #[inline]
-    pub fn as_slice(&self) -> (&[u8], usize, usize) {
-        let start = self.offset / 8;
-        let len = (self.offset % 8 + self.length).saturating_add(7) / 8;
-        (
-            &self.bytes[start..start + len],
-            self.offset % 8,
-            self.length,
-        )
-    }
-}
-
 impl<'a> IntoIterator for &'a Bitmap {
     type Item = bool;
     type IntoIter = BitmapIter<'a>;
@@ -286,10 +334,3 @@ impl<'a> IntoIterator for &'a Bitmap {
         BitmapIter::<'a>::new(&self.bytes, self.offset, self.length)
     }
 }
-
-impl<'a> Bitmap {
-    /// constructs a new iterator
-    pub fn iter(&'a self) -> BitmapIter<'a> {
-        BitmapIter::<'a>::new(&self.bytes, self.offset, self.length)
-    }
-}
diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs
index c8e5e04649..4cc76a897c 100644
--- a/src/bitmap/mutable.rs
+++ b/src/bitmap/mutable.rs
@@ -264,7 +264,7 @@ impl MutableBitmap {
 impl From<MutableBitmap> for Bitmap {
     #[inline]
     fn from(buffer: MutableBitmap) -> Self {
-        Bitmap::from_bytes(buffer.buffer.into(), buffer.length)
+        Bitmap::try_new(buffer.buffer, buffer.length).unwrap()
     }
 }
 
@@ -272,7 +272,7 @@ impl From<MutableBitmap> for Option<Bitmap> {
     #[inline]
     fn from(buffer: MutableBitmap) -> Self {
         if buffer.null_count() > 0 {
-            Some(Bitmap::from_bytes(buffer.buffer.into(), buffer.length))
+            Some(Bitmap::try_new(buffer.buffer, buffer.length).unwrap())
         } else {
             None
         }
diff --git a/src/bitmap/utils/mod.rs b/src/bitmap/utils/mod.rs
index aba94cca00..55ca42ba14 100644
--- a/src/bitmap/utils/mod.rs
+++ b/src/bitmap/utils/mod.rs
@@ -59,10 +59,12 @@ pub unsafe fn set_bit_unchecked(data: &mut [u8], i: usize, value: bool) {
     *byte = set(*byte, i % 8, value);
 }
 
-/// Returns whether bit at position `i` in `data` is set or not
+/// Returns whether bit at position `i` in `data` is set
+/// # Panic
+/// This function panics iff `i / 8 >= bytes.len()`
 #[inline]
-pub fn get_bit(data: &[u8], i: usize) -> bool {
-    is_set(data[i / 8], i % 8)
+pub fn get_bit(bytes: &[u8], i: usize) -> bool {
+    is_set(bytes[i / 8], i % 8)
 }
 
 /// Returns whether bit at position `i` in `data` is set or not.
diff --git a/src/compute/sort/lex_sort.rs b/src/compute/sort/lex_sort.rs
index 114093f49a..c598cfd3ab 100644
--- a/src/compute/sort/lex_sort.rs
+++ b/src/compute/sort/lex_sort.rs
@@ -198,7 +198,12 @@ pub fn lexsort_to_indices_impl<I: Index>(
 
     if let Some(limit) = limit {
         let limit = limit.min(row_count);
-        let (before, _, _) = values.select_nth_unstable_by(limit, lex_comparator);
+        let before = if limit < row_count {
+            let (before, _, _) = values.select_nth_unstable_by(limit, lex_comparator);
+            before
+        } else {
+            &mut values[..]
+        };
         before.sort_unstable_by(lex_comparator);
         values.truncate(limit);
         values.shrink_to_fit();
diff --git a/src/io/ipc/read/read_basic.rs b/src/io/ipc/read/read_basic.rs
index 5f449917ad..893d62fc4d 100644
--- a/src/io/ipc/read/read_basic.rs
+++ b/src/io/ipc/read/read_basic.rs
@@ -209,7 +209,7 @@ pub fn read_bitmap<R: Read + Seek>(
         read_uncompressed_bitmap(length, bytes, reader)
     }?;
 
-    Ok(Bitmap::from_bytes(buffer.into(), length))
+    Bitmap::try_new(buffer, length)
 }
 
 pub fn read_validity<R: Read + Seek>(
diff --git a/src/types/bit_chunk.rs b/src/types/bit_chunk.rs
index 6466d7cbbb..5bc4e398fd 100644
--- a/src/types/bit_chunk.rs
+++ b/src/types/bit_chunk.rs
@@ -8,7 +8,7 @@ use num_traits::PrimInt;
 use super::NativeType;
 
 /// A chunk of bits. This is used to create masks of a given length
-/// whose width is `1` bit. In `simd_packed` notation, this corresponds to `m1xY`.
+/// whose width is `1` bit. In `portable_simd` notation, this corresponds to `m1xY`.
 pub trait BitChunk:
     super::private::Sealed
     + PrimInt


Test patch:
diff --git a/tests/it/compute/sort/lex_sort.rs b/tests/it/compute/sort/lex_sort.rs
index bf0b146a8f..a24f116dc8 100644
--- a/tests/it/compute/sort/lex_sort.rs
+++ b/tests/it/compute/sort/lex_sort.rs
@@ -5,6 +5,13 @@ fn test_lex_sort_arrays(input: Vec<SortColumn>, expected: Vec<Box<dyn Array>>) {
     let sorted = lexsort::<i32>(&input, None).unwrap();
     assert_eq!(sorted, expected);
 
+    let sorted = lexsort::<i32>(&input, Some(4)).unwrap();
+    let expected = expected
+        .into_iter()
+        .map(|x| x.slice(0, 4))
+        .collect::<Vec<_>>();
+    assert_eq!(sorted, expected);
+
     let sorted = lexsort::<i32>(&input, Some(2)).unwrap();
     let expected = expected
         .into_iter()