Repository: jorgecarleitao/arrow2
Pull request #1067 (2022-06-10 11:51:08)

Improved performance in lexical write (~5%)
It might be interesting to explore a slightly adapted version of `BufStreamingIterator`.

If I am correct it currently writes to an internal buffer, and then copies that buffer to the `W: Write` buffer. If we would accept that final buffer in the `next()` method as argument, we could directly write to that.

Will explore that idea later.

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1067](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (0506acf) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/6919cc3281c3ee7dd3e973a7b303c67277458f20?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (6919cc3) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##             main    #1067   +/-   ##
=======================================
  Coverage   81.30%   81.31%           
=======================================
  Files         363      363           
  Lines       34689    34692    +3     
=======================================
+ Hits        28205    28209    +4     
+ Misses       6484     6483    -1     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/util/lexical.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL3V0aWwvbGV4aWNhbC5ycw==) | `69.23% <100.00%> (+4.01%)` | :arrow_up: |
| [src/bitmap/utils/slice\_iterator.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2JpdG1hcC91dGlscy9zbGljZV9pdGVyYXRvci5ycw==) | `97.56% <0.00%> (-1.22%)` | :arrow_down: |
| [src/io/ipc/read/stream\_async.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3N0cmVhbV9hc3luYy5ycw==) | `74.38% <0.00%> (+0.82%)` | :arrow_up: |
| [src/array/primitive/fmt.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3ByaW1pdGl2ZS9mbXQucnM=) | `86.07% <0.00%> (+1.26%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067?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/1067?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [6919cc3...0506acf](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1067?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/Cargo.toml b/Cargo.toml
index 921facece5..17ddddab43 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -321,3 +321,11 @@ harness = false
 [[bench]]
 name = "slices_iterator"
 harness = false
+
+[[bench]]
+name = "bitmap_assign_ops"
+harness = false
+
+[[bench]]
+name = "assign_ops"
+harness = false
diff --git a/benches/assign_ops.rs b/benches/assign_ops.rs
new file mode 100644
index 0000000000..0376edccf0
--- /dev/null
+++ b/benches/assign_ops.rs
@@ -0,0 +1,29 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use arrow2::{compute::arithmetics::basic::mul_scalar, util::bench_util::create_primitive_array};
+
+fn add_benchmark(c: &mut Criterion) {
+    (10..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+
+        let mut arr_a = create_primitive_array::<f32>(size, 0.2);
+        c.bench_function(&format!("apply_mul 2^{}", log2_size), |b| {
+            b.iter(|| {
+                criterion::black_box(&mut arr_a)
+                    .apply_values(|x| x.iter_mut().for_each(|x| *x *= 1.01));
+                assert!(!arr_a.value(10).is_nan());
+            })
+        });
+
+        let arr_a = create_primitive_array::<f32>(size, 0.2);
+        c.bench_function(&format!("mul 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let a = mul_scalar(criterion::black_box(&arr_a), &1.01f32);
+                assert!(!a.value(10).is_nan());
+            })
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);
diff --git a/benches/bitmap_assign_ops.rs b/benches/bitmap_assign_ops.rs
new file mode 100644
index 0000000000..c962e8ed95
--- /dev/null
+++ b/benches/bitmap_assign_ops.rs
@@ -0,0 +1,47 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+use arrow2::bitmap::{binary_assign, unary_assign};
+use arrow2::bitmap::{Bitmap, MutableBitmap};
+
+fn add_benchmark(c: &mut Criterion) {
+    (10..=20).step_by(2).for_each(|log2_size| {
+        let size = 2usize.pow(log2_size);
+
+        let mut bitmap: MutableBitmap = (0..size).into_iter().map(|x| x % 3 == 0).collect();
+        c.bench_function(&format!("mutablebitmap not 2^{}", log2_size), |b| {
+            b.iter(|| {
+                unary_assign(criterion::black_box(&mut bitmap), |x: u64| !x);
+                assert!(!bitmap.is_empty());
+            })
+        });
+
+        let bitmap: Bitmap = (0..size).into_iter().map(|x| x % 3 == 0).collect();
+        c.bench_function(&format!("bitmap not 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let r = !criterion::black_box(&bitmap);
+                assert!(!r.is_empty());
+            })
+        });
+
+        let mut lhs: MutableBitmap = (0..size).into_iter().map(|x| x % 3 == 0).collect();
+        let rhs: Bitmap = (0..size).into_iter().map(|x| x % 4 == 0).collect();
+        c.bench_function(&format!("mutablebitmap and 2^{}", log2_size), |b| {
+            b.iter(|| {
+                binary_assign(criterion::black_box(&mut lhs), &rhs, |x: u64, y| x & y);
+                assert!(!bitmap.is_empty());
+            })
+        });
+
+        let lhs: Bitmap = (0..size).into_iter().map(|x| x % 3 == 0).collect();
+        let rhs: Bitmap = (0..size).into_iter().map(|x| x % 4 == 0).collect();
+        c.bench_function(&format!("bitmap and 2^{}", log2_size), |b| {
+            b.iter(|| {
+                let r = criterion::black_box(&lhs) & &rhs;
+                assert!(!r.is_empty());
+            })
+        });
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);
diff --git a/src/array/boolean/mod.rs b/src/array/boolean/mod.rs
index f25c044ae3..977529cf5c 100644
--- a/src/array/boolean/mod.rs
+++ b/src/array/boolean/mod.rs
@@ -1,5 +1,5 @@
 use crate::{
-    bitmap::Bitmap,
+    bitmap::{Bitmap, MutableBitmap},
     datatypes::{DataType, PhysicalType},
     error::Error,
 };
@@ -92,6 +92,44 @@ impl BooleanArray {
     pub fn arced(self) -> std::sync::Arc<dyn Array> {
         std::sync::Arc::new(self)
     }
+
+    /// Applies a function `f` to the values of this array, cloning the values
+    /// iff they are being shared with others
+    ///
+    /// This is an API to use clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    /// # Panics
+    /// This function panics if the function modifies the length of the [`MutableBitmap`].
+    pub fn apply_values<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+        let values = std::mem::take(&mut self.values);
+        let mut values = values.make_mut();
+        f(&mut values);
+        if let Some(validity) = &self.validity {
+            assert_eq!(validity.len(), values.len());
+        }
+        self.values = values.into();
+    }
+
+    /// Applies a function `f` to the validity of this array, cloning it
+    /// iff it is being shared.
+    ///
+    /// This is an API to leverage clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    /// # Panics
+    /// This function panics if the function modifies the length of the [`MutableBitmap`].
+    pub fn apply_validity<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+        if let Some(validity) = self.validity.as_mut() {
+            let values = std::mem::take(validity);
+            let mut bitmap = values.make_mut();
+            f(&mut bitmap);
+            assert_eq!(bitmap.len(), self.values.len());
+            *validity = bitmap.into();
+        }
+    }
 }
 
 // must use
diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
index 5270be8ea8..cc419c0645 100644
--- a/src/array/primitive/mod.rs
+++ b/src/array/primitive/mod.rs
@@ -1,7 +1,7 @@
 use crate::{
     bitmap::{
         utils::{zip_validity, ZipValidity},
-        Bitmap,
+        Bitmap, MutableBitmap,
     },
     buffer::Buffer,
     datatypes::*,
@@ -252,6 +252,39 @@ impl<T: NativeType> PrimitiveArray<T> {
         arr
     }
 
+    /// Applies a function `f` to the values of this array, cloning the values
+    /// iff they are being shared with others
+    ///
+    /// This is an API to use clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    pub fn apply_values<F: Fn(&mut [T])>(&mut self, f: F) {
+        let values = std::mem::take(&mut self.values);
+        let mut values = values.make_mut();
+        f(&mut values);
+        self.values = values.into();
+    }
+
+    /// Applies a function `f` to the validity of this array, cloning it
+    /// iff it is being shared.
+    ///
+    /// This is an API to leverage clone-on-write
+    /// # Implementation
+    /// This function is `O(f)` if the data is not being shared, and `O(N) + O(f)`
+    /// if it is being shared (since it results in a `O(N)` memcopy).
+    /// # Panics
+    /// This function panics if the function modifies the length of the [`MutableBitmap`].
+    pub fn apply_validity<F: Fn(&mut MutableBitmap)>(&mut self, f: F) {
+        if let Some(validity) = self.validity.as_mut() {
+            let values = std::mem::take(validity);
+            let mut bitmap = values.make_mut();
+            f(&mut bitmap);
+            assert_eq!(bitmap.len(), self.values.len());
+            *validity = bitmap.into();
+        }
+    }
+
     /// Try to convert this [`PrimitiveArray`] to a [`MutablePrimitiveArray`] via copy-on-write semantics.
     ///
     /// A [`PrimitiveArray`] is backed by a [`Buffer`] and [`Bitmap`] which are essentially `Arc<Vec<_>>`.
diff --git a/src/bitmap/assign_ops.rs b/src/bitmap/assign_ops.rs
new file mode 100644
index 0000000000..31bfa24e2e
--- /dev/null
+++ b/src/bitmap/assign_ops.rs
@@ -0,0 +1,191 @@
+use crate::bitmap::{Bitmap, MutableBitmap};
+
+use super::utils::{BitChunk, BitChunkIterExact, BitChunksExact};
+
+/// Applies a function to every bit of this [`MutableBitmap`] in chunks
+///
+/// This function can be for operations like `!` to a [`MutableBitmap`].
+pub fn unary_assign<T: BitChunk, F: Fn(T) -> T>(bitmap: &mut MutableBitmap, op: F) {
+    let mut chunks = bitmap.bitchunks_exact_mut::<T>();
+
+    chunks.by_ref().for_each(|chunk| {
+        let new_chunk: T = match (chunk as &[u8]).try_into() {
+            Ok(a) => T::from_ne_bytes(a),
+            Err(_) => unreachable!(),
+        };
+        let new_chunk = op(new_chunk);
+        chunk.copy_from_slice(new_chunk.to_ne_bytes().as_ref());
+    });
+
+    if chunks.remainder().is_empty() {
+        return;
+    }
+    let mut new_remainder = T::zero().to_ne_bytes();
+    chunks
+        .remainder()
+        .iter()
+        .enumerate()
+        .for_each(|(index, b)| new_remainder[index] = *b);
+    new_remainder = op(T::from_ne_bytes(new_remainder)).to_ne_bytes();
+
+    let len = chunks.remainder().len();
+    chunks
+        .remainder()
+        .copy_from_slice(&new_remainder.as_ref()[..len]);
+}
+
+impl std::ops::Not for MutableBitmap {
+    type Output = Self;
+
+    #[inline]
+    fn not(mut self) -> Self {
+        unary_assign(&mut self, |a: u64| !a);
+        self
+    }
+}
+
+fn binary_assign_impl<I, T, F>(lhs: &mut MutableBitmap, mut rhs: I, op: F)
+where
+    I: BitChunkIterExact<T>,
+    T: BitChunk,
+    F: Fn(T, T) -> T,
+{
+    let mut lhs_chunks = lhs.bitchunks_exact_mut::<T>();
+
+    lhs_chunks
+        .by_ref()
+        .zip(rhs.by_ref())
+        .for_each(|(lhs, rhs)| {
+            let new_chunk: T = match (lhs as &[u8]).try_into() {
+                Ok(a) => T::from_ne_bytes(a),
+                Err(_) => unreachable!(),
+            };
+            let new_chunk = op(new_chunk, rhs);
+            lhs.copy_from_slice(new_chunk.to_ne_bytes().as_ref());
+        });
+
+    let rem_lhs = lhs_chunks.remainder();
+    let rem_rhs = rhs.remainder();
+    if rem_lhs.is_empty() {
+        return;
+    }
+    let mut new_remainder = T::zero().to_ne_bytes();
+    lhs_chunks
+        .remainder()
+        .iter()
+        .enumerate()
+        .for_each(|(index, b)| new_remainder[index] = *b);
+    new_remainder = op(T::from_ne_bytes(new_remainder), rem_rhs).to_ne_bytes();
+
+    let len = lhs_chunks.remainder().len();
+    lhs_chunks
+        .remainder()
+        .copy_from_slice(&new_remainder.as_ref()[..len]);
+}
+
+/// Apply a bitwise binary operation to a [`MutableBitmap`].
+///
+/// This function can be used for operations like `&=` to a [`MutableBitmap`].
+/// # Panics
+/// This function panics iff `lhs.len() != `rhs.len()`
+pub fn binary_assign<T: BitChunk, F>(lhs: &mut MutableBitmap, rhs: &Bitmap, op: F)
+where
+    F: Fn(T, T) -> T,
+{
+    assert_eq!(lhs.len(), rhs.len());
+
+    let (slice, offset, length) = rhs.as_slice();
+    if offset == 0 {
+        let iter = BitChunksExact::<T>::new(slice, length);
+        binary_assign_impl(lhs, iter, op)
+    } else {
+        let rhs_chunks = rhs.chunks::<T>();
+        binary_assign_impl(lhs, rhs_chunks, op)
+    }
+}
+
+#[inline]
+/// Compute bitwise OR operation in-place
+fn or_assign<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &Bitmap) {
+    if rhs.null_count() == 0 {
+        assert_eq!(lhs.len(), rhs.len());
+        lhs.clear();
+        lhs.extend_constant(rhs.len(), true);
+    } else if rhs.null_count() == rhs.len() {
+        // bitmap remains
+    } else {
+        binary_assign(lhs, rhs, |x: T, y| x | y)
+    }
+}
+
+impl<'a, 'b> std::ops::BitOrAssign<&'a Bitmap> for &'b mut MutableBitmap {
+    #[inline]
+    fn bitor_assign(&mut self, rhs: &'a Bitmap) {
+        or_assign::<u64>(self, rhs)
+    }
+}
+
+impl<'a, 'b> std::ops::BitOr<&'a Bitmap> for MutableBitmap {
+    type Output = Self;
+
+    #[inline]
+    fn bitor(mut self, rhs: &'a Bitmap) -> Self {
+        or_assign::<u64>(&mut self, rhs);
+        self
+    }
+}
+
+#[inline]
+/// Compute bitwise `&` between `lhs` and `rhs`, assigning it to `lhs`
+fn and_assign<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &Bitmap) {
+    if rhs.null_count() == 0 {
+        // bitmap remains
+    }
+    if rhs.null_count() == rhs.len() {
+        assert_eq!(lhs.len(), rhs.len());
+        lhs.clear();
+        lhs.extend_constant(rhs.len(), false);
+    } else {
+        binary_assign(lhs, rhs, |x: T, y| x & y)
+    }
+}
+
+impl<'a, 'b> std::ops::BitAndAssign<&'a Bitmap> for &'b mut MutableBitmap {
+    #[inline]
+    fn bitand_assign(&mut self, rhs: &'a Bitmap) {
+        and_assign::<u64>(self, rhs)
+    }
+}
+
+impl<'a, 'b> std::ops::BitAnd<&'a Bitmap> for MutableBitmap {
+    type Output = Self;
+
+    #[inline]
+    fn bitand(mut self, rhs: &'a Bitmap) -> Self {
+        and_assign::<u64>(&mut self, rhs);
+        self
+    }
+}
+
+#[inline]
+/// Compute bitwise XOR operation
+fn xor_assign<T: BitChunk>(lhs: &mut MutableBitmap, rhs: &Bitmap) {
+    binary_assign(lhs, rhs, |x: T, y| x ^ y)
+}
+
+impl<'a, 'b> std::ops::BitXorAssign<&'a Bitmap> for &'b mut MutableBitmap {
+    #[inline]
+    fn bitxor_assign(&mut self, rhs: &'a Bitmap) {
+        xor_assign::<u64>(self, rhs)
+    }
+}
+
+impl<'a, 'b> std::ops::BitXor<&'a Bitmap> for MutableBitmap {
+    type Output = Self;
+
+    #[inline]
+    fn bitxor(mut self, rhs: &'a Bitmap) -> Self {
+        xor_assign::<u64>(&mut self, rhs);
+        self
+    }
+}
diff --git a/src/bitmap/immutable.rs b/src/bitmap/immutable.rs
index bee9e62c63..45d8e89e4e 100644
--- a/src/bitmap/immutable.rs
+++ b/src/bitmap/immutable.rs
@@ -234,6 +234,17 @@ impl Bitmap {
         }
     }
 
