diff --git a/crates/polars-arrow/src/array/growable/binary.rs b/crates/polars-arrow/src/array/growable/binary.rs index bfcbe9927158..b0063bbb1d84 100644 --- a/crates/polars-arrow/src/array/growable/binary.rs +++ b/crates/polars-arrow/src/array/growable/binary.rs @@ -4,7 +4,7 @@ use super::utils::extend_offset_values; use super::Growable; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, BinaryArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::datatypes::ArrowDataType; use crate::offset::{Offset, Offsets}; @@ -12,7 +12,7 @@ use crate::offset::{Offset, Offsets}; pub struct GrowableBinary<'a, O: Offset> { arrays: Vec<&'a BinaryArray>, dtype: ArrowDataType, - validity: Option, + validity: Option, values: Vec, offsets: Offsets, } @@ -49,7 +49,7 @@ impl<'a, O: Offset> GrowableBinary<'a, O> { dtype, offsets.into(), values.into(), - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } @@ -97,7 +97,7 @@ impl<'a, O: Offset> From> for BinaryArray { val.dtype, val.offsets.into(), val.values.into(), - val.validity.map(|v| v.into()), + val.validity.map(|v| v.freeze()), ) } } diff --git a/crates/polars-arrow/src/array/growable/binview.rs b/crates/polars-arrow/src/array/growable/binview.rs index 6c974510fc46..c2f6ed6f5676 100644 --- a/crates/polars-arrow/src/array/growable/binview.rs +++ b/crates/polars-arrow/src/array/growable/binview.rs @@ -8,7 +8,7 @@ use super::Growable; use crate::array::binview::{BinaryViewArrayGeneric, ViewType}; use crate::array::growable::utils::{extend_validity, extend_validity_copies, prepare_validity}; use crate::array::{Array, MutableBinaryViewArray, View}; -use crate::bitmap::{Bitmap, MutableBitmap}; +use crate::bitmap::BitmapBuilder; use crate::buffer::Buffer; use crate::datatypes::ArrowDataType; @@ -16,7 +16,7 @@ use crate::datatypes::ArrowDataType; pub struct GrowableBinaryViewArray<'a, T: ViewType + ?Sized> { arrays: Vec<&'a BinaryViewArrayGeneric>, dtype: ArrowDataType, - validity: Option, + validity: Option, inner: MutableBinaryViewArray, same_buffers: Option<&'a Arc<[Buffer]>>, total_same_buffers_len: usize, // Only valid if same_buffers is Some. @@ -81,14 +81,14 @@ impl<'a, T: ViewType + ?Sized> GrowableBinaryViewArray<'a, T> { self.dtype.clone(), arr.views.into(), buffers.clone(), - self.validity.take().map(Bitmap::from), + self.validity.take().map(BitmapBuilder::freeze), arr.total_bytes_len, self.total_same_buffers_len, ) } } else { arr.freeze_with_dtype(self.dtype.clone()) - .with_validity(self.validity.take().map(Bitmap::from)) + .with_validity(self.validity.take().map(BitmapBuilder::freeze)) } } } diff --git a/crates/polars-arrow/src/array/growable/boolean.rs b/crates/polars-arrow/src/array/growable/boolean.rs index 739622fecd91..20aa0740a298 100644 --- a/crates/polars-arrow/src/array/growable/boolean.rs +++ b/crates/polars-arrow/src/array/growable/boolean.rs @@ -3,15 +3,15 @@ use std::sync::Arc; use super::Growable; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, BooleanArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::datatypes::ArrowDataType; /// Concrete [`Growable`] for the [`BooleanArray`]. pub struct GrowableBoolean<'a> { arrays: Vec<&'a BooleanArray>, dtype: ArrowDataType, - validity: Option, - values: MutableBitmap, + validity: Option, + values: BitmapBuilder, } impl<'a> GrowableBoolean<'a> { @@ -30,7 +30,7 @@ impl<'a> GrowableBoolean<'a> { Self { arrays, dtype, - values: MutableBitmap::with_capacity(capacity), + values: BitmapBuilder::with_capacity(capacity), validity: prepare_validity(use_validity, capacity), } } @@ -41,8 +41,8 @@ impl<'a> GrowableBoolean<'a> { BooleanArray::new( self.dtype.clone(), - values.into(), - validity.map(|v| v.into()), + values.freeze(), + validity.map(|v| v.freeze()), ) } } @@ -55,11 +55,7 @@ impl<'a> Growable<'a> for GrowableBoolean<'a> { let values = array.values(); let (slice, offset, _) = values.as_slice(); - // SAFETY: invariant offset + length <= slice.len() - unsafe { - self.values - .extend_from_slice_unchecked(slice, start + offset, len); - } + self.values.extend_from_slice(slice, start + offset, len); } fn extend_validity(&mut self, additional: usize) { @@ -85,6 +81,10 @@ impl<'a> Growable<'a> for GrowableBoolean<'a> { impl<'a> From> for BooleanArray { fn from(val: GrowableBoolean<'a>) -> Self { - BooleanArray::new(val.dtype, val.values.into(), val.validity.map(|v| v.into())) + BooleanArray::new( + val.dtype, + val.values.freeze(), + val.validity.map(|v| v.freeze()), + ) } } diff --git a/crates/polars-arrow/src/array/growable/dictionary.rs b/crates/polars-arrow/src/array/growable/dictionary.rs index d503e35a3159..2d7f7fe24794 100644 --- a/crates/polars-arrow/src/array/growable/dictionary.rs +++ b/crates/polars-arrow/src/array/growable/dictionary.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use super::{make_growable, Growable}; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, DictionaryArray, DictionaryKey, PrimitiveArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::datatypes::ArrowDataType; /// Concrete [`Growable`] for the [`DictionaryArray`]. @@ -14,7 +14,7 @@ pub struct GrowableDictionary<'a, K: DictionaryKey> { dtype: ArrowDataType, keys: Vec<&'a PrimitiveArray>, key_values: Vec, - validity: Option, + validity: Option, offsets: Vec, values: Box, } @@ -77,7 +77,7 @@ impl<'a, T: DictionaryKey> GrowableDictionary<'a, T> { let keys = PrimitiveArray::::new( T::PRIMITIVE.into(), key_values.into(), - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ); // SAFETY: the invariant of this struct ensures that this is up-held diff --git a/crates/polars-arrow/src/array/growable/fixed_binary.rs b/crates/polars-arrow/src/array/growable/fixed_binary.rs index 3c483121e8b8..da1e598d4be3 100644 --- a/crates/polars-arrow/src/array/growable/fixed_binary.rs +++ b/crates/polars-arrow/src/array/growable/fixed_binary.rs @@ -3,12 +3,12 @@ use std::sync::Arc; use super::Growable; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, FixedSizeBinaryArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; /// Concrete [`Growable`] for the [`FixedSizeBinaryArray`]. pub struct GrowableFixedSizeBinary<'a> { arrays: Vec<&'a FixedSizeBinaryArray>, - validity: Option, + validity: Option, values: Vec, size: usize, // just a cache } @@ -44,7 +44,7 @@ impl<'a> GrowableFixedSizeBinary<'a> { FixedSizeBinaryArray::new( self.arrays[0].dtype().clone(), values.into(), - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } @@ -88,7 +88,7 @@ impl<'a> From> for FixedSizeBinaryArray { FixedSizeBinaryArray::new( val.arrays[0].dtype().clone(), val.values.into(), - val.validity.map(|v| v.into()), + val.validity.map(|v| v.freeze()), ) } } diff --git a/crates/polars-arrow/src/array/growable/fixed_size_list.rs b/crates/polars-arrow/src/array/growable/fixed_size_list.rs index ceee42f49399..5972e24e4ee9 100644 --- a/crates/polars-arrow/src/array/growable/fixed_size_list.rs +++ b/crates/polars-arrow/src/array/growable/fixed_size_list.rs @@ -3,12 +3,12 @@ use std::sync::Arc; use super::{make_growable, Growable}; use crate::array::growable::utils::{extend_validity, extend_validity_copies, prepare_validity}; use crate::array::{Array, FixedSizeListArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; /// Concrete [`Growable`] for the [`FixedSizeListArray`]. pub struct GrowableFixedSizeList<'a> { arrays: Vec<&'a FixedSizeListArray>, - validity: Option, + validity: Option, values: Box + 'a>, size: usize, length: usize, @@ -61,7 +61,7 @@ impl<'a> GrowableFixedSizeList<'a> { self.arrays[0].dtype().clone(), self.length, values, - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } @@ -122,7 +122,7 @@ impl<'a> From> for FixedSizeListArray { val.arrays[0].dtype().clone(), val.length, values, - val.validity.map(|v| v.into()), + val.validity.map(|v| v.freeze()), ) } } diff --git a/crates/polars-arrow/src/array/growable/list.rs b/crates/polars-arrow/src/array/growable/list.rs index 211dfe32563b..8dd070a003e6 100644 --- a/crates/polars-arrow/src/array/growable/list.rs +++ b/crates/polars-arrow/src/array/growable/list.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use super::{make_growable, Growable}; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, ListArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::offset::{Offset, Offsets}; unsafe fn extend_offset_values( @@ -29,7 +29,7 @@ unsafe fn extend_offset_values( /// Concrete [`Growable`] for the [`ListArray`]. pub struct GrowableList<'a, O: Offset> { arrays: Vec<&'a ListArray>, - validity: Option, + validity: Option, values: Box + 'a>, offsets: Offsets, } @@ -68,7 +68,7 @@ impl<'a, O: Offset> GrowableList<'a, O> { self.arrays[0].dtype().clone(), offsets.into(), values, - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } diff --git a/crates/polars-arrow/src/array/growable/primitive.rs b/crates/polars-arrow/src/array/growable/primitive.rs index dc27f6f029a9..c2f3188634d8 100644 --- a/crates/polars-arrow/src/array/growable/primitive.rs +++ b/crates/polars-arrow/src/array/growable/primitive.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use super::Growable; use crate::array::growable::utils::{extend_validity, extend_validity_copies, prepare_validity}; use crate::array::{Array, PrimitiveArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::datatypes::ArrowDataType; use crate::types::NativeType; @@ -11,7 +11,7 @@ use crate::types::NativeType; pub struct GrowablePrimitive<'a, T: NativeType> { dtype: ArrowDataType, arrays: Vec<&'a PrimitiveArray>, - validity: Option, + validity: Option, values: Vec, } @@ -48,7 +48,7 @@ impl<'a, T: NativeType> GrowablePrimitive<'a, T> { PrimitiveArray::::new( self.dtype.clone(), values.into(), - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } @@ -105,6 +105,10 @@ impl<'a, T: NativeType> Growable<'a> for GrowablePrimitive<'a, T> { impl<'a, T: NativeType> From> for PrimitiveArray { #[inline] fn from(val: GrowablePrimitive<'a, T>) -> Self { - PrimitiveArray::::new(val.dtype, val.values.into(), val.validity.map(|v| v.into())) + PrimitiveArray::::new( + val.dtype, + val.values.into(), + val.validity.map(|v| v.freeze()), + ) } } diff --git a/crates/polars-arrow/src/array/growable/structure.rs b/crates/polars-arrow/src/array/growable/structure.rs index 127d5673f660..386c3cc0d470 100644 --- a/crates/polars-arrow/src/array/growable/structure.rs +++ b/crates/polars-arrow/src/array/growable/structure.rs @@ -3,13 +3,13 @@ use std::sync::Arc; use super::{make_growable, Growable}; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, StructArray}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; /// Concrete [`Growable`] for the [`StructArray`]. pub struct GrowableStruct<'a> { arrays: Vec<&'a StructArray>, length: usize, - validity: Option, + validity: Option, values: Vec + 'a>>, } @@ -62,7 +62,7 @@ impl<'a> GrowableStruct<'a> { self.arrays[0].dtype().clone(), self.length, values, - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } @@ -129,7 +129,7 @@ impl<'a> From> for StructArray { val.arrays[0].dtype().clone(), val.length, values, - val.validity.map(|v| v.into()), + val.validity.map(|v| v.freeze()), ) } } diff --git a/crates/polars-arrow/src/array/growable/utf8.rs b/crates/polars-arrow/src/array/growable/utf8.rs index 03c0a8f3f814..070cdc9d85b8 100644 --- a/crates/polars-arrow/src/array/growable/utf8.rs +++ b/crates/polars-arrow/src/array/growable/utf8.rs @@ -4,13 +4,13 @@ use super::utils::extend_offset_values; use super::Growable; use crate::array::growable::utils::{extend_validity, prepare_validity}; use crate::array::{Array, Utf8Array}; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::offset::{Offset, Offsets}; /// Concrete [`Growable`] for the [`Utf8Array`]. pub struct GrowableUtf8<'a, O: Offset> { arrays: Vec<&'a Utf8Array>, - validity: Option, + validity: Option, values: Vec, offsets: Offsets, } @@ -49,7 +49,7 @@ impl<'a, O: Offset> GrowableUtf8<'a, O> { self.arrays[0].dtype().clone(), offsets.into(), values.into(), - validity.map(|v| v.into()), + validity.map(|v| v.freeze()), ) } } diff --git a/crates/polars-arrow/src/array/growable/utils.rs b/crates/polars-arrow/src/array/growable/utils.rs index 0b115fda95ca..cb7a8a47fb6f 100644 --- a/crates/polars-arrow/src/array/growable/utils.rs +++ b/crates/polars-arrow/src/array/growable/utils.rs @@ -1,5 +1,5 @@ use crate::array::Array; -use crate::bitmap::MutableBitmap; +use crate::bitmap::BitmapBuilder; use crate::offset::Offset; #[inline] @@ -16,16 +16,16 @@ pub(super) unsafe fn extend_offset_values( buffer.extend_from_slice(new_values); } -pub(super) fn prepare_validity(use_validity: bool, capacity: usize) -> Option { +pub(super) fn prepare_validity(use_validity: bool, capacity: usize) -> Option { if use_validity { - Some(MutableBitmap::with_capacity(capacity)) + Some(BitmapBuilder::with_capacity(capacity)) } else { None } } pub(super) fn extend_validity( - mutable_validity: &mut Option, + mutable_validity: &mut Option, array: &dyn Array, start: usize, len: usize, @@ -36,17 +36,14 @@ pub(super) fn extend_validity( Some(validity) => { debug_assert!(start + len <= validity.len()); let (slice, offset, _) = validity.as_slice(); - // SAFETY: invariant offset + length <= slice.len() - unsafe { - mutable_validity.extend_from_slice_unchecked(slice, start + offset, len); - } + mutable_validity.extend_from_slice(slice, start + offset, len); }, } } } pub(super) fn extend_validity_copies( - mutable_validity: &mut Option, + mutable_validity: &mut Option, array: &dyn Array, start: usize, len: usize, @@ -58,11 +55,8 @@ pub(super) fn extend_validity_copies( Some(validity) => { debug_assert!(start + len <= validity.len()); let (slice, offset, _) = validity.as_slice(); - // SAFETY: invariant offset + length <= slice.len() for _ in 0..copies { - unsafe { - mutable_validity.extend_from_slice_unchecked(slice, start + offset, len); - } + mutable_validity.extend_from_slice(slice, start + offset, len); } }, } diff --git a/crates/polars-arrow/src/bitmap/builder.rs b/crates/polars-arrow/src/bitmap/builder.rs index c507df97c5ba..c16e17b03db0 100644 --- a/crates/polars-arrow/src/bitmap/builder.rs +++ b/crates/polars-arrow/src/bitmap/builder.rs @@ -1,3 +1,5 @@ +use polars_utils::slice::load_padded_le_u64; + use crate::bitmap::{Bitmap, MutableBitmap}; use crate::storage::SharedStorage; @@ -5,8 +7,8 @@ use crate::storage::SharedStorage; #[derive(Default, Clone)] pub struct BitmapBuilder { buf: u64, - len: usize, - cap: usize, + len: usize, // Length in bits. + cap: usize, // Capacity in bits. set_bits: usize, bytes: Vec, } @@ -47,7 +49,7 @@ impl BitmapBuilder { #[inline(never)] fn reserve_slow(&mut self, additional: usize) { let bytes_needed = (self.len + additional).div_ceil(64) * 8; - self.bytes.reserve(bytes_needed - self.bytes.capacity()); + self.bytes.reserve(bytes_needed - self.bytes.len()); let words_available = self.bytes.capacity() / 8; self.cap = words_available * 64; } @@ -58,6 +60,17 @@ impl BitmapBuilder { unsafe { self.push_unchecked(x) } } + /// Does not update len/set_bits, simply writes to the output buffer. + /// # Safety + /// self.bytes.len() + 8 <= self.bytes.capacity() must hold. + #[inline(always)] + unsafe fn flush_word_unchecked(&mut self, w: u64) { + let cur_len = self.bytes.len(); + let p = self.bytes.as_mut_ptr().add(cur_len).cast::(); + p.write_unaligned(w.to_le()); + self.bytes.set_len(cur_len + 8); + } + /// # Safety /// self.len() < self.capacity() must hold. #[inline(always)] @@ -66,14 +79,133 @@ impl BitmapBuilder { self.buf |= (x as u64) << (self.len % 64); self.len += 1; if self.len % 64 == 0 { - let p = self.bytes.as_mut_ptr().add(self.bytes.len()).cast::(); - p.write_unaligned(self.buf.to_le()); - self.bytes.set_len(self.bytes.len() + 8); + self.flush_word_unchecked(self.buf); self.set_bits += self.buf.count_ones() as usize; self.buf = 0; } } + pub fn extend_constant(&mut self, length: usize, value: bool) { + // Fast path if the extension still fits in buf with room left to spare. + let bits_in_buf = self.len % 64; + if bits_in_buf + length < 64 { + let bit_block = ((value as u64) << length) - (value as u64); + self.buf |= bit_block << bits_in_buf; + self.len += length; + return; + } + + unsafe { + let value_spread = if value { u64::MAX } else { 0 }; // Branchless neg. + + // Extend and flush current buf. + self.reserve(length); + let ext_buf = self.buf | (value_spread << bits_in_buf); + self.flush_word_unchecked(ext_buf); + self.set_bits += ext_buf.count_ones() as usize; + + // Write complete words. + let remaining_bits = length - (64 - bits_in_buf); + let remaining_words = remaining_bits / 64; + for _ in 0..remaining_words { + self.flush_word_unchecked(value_spread); + } + self.set_bits += (remaining_words * 64) & value_spread as usize; + + // Put remainder in buf and update length. + self.buf = ((value as u64) << (remaining_bits % 64)) - (value as u64); + self.len += length; + } + } + + /// Pushes the first length bits from the given word, assuming the rest of + /// the bits are zero. + /// # Safety + /// self.len + length <= self.cap and length <= 64 must hold. + pub unsafe fn push_word_with_len_unchecked(&mut self, word: u64, length: usize) { + debug_assert!(self.len + length <= self.cap); + debug_assert!(length <= 64); + debug_assert!(length == 64 || (word >> length) == 0); + let bits_in_buf = self.len % 64; + self.buf |= word << bits_in_buf; + if bits_in_buf + length >= 64 { + self.flush_word_unchecked(self.buf); + self.set_bits += self.buf.count_ones() as usize; + self.buf = if bits_in_buf > 0 { + word >> (64 - bits_in_buf) + } else { + 0 + }; + } + self.len += length; + } + + /// # Safety + /// self.len() + length <= self.capacity() must hold, as well as + /// offset + length <= 8 * slice.len(). + unsafe fn extend_from_slice_unchecked( + &mut self, + mut slice: &[u8], + mut offset: usize, + mut length: usize, + ) { + if length == 0 { + return; + } + + // Deal with slice offset so it's aligned to bytes. + let slice_bit_offset = offset % 8; + if slice_bit_offset > 0 { + let bits_in_first_byte = (8 - slice_bit_offset).min(length); + let first_byte = *slice.get_unchecked(offset / 8) >> slice_bit_offset; + self.push_word_with_len_unchecked( + first_byte as u64 & ((1 << bits_in_first_byte) - 1), + bits_in_first_byte, + ); + length -= bits_in_first_byte; + offset += bits_in_first_byte; + } + slice = slice.get_unchecked(offset / 8..); + + // Write word-by-word. + let bits_in_buf = self.len % 64; + if bits_in_buf > 0 { + while length >= 64 { + let word = u64::from_le_bytes(slice.get_unchecked(0..8).try_into().unwrap()); + self.buf |= word << bits_in_buf; + self.flush_word_unchecked(self.buf); + self.set_bits += self.buf.count_ones() as usize; + self.buf = word >> (64 - bits_in_buf); + self.len += 64; + length -= 64; + slice = slice.get_unchecked(8..); + } + } else { + while length >= 64 { + let word = u64::from_le_bytes(slice.get_unchecked(0..8).try_into().unwrap()); + self.flush_word_unchecked(word); + self.set_bits += word.count_ones() as usize; + self.len += 64; + length -= 64; + slice = slice.get_unchecked(8..); + } + } + + // Just the last word left. + if length > 0 { + let word = load_padded_le_u64(slice); + self.push_word_with_len_unchecked(word & ((1 << length) - 1), length); + } + } + + pub fn extend_from_slice(&mut self, slice: &[u8], offset: usize, length: usize) { + assert!(8 * slice.len() >= offset + length); + self.reserve(length); + unsafe { + self.extend_from_slice_unchecked(slice, offset, length); + } + } + /// # Safety /// May only be called once at the end. unsafe fn finish(&mut self) {