+    /// Converts this [`Bitmap`] into a [`MutableBitmap`], cloning its internal
+    /// buffer if required (clone-on-write).
+    pub fn make_mut(self) -> MutableBitmap {
+        match self.into_mut() {
+            Either::Left(data) => {
+                MutableBitmap::from_vec(data.bytes.as_ref().to_vec(), data.length)
+            }
+            Either::Right(data) => data,
+        }
+    }
+
     /// Initializes an new [`Bitmap`] filled with unset values.
     #[inline]
     pub fn new_zeroed(length: usize) -> Self {
diff --git a/src/bitmap/mod.rs b/src/bitmap/mod.rs
index 1ee5c3f301..a9325286f7 100644
--- a/src/bitmap/mod.rs
+++ b/src/bitmap/mod.rs
@@ -8,4 +8,7 @@ pub use mutable::MutableBitmap;
 mod bitmap_ops;
 pub use bitmap_ops::*;
 
+mod assign_ops;
+pub use assign_ops::*;
+
 pub mod utils;
diff --git a/src/bitmap/mutable.rs b/src/bitmap/mutable.rs
index b9056164bb..d490a74eac 100644
--- a/src/bitmap/mutable.rs
+++ b/src/bitmap/mutable.rs
@@ -5,7 +5,9 @@ use crate::bitmap::utils::{merge_reversed, set_bit_unchecked};
 use crate::error::Error;
 use crate::trusted_len::TrustedLen;
 
-use super::utils::{count_zeros, fmt, get_bit, set, set_bit, BitmapIter};
+use super::utils::{
+    count_zeros, fmt, get_bit, set, set_bit, BitChunk, BitChunksExactMut, BitmapIter,
+};
 use super::Bitmap;
 
 /// A container of booleans. [`MutableBitmap`] is semantically equivalent
@@ -149,7 +151,7 @@ impl MutableBitmap {
         set_bit(self.buffer.as_mut_slice(), index, value)
     }
 
-    /// constructs a new iterator over the values of [`MutableBitmap`].
+    /// constructs a new iterator over the bits of [`MutableBitmap`].
     pub fn iter(&self) -> BitmapIter {
         BitmapIter::new(&self.buffer, 0, self.length)
     }
@@ -308,6 +310,11 @@ impl MutableBitmap {
     pub fn shrink_to_fit(&mut self) {
         self.buffer.shrink_to_fit();
     }
+
+    /// Returns an iterator over mutable slices, [`BitChunksExactMut`]
+    pub(crate) fn bitchunks_exact_mut<T: BitChunk>(&mut self) -> BitChunksExactMut<T> {
+        BitChunksExactMut::new(&mut self.buffer, self.length)
+    }
 }
 
 impl From<MutableBitmap> for Bitmap {
diff --git a/src/bitmap/utils/chunks_exact_mut.rs b/src/bitmap/utils/chunks_exact_mut.rs
new file mode 100644
index 0000000000..afc55ee093
--- /dev/null
+++ b/src/bitmap/utils/chunks_exact_mut.rs
@@ -0,0 +1,60 @@
+use super::BitChunk;
+
+/// An iterator over mutable slices of bytes of exact size.
+///
+/// # Safety
+/// The slices returned by this iterator are guaranteed to have length equal to
+/// `std::mem::size_of::<T>()`.
+#[derive(Debug)]
+pub struct BitChunksExactMut<'a, T: BitChunk> {
+    chunks: std::slice::ChunksExactMut<'a, u8>,
+    remainder: &'a mut [u8],
+    remainder_len: usize,
+    marker: std::marker::PhantomData<T>,
+}
+
+impl<'a, T: BitChunk> BitChunksExactMut<'a, T> {
+    /// Returns a new [`BitChunksExactMut`]
+    #[inline]
+    pub fn new(bitmap: &'a mut [u8], length: usize) -> Self {
+        let size_of = std::mem::size_of::<T>();
+
+        let split = (length / 8 / size_of) * size_of;
+        let (chunks, remainder) = bitmap.split_at_mut(split);
+        let remainder_len = length - chunks.len() * 8;
+
+        let chunks = chunks.chunks_exact_mut(size_of);
+        Self {
+            chunks,
+            remainder,
+            remainder_len,
+            marker: std::marker::PhantomData,
+        }
+    }
+
+    /// The remainder slice
+    #[inline]
+    pub fn remainder(&mut self) -> &mut [u8] {
+        self.remainder
+    }
+
+    /// The length of the remainder slice in bits.
+    #[inline]
+    pub fn remainder_len(&mut self) -> usize {
+        self.remainder_len
+    }
+}
+
+impl<'a, T: BitChunk> Iterator for BitChunksExactMut<'a, T> {
+    type Item = &'a mut [u8];
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.chunks.next()
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.chunks.size_hint()
+    }
+}
diff --git a/src/bitmap/utils/mod.rs b/src/bitmap/utils/mod.rs
index 55ca42ba14..e2a5fdcb98 100644
--- a/src/bitmap/utils/mod.rs
+++ b/src/bitmap/utils/mod.rs
@@ -1,5 +1,6 @@
 //! General utilities for bitmaps representing items where LSB is the first item.
 mod chunk_iterator;
+mod chunks_exact_mut;
 mod fmt;
 mod iterator;
 mod slice_iterator;
@@ -9,6 +10,7 @@ use std::convert::TryInto;
 
 pub(crate) use chunk_iterator::merge_reversed;
 pub use chunk_iterator::{BitChunk, BitChunkIterExact, BitChunks, BitChunksExact};
+pub use chunks_exact_mut::BitChunksExactMut;
 pub use fmt::fmt;
 pub use iterator::BitmapIter;
 pub use slice_iterator::SlicesIterator;
diff --git a/src/buffer/immutable.rs b/src/buffer/immutable.rs
index 758de74f29..726099b14d 100644
--- a/src/buffer/immutable.rs
+++ b/src/buffer/immutable.rs
@@ -157,6 +157,17 @@ impl<T: NativeType> Buffer<T> {
             }
         }
     }
+
+    /// Converts this [`Buffer`] to a [`Vec`], cloning the data if needed, also
+    /// known as clone-on-write semantics.
+    ///
+    /// This function is O(1) under the same conditions that [`Self::into_mut`] returns `Vec`.
+    pub fn make_mut(self) -> Vec<T> {
+        match self.into_mut() {
+            Either::Left(data) => data.as_ref().to_vec(),
+            Either::Right(data) => data,
+        }
+    }
 }
 
 impl<T: NativeType> From<Vec<T>> for Buffer<T> {
diff --git a/src/compute/cast/mod.rs b/src/compute/cast/mod.rs
index f236473e88..2cae7a2192 100644
--- a/src/compute/cast/mod.rs
+++ b/src/compute/cast/mod.rs
@@ -79,16 +79,7 @@ pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
     }
 
     match (from_type, to_type) {
-        (
-            Null,
-            Boolean | Int8 | UInt8 | Int16 | UInt16 | Int32 | UInt32 | Float32 | Date32 | Time32(_)
-            | Int64 | UInt64 | Float64 | Date64 | List(_) | Dictionary(..),
-        )
-        | (
-            Boolean | Int8 | UInt8 | Int16 | UInt16 | Int32 | UInt32 | Float32 | Date32 | Time32(_)
-            | Int64 | UInt64 | Float64 | Date64 | List(_) | Dictionary(..),
-            Null,
-        ) => true,
+        (Null, _) | (_, Null) => true,
         (Struct(_), _) => false,
         (_, Struct(_)) => false,
         (List(list_from), List(list_to)) => {
@@ -379,16 +370,7 @@ pub fn cast(array: &dyn Array, to_type: &DataType, options: CastOptions) -> Resu
 
     let as_options = options.with_wrapped(true);
     match (from_type, to_type) {
-        (
-            Null,
-            Boolean | Int8 | UInt8 | Int16 | UInt16 | Int32 | UInt32 | Float32 | Date32 | Time32(_)
-            | Int64 | UInt64 | Float64 | Date64 | List(_) | Dictionary(..),
-        )
-        | (
-            Boolean | Int8 | UInt8 | Int16 | UInt16 | Int32 | UInt32 | Float32 | Date32 | Time32(_)
-            | Int64 | UInt64 | Float64 | Date64 | List(_) | Dictionary(..),
-            Null,
-        ) => Ok(new_null_array(to_type.clone(), array.len())),
+        (Null, _) | (_, Null) => Ok(new_null_array(to_type.clone(), array.len())),
         (Struct(_), _) => Err(Error::NotYetImplemented(
             "Cannot cast from struct to other types".to_string(),
         )),
diff --git a/src/io/parquet/write/file.rs b/src/io/parquet/write/file.rs
index 42153a93a2..e93aa9fa67 100644
--- a/src/io/parquet/write/file.rs
+++ b/src/io/parquet/write/file.rs
@@ -8,7 +8,7 @@ use parquet2::write::WriteOptions as FileWriteOptions;
 use crate::datatypes::Schema;
 use crate::error::{Error, Result};
 
-use super::{schema::schema_to_metadata_key, to_parquet_schema, WriteOptions};
+use super::{schema::schema_to_metadata_key, to_parquet_schema, ThriftFileMetaData, WriteOptions};
 
 /// Attaches [`Schema`] to `key_value_metadata`
 pub fn add_arrow_schema(
@@ -87,4 +87,11 @@ impl<W: Write> FileWriter<W> {
     pub fn into_inner(self) -> W {
         self.writer.into_inner()
     }
+
+    /// Returns the underlying writer and [`ThriftFileMetaData`]
+    /// # Panics
+    /// This function panics if [`Self::end`] has not yet been called
+    pub fn into_inner_and_metadata(self) -> (W, ThriftFileMetaData) {
+        self.writer.into_inner_and_metadata()
+    }
 }
diff --git a/src/io/parquet/write/mod.rs b/src/io/parquet/write/mod.rs
index 72a90ebb7c..5dfd8ae28b 100644
--- a/src/io/parquet/write/mod.rs
+++ b/src/io/parquet/write/mod.rs
@@ -24,10 +24,13 @@ pub use parquet2::{
     compression::{BrotliLevel, CompressionLevel, CompressionOptions, GzipLevel, ZstdLevel},
     encoding::Encoding,
     fallible_streaming_iterator,
-    metadata::{Descriptor, KeyValue, SchemaDescriptor},
+    metadata::{Descriptor, FileMetaData, KeyValue, SchemaDescriptor, ThriftFileMetaData},
     page::{CompressedDataPage, CompressedPage, EncodedPage},
     schema::types::{FieldInfo, ParquetType, PhysicalType as ParquetPhysicalType},
-    write::{compress, Compressor, DynIter, DynStreamingIterator, RowGroupIter, Version},
+    write::{
+        compress, write_metadata_sidecar, Compressor, DynIter, DynStreamingIterator, RowGroupIter,
+        Version,
+    },
     FallibleStreamingIterator,
 };
 
diff --git a/src/types/bit_chunk.rs b/src/types/bit_chunk.rs
index 5bc4e398fd..0937b50058 100644
--- a/src/types/bit_chunk.rs
+++ b/src/types/bit_chunk.rs
@@ -9,6 +9,8 @@ use super::NativeType;
 
 /// A chunk of bits. This is used to create masks of a given length
 /// whose width is `1` bit. In `portable_simd` notation, this corresponds to `m1xY`.
+///
+/// This (sealed) trait is implemented for [`u8`], [`u16`], [`u32`] and [`u64`].
 pub trait BitChunk:
     super::private::Sealed
     + PrimInt
diff --git a/src/util/lexical.rs b/src/util/lexical.rs
index 970c972e0a..047986cbbe 100644
--- a/src/util/lexical.rs
+++ b/src/util/lexical.rs
@@ -19,8 +19,19 @@ pub fn lexical_to_bytes_mut<N: lexical_core::ToLexical>(n: N, buf: &mut Vec<u8>)
         //      Length of buf is set as written length afterwards. lexical_core
         //      creates a valid string, so doesn't need to be checked.
         let slice = std::slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.capacity());
-        let len = lexical_core::write(n, slice).len();
-        buf.set_len(len);
+
+        //  Safety:
+        //  Omits an unneeded bound check as we just ensured that we reserved `N::FORMATTED_SIZE_DECIMAL`
+        #[cfg(debug_assertions)]
+        {
+            let len = lexical_core::write(n, slice).len();
+            buf.set_len(len);
+        }
+        #[cfg(not(debug_assertions))]
+        {
+            let len = lexical_core::write_unchecked(n, slice).len();
+            buf.set_len(len);
+        }
     }
 }
 


Test patch:
diff --git a/tests/it/array/boolean/mod.rs b/tests/it/array/boolean/mod.rs
index 2dd47d480d..786fdb1dbc 100644
--- a/tests/it/array/boolean/mod.rs
+++ b/tests/it/array/boolean/mod.rs
@@ -131,3 +131,27 @@ fn from_iter() {
     let a: BooleanArray = iter.collect();
     assert_eq!(a.len(), 2);
 }
+
+#[test]
+fn apply_values() {
+    let mut a = BooleanArray::from([Some(true), Some(false), None]);
+    a.apply_values(|x| {
+        let mut a = std::mem::take(x);
+        a = !a;
+        *x = a;
+    });
+    let expected = BooleanArray::from([Some(false), Some(true), None]);
+    assert_eq!(a, expected);
+}
+
+#[test]
+fn apply_validity() {
+    let mut a = BooleanArray::from([Some(true), Some(false), None]);
+    a.apply_validity(|x| {
+        let mut a = std::mem::take(x);
+        a = !a;
+        *x = a;
+    });
+    let expected = BooleanArray::from([None, None, Some(false)]);
+    assert_eq!(a, expected);
+}
diff --git a/tests/it/array/primitive/mod.rs b/tests/it/array/primitive/mod.rs
index 058d503f47..2af71d504a 100644
--- a/tests/it/array/primitive/mod.rs
+++ b/tests/it/array/primitive/mod.rs
@@ -124,3 +124,25 @@ fn into_mut_3() {
     let array = PrimitiveArray::new(DataType::Int32, values, validity);
     assert!(array.into_mut().is_right());
 }
+
+#[test]
+fn apply_values() {
+    let mut a = PrimitiveArray::from([Some(1), Some(2), None]);
+    a.apply_values(|x| {
+        x[0] = 10;
+    });
+    let expected = PrimitiveArray::from([Some(10), Some(2), None]);
+    assert_eq!(a, expected);
+}
+
+#[test]
+fn apply_validity() {
+    let mut a = PrimitiveArray::from([Some(1), Some(2), None]);
+    a.apply_validity(|x| {
+        let mut a = std::mem::take(x);
+        a = !a;
+        *x = a;
+    });
+    let expected = PrimitiveArray::from([None, None, Some(0)]);
+    assert_eq!(a, expected);
+}
diff --git a/tests/it/bitmap/assign_ops.rs b/tests/it/bitmap/assign_ops.rs
new file mode 100644
index 0000000000..bea2b0dcbf
--- /dev/null
+++ b/tests/it/bitmap/assign_ops.rs
@@ -0,0 +1,42 @@
+use arrow2::bitmap::{binary_assign, unary_assign, Bitmap, MutableBitmap};
+
+#[test]
+fn basics() {
+    let mut b = MutableBitmap::from_iter(std::iter::repeat(true).take(10));
+    unary_assign(&mut b, |x: u8| !x);
+    assert_eq!(
+        b,
+        MutableBitmap::from_iter(std::iter::repeat(false).take(10))
+    );
+
+    let mut b = MutableBitmap::from_iter(std::iter::repeat(true).take(10));
+    let c = Bitmap::from_iter(std::iter::repeat(true).take(10));
+    binary_assign(&mut b, &c, |x: u8, y| x | y);
+    assert_eq!(
+        b,
+        MutableBitmap::from_iter(std::iter::repeat(true).take(10))
+    );
+}
+
+#[test]
+fn fast_paths() {
+    let b = MutableBitmap::from([true, false]);
+    let c = Bitmap::from_iter([true, true]);
+    let b = b & &c;
+    assert_eq!(b, MutableBitmap::from_iter([true, false]));
+
+    let b = MutableBitmap::from([true, false]);
+    let c = Bitmap::from_iter([false, false]);
+    let b = b & &c;
+    assert_eq!(b, MutableBitmap::from_iter([false, false]));
+
+    let b = MutableBitmap::from([true, false]);
+    let c = Bitmap::from_iter([true, true]);
+    let b = b | &c;
+    assert_eq!(b, MutableBitmap::from_iter([true, true]));
+
+    let b = MutableBitmap::from([true, false]);
+    let c = Bitmap::from_iter([false, false]);
+    let b = b | &c;
+    assert_eq!(b, MutableBitmap::from_iter([true, false]));
+}
diff --git a/tests/it/bitmap/mod.rs b/tests/it/bitmap/mod.rs
index af0c030f5c..c04e2c009d 100644
--- a/tests/it/bitmap/mod.rs
+++ b/tests/it/bitmap/mod.rs
@@ -1,3 +1,4 @@
+mod assign_ops;
 mod bitmap_ops;
 mod immutable;
 mod mutable;