From ed3cd1ad363dd27964b2dbf8a04332fbc419c083 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 11 Oct 2024 10:26:54 +0200 Subject: [PATCH 01/32] perf: Branchless Parquet Prefiltering --- .../read/deserialize/primitive/integer.rs | 257 +++++++++++- .../read/deserialize/utils/array_chunks.rs | 4 + .../read/deserialize/utils/dict_encoded.rs | 375 ++++++++++++++++++ .../arrow/read/deserialize/utils/filter.rs | 7 + .../src/arrow/read/deserialize/utils/mod.rs | 177 +++++---- 5 files changed, 742 insertions(+), 78 deletions(-) create mode 100644 crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index ff9c0b014b08..5b18935d6855 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -1,7 +1,8 @@ use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::types::NativeType; +use polars_compute::filter::filter_boolean_kernel; use super::super::utils; use super::{ @@ -9,6 +10,7 @@ use super::{ DeltaTranslator, IntoDecoderFunction, PlainDecoderFnCollector, PrimitiveDecoder, UnitDecoderFunction, }; +use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; use crate::parquet::encoding::hybrid_rle::{self, DictionaryTranslator}; use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, Encoding}; use crate::parquet::error::ParquetResult; @@ -19,6 +21,7 @@ use crate::read::deserialize::utils::{ dict_indices_decoder, freeze_validity, BatchableCollector, Decoder, PageValidity, TranslatedHybridRle, }; +use crate::read::Filter; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -358,6 +361,258 @@ where let validity = freeze_validity(validity); Ok(PrimitiveArray::try_new(dtype, values.into(), validity).unwrap()) } + + fn extend_filtered_with_state<'a>( + &mut self, + state: &mut utils::State<'a, Self>, + decoded: &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + struct BitmapGatherer; + + impl HybridRleGatherer for BitmapGatherer { + type Target = MutableBitmap; + + fn target_reserve(&self, target: &mut Self::Target, n: usize) { + target.reserve(n); + } + + fn target_num_elements(&self, target: &Self::Target) -> usize { + target.len() + } + + fn hybridrle_to_target(&self, value: u32) -> ParquetResult { + Ok(value != 0) + } + + fn gather_one(&self, target: &mut Self::Target, value: bool) -> ParquetResult<()> { + target.push(value); + Ok(()) + } + + fn gather_repeated( + &self, + target: &mut Self::Target, + value: bool, + n: usize, + ) -> ParquetResult<()> { + target.extend_constant(n, value); + Ok(()) + } + } + + let num_rows = state.len(); + let mut max_offset = num_rows; + + if let Some(ref filter) = filter { + max_offset = filter.max_offset(); + assert!(filter.max_offset() <= num_rows); + } + + match state.translation { + StateTranslation::Plain(ref mut chunks) => { + let page_validity = state + .page_validity + .take() + .map(|v| { + let mut bm = MutableBitmap::with_capacity(max_offset); + + let gatherer = BitmapGatherer; + + v.clone().gather_n_into(&mut bm, max_offset, &gatherer)?; + + ParquetResult::Ok(bm.freeze()) + }) + .transpose()?; + + let filter = match filter { + None => Bitmap::new_with_value(true, max_offset), + Some(Filter::Range(rng)) => { + let mut bm = MutableBitmap::with_capacity(max_offset); + + bm.extend_constant(rng.start, false); + bm.extend_constant(rng.len(), true); + + bm.freeze() + }, + Some(Filter::Mask(bm)) => bm, + }; + let validity = + page_validity.unwrap_or_else(|| Bitmap::new_with_value(true, max_offset)); + + let start_len = decoded.0.len(); + decode_masked_plain(*chunks, &filter, &validity, &mut decoded.0, self.0.decoder); + debug_assert_eq!(decoded.0.len() - start_len, filter.set_bits()); + if state.is_optional { + let validity = filter_boolean_kernel(&validity, &filter); + decoded.1.extend_from_bitmap(&validity); + } + + chunks.skip_in_place(max_offset); + + Ok(()) + }, + StateTranslation::Dictionary(ref mut indexes) => { + utils::dict_encoded::decode_dict( + indexes.clone(), + state.dict.unwrap(), + state.is_optional, + state.page_validity.clone(), + filter, + &mut decoded.1, + &mut decoded.0, + )?; + + // @NOTE: Needed for compatibility now. + indexes.skip_in_place(max_offset)?; + if let Some(ref mut page_validity) = state.page_validity { + page_validity.skip_in_place(max_offset)?; + } + + Ok(()) + }, + _ => self.extend_filtered_with_state_default(state, decoded, filter), + } + } +} + +fn decode_masked_plain( + values: ArrayChunks<'_, P>, + filter: &Bitmap, + validity: &Bitmap, + target: &mut Vec, + dfn: D, +) where + T: NativeType, + P: ParquetNativeType, + i64: num_traits::AsPrimitive

, + D: DecoderFunction, +{ + /// # Safety + /// + /// All of the below need to hold: + /// - `values.len() >= validity.count_ones()` + /// - it needs to be safe to write to out[..filter.count_ones()] + unsafe fn decode_iter_allow_last( + values: ArrayChunks<'_, P>, + mut filter: u64, + mut validity: u64, + out: *mut T, + dfn: D, + ) -> (usize, usize) + where + T: NativeType, + P: ParquetNativeType, + i64: num_traits::AsPrimitive

, + D: DecoderFunction, + { + let mut num_read = 0; + let mut num_written = 0; + + while filter != 0 { + let offset = filter.trailing_zeros(); + + num_read += (validity & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + validity >>= offset; + + let v = if validity & 1 != 0 { + dfn.decode(unsafe { values.get_unchecked(num_read) }) + } else { + T::zeroed() + }; + + *out.add(num_written) = v; + + num_written += 1; + num_read += (validity & 1) as usize; + + filter >>= offset + 1; + validity >>= 1; + } + + num_read += validity.count_ones() as usize; + + (num_read, num_written) + } + + let start_length = target.len(); + let num_rows = filter.set_bits(); + let num_values = validity.set_bits(); + + assert_eq!(filter.len(), validity.len()); + assert!(values.len() >= num_values); + + let mut filter_iter = filter.fast_iter_u56(); + let mut validity_iter = validity.fast_iter_u56(); + + let mut values = values; + + target.reserve(num_rows); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut num_rows_done = 0; + let mut num_values_done = 0; + + for (f, v) in filter_iter.by_ref().zip(validity_iter.by_ref()) { + let nv = v.count_ones() as usize; + let nr = f.count_ones() as usize; + + // @NOTE: + // If we cannot guarantee that there are more values than we need, we run the branching + // variant and know we can write the remaining values as dummy values because they will be + // None. + if num_values_done + nv == num_values { + // We don't really need to update `values` or `target_ptr` because we won't be + // using them after this. + _ = unsafe { decode_iter_allow_last::(values, f, v, target_ptr, dfn) }; + + num_rows_done += nr; + unsafe { target.set_len(start_length + num_rows_done) }; + target.resize(start_length + num_rows, T::zeroed()); + return; + } + + let mut v = v; + let mut f = f; + let mut num_read = 0; + let mut num_written = 0; + + while f != 0 { + let offset = f.trailing_zeros(); + + num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + v >>= offset; + + unsafe { + *target_ptr.add(num_written) = dfn.decode(values.get_unchecked(num_read)); + } + + num_written += 1; + num_read += (v & 1) as usize; + + f >>= offset + 1; + v >>= 1; + } + + num_read += v.count_ones() as usize; + + values = ArrayChunks { + bytes: unsafe { values.bytes.get_unchecked(num_read..) }, + }; + target_ptr = unsafe { target_ptr.add(num_written) }; + + num_values_done += nv; + num_rows_done += nr; + } + + let (f, fl) = filter_iter.remainder(); + let (v, vl) = validity_iter.remainder(); + + assert_eq!(fl, vl); + + _ = unsafe { decode_iter_allow_last::(values, f, v, target_ptr, dfn) }; + + unsafe { target.set_len(start_length + num_rows) }; } impl utils::DictDecodable for IntDecoder diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs index 0882cbed5cb0..eede2f910ffa 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs @@ -25,6 +25,10 @@ impl<'a, P: ParquetNativeType> ArrayChunks<'a, P> { Some(Self { bytes }) } + pub(crate) unsafe fn get_unchecked(&self, at: usize) -> P { + P::from_le_bytes(*unsafe { self.bytes.get_unchecked(at) }) + } + pub(crate) fn skip_in_place(&mut self, n: usize) { let n = usize::min(self.bytes.len(), n); self.bytes = &self.bytes[n..]; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs new file mode 100644 index 000000000000..7fc3f96d96ae --- /dev/null +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -0,0 +1,375 @@ +use std::ops::Range; + +use arrow::bitmap::{Bitmap, MutableBitmap}; +use bytemuck::Zeroable; +use polars_compute::filter::filter_boolean_kernel; + +use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; +use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; +use crate::parquet::error::ParquetResult; +use crate::read::Filter; + +fn filter_from_range(rng: Range) -> Bitmap { + let mut bm = MutableBitmap::with_capacity(rng.end); + + bm.extend_constant(rng.start, false); + bm.extend_constant(rng.len(), true); + + bm.freeze() +} + +fn decode_page_validity( + mut page_validity: HybridRleDecoder<'_>, + limit: Option, +) -> ParquetResult { + struct BitmapGatherer; + + impl HybridRleGatherer for BitmapGatherer { + type Target = MutableBitmap; + + fn target_reserve(&self, target: &mut Self::Target, n: usize) { + target.reserve(n); + } + + fn target_num_elements(&self, target: &Self::Target) -> usize { + target.len() + } + + fn hybridrle_to_target(&self, value: u32) -> ParquetResult { + Ok(value != 0) + } + + fn gather_one(&self, target: &mut Self::Target, value: bool) -> ParquetResult<()> { + target.push(value); + Ok(()) + } + + fn gather_repeated( + &self, + target: &mut Self::Target, + value: bool, + n: usize, + ) -> ParquetResult<()> { + target.extend_constant(n, value); + Ok(()) + } + } + + let mut bm = MutableBitmap::with_capacity(limit.unwrap_or(page_validity.len())); + + let gatherer = BitmapGatherer; + + match limit { + None => page_validity.gather_into(&mut bm, &gatherer)?, + Some(limit) => page_validity.gather_n_into(&mut bm, limit, &gatherer)?, + } + + ParquetResult::Ok(bm.freeze()) +} + +pub fn decode_dict( + values: HybridRleDecoder<'_>, + dict: &[T], + is_optional: bool, + page_validity: Option>, + filter: Option, + validity: &mut MutableBitmap, + target: &mut Vec, +) -> ParquetResult<()> { + let page_validity = match (page_validity, filter.as_ref()) { + (None, _) => None, + (Some(page_validity), None) => Some(decode_page_validity(page_validity, None)?), + (Some(page_validity), Some(filter)) => Some(decode_page_validity( + page_validity, + Some(filter.max_offset()), + )?), + }; + + if is_optional { + match (page_validity.as_ref(), filter.as_ref()) { + (None, None) => validity.extend_constant(values.len(), true), + (None, Some(f)) => validity.extend_constant(f.num_rows(), true), + (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), + (Some(page_validity), Some(Filter::Range(rng))) => { + validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) + }, + (Some(page_validity), Some(Filter::Mask(mask))) => { + validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, &mask)) + }, + } + } + + match (filter, page_validity) { + (None, None) => decode_non_optional_dict(values, dict, None, target), + (Some(Filter::Range(rng)), None) if rng.start == 0 => { + decode_non_optional_dict(values, dict, Some(rng.end), target) + }, + (None, Some(page_validity)) => decode_optional_dict(values, dict, &page_validity, target), + (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { + decode_optional_dict(values, dict, &page_validity, target) + }, + (Some(Filter::Mask(filter)), None) => decode_masked_dict( + values, + dict, + &filter, + &Bitmap::new_with_value(true, filter.len()), + target, + ), + (Some(Filter::Mask(filter)), Some(page_validity)) => { + decode_masked_dict(values, dict, &filter, &page_validity, target) + }, + (Some(Filter::Range(rng)), None) => decode_masked_dict( + values, + dict, + &filter_from_range(rng.clone()), + &Bitmap::new_with_value(true, rng.end), + target, + ), + (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_dict( + values, + dict, + &filter_from_range(rng.clone()), + &page_validity, + target, + ), + } +} + +pub fn decode_non_optional_dict( + values: HybridRleDecoder<'_>, + dict: &[T], + limit: Option, + target: &mut Vec, +) -> ParquetResult<()> { + let limit = limit.unwrap_or(values.len()); + assert!(limit <= values.len()); + let start_length = target.len(); + let end_length = start_length + limit; + + target.reserve(limit); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut limit = limit; + let mut values = values; + let mut intermediate_buffer = Vec::with_capacity(32); + + while limit > 0 { + let num_added_values = limit.min(32); + values.collect_n_into(&mut intermediate_buffer, num_added_values)?; + + // @TODO: Make this into a chunk by chunk operation. + // let (added_buffer, num_added_values) = values.next_chunk().unwrap(); + + let highest_idx = (0..num_added_values) + .map(|i| intermediate_buffer[i]) + .max() + .unwrap(); + assert!((highest_idx as usize) < dict.len()); + + for (i, v) in intermediate_buffer.iter().enumerate() { + let v = unsafe { dict.get_unchecked(*v as usize) }.clone(); + unsafe { target_ptr.add(i).write(v) }; + } + + unsafe { + target_ptr = target_ptr.add(num_added_values); + } + limit -= num_added_values; + } + + unsafe { + target.set_len(end_length); + } + + Ok(()) +} + +pub fn decode_optional_dict( + values: HybridRleDecoder<'_>, + dict: &[T], + validity: &Bitmap, + target: &mut Vec, +) -> ParquetResult<()> { + let num_valid_values = validity.set_bits(); + + assert!(num_valid_values <= values.len()); + let start_length = target.len(); + + target.reserve(validity.len()); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut validity_iter = validity.fast_iter_u56(); + let mut values = values; + let mut values_buffer = [0u32; 64]; + let mut values_offset = 0; + let mut num_buffered = 0; + let mut buffer_side = false; + let mut intermediate_buffer = Vec::with_capacity(32); + + let mut iter_u32 = |v: u32| { + while num_buffered < v.count_ones() as usize { + intermediate_buffer.clear(); + let num_added_values = values.len().min(32); + values.collect_n_into(&mut intermediate_buffer, num_added_values)?; + + // @TODO: Make this into a chunk by chunk operation. + // let (added_buffer, num_added_values) = values.next_chunk().unwrap(); + + values_buffer[usize::from(buffer_side) * 32..][..num_added_values] + .copy_from_slice(&intermediate_buffer[..num_added_values]); + + let highest_idx = (0..num_added_values) + .map(|i| values_buffer[(values_offset + i) % 64]) + .max() + .unwrap(); + assert!((highest_idx as usize) < dict.len()); + + buffer_side = !buffer_side; + num_buffered += num_added_values; + } + + let mut num_read = 0; + + for i in 0..32 { + let idx = values_buffer[(values_offset + num_read) % 64]; + let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + unsafe { *target_ptr.add(i) = value }; + num_read += ((v >> i) & 1) as usize; + } + + values_offset = (values_offset + num_read) % 64; + num_buffered -= num_read; + target_ptr = unsafe { target_ptr.add(32) }; + + ParquetResult::Ok(()) + }; + let mut iter = |v: u64| { + iter_u32((v & 0xFFFF_FFFF) as u32)?; + iter_u32((v >> 32) as u32)?; + + ParquetResult::Ok(()) + }; + + for v in validity_iter.by_ref() { + iter(v)?; + } + + let (v, _) = validity_iter.remainder(); + iter(v)?; + + unsafe { target.set_len(start_length + validity.len()) }; + + Ok(()) +} + +pub fn decode_masked_dict( + values: HybridRleDecoder<'_>, + dict: &[T], + filter: &Bitmap, + validity: &Bitmap, + target: &mut Vec, +) -> ParquetResult<()> { + let start_length = target.len(); + let num_rows = filter.set_bits(); + let num_valid_values = validity.set_bits(); + + if dict.is_empty() { + assert_eq!(num_valid_values, 0); + target.resize(start_length + num_rows, T::zeroed()); + } + + assert_eq!(filter.len(), validity.len()); + assert!(values.len() >= num_valid_values); + + let mut filter_iter = filter.fast_iter_u56(); + let mut validity_iter = validity.fast_iter_u56(); + + let mut values = values; + let mut values_offset = 0; + // @NOTE: + // Since we asserted before that dict is not empty and every time we copy values into this + // buffer we verify the indexes. These will now always be valid offsets into the dict array. + let mut values_buffer = [0u32; 64]; + + target.reserve(num_rows); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut num_buffered = 0; + let mut buffer_side = false; + + let mut intermediate_buffer = Vec::with_capacity(32); + + let mut iter_u32 = |f: u32, v: u32| { + while num_buffered < v.count_ones() as usize { + intermediate_buffer.clear(); + let num_added_values = values.len().min(32); + values.collect_n_into(&mut intermediate_buffer, num_added_values)?; + + // @TODO: Make this into a chunk by chunk operation. + // let (added_buffer, num_added_values) = values.next_chunk().unwrap(); + + values_buffer[usize::from(buffer_side) * 32..][..num_added_values] + .copy_from_slice(&intermediate_buffer[..num_added_values]); + + let highest_idx = (0..num_added_values) + .map(|i| values_buffer[(values_offset + i) % 64]) + .max() + .unwrap(); + assert!((highest_idx as usize) < dict.len()); + + buffer_side = !buffer_side; + num_buffered += num_added_values; + } + + // @NOTE: We have to cast to u64 here to avoid the `shr` overflow on the filter. + let mut f = f as u64; + let mut v = v as u64; + let mut num_read = 0; + let mut num_written = 0; + + while f != 0 { + let offset = f.trailing_zeros(); + + num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + v >>= offset; + + let idx = values_buffer[(values_offset + num_read) % 64]; + let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + unsafe { *target_ptr.add(num_written) = value }; + + num_written += 1; + num_read += (v & 1) as usize; + + f >>= offset + 1; + v >>= 1; + } + + num_read += v.count_ones() as usize; + + values_offset = (values_offset + num_read) % 64; + num_buffered -= num_read; + target_ptr = unsafe { target_ptr.add(num_written) }; + + ParquetResult::Ok(()) + }; + let mut iter = |f: u64, v: u64| { + iter_u32((f & 0xFFFF_FFFF) as u32, (v & 0xFFFF_FFFF) as u32)?; + iter_u32((f >> 32) as u32, (v >> 32) as u32)?; + + ParquetResult::Ok(()) + }; + + for (f, v) in filter_iter.by_ref().zip(validity_iter.by_ref()) { + iter(f, v)?; + } + + let (f, fl) = filter_iter.remainder(); + let (v, vl) = validity_iter.remainder(); + + assert_eq!(fl, vl); + + iter(f, v)?; + + unsafe { target.set_len(start_length + num_rows) }; + Ok(()) +} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/filter.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/filter.rs index a9f0f7b3ef87..bc37d79fc868 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/filter.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/filter.rs @@ -36,6 +36,13 @@ impl Filter { } } + pub fn max_offset(&self) -> usize { + match self { + Filter::Range(range) => range.end, + Filter::Mask(bitmap) => bitmap.len(), + } + } + pub(crate) fn split_at(&self, at: usize) -> (Filter, Filter) { use Filter as F; match self { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index dba00fc97930..5b9213afa152 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod array_chunks; pub(crate) mod filter; +pub(crate) mod dict_encoded; use arrow::array::{DictionaryArray, DictionaryKey, MutableBinaryViewArray, PrimitiveArray, View}; use arrow::bitmap::{Bitmap, MutableBitmap}; @@ -125,81 +126,7 @@ impl<'a, D: Decoder> State<'a, D> { decoded: &mut D::DecodedState, filter: Option, ) -> ParquetResult<()> { - match filter { - None => { - let num_rows = self.len(); - - if num_rows == 0 { - return Ok(()); - } - - self.translation.extend_from_state( - decoder, - decoded, - self.is_optional, - &mut self.page_validity, - self.dict, - num_rows, - ) - }, - Some(filter) => match filter { - Filter::Range(range) => { - let start = range.start; - let end = range.end; - - self.skip_in_place(start)?; - debug_assert!(end - start <= self.len()); - - if end - start > 0 { - self.translation.extend_from_state( - decoder, - decoded, - self.is_optional, - &mut self.page_validity, - self.dict, - end - start, - )?; - } - - Ok(()) - }, - Filter::Mask(bitmap) => { - debug_assert!(bitmap.len() == self.len()); - - let mut iter = bitmap.iter(); - while iter.num_remaining() > 0 && self.len() > 0 { - let prev_state_len = self.len(); - - let num_ones = iter.take_leading_ones(); - - if num_ones > 0 { - self.translation.extend_from_state( - decoder, - decoded, - self.is_optional, - &mut self.page_validity, - self.dict, - num_ones, - )?; - } - - if iter.num_remaining() == 0 || self.len() == 0 { - break; - } - - let num_zeros = iter.take_leading_zeros(); - self.skip_in_place(num_zeros)?; - - assert!( - prev_state_len != self.len(), - "No forward progress was booked in a filtered parquet file." - ); - } - - Ok(()) - }, - }, - } + decoder.extend_filtered_with_state(self, decoded, filter) } } @@ -311,7 +238,8 @@ pub(crate) fn page_validity_decoder(page: &DataPage) -> ParquetResult>( +#[derive(Default)] +pub(crate) struct BatchGatherer<'a, I, T, C: BatchableCollector>( std::marker::PhantomData<&'a (I, T, C)>, ); impl<'a, I, T, C: BatchableCollector> HybridRleGatherer for BatchGatherer<'a, I, T, C> { @@ -607,6 +535,98 @@ pub(super) trait Decoder: Sized { /// Deserializes a [`DictPage`] into [`Self::Dict`]. fn deserialize_dict(&self, page: DictPage) -> ParquetResult; + fn extend_filtered_with_state<'a>( + &mut self, + state: &mut State<'a, Self>, + decoded: &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + self.extend_filtered_with_state_default(state, decoded, filter) + } + + fn extend_filtered_with_state_default<'a>( + &mut self, + state: &mut State<'a, Self>, + decoded: &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + match filter { + None => { + let num_rows = state.len(); + + if num_rows == 0 { + return Ok(()); + } + + state.translation.extend_from_state( + self, + decoded, + state.is_optional, + &mut state.page_validity, + state.dict, + num_rows, + ) + }, + Some(filter) => match filter { + Filter::Range(range) => { + let start = range.start; + let end = range.end; + + state.skip_in_place(start)?; + debug_assert!(end - start <= state.len()); + + if end - start > 0 { + state.translation.extend_from_state( + self, + decoded, + state.is_optional, + &mut state.page_validity, + state.dict, + end - start, + )?; + } + + Ok(()) + }, + Filter::Mask(bitmap) => { + debug_assert!(bitmap.len() == state.len()); + + let mut iter = bitmap.iter(); + while iter.num_remaining() > 0 && state.len() > 0 { + let prev_state_len = state.len(); + + let num_ones = iter.take_leading_ones(); + + if num_ones > 0 { + state.translation.extend_from_state( + self, + decoded, + state.is_optional, + &mut state.page_validity, + state.dict, + num_ones, + )?; + } + + if iter.num_remaining() == 0 || state.len() == 0 { + break; + } + + let num_zeros = iter.take_leading_zeros(); + state.skip_in_place(num_zeros)?; + + assert!( + prev_state_len != state.len(), + "No forward progress was booked in a filtered parquet file." + ); + } + + Ok(()) + }, + }, + } + } + fn apply_dictionary( &mut self, _decoded: &mut Self::DecodedState, @@ -725,7 +745,7 @@ impl PageDecoder { (state_filter, filter) = Filter::opt_split_at(&filter, page.num_values()); // Skip the whole page if we don't need any rows from it - if state_filter.as_ref().is_some_and(|f| f.num_rows() == 0) { + if state_filter.as_ref().is_some_and(|f| dbg!(f.num_rows()) == 0) { continue; } @@ -737,6 +757,9 @@ impl PageDecoder { state.extend_from_state(&mut self.decoder, &mut target, state_filter)?; let end_length = target.len(); + dbg!(num_rows_remaining); + dbg!(start_length); + dbg!(end_length); num_rows_remaining -= end_length - start_length; debug_assert!(state.len() == 0 || num_rows_remaining == 0); From 4c06099ec48cc4a1dbcba33fcc21b31c94727a42 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 11 Oct 2024 11:13:02 +0200 Subject: [PATCH 02/32] add filter_boolean_kernel --- crates/polars-compute/src/filter/mod.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/polars-compute/src/filter/mod.rs b/crates/polars-compute/src/filter/mod.rs index 6de1afbab2ed..a90ac3aa2b15 100644 --- a/crates/polars-compute/src/filter/mod.rs +++ b/crates/polars-compute/src/filter/mod.rs @@ -11,6 +11,7 @@ use arrow::array::{new_empty_array, Array, BinaryViewArray, BooleanArray, Primit use arrow::bitmap::utils::SlicesIterator; use arrow::bitmap::Bitmap; use arrow::with_match_primitive_type_full; +pub use boolean::filter_boolean_kernel; pub fn filter(array: &dyn Array, mask: &BooleanArray) -> Box { assert_eq!(array.len(), mask.len()); From dfc3600809527a0cbc930dc965898381aa5e2b43 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Sat, 12 Oct 2024 20:31:50 +0200 Subject: [PATCH 03/32] work --- crates/polars-arrow/src/bitmap/iterator.rs | 77 ++++++ .../read/deserialize/utils/dict_encoded.rs | 236 ++++++++++++------ .../src/arrow/read/deserialize/utils/mod.rs | 5 +- .../src/parquet/encoding/bitpacked/decode.rs | 17 +- .../src/parquet/encoding/bitpacked/mod.rs | 2 +- .../src/parquet/encoding/hybrid_rle/mod.rs | 70 ++++++ crates/polars-utils/src/lib.rs | 1 + 7 files changed, 330 insertions(+), 78 deletions(-) diff --git a/crates/polars-arrow/src/bitmap/iterator.rs b/crates/polars-arrow/src/bitmap/iterator.rs index bd48fb706c0a..cab8454f0e19 100644 --- a/crates/polars-arrow/src/bitmap/iterator.rs +++ b/crates/polars-arrow/src/bitmap/iterator.rs @@ -1,3 +1,4 @@ +use polars_utils::bit_fns::nth_setbit; use polars_utils::slice::load_padded_le_u64; use super::bitmask::BitMask; @@ -173,6 +174,7 @@ impl<'a> Iterator for FastU32BitmapIter<'a> { unsafe impl<'a> TrustedLen for FastU32BitmapIter<'a> {} +#[derive(Clone)] pub struct FastU56BitmapIter<'a> { bytes: &'a [u8], shift: u32, @@ -220,6 +222,81 @@ impl<'a> FastU56BitmapIter<'a> { let hi = self.next_remainder().unwrap_or(0); ((hi << 56) | lo, bits_left) } + + /// Return how many bits need to be consumed until we have seen `num_ones` 1s. + /// + /// It will return `self.bits_left()` if there are less 1s in the remaining values that + /// `num_ones`. + pub fn num_bits_until_n_ones(&self, mut num_ones: usize) -> usize { + let mut num_skipped = 0; + let mut slf = self.clone(); + + loop { + if num_ones < 56 { + break; + } + + let Some(v) = slf.next() else { + break; + }; + + num_ones -= v.count_ones() as usize; + num_skipped += 56; + } + + for v in slf.by_ref() { + let vco = v.count_ones() as usize; + + if vco <= num_ones { + return num_skipped + nth_setbit(v, num_ones as u64) as usize; + }; + + num_ones -= vco; + num_skipped += 56; + } + + let (v, _) = slf.remainder(); + + let vco = v.count_ones() as usize; + + if vco > num_ones { + self.bits_left as usize + } else { + num_skipped + nth_setbit(v, num_ones as u64) as usize + } + } + + /// Limit the iterator `num_bits`. + pub fn limit_to_bits(&self, num_bits: usize) -> Self { + let mut new = self.clone(); + new.bits_left = new.bits_left.min(num_bits); + new + } + + pub fn count_ones(&self) -> usize { + let mut slf = self.clone(); + + let base = slf.by_ref().map(|v| v.count_ones() as usize).sum::(); + let rem = slf.remainder().0.count_ones() as usize; + + base + rem + } + + /// Advance the iterator by `num_bits` returning the amount. + pub fn advance_by_bits(&mut self, num_bits: usize) { + if self.bits_left <= num_bits { + self.bytes = &[]; + self.bits_left = 0; + return; + } + + self.bits_left -= num_bits; + + self.shift += (num_bits % 8) as u32; + self.shift %= 8; + + self.bytes = unsafe { self.bytes.get_unchecked(num_bits / 8..) }; + } } impl<'a> Iterator for FastU56BitmapIter<'a> { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index 7fc3f96d96ae..c00f04da6bba 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -5,7 +5,7 @@ use bytemuck::Zeroable; use polars_compute::filter::filter_boolean_kernel; use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; -use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; +use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::read::Filter; @@ -135,7 +135,7 @@ pub fn decode_dict( } } -pub fn decode_non_optional_dict( +fn decode_non_optional_dict( values: HybridRleDecoder<'_>, dict: &[T], limit: Option, @@ -151,30 +151,67 @@ pub fn decode_non_optional_dict( let mut limit = limit; let mut values = values; - let mut intermediate_buffer = Vec::with_capacity(32); while limit > 0 { - let num_added_values = limit.min(32); - values.collect_n_into(&mut intermediate_buffer, num_added_values)?; - - // @TODO: Make this into a chunk by chunk operation. - // let (added_buffer, num_added_values) = values.next_chunk().unwrap(); + let chunk = values.next_chunk()?.unwrap(); - let highest_idx = (0..num_added_values) - .map(|i| intermediate_buffer[i]) - .max() - .unwrap(); - assert!((highest_idx as usize) < dict.len()); + match chunk { + HybridRleChunk::Rle(value, length) => { + let length = length.min(limit); - for (i, v) in intermediate_buffer.iter().enumerate() { - let v = unsafe { dict.get_unchecked(*v as usize) }.clone(); - unsafe { target_ptr.add(i).write(v) }; - } + let v = dict.get(value as usize).unwrap(); + let target_slice = unsafe { std::slice::from_raw_parts_mut(target_ptr, length) }; + target_slice.fill(v.clone()); - unsafe { - target_ptr = target_ptr.add(num_added_values); + unsafe { + target_ptr = target_ptr.add(length); + } + limit -= length; + }, + HybridRleChunk::Bitpacked(mut decoder) => { + let mut chunked = decoder.chunked(); + loop { + if limit < 32 { + break; + } + + let Some(chunk) = chunked.next() else { + break; + }; + + let highest_idx = chunk.iter().copied().max().unwrap(); + assert!((highest_idx as usize) < dict.len()); + + for i in 0..32 { + let v = unsafe { dict.get_unchecked(chunk[i] as usize) }.clone(); + unsafe { target_ptr.add(i).write(v) }; + } + + unsafe { + target_ptr = target_ptr.add(32); + } + limit -= 32; + } + + if let Some((chunk, chunk_size)) = chunked.next_inexact() { + let chunk_size = chunk_size.min(limit); + + let highest_idx = chunk[..chunk_size].iter().copied().max().unwrap(); + assert!((highest_idx as usize) < dict.len()); + + for i in 0..chunk_size { + let v = unsafe { dict.get_unchecked(chunk[i] as usize) }.clone(); + unsafe { target_ptr.add(i).write(v) }; + } + + unsafe { + target_ptr = target_ptr.add(chunk_size); + } + + limit -= chunk_size; + } + }, } - limit -= num_added_values; } unsafe { @@ -184,12 +221,13 @@ pub fn decode_non_optional_dict( Ok(()) } -pub fn decode_optional_dict( +fn decode_optional_dict( values: HybridRleDecoder<'_>, dict: &[T], validity: &Bitmap, target: &mut Vec, ) -> ParquetResult<()> { + let mut limit = validity.len(); let num_valid_values = validity.set_bits(); assert!(num_valid_values <= values.len()); @@ -198,71 +236,129 @@ pub fn decode_optional_dict( target.reserve(validity.len()); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + let mut chunk_iter = values.into_chunk_iter(); let mut validity_iter = validity.fast_iter_u56(); - let mut values = values; - let mut values_buffer = [0u32; 64]; + let values_buffer = [0u32; 128]; let mut values_offset = 0; - let mut num_buffered = 0; - let mut buffer_side = false; - let mut intermediate_buffer = Vec::with_capacity(32); + let mut num_buffered: usize = 0; + let mut rridx = 0; - let mut iter_u32 = |v: u32| { - while num_buffered < v.count_ones() as usize { - intermediate_buffer.clear(); - let num_added_values = values.len().min(32); - values.collect_n_into(&mut intermediate_buffer, num_added_values)?; - - // @TODO: Make this into a chunk by chunk operation. - // let (added_buffer, num_added_values) = values.next_chunk().unwrap(); - - values_buffer[usize::from(buffer_side) * 32..][..num_added_values] - .copy_from_slice(&intermediate_buffer[..num_added_values]); - - let highest_idx = (0..num_added_values) - .map(|i| values_buffer[(values_offset + i) % 64]) - .max() - .unwrap(); - assert!((highest_idx as usize) < dict.len()); - - buffer_side = !buffer_side; - num_buffered += num_added_values; + loop { + if limit == 0 { + break; } - let mut num_read = 0; + let Some(chunk) = chunk_iter.next() else { + break; + }; - for i in 0..32 { - let idx = values_buffer[(values_offset + num_read) % 64]; - let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); - unsafe { *target_ptr.add(i) = value }; - num_read += ((v >> i) & 1) as usize; - } + let chunk = chunk?; - values_offset = (values_offset + num_read) % 64; - num_buffered -= num_read; - target_ptr = unsafe { target_ptr.add(32) }; + match chunk { + HybridRleChunk::Rle(value, size) => { + let size = size; - ParquetResult::Ok(()) - }; - let mut iter = |v: u64| { - iter_u32((v & 0xFFFF_FFFF) as u32)?; - iter_u32((v >> 32) as u32)?; + let num_rows = validity_iter.num_bits_until_n_ones(size); + let num_rows = num_rows.min(limit); - ParquetResult::Ok(()) - }; + let value = unsafe { dict.get_unchecked(value as usize) }.clone(); + let target_slice = unsafe { std::slice::from_raw_parts_mut(target_ptr, num_rows) }; + target_slice.fill(value); - for v in validity_iter.by_ref() { - iter(v)?; + unsafe { + target_ptr = target_ptr.add(num_rows); + } + validity_iter.advance_by_bits(num_rows); + limit -= num_rows; + }, + HybridRleChunk::Bitpacked(mut decoder) => { + let mut chunked = decoder.chunked(); + + 'outer: while limit > 56 { + let v = validity_iter.next().unwrap(); + + while num_buffered < v.count_ones() as usize { + rridx %= 4; + + let buffer_part = + &mut values_buffer[rridx * 32..][..32].try_into().unwrap(); + let Some(num_added) = chunked.next_into(buffer_part) else { + break 'outer; + }; + + let highest_idx = buffer_part.iter().copied().max().unwrap(); + assert!((highest_idx as usize) < dict.len()); + + num_buffered += num_added; + rridx += 1; + } + + let mut num_read = 0; + + for i in 0..56 { + let idx = values_buffer[(values_offset + num_read) % 128]; + let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + unsafe { *target_ptr.add(i) = value }; + num_read += ((v >> i) & 1) as usize; + } + + values_offset += num_read; + values_offset %= 128; + num_buffered -= num_read; + unsafe { + target_ptr = target_ptr.add(56); + } + limit -= 56; + } + + let num_remaining = limit.min(num_buffered + chunked.decoder.len()); + debug_assert!(num_remaining < 56); + let (v, _) = validity_iter.limit_to_bits(num_remaining).remainder(); + validity_iter.advance_by_bits(num_remaining); + + while num_buffered < v.count_ones() as usize { + rridx %= 4; + + let buffer_part = &mut values_buffer[rridx * 32..][..32].try_into().unwrap(); + let num_added = chunked.next_into(buffer_part).unwrap(); + + let highest_idx = buffer_part.iter().copied().max().unwrap(); + assert!((highest_idx as usize) < dict.len()); + + num_buffered += num_added; + rridx += 1; + } + + let mut num_read = 0; + + for i in 0..num_remaining { + let idx = values_buffer[(values_offset + num_read) % 128]; + let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + unsafe { *target_ptr.add(i) = value }; + num_read += ((v >> i) & 1) as usize; + } + + values_offset += num_read; + values_offset %= 128; + num_buffered -= num_read; + unsafe { + target_ptr = target_ptr.add(num_remaining); + } + limit -= num_remaining; + + debug_assert!(limit == 0 || num_buffered == 0); + }, + } } - let (v, _) = validity_iter.remainder(); - iter(v)?; - - unsafe { target.set_len(start_length + validity.len()) }; + unsafe { + target.set_len(start_length + validity.len()); + } Ok(()) } -pub fn decode_masked_dict( +fn decode_masked_dict( values: HybridRleDecoder<'_>, dict: &[T], filter: &Bitmap, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 5b9213afa152..0e68863abb47 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -745,7 +745,7 @@ impl PageDecoder { (state_filter, filter) = Filter::opt_split_at(&filter, page.num_values()); // Skip the whole page if we don't need any rows from it - if state_filter.as_ref().is_some_and(|f| dbg!(f.num_rows()) == 0) { + if state_filter.as_ref().is_some_and(|f| f.num_rows() == 0) { continue; } @@ -757,9 +757,6 @@ impl PageDecoder { state.extend_from_state(&mut self.decoder, &mut target, state_filter)?; let end_length = target.len(); - dbg!(num_rows_remaining); - dbg!(start_length); - dbg!(end_length); num_rows_remaining -= end_length - start_length; debug_assert!(state.len() == 0 || num_rows_remaining == 0); diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs index b453c7b6e0f6..dd2feaba44d6 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs @@ -124,9 +124,7 @@ impl<'a, 'b, T: Unpackable> Iterator for ChunkedDecoder<'a, 'b, T> { } let mut unpacked = T::Unpacked::zero(); - let packed = self.decoder.packed.next()?; - decode_pack::(packed, self.decoder.num_bits, &mut unpacked); - self.decoder.length -= T::Unpacked::LENGTH; + self.next_into(&mut unpacked)?; Some(unpacked) } @@ -162,6 +160,19 @@ impl<'a, 'b, T: Unpackable> ChunkedDecoder<'a, 'b, T> { self.remainder() } } + + pub fn next_into(&mut self, unpacked: &mut T::Unpacked) -> Option { + if self.decoder.len() == 0 { + return None; + } + + let unpacked_len = self.decoder.len().min(T::Unpacked::LENGTH); + let packed = self.decoder.packed.next()?; + decode_pack::(packed, self.decoder.num_bits, unpacked); + self.decoder.length -= unpacked_len; + + Some(unpacked_len) + } } impl<'a, T: Unpackable> Decoder<'a, T> { diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs index 94f310d28f14..0a9045ef40d5 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs @@ -57,7 +57,7 @@ mod encode; mod pack; mod unpack; -pub use decode::Decoder; +pub use decode::{Decoder, ChunkedDecoder}; pub use encode::{encode, encode_pack}; /// A byte slice (e.g. `[u8; 8]`) denoting types that represent complete packs. diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs index f721274d00cd..49bc7682da1d 100644 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs @@ -51,6 +51,23 @@ pub struct HybridRleDecoder<'a> { buffered: Option>, } +pub struct HybridRleChunkIter<'a> { + decoder: HybridRleDecoder<'a> +} + +pub enum HybridRleChunk<'a> { + Rle(u32, usize), + Bitpacked(bitpacked::Decoder<'a, u32>), +} + +impl<'a> Iterator for HybridRleChunkIter<'a> { + type Item = ParquetResult>; + + fn next(&mut self) -> Option { + self.decoder.next_chunk().transpose() + } +} + impl<'a> HybridRleDecoder<'a> { /// Returns a new [`HybridRleDecoder`] pub fn new(data: &'a [u8], num_bits: u32, num_values: usize) -> Self { @@ -67,6 +84,59 @@ impl<'a> HybridRleDecoder<'a> { self.num_values } + pub fn into_chunk_iter(self) -> HybridRleChunkIter<'a> { + HybridRleChunkIter { decoder: self } + } + + pub fn next_chunk(&mut self) -> ParquetResult>> { + if self.len() == 0 { + return Ok(None); + } + + if self.num_bits == 0 { + let num_values = self.num_values; + self.num_values = 0; + return Ok(Some(HybridRleChunk::Rle(0, num_values))); + } + + let (indicator, consumed) = uleb128::decode(self.data); + self.data = unsafe { self.data.get_unchecked_release(consumed..) }; + + Ok(Some(if indicator & 1 == 1 { + // is bitpacking + let bytes = (indicator as usize >> 1) * self.num_bits; + let bytes = std::cmp::min(bytes, self.data.len()); + let (packed, remaining) = self.data.split_at(bytes); + self.data = remaining; + + let length = std::cmp::min(packed.len() * 8 / self.num_bits, self.num_values); + let decoder = bitpacked::Decoder::::try_new(packed, self.num_bits, length)?; + + self.num_values -= length; + + HybridRleChunk::Bitpacked(decoder) + } else { + // is rle + let run_length = indicator as usize >> 1; + // repeated-value := value that is repeated, using a fixed-width of round-up-to-next-byte(bit-width) + let rle_bytes = self.num_bits.div_ceil(8); + let (pack, remaining) = self.data.split_at(rle_bytes); + self.data = remaining; + + let mut bytes = [0u8; std::mem::size_of::()]; + pack.iter().zip(bytes.iter_mut()).for_each(|(src, dst)| { + *dst = *src; + }); + let value = u32::from_le_bytes(bytes); + + let length = std::cmp::min(run_length, self.num_values); + + self.num_values -= length; + + HybridRleChunk::Rle(value, length) + })) + } + fn gather_limited_once>( &mut self, target: &mut G::Target, diff --git a/crates/polars-utils/src/lib.rs b/crates/polars-utils/src/lib.rs index 5c302067e146..ea80dd169671 100644 --- a/crates/polars-utils/src/lib.rs +++ b/crates/polars-utils/src/lib.rs @@ -7,6 +7,7 @@ pub mod abs_diff; pub mod arena; pub mod atomic; pub mod binary_search; +pub mod bit_fns; pub mod cache; pub mod cell; pub mod clmul; From 97c9b04037fd1f13302f720873c43c322b8e7760 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Mon, 14 Oct 2024 15:19:01 +0200 Subject: [PATCH 04/32] working dictionary --- crates/polars-arrow/src/bitmap/iterator.rs | 31 +- .../src/arrow/read/deserialize/binview.rs | 173 ++-- .../arrow/read/deserialize/primitive/float.rs | 66 +- .../read/deserialize/primitive/integer.rs | 118 +-- .../read/deserialize/utils/dict_encoded.rs | 766 +++++++++++++----- .../src/arrow/read/deserialize/utils/mod.rs | 102 +-- .../src/parquet/encoding/bitpacked/decode.rs | 19 +- crates/polars-utils/src/bit_fns.rs | 30 + crates/polars-utils/src/chunks.rs | 67 ++ crates/polars-utils/src/lib.rs | 1 + 10 files changed, 781 insertions(+), 592 deletions(-) create mode 100644 crates/polars-utils/src/bit_fns.rs create mode 100644 crates/polars-utils/src/chunks.rs diff --git a/crates/polars-arrow/src/bitmap/iterator.rs b/crates/polars-arrow/src/bitmap/iterator.rs index cab8454f0e19..66db8a6260c7 100644 --- a/crates/polars-arrow/src/bitmap/iterator.rs +++ b/crates/polars-arrow/src/bitmap/iterator.rs @@ -227,7 +227,7 @@ impl<'a> FastU56BitmapIter<'a> { /// /// It will return `self.bits_left()` if there are less 1s in the remaining values that /// `num_ones`. - pub fn num_bits_until_n_ones(&self, mut num_ones: usize) -> usize { + pub fn num_bits_before_nth_one(&self, mut num_ones: usize) -> usize { let mut num_skipped = 0; let mut slf = self.clone(); @@ -245,24 +245,23 @@ impl<'a> FastU56BitmapIter<'a> { } for v in slf.by_ref() { - let vco = v.count_ones() as usize; + let v_num_ones = v.count_ones() as usize; - if vco <= num_ones { - return num_skipped + nth_setbit(v, num_ones as u64) as usize; - }; + if num_ones < v_num_ones { + return num_skipped + nth_setbit(v, num_ones as u32) as usize; + } - num_ones -= vco; + num_ones -= v_num_ones; num_skipped += 56; } let (v, _) = slf.remainder(); + let v_num_ones = v.count_ones() as usize; - let vco = v.count_ones() as usize; - - if vco > num_ones { - self.bits_left as usize + if num_ones < v_num_ones { + num_skipped + nth_setbit(v, num_ones as u32) as usize } else { - num_skipped + nth_setbit(v, num_ones as u64) as usize + self.bits_left } } @@ -275,10 +274,10 @@ impl<'a> FastU56BitmapIter<'a> { pub fn count_ones(&self) -> usize { let mut slf = self.clone(); - + let base = slf.by_ref().map(|v| v.count_ones() as usize).sum::(); let rem = slf.remainder().0.count_ones() as usize; - + base + rem } @@ -293,9 +292,13 @@ impl<'a> FastU56BitmapIter<'a> { self.bits_left -= num_bits; self.shift += (num_bits % 8) as u32; + let extra = self.shift >= 8; self.shift %= 8; - self.bytes = unsafe { self.bytes.get_unchecked(num_bits / 8..) }; + self.bytes = unsafe { + self.bytes + .get_unchecked(num_bits / 8 + usize::from(extra)..) + }; } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index 6777f7e639c9..d987419f8237 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -11,7 +11,6 @@ use arrow::datatypes::{ArrowDataType, PhysicalType}; use super::utils::{dict_indices_decoder, freeze_validity, BatchableCollector}; use crate::parquet::encoding::delta_bitpacked::{lin_natural_sum, DeltaGatherer}; -use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; use crate::parquet::encoding::{delta_byte_array, delta_length_byte_array, hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; @@ -700,145 +699,69 @@ impl utils::Decoder for BinViewDecoder { fn decode_dictionary_encoded<'a>( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut hybrid_rle::HybridRleDecoder<'a>, - is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, - dict: &Self::Dict, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _is_optional: bool, + _page_validity: Option<&mut PageValidity<'a>>, + _dict: &Self::Dict, + _limit: usize, ) -> ParquetResult<()> { - struct DictionaryTranslator<'a>(&'a [View]); + unreachable!() + } - impl<'a> HybridRleGatherer for DictionaryTranslator<'a> { - type Target = MutableBinaryViewArray<[u8]>; + fn extend_filtered_with_state<'a>( + &mut self, + state: &mut utils::State<'a, Self>, + decoded: &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + let num_rows = state.len(); + let mut max_offset = num_rows; - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } + if let Some(ref filter) = filter { + max_offset = filter.max_offset(); + assert!(filter.max_offset() <= num_rows); + } - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.len() - } + match state.translation { + StateTranslation::Dictionary(ref mut indexes) => { + let (dict, _) = state.dict.unwrap(); - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - self.0 - .get(value as usize) - .cloned() - .ok_or(ParquetError::oos("Dictionary index is out of range")) - } - - fn gather_one(&self, target: &mut Self::Target, value: View) -> ParquetResult<()> { - // SAFETY: - // - All the dictionary values are already buffered - // - We keep the `total_bytes_len` in-sync with the views - unsafe { - target.views_mut().push(value); - target.set_total_bytes_len(target.total_bytes_len() + value.length as usize); - } + let start_length = decoded.0.views().len(); - Ok(()) - } + utils::dict_encoded::decode_dict( + indexes.clone(), + &dict, + state.is_optional, + state.page_validity.clone(), + filter, + &mut decoded.1, + unsafe { decoded.0.views_mut() }, + )?; - fn gather_repeated( - &self, - target: &mut Self::Target, - value: View, - n: usize, - ) -> ParquetResult<()> { - // SAFETY: - // - All the dictionary values are already buffered - // - We keep the `total_bytes_len` in-sync with the views + let total_length: usize = decoded + .0 + .views() + .iter() + .skip(start_length) + .map(|view| view.length as usize) + .sum(); unsafe { - let length = target.views_mut().len(); - target.views_mut().resize(length + n, value); - target - .set_total_bytes_len(target.total_bytes_len() + n * value.length as usize); - } - - Ok(()) - } - - fn gather_slice(&self, target: &mut Self::Target, source: &[u32]) -> ParquetResult<()> { - let Some(source_max) = source.iter().copied().max() else { - return Ok(()); - }; - - if source_max as usize >= self.0.len() { - return Err(ParquetError::oos("Dictionary index is out of range")); + decoded + .0 + .set_total_bytes_len(decoded.0.total_bytes_len() + total_length); } - let mut view_length_sum = 0usize; - // Safety: We have checked before that source only has indexes that are smaller than the - // dictionary length. - // - // Safety: - // - All the dictionary values are already buffered - // - We keep the `total_bytes_len` in-sync with the views - unsafe { - target.views_mut().extend(source.iter().map(|&src_idx| { - let v = *self.0.get_unchecked(src_idx as usize); - view_length_sum += v.length as usize; - v - })); - target.set_total_bytes_len(target.total_bytes_len() + view_length_sum); + // @NOTE: Needed for compatibility now. + indexes.skip_in_place(max_offset)?; + if let Some(ref mut page_validity) = state.page_validity { + page_validity.skip_in_place(max_offset)?; } Ok(()) - } - } - - let translator = DictionaryTranslator(&dict.0); - - match page_validity { - None => { - page_values.gather_n_into(values, limit, &translator)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - struct Collector<'a, 'b> { - decoder: &'b mut hybrid_rle::HybridRleDecoder<'a>, - translator: DictionaryTranslator<'b>, - } - - impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for Collector<'a, 'b> { - fn reserve(target: &mut MutableBinaryViewArray<[u8]>, n: usize) { - target.reserve(n); - } - - fn push_n( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - self.decoder.gather_n_into(target, n, &self.translator)?; - Ok(()) - } - - fn push_n_nulls( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - target.extend_constant(n, >::None); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } - } - let collector = Collector { - decoder: page_values, - translator, - }; - extend_from_decoder(validity, page_validity, Some(limit), values, collector)?; }, + _ => self.extend_filtered_with_state_default(state, decoded, filter), } - - Ok(()) } fn finalize( diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index eb4815b6fbfe..2b15709fb6f4 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -8,7 +8,6 @@ use super::{ deserialize_plain, AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, PlainDecoderFnCollector, PrimitiveDecoder, UnitDecoderFunction, }; -use crate::parquet::encoding::hybrid_rle::DictionaryTranslator; use crate::parquet::encoding::{byte_stream_split, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; @@ -16,8 +15,8 @@ use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{ dict_indices_decoder, freeze_validity, BatchableCollector, Decoder, PageValidity, - TranslatedHybridRle, }; +use crate::read::Filter; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -261,37 +260,52 @@ where fn decode_dictionary_encoded<'a>( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut hybrid_rle::HybridRleDecoder<'a>, - is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, - dict: &Self::Dict, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _is_optional: bool, + _page_validity: Option<&mut PageValidity<'a>>, + _dict: &Self::Dict, + _limit: usize, ) -> ParquetResult<()> { - let translator = DictionaryTranslator(dict); + unreachable!() + } - match page_validity { - None => { - page_values.translate_and_collect_n_into(values, limit, &translator)?; + fn extend_filtered_with_state<'a>( + &mut self, + state: &mut utils::State<'a, Self>, + decoded: &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + let num_rows = state.len(); + let mut max_offset = num_rows; - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - let translated_hybridrle = TranslatedHybridRle::new(page_values, &translator); + if let Some(ref filter) = filter { + max_offset = filter.max_offset(); + assert!(filter.max_offset() <= num_rows); + } - utils::extend_from_decoder( - validity, - page_validity, - Some(limit), - values, - translated_hybridrle, + match state.translation { + StateTranslation::Dictionary(ref mut indexes) => { + utils::dict_encoded::decode_dict( + indexes.clone(), + state.dict.unwrap(), + state.is_optional, + state.page_validity.clone(), + filter, + &mut decoded.1, + &mut decoded.0, )?; + + // @NOTE: Needed for compatibility now. + indexes.skip_in_place(max_offset)?; + if let Some(ref mut page_validity) = state.page_validity { + page_validity.skip_in_place(max_offset)?; + } + + Ok(()) }, + _ => self.extend_filtered_with_state_default(state, decoded, filter), } - - Ok(()) } fn finalize( diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 5b18935d6855..925f4a7d3d97 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -2,7 +2,6 @@ use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::types::NativeType; -use polars_compute::filter::filter_boolean_kernel; use super::super::utils; use super::{ @@ -10,16 +9,13 @@ use super::{ DeltaTranslator, IntoDecoderFunction, PlainDecoderFnCollector, PrimitiveDecoder, UnitDecoderFunction, }; -use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; -use crate::parquet::encoding::hybrid_rle::{self, DictionaryTranslator}; -use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, Encoding}; +use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{ dict_indices_decoder, freeze_validity, BatchableCollector, Decoder, PageValidity, - TranslatedHybridRle, }; use crate::read::Filter; @@ -319,37 +315,14 @@ where fn decode_dictionary_encoded<'a>( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut hybrid_rle::HybridRleDecoder<'a>, - is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, - dict: &Self::Dict, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _is_optional: bool, + _page_validity: Option<&mut PageValidity<'a>>, + _dict: &Self::Dict, + _limit: usize, ) -> ParquetResult<()> { - match page_validity { - None => { - let translator = DictionaryTranslator(dict); - page_values.translate_and_collect_n_into(values, limit, &translator)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - let translator = DictionaryTranslator(dict); - let translated_hybridrle = TranslatedHybridRle::new(page_values, &translator); - - utils::extend_from_decoder( - validity, - page_validity, - Some(limit), - values, - translated_hybridrle, - )?; - }, - } - - Ok(()) + unreachable!() } fn finalize( @@ -368,39 +341,6 @@ where decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { - struct BitmapGatherer; - - impl HybridRleGatherer for BitmapGatherer { - type Target = MutableBitmap; - - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.len() - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(value != 0) - } - - fn gather_one(&self, target: &mut Self::Target, value: bool) -> ParquetResult<()> { - target.push(value); - Ok(()) - } - - fn gather_repeated( - &self, - target: &mut Self::Target, - value: bool, - n: usize, - ) -> ParquetResult<()> { - target.extend_constant(n, value); - Ok(()) - } - } - let num_rows = state.len(); let mut max_offset = num_rows; @@ -410,48 +350,6 @@ where } match state.translation { - StateTranslation::Plain(ref mut chunks) => { - let page_validity = state - .page_validity - .take() - .map(|v| { - let mut bm = MutableBitmap::with_capacity(max_offset); - - let gatherer = BitmapGatherer; - - v.clone().gather_n_into(&mut bm, max_offset, &gatherer)?; - - ParquetResult::Ok(bm.freeze()) - }) - .transpose()?; - - let filter = match filter { - None => Bitmap::new_with_value(true, max_offset), - Some(Filter::Range(rng)) => { - let mut bm = MutableBitmap::with_capacity(max_offset); - - bm.extend_constant(rng.start, false); - bm.extend_constant(rng.len(), true); - - bm.freeze() - }, - Some(Filter::Mask(bm)) => bm, - }; - let validity = - page_validity.unwrap_or_else(|| Bitmap::new_with_value(true, max_offset)); - - let start_len = decoded.0.len(); - decode_masked_plain(*chunks, &filter, &validity, &mut decoded.0, self.0.decoder); - debug_assert_eq!(decoded.0.len() - start_len, filter.set_bits()); - if state.is_optional { - let validity = filter_boolean_kernel(&validity, &filter); - decoded.1.extend_from_bitmap(&validity); - } - - chunks.skip_in_place(max_offset); - - Ok(()) - }, StateTranslation::Dictionary(ref mut indexes) => { utils::dict_encoded::decode_dict( indexes.clone(), diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index c00f04da6bba..b578f55892ca 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -1,13 +1,13 @@ use std::ops::Range; use arrow::bitmap::{Bitmap, MutableBitmap}; -use bytemuck::Zeroable; +use arrow::types::NativeType; +use bytemuck::Pod; use polars_compute::filter::filter_boolean_kernel; -use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; -use crate::read::Filter; +use crate::read::{Filter, ParquetError}; fn filter_from_range(rng: Range) -> Bitmap { let mut bm = MutableBitmap::with_capacity(rng.end); @@ -22,52 +22,57 @@ fn decode_page_validity( mut page_validity: HybridRleDecoder<'_>, limit: Option, ) -> ParquetResult { - struct BitmapGatherer; + let mut limit = limit.unwrap_or(page_validity.len()); + let mut bm = MutableBitmap::with_capacity(limit); - impl HybridRleGatherer for BitmapGatherer { - type Target = MutableBitmap; - - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.len() - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(value != 0) - } - - fn gather_one(&self, target: &mut Self::Target, value: bool) -> ParquetResult<()> { - target.push(value); - Ok(()) + while let Some(chunk) = page_validity.next_chunk()? { + if limit == 0 { + break; } - fn gather_repeated( - &self, - target: &mut Self::Target, - value: bool, - n: usize, - ) -> ParquetResult<()> { - target.extend_constant(n, value); - Ok(()) + match chunk { + HybridRleChunk::Rle(value, size) => { + let size = size.min(limit); + bm.extend_constant(size, value != 0); + limit -= size; + }, + HybridRleChunk::Bitpacked(decoder) => { + let len = decoder.len().min(limit); + bm.extend_from_slice(decoder.as_slice(), 0, len); + limit -= len; + }, } } - let mut bm = MutableBitmap::with_capacity(limit.unwrap_or(page_validity.len())); - - let gatherer = BitmapGatherer; - - match limit { - None => page_validity.gather_into(&mut bm, &gatherer)?, - Some(limit) => page_validity.gather_n_into(&mut bm, limit, &gatherer)?, - } + Ok(bm.freeze()) +} - ParquetResult::Ok(bm.freeze()) +pub fn decode_dict( + values: HybridRleDecoder<'_>, + dict: &[T], + is_optional: bool, + page_validity: Option>, + filter: Option, + validity: &mut MutableBitmap, + target: &mut Vec, +) -> ParquetResult<()> { + // @TODO: It would be really nice to reduce monomorphizations here. All decode kernels only + // dependent on the alignment and size of `T` so we could make it downcast here to types that + // are `Pod` and have the same alignment and size. + + decode_dict_dispatch( + values, + dict, + is_optional, + page_validity, + filter, + validity, + target, + ) } -pub fn decode_dict( +#[inline(never)] +fn decode_dict_dispatch( values: HybridRleDecoder<'_>, dict: &[T], is_optional: bool, @@ -100,32 +105,24 @@ pub fn decode_dict( } match (filter, page_validity) { - (None, None) => decode_non_optional_dict(values, dict, None, target), + (None, None) => decode_required_dict(values, dict, None, target), (Some(Filter::Range(rng)), None) if rng.start == 0 => { - decode_non_optional_dict(values, dict, Some(rng.end), target) + decode_required_dict(values, dict, Some(rng.end), target) }, (None, Some(page_validity)) => decode_optional_dict(values, dict, &page_validity, target), (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { decode_optional_dict(values, dict, &page_validity, target) }, - (Some(Filter::Mask(filter)), None) => decode_masked_dict( - values, - dict, - &filter, - &Bitmap::new_with_value(true, filter.len()), - target, - ), + (Some(Filter::Mask(filter)), None) => { + decode_masked_required_dict(values, dict, &filter, target) + }, (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_dict(values, dict, &filter, &page_validity, target) + decode_masked_optional_dict(values, dict, &filter, &page_validity, target) }, - (Some(Filter::Range(rng)), None) => decode_masked_dict( - values, - dict, - &filter_from_range(rng.clone()), - &Bitmap::new_with_value(true, rng.end), - target, - ), - (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_dict( + (Some(Filter::Range(rng)), None) => { + decode_masked_required_dict(values, dict, &filter_from_range(rng.clone()), target) + }, + (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional_dict( values, dict, &filter_from_range(rng.clone()), @@ -135,13 +132,38 @@ pub fn decode_dict( } } -fn decode_non_optional_dict( - values: HybridRleDecoder<'_>, +#[cold] +#[inline(always)] +fn oob_dict_idx() -> ParquetError { + ParquetError::oos("Dictionary Index is out-of-bounds") +} + +#[inline(always)] +fn verify_dict_indices(indices: &[u32; 32], dict_size: usize) -> ParquetResult<()> { + let mut is_valid = true; + for &idx in indices { + is_valid &= (idx as usize) < dict_size; + } + + if is_valid { + return Ok(()); + } + + Err(oob_dict_idx()) +} + +#[inline(never)] +pub fn decode_required_dict( + mut values: HybridRleDecoder<'_>, dict: &[T], limit: Option, target: &mut Vec, ) -> ParquetResult<()> { - let limit = limit.unwrap_or(values.len()); + if dict.is_empty() && values.len() > 0 { + return Err(oob_dict_idx()); + } + + let mut limit = limit.unwrap_or(values.len()); assert!(limit <= values.len()); let start_length = target.len(); let end_length = start_length + limit; @@ -149,9 +171,6 @@ fn decode_non_optional_dict( target.reserve(limit); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - let mut limit = limit; - let mut values = values; - while limit > 0 { let chunk = values.next_chunk()?.unwrap(); @@ -159,13 +178,19 @@ fn decode_non_optional_dict( HybridRleChunk::Rle(value, length) => { let length = length.min(limit); - let v = dict.get(value as usize).unwrap(); - let target_slice = unsafe { std::slice::from_raw_parts_mut(target_ptr, length) }; - target_slice.fill(v.clone()); - + let Some(&value) = dict.get(value as usize) else { + return Err(oob_dict_idx()); + }; + let target_slice; + // SAFETY: + // 1. `target_ptr..target_ptr + limit` is allocated + // 2. `length <= limit` unsafe { + target_slice = std::slice::from_raw_parts_mut(target_ptr, length); target_ptr = target_ptr.add(length); } + + target_slice.fill(value); limit -= length; }, HybridRleChunk::Bitpacked(mut decoder) => { @@ -179,12 +204,11 @@ fn decode_non_optional_dict( break; }; - let highest_idx = chunk.iter().copied().max().unwrap(); - assert!((highest_idx as usize) < dict.len()); + verify_dict_indices(&chunk, dict.len())?; for i in 0..32 { - let v = unsafe { dict.get_unchecked(chunk[i] as usize) }.clone(); - unsafe { target_ptr.add(i).write(v) }; + let v = unsafe { dict.get_unchecked(chunk[i] as usize) }; + unsafe { target_ptr.add(i).write(*v) }; } unsafe { @@ -200,8 +224,8 @@ fn decode_non_optional_dict( assert!((highest_idx as usize) < dict.len()); for i in 0..chunk_size { - let v = unsafe { dict.get_unchecked(chunk[i] as usize) }.clone(); - unsafe { target_ptr.add(i).write(v) }; + let v = unsafe { dict.get_unchecked(chunk[i] as usize) }; + unsafe { target_ptr.add(i).write(*v) }; } unsafe { @@ -221,8 +245,9 @@ fn decode_non_optional_dict( Ok(()) } -fn decode_optional_dict( - values: HybridRleDecoder<'_>, +#[inline(never)] +pub fn decode_optional_dict( + mut values: HybridRleDecoder<'_>, dict: &[T], validity: &Bitmap, target: &mut Vec, @@ -230,103 +255,144 @@ fn decode_optional_dict( let mut limit = validity.len(); let num_valid_values = validity.set_bits(); + // Dispatch to the required kernel if all rows are valid anyway. + if num_valid_values == validity.len() { + return decode_required_dict(values, dict, None, target); + } + + if dict.is_empty() && num_valid_values > 0 { + return Err(oob_dict_idx()); + } + assert!(num_valid_values <= values.len()); let start_length = target.len(); target.reserve(validity.len()); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - let mut chunk_iter = values.into_chunk_iter(); let mut validity_iter = validity.fast_iter_u56(); - let values_buffer = [0u32; 128]; - let mut values_offset = 0; - let mut num_buffered: usize = 0; - let mut rridx = 0; + let mut values_buffer = [0u32; 128]; + let values_buffer = &mut values_buffer; loop { if limit == 0 { break; } - let Some(chunk) = chunk_iter.next() else { + let Some(chunk) = values.next_chunk()? else { break; }; - let chunk = chunk?; - match chunk { HybridRleChunk::Rle(value, size) => { - let size = size; - - let num_rows = validity_iter.num_bits_until_n_ones(size); - let num_rows = num_rows.min(limit); - - let value = unsafe { dict.get_unchecked(value as usize) }.clone(); - let target_slice = unsafe { std::slice::from_raw_parts_mut(target_ptr, num_rows) }; - target_slice.fill(value); - + // If we know that we have `size` times `value` that we can append, but there might + // be nulls in between those values. + // + // 1. See how many `num_rows = valid + invalid` values `size` would entail. This is + // done with `num_bits_before_nth_one` on the validity mask. + // 2. Fill `num_rows` values into the target buffer. + // 3. Advance the validity mask by `num_rows` values. + + let num_chunk_rows = validity_iter.num_bits_before_nth_one(size); + + validity_iter.advance_by_bits(num_chunk_rows); + + let Some(&value) = dict.get(value as usize) else { + return Err(oob_dict_idx()); + }; + let target_slice; + // SAFETY: + // Given `validity_iter` before the `advance_by_bits` + // + // 1. `target_ptr..target_ptr + validity_iter.bits_left()` is allocated + // 2. `num_chunk_rows <= validity_iter.bits_left()` unsafe { - target_ptr = target_ptr.add(num_rows); + target_slice = std::slice::from_raw_parts_mut(target_ptr, num_chunk_rows); + target_ptr = target_ptr.add(num_chunk_rows); } - validity_iter.advance_by_bits(num_rows); - limit -= num_rows; + + target_slice.fill(value); + limit -= num_chunk_rows; }, HybridRleChunk::Bitpacked(mut decoder) => { let mut chunked = decoder.chunked(); - 'outer: while limit > 56 { - let v = validity_iter.next().unwrap(); - - while num_buffered < v.count_ones() as usize { - rridx %= 4; - - let buffer_part = - &mut values_buffer[rridx * 32..][..32].try_into().unwrap(); - let Some(num_added) = chunked.next_into(buffer_part) else { - break 'outer; - }; - - let highest_idx = buffer_part.iter().copied().max().unwrap(); - assert!((highest_idx as usize) < dict.len()); - - num_buffered += num_added; - rridx += 1; - } - - let mut num_read = 0; - - for i in 0..56 { - let idx = values_buffer[(values_offset + num_read) % 128]; - let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); - unsafe { *target_ptr.add(i) = value }; - num_read += ((v >> i) & 1) as usize; + let mut buffer_part_idx = 0; + let mut values_offset = 0; + let mut num_buffered: usize = 0; + + { + let mut num_done = 0; + let mut cpy_validity_iter = validity_iter.clone(); + + 'outer: while limit >= 64 { + let v = cpy_validity_iter.next().unwrap(); + + while num_buffered < v.count_ones() as usize { + let buffer_part = <&mut [u32; 32]>::try_from( + &mut values_buffer[buffer_part_idx * 32..][..32], + ) + .unwrap(); + let Some(num_added) = chunked.next_into(buffer_part) else { + break 'outer; + }; + + verify_dict_indices(&buffer_part, dict.len())?; + + num_buffered += num_added; + + buffer_part_idx += 1; + buffer_part_idx %= 4; + } + + let mut num_read = 0; + + for i in 0..56 { + let idx = values_buffer[(values_offset + num_read) % 128]; + + // SAFETY: + // 1. `values_buffer` starts out as only zeros, which we know is in the + // dictionary following the original `dict.is_empty` check. + // 2. Each time we write to `values_buffer`, it is followed by a + // `verify_dict_indices`. + let value = unsafe { dict.get_unchecked(idx as usize) }; + unsafe { target_ptr.add(i).write(*value) }; + num_read += ((v >> i) & 1) as usize; + } + + values_offset += num_read; + values_offset %= 128; + num_buffered -= num_read; + unsafe { + target_ptr = target_ptr.add(56); + } + num_done += 56; + limit -= 56; } - values_offset += num_read; - values_offset %= 128; - num_buffered -= num_read; - unsafe { - target_ptr = target_ptr.add(56); - } - limit -= 56; + validity_iter.advance_by_bits(num_done); } - let num_remaining = limit.min(num_buffered + chunked.decoder.len()); - debug_assert!(num_remaining < 56); + let num_decoder_remaining = num_buffered + chunked.decoder.len(); + let decoder_limit = validity_iter.num_bits_before_nth_one(num_decoder_remaining); + + let num_remaining = limit.min(decoder_limit); let (v, _) = validity_iter.limit_to_bits(num_remaining).remainder(); validity_iter.advance_by_bits(num_remaining); while num_buffered < v.count_ones() as usize { - rridx %= 4; - - let buffer_part = &mut values_buffer[rridx * 32..][..32].try_into().unwrap(); + let buffer_part = <&mut [u32; 32]>::try_from( + &mut values_buffer[buffer_part_idx * 32..][..32], + ) + .unwrap(); let num_added = chunked.next_into(buffer_part).unwrap(); - let highest_idx = buffer_part.iter().copied().max().unwrap(); - assert!((highest_idx as usize) < dict.len()); + verify_dict_indices(&buffer_part, dict.len())?; num_buffered += num_added; - rridx += 1; + + buffer_part_idx += 1; + buffer_part_idx %= 4; } let mut num_read = 0; @@ -338,15 +404,10 @@ fn decode_optional_dict( num_read += ((v >> i) & 1) as usize; } - values_offset += num_read; - values_offset %= 128; - num_buffered -= num_read; unsafe { target_ptr = target_ptr.add(num_remaining); } limit -= num_remaining; - - debug_assert!(limit == 0 || num_buffered == 0); }, } } @@ -358,114 +419,395 @@ fn decode_optional_dict( Ok(()) } -fn decode_masked_dict( - values: HybridRleDecoder<'_>, +#[inline(never)] +pub fn decode_masked_optional_dict( + mut values: HybridRleDecoder<'_>, dict: &[T], filter: &Bitmap, validity: &Bitmap, target: &mut Vec, ) -> ParquetResult<()> { - let start_length = target.len(); let num_rows = filter.set_bits(); let num_valid_values = validity.set_bits(); - if dict.is_empty() { - assert_eq!(num_valid_values, 0); - target.resize(start_length + num_rows, T::zeroed()); + // Dispatch to the non-filter kernel if all rows are needed anyway. + if num_rows == filter.len() { + return decode_optional_dict(values, dict, validity, target); } - assert_eq!(filter.len(), validity.len()); - assert!(values.len() >= num_valid_values); + // Dispatch to the required kernel if all rows are valid anyway. + if num_valid_values == validity.len() { + return decode_masked_required_dict(values, dict, filter, target); + } - let mut filter_iter = filter.fast_iter_u56(); - let mut validity_iter = validity.fast_iter_u56(); + if dict.is_empty() && num_valid_values > 0 { + return Err(oob_dict_idx()); + } - let mut values = values; - let mut values_offset = 0; - // @NOTE: - // Since we asserted before that dict is not empty and every time we copy values into this - // buffer we verify the indexes. These will now always be valid offsets into the dict array. - let mut values_buffer = [0u32; 64]; + debug_assert_eq!(filter.len(), validity.len()); + assert!(num_valid_values <= values.len()); + let start_length = target.len(); target.reserve(num_rows); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - let mut num_buffered = 0; - let mut buffer_side = false; + let mut filter_iter = filter.fast_iter_u56(); + let mut validity_iter = validity.fast_iter_u56(); - let mut intermediate_buffer = Vec::with_capacity(32); + let mut values_buffer = [0u32; 128]; + let values_buffer = &mut values_buffer; - let mut iter_u32 = |f: u32, v: u32| { - while num_buffered < v.count_ones() as usize { - intermediate_buffer.clear(); - let num_added_values = values.len().min(32); - values.collect_n_into(&mut intermediate_buffer, num_added_values)?; + let mut num_rows_left = num_rows; - // @TODO: Make this into a chunk by chunk operation. - // let (added_buffer, num_added_values) = values.next_chunk().unwrap(); + loop { + // Early stop if we have no more rows to load. + if num_rows_left == 0 { + break; + } - values_buffer[usize::from(buffer_side) * 32..][..num_added_values] - .copy_from_slice(&intermediate_buffer[..num_added_values]); + let Some(chunk) = values.next_chunk()? else { + break; + }; - let highest_idx = (0..num_added_values) - .map(|i| values_buffer[(values_offset + i) % 64]) - .max() - .unwrap(); - assert!((highest_idx as usize) < dict.len()); + match chunk { + HybridRleChunk::Rle(value, size) => { + if value as usize >= dict.len() { + return Err(oob_dict_idx()); + } - buffer_side = !buffer_side; - num_buffered += num_added_values; - } + // If we know that we have `size` times `value` that we can append, but there might + // be nulls in between those values. + // + // 1. See how many `num_rows = valid + invalid` values `size` would entail. This is + // done with `num_bits_before_nth_one` on the validity mask. + // 2. Fill `num_rows` values into the target buffer. + // 3. Advance the validity mask by `num_rows` values. + + let num_chunk_values = validity_iter.num_bits_before_nth_one(size); + let num_chunk_rows = filter_iter.limit_to_bits(num_chunk_values).count_ones(); + + validity_iter.advance_by_bits(num_chunk_values); + filter_iter.advance_by_bits(num_chunk_values); + + if num_chunk_rows > 0 { + // SAFETY: Bounds check done before. + let value = unsafe { dict.get_unchecked(value as usize) }; + + let target_slice; + // SAFETY: + // Given `filter_iter` before the `advance_by_bits`. + // + // 1. `target_ptr..target_ptr + filter_iter.count_ones()` is allocated + // 2. `num_chunk_rows < filter_iter.count_ones()` + unsafe { + target_slice = std::slice::from_raw_parts_mut(target_ptr, num_chunk_rows); + target_ptr = target_ptr.add(num_chunk_rows); + } + + target_slice.fill(*value); + num_rows_left -= num_chunk_rows; + } + }, + HybridRleChunk::Bitpacked(mut decoder) => { + // For bitpacked we do the following: + // 1. See how many rows are encoded by this `decoder`. + // 2. Go through the filter and validity 56 bits at a time and: + // 0. If filter bits are 0, skip the chunk entirely. + // 1. Buffer enough values so that we can branchlessly decode with the filter + // and validity. + // 2. Decode with filter and validity. + // 3. Decode remainder. + + let size = decoder.len(); + let mut chunked = decoder.chunked(); - // @NOTE: We have to cast to u64 here to avoid the `shr` overflow on the filter. - let mut f = f as u64; - let mut v = v as u64; - let mut num_read = 0; - let mut num_written = 0; + let num_chunk_values = validity_iter.num_bits_before_nth_one(size); - while f != 0 { - let offset = f.trailing_zeros(); + let mut buffer_part_idx = 0; + let mut values_offset = 0; + let mut num_buffered: usize = 0; + let mut skip_values = 0; - num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; - v >>= offset; + let mut f_cpy = filter_iter.limit_to_bits(num_chunk_values); + let mut v_cpy = validity_iter.limit_to_bits(num_chunk_values); - let idx = values_buffer[(values_offset + num_read) % 64]; - let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); - unsafe { *target_ptr.add(num_written) = value }; + filter_iter.advance_by_bits(num_chunk_values); + validity_iter.advance_by_bits(num_chunk_values); - num_written += 1; - num_read += (v & 1) as usize; + let mut iter = |mut f: u64, mut v: u64| { + // Skip chunk if we don't any values from here. + if f == 0 { + skip_values += v.count_ones() as usize; + return ParquetResult::Ok(()); + } + + // Skip over already buffered items. + let num_buffered_skipped = skip_values.min(num_buffered); + values_offset += num_buffered_skipped; + num_buffered -= num_buffered_skipped; + skip_values -= num_buffered_skipped; + + // If we skipped plenty already, just skip decoding those chunks instead of + // decoding them and throwing them away. + chunked.decoder.skip_chunks((skip_values / 32) as usize); + // The leftovers we have to decode but we can also just skip. + skip_values %= 32; + + while num_buffered < v.count_ones() as usize { + let buffer_part = <&mut [u32; 32]>::try_from( + &mut values_buffer[buffer_part_idx * 32..][..32], + ) + .unwrap(); + let num_added = chunked.next_into(buffer_part).unwrap(); + + verify_dict_indices(&buffer_part, dict.len())?; + + let skip_chunk_values = (skip_values as usize).min(num_added); + + values_offset += skip_chunk_values; + num_buffered += num_added - skip_chunk_values; + skip_values -= skip_chunk_values; + + buffer_part_idx += 1; + buffer_part_idx %= 4; + } + + let mut num_read = 0; + let mut num_written = 0; + + while f != 0 { + let offset = f.trailing_zeros(); + + num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + v >>= offset; + + let idx = values_buffer[(values_offset + num_read) % 128]; + // SAFETY: + // 1. `values_buffer` starts out as only zeros, which we know is in the + // dictionary following the original `dict.is_empty` check. + // 2. Each time we write to `values_buffer`, it is followed by a + // `verify_dict_indices`. + let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + unsafe { target_ptr.add(num_written).write(value) }; + + num_written += 1; + num_read += (v & 1) as usize; + + f >>= offset + 1; // Clear least significant bit. + v >>= 1; + } + + num_read += v.count_ones() as usize; + + values_offset += num_read; + values_offset %= 128; + num_buffered -= num_read; + unsafe { + target_ptr = target_ptr.add(num_written); + } + num_rows_left -= num_written; + + ParquetResult::Ok(()) + }; - f >>= offset + 1; - v >>= 1; + for (f, v) in f_cpy.by_ref().zip(v_cpy.by_ref()) { + iter(f, v)?; + } + + let (f, fl) = f_cpy.remainder(); + let (v, vl) = v_cpy.remainder(); + + assert_eq!(fl, vl); + + iter(f, v)?; + }, } + } - num_read += v.count_ones() as usize; + unsafe { + target.set_len(start_length + num_rows); + } - values_offset = (values_offset + num_read) % 64; - num_buffered -= num_read; - target_ptr = unsafe { target_ptr.add(num_written) }; + Ok(()) +} - ParquetResult::Ok(()) - }; - let mut iter = |f: u64, v: u64| { - iter_u32((f & 0xFFFF_FFFF) as u32, (v & 0xFFFF_FFFF) as u32)?; - iter_u32((f >> 32) as u32, (v >> 32) as u32)?; +#[inline(never)] +pub fn decode_masked_required_dict( + mut values: HybridRleDecoder<'_>, + dict: &[T], + filter: &Bitmap, + target: &mut Vec, +) -> ParquetResult<()> { + let num_rows = filter.set_bits(); - ParquetResult::Ok(()) - }; + // Dispatch to the non-filter kernel if all rows are needed anyway. + if num_rows == filter.len() { + return decode_required_dict(values, dict, None, target); + } - for (f, v) in filter_iter.by_ref().zip(validity_iter.by_ref()) { - iter(f, v)?; + if dict.is_empty() && values.len() > 0 { + return Err(oob_dict_idx()); } - let (f, fl) = filter_iter.remainder(); - let (v, vl) = validity_iter.remainder(); + let start_length = target.len(); + + target.reserve(num_rows); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut filter_iter = filter.fast_iter_u56(); + + let mut values_buffer = [0u32; 128]; + let values_buffer = &mut values_buffer; - assert_eq!(fl, vl); + let mut num_rows_left = num_rows; - iter(f, v)?; + loop { + if num_rows_left == 0 { + break; + } + + let Some(chunk) = values.next_chunk()? else { + break; + }; + + match chunk { + HybridRleChunk::Rle(value, size) => { + if value as usize >= dict.len() { + return Err(oob_dict_idx()); + } + + // If we know that we have `size` times `value` that we can append, but there might + // be nulls in between those values. + // + // 1. See how many `num_rows = valid + invalid` values `size` would entail. This is + // done with `num_bits_before_nth_one` on the validity mask. + // 2. Fill `num_rows` values into the target buffer. + // 3. Advance the validity mask by `num_rows` values. + + let num_chunk_rows = filter_iter.limit_to_bits(size).count_ones(); + + filter_iter.advance_by_bits(size); + + if num_chunk_rows > 0 { + // SAFETY: Bounds check done before. + let value = unsafe { dict.get_unchecked(value as usize) }; + + let target_slice; + // SAFETY: + // Given `filter_iter` before the `advance_by_bits`. + // + // 1. `target_ptr..target_ptr + filter_iter.count_ones()` is allocated + // 2. `num_chunk_rows < filter_iter.count_ones()` + unsafe { + target_slice = std::slice::from_raw_parts_mut(target_ptr, num_chunk_rows); + target_ptr = target_ptr.add(num_chunk_rows); + } + + target_slice.fill(*value); + num_rows_left -= num_chunk_rows; + } + }, + HybridRleChunk::Bitpacked(mut decoder) => { + let size = decoder.len(); + let mut chunked = decoder.chunked(); + + let mut buffer_part_idx = 0; + let mut values_offset = 0; + let mut num_buffered: usize = 0; + let mut skip_values = 0; + + let mut f_cpy = filter_iter.limit_to_bits(size); + + filter_iter.advance_by_bits(size); + + let mut iter = |mut f: u64, len: usize| { + debug_assert!(len <= 64); + + // Skip chunk if we don't any values from here. + if f == 0 { + skip_values += len; + return ParquetResult::Ok(()); + } + + // Skip over already buffered items. + let num_buffered_skipped = skip_values.min(num_buffered); + values_offset += num_buffered_skipped; + num_buffered -= num_buffered_skipped; + skip_values -= num_buffered_skipped; + + // If we skipped plenty already, just skip decoding those chunks instead of + // decoding them and throwing them away. + chunked.decoder.skip_chunks((skip_values / 32) as usize); + // The leftovers we have to decode but we can also just skip. + skip_values %= 32; + + while num_buffered < len { + let buffer_part = <&mut [u32; 32]>::try_from( + &mut values_buffer[buffer_part_idx * 32..][..32], + ) + .unwrap(); + let num_added = chunked.next_into(buffer_part).unwrap(); + + verify_dict_indices(&buffer_part, dict.len())?; + + let skip_chunk_values = (skip_values as usize).min(num_added); + + values_offset += skip_chunk_values; + num_buffered += num_added - skip_chunk_values; + skip_values -= skip_chunk_values; + + buffer_part_idx += 1; + buffer_part_idx %= 4; + } + + let mut num_read = 0; + let mut num_written = 0; + + while f != 0 { + let offset = f.trailing_zeros() as usize; + + num_read += offset; + + let idx = values_buffer[(values_offset + num_read) % 128]; + // SAFETY: + // 1. `values_buffer` starts out as only zeros, which we know is in the + // dictionary following the original `dict.is_empty` check. + // 2. Each time we write to `values_buffer`, it is followed by a + // `verify_dict_indices`. + let value = *unsafe { dict.get_unchecked(idx as usize) }; + unsafe { target_ptr.add(num_written).write(value) }; + + num_written += 1; + num_read += 1; + + f >>= offset + 1; // Clear least significant bit. + } + + values_offset += len; + values_offset %= 128; + num_buffered -= len; + unsafe { + target_ptr = target_ptr.add(num_written); + } + num_rows_left -= num_written; + + ParquetResult::Ok(()) + }; + + for f in f_cpy.by_ref() { + iter(f, 56)?; + } + + let (f, fl) = f_cpy.remainder(); + + iter(f, fl)?; + }, + } + } + + unsafe { + target.set_len(start_length + num_rows); + } - unsafe { target.set_len(start_length + num_rows) }; Ok(()) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 0e68863abb47..16e5b5cb6aa6 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -1,8 +1,8 @@ pub(crate) mod array_chunks; -pub(crate) mod filter; pub(crate) mod dict_encoded; +pub(crate) mod filter; -use arrow::array::{DictionaryArray, DictionaryKey, MutableBinaryViewArray, PrimitiveArray, View}; +use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::pushable::Pushable; @@ -12,7 +12,7 @@ use super::BasicDecompressor; use crate::parquet::encoding::hybrid_rle::gatherer::{ HybridRleGatherer, ZeroCount, ZeroCountGatherer, }; -use crate::parquet::encoding::hybrid_rle::{self, HybridRleDecoder, Translator}; +use crate::parquet::encoding::hybrid_rle::{self, HybridRleDecoder}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::schema::Repetition; @@ -334,61 +334,6 @@ pub(super) fn extend_from_decoder>( Ok(()) } -/// This translates and collects items from a [`HybridRleDecoder`] into a target [`Vec`]. -/// -/// This batches sequential collect operations to try and prevent unnecessary buffering. -pub struct TranslatedHybridRle<'a, 'b, 'c, O, T> -where - O: Clone + Default, - T: Translator, -{ - decoder: &'a mut HybridRleDecoder<'b>, - translator: &'c T, - _pd: std::marker::PhantomData, -} - -impl<'a, 'b, 'c, O, T> TranslatedHybridRle<'a, 'b, 'c, O, T> -where - O: Clone + Default, - T: Translator, -{ - pub fn new(decoder: &'a mut HybridRleDecoder<'b>, translator: &'c T) -> Self { - Self { - decoder, - translator, - _pd: Default::default(), - } - } -} - -impl<'a, 'b, 'c, O, T> BatchableCollector> for TranslatedHybridRle<'a, 'b, 'c, O, T> -where - O: Clone + Default, - T: Translator, -{ - #[inline] - fn reserve(target: &mut Vec, n: usize) { - target.reserve(n); - } - - #[inline] - fn push_n(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - self.decoder - .translate_and_collect_n_into(target, n, self.translator) - } - - #[inline] - fn push_n_nulls(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - target.resize(target.len() + n, O::default()); - Ok(()) - } - - #[inline] - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } -} - pub struct GatheredHybridRle<'a, 'b, 'c, O, G> where O: Clone, @@ -444,47 +389,6 @@ where } } -impl<'a, 'b, 'c, T> BatchableCollector> - for TranslatedHybridRle<'a, 'b, 'c, View, T> -where - T: Translator, -{ - #[inline] - fn reserve(target: &mut MutableBinaryViewArray<[u8]>, n: usize) { - target.reserve(n); - } - - #[inline] - fn push_n(&mut self, target: &mut MutableBinaryViewArray<[u8]>, n: usize) -> ParquetResult<()> { - self.decoder.translate_and_collect_n_into( - unsafe { target.views_mut() }, - n, - self.translator, - )?; - - if let Some(validity) = target.validity() { - validity.extend_constant(n, true); - } - - Ok(()) - } - - #[inline] - fn push_n_nulls( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - target.extend_null(n); - Ok(()) - } - - #[inline] - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } -} - impl, I: Iterator> BatchableCollector for I { #[inline] fn reserve(target: &mut P, n: usize) { diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs index dd2feaba44d6..f860989689e6 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs @@ -1,3 +1,5 @@ +use polars_utils::chunks::Chunks; + use super::{Packed, Unpackable, Unpacked}; use crate::parquet::error::{ParquetError, ParquetResult}; @@ -6,7 +8,7 @@ use crate::parquet::error::{ParquetError, ParquetResult}; /// This iterator unpacks bytes in chunks and does not allocate. #[derive(Debug, Clone)] pub struct Decoder<'a, T: Unpackable> { - packed: std::slice::Chunks<'a, u8>, + packed: Chunks<'a, u8>, num_bits: usize, /// number of items pub(crate) length: usize, @@ -16,7 +18,7 @@ pub struct Decoder<'a, T: Unpackable> { impl<'a, T: Unpackable> Default for Decoder<'a, T> { fn default() -> Self { Self { - packed: [].chunks(1), + packed: Chunks::new(&[], 1), num_bits: 0, length: 0, _pd: std::marker::PhantomData, @@ -66,7 +68,8 @@ impl<'a, T: Unpackable> Decoder<'a, T> { } debug_assert!(num_bits != 0 || packed.is_empty()); - let packed = packed.chunks(block_size.max(1)); + let block_size = block_size.max(1); + let packed = Chunks::new(packed, block_size); Ok(Self { length, @@ -91,7 +94,7 @@ impl<'a, T: Unpackable> Decoder<'a, T> { ))); } - let packed = packed.chunks(block_size); + let packed = Chunks::new(packed, block_size); Ok(Self { length, @@ -104,6 +107,10 @@ impl<'a, T: Unpackable> Decoder<'a, T> { pub fn num_bits(&self) -> usize { self.num_bits } + + pub fn as_slice(&self) -> &[u8] { + self.packed.as_slice() + } } /// A iterator over the exact chunks in a [`Decoder`]. @@ -192,8 +199,8 @@ impl<'a, T: Unpackable> Decoder<'a, T> { } pub fn take(&mut self) -> Self { - let block_size = size_of::() * self.num_bits; - let packed = std::mem::replace(&mut self.packed, [].chunks(block_size)); + let block_size = self.packed.chunk_size(); + let packed = std::mem::replace(&mut self.packed, Chunks::new(&[], block_size)); let length = self.length; self.length = 0; diff --git a/crates/polars-utils/src/bit_fns.rs b/crates/polars-utils/src/bit_fns.rs new file mode 100644 index 000000000000..50d683f347f8 --- /dev/null +++ b/crates/polars-utils/src/bit_fns.rs @@ -0,0 +1,30 @@ +/// Get the index of the `n`th set bit in `x`. +/// +/// This is starting from `0`. +pub fn nth_setbit(x: u64, n: u32) -> u32 { + #[cfg(target_feature = "bmi2")] + { + use std::arch::x86_64::*; + + return unsafe { _tzcnt_u64(_pdep_u64(1u64 << n, x)) as u32 }; + } + + + #[cfg(not(target_feature = "bmi2"))] + { + let mut x = x; + let mut n = n; + + assert!(x.count_ones() > n); + + let mut idx = 0; + + while x & 1 == 0 || n > 0 { + n -= (x & 1) as u32; + x >>= 1; + idx += 1; + } + + return idx; + } +} diff --git a/crates/polars-utils/src/chunks.rs b/crates/polars-utils/src/chunks.rs new file mode 100644 index 000000000000..1c5d0e8c4f8e --- /dev/null +++ b/crates/polars-utils/src/chunks.rs @@ -0,0 +1,67 @@ +/// A copy of the [`std::slice::Chunks`] that exposes the inner `slice` and `chunk_size`. +#[derive(Clone, Debug)] +pub struct Chunks<'a, T> { + slice: &'a [T], + chunk_size: usize, +} + +impl<'a, T> Iterator for Chunks<'a, T> { + type Item = &'a [T]; + + fn next(&mut self) -> Option { + if self.slice.is_empty() { + return None; + } + + let item; + (item, self.slice) = self.slice.split_at(self.chunk_size.min(self.slice.len())); + + Some(item) + } + + fn size_hint(&self) -> (usize, Option) { + let len = self.slice.len().div_ceil(self.chunk_size); + (len, Some(len)) + } +} + +impl<'a, T> DoubleEndedIterator for Chunks<'a, T> { + fn next_back(&mut self) -> Option { + if self.slice.is_empty() { + return None; + } + + let rem = self.slice.len() % self.chunk_size; + let offset = if rem == 0 { + self.chunk_size + } else { + rem + }; + + let item; + (self.slice, item) = self.slice.split_at(self.slice.len() - offset); + + Some(item) + } +} + +impl<'a, T> ExactSizeIterator for Chunks<'a, T> {} + +impl<'a, T> Chunks<'a, T> { + pub const fn new(slice: &'a [T], chunk_size: usize) -> Self { + Self { slice, chunk_size } + } + + pub const fn as_slice(&self) -> &'a [T] { + self.slice + } + + pub const fn chunk_size(&self) -> usize { + self.chunk_size + } + + pub fn skip_in_place(&mut self, n: usize) { + let n = n * self.chunk_size; + self.slice = &self.slice[n.min(self.slice.len())..]; + } +} diff --git a/crates/polars-utils/src/lib.rs b/crates/polars-utils/src/lib.rs index ea80dd169671..109414da657a 100644 --- a/crates/polars-utils/src/lib.rs +++ b/crates/polars-utils/src/lib.rs @@ -10,6 +10,7 @@ pub mod binary_search; pub mod bit_fns; pub mod cache; pub mod cell; +pub mod chunks; pub mod clmul; pub mod contention_pool; pub mod cpuid; From e45d5e4b58d781113bb9730b2760ee1231538707 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Mon, 14 Oct 2024 17:37:25 +0200 Subject: [PATCH 05/32] working plain --- .../arrow/read/deserialize/primitive/float.rs | 87 ++-- .../read/deserialize/primitive/integer.rs | 229 +++-------- .../arrow/read/deserialize/primitive/mod.rs | 63 +-- .../arrow/read/deserialize/primitive/plain.rs | 376 ++++++++++++++++++ .../read/deserialize/utils/dict_encoded.rs | 40 +- .../src/arrow/read/deserialize/utils/mod.rs | 42 +- 6 files changed, 513 insertions(+), 324 deletions(-) create mode 100644 crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index 2b15709fb6f4..73adb50a4d88 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -5,8 +5,8 @@ use arrow::types::NativeType; use super::super::utils; use super::{ - deserialize_plain, AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, - PlainDecoderFnCollector, PrimitiveDecoder, UnitDecoderFunction, + AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, PrimitiveDecoder, + UnitDecoderFunction, }; use crate::parquet::encoding::{byte_stream_split, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; @@ -14,9 +14,9 @@ use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{ - dict_indices_decoder, freeze_validity, BatchableCollector, Decoder, PageValidity, + dict_indices_decoder, freeze_validity, Decoder, PageValidity, }; -use crate::read::Filter; +use crate::read::{Filter, ParquetError}; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -214,48 +214,34 @@ where } fn deserialize_dict(&self, page: DictPage) -> ParquetResult { - Ok(deserialize_plain::(&page.buffer, self.0.decoder)) + let Some(values) = ArrayChunks::

::new(page.buffer.as_ref()) else { + return Err(ParquetError::oos( + "Primitive dictionary page size is not a multiple of primitive size", + )); + }; + + let mut target = Vec::new(); + super::plain::decode( + values, + false, + None, + None, + &mut MutableBitmap::new(), + &mut target, + self.0.decoder, + )?; + Ok(target) } fn decode_plain_encoded<'a>( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, + _is_optional: bool, + _page_validity: Option<&mut PageValidity<'a>>, + _limit: usize, ) -> ParquetResult<()> { - match page_validity { - None => { - PlainDecoderFnCollector { - chunks: page_values, - decoder: self.0.decoder, - _pd: std::marker::PhantomData, - } - .push_n(values, limit)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - let collector = PlainDecoderFnCollector { - chunks: page_values, - decoder: self.0.decoder, - _pd: std::marker::PhantomData, - }; - - utils::extend_from_decoder( - validity, - page_validity, - Some(limit), - values, - collector, - )?; - }, - } - - Ok(()) + unreachable!() } fn decode_dictionary_encoded<'a>( @@ -285,6 +271,25 @@ where } match state.translation { + StateTranslation::Plain(ref mut values) => { + super::plain::decode( + values.clone(), + state.is_optional, + state.page_validity.clone(), + filter, + &mut decoded.1, + &mut decoded.0, + self.0.decoder, + )?; + + // @NOTE: Needed for compatibility now. + values.skip_in_place(max_offset); + if let Some(ref mut page_validity) = state.page_validity { + page_validity.skip_in_place(max_offset)?; + } + + Ok(()) + }, StateTranslation::Dictionary(ref mut indexes) => { utils::dict_encoded::decode_dict( indexes.clone(), diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 925f4a7d3d97..17d40c9c2998 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -1,13 +1,12 @@ use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; -use arrow::bitmap::{Bitmap, MutableBitmap}; +use arrow::bitmap::MutableBitmap; use arrow::datatypes::ArrowDataType; use arrow::types::NativeType; use super::super::utils; use super::{ - deserialize_plain, AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, DeltaCollector, - DeltaTranslator, IntoDecoderFunction, PlainDecoderFnCollector, PrimitiveDecoder, - UnitDecoderFunction, + AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, DeltaCollector, DeltaTranslator, + IntoDecoderFunction, PrimitiveDecoder, UnitDecoderFunction, }; use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; @@ -15,9 +14,9 @@ use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{ - dict_indices_decoder, freeze_validity, BatchableCollector, Decoder, PageValidity, + dict_indices_decoder, freeze_validity, Decoder, PageValidity, }; -use crate::read::Filter; +use crate::read::{Filter, ParquetError}; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -269,48 +268,34 @@ where } fn deserialize_dict(&self, page: DictPage) -> ParquetResult { - Ok(deserialize_plain::(&page.buffer, self.0.decoder)) + let Some(values) = ArrayChunks::

::new(page.buffer.as_ref()) else { + return Err(ParquetError::oos( + "Primitive dictionary page size is not a multiple of primitive size", + )); + }; + + let mut target = Vec::new(); + super::plain::decode( + values, + false, + None, + None, + &mut MutableBitmap::new(), + &mut target, + self.0.decoder, + )?; + Ok(target) } fn decode_plain_encoded<'a>( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, + _is_optional: bool, + _page_validity: Option<&mut PageValidity<'a>>, + _limit: usize, ) -> ParquetResult<()> { - match page_validity { - None => { - PlainDecoderFnCollector { - chunks: page_values, - decoder: self.0.decoder, - _pd: Default::default(), - } - .push_n(values, limit)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - let collector = PlainDecoderFnCollector { - chunks: page_values, - decoder: self.0.decoder, - _pd: Default::default(), - }; - - utils::extend_from_decoder( - validity, - page_validity, - Some(limit), - values, - collector, - )?; - }, - } - - Ok(()) + unreachable!() } fn decode_dictionary_encoded<'a>( @@ -350,6 +335,25 @@ where } match state.translation { + StateTranslation::Plain(ref mut values) => { + super::plain::decode( + values.clone(), + state.is_optional, + state.page_validity.clone(), + filter, + &mut decoded.1, + &mut decoded.0, + self.0.decoder, + )?; + + // @NOTE: Needed for compatibility now. + values.skip_in_place(max_offset); + if let Some(ref mut page_validity) = state.page_validity { + page_validity.skip_in_place(max_offset)?; + } + + Ok(()) + }, StateTranslation::Dictionary(ref mut indexes) => { utils::dict_encoded::decode_dict( indexes.clone(), @@ -374,145 +378,6 @@ where } } -fn decode_masked_plain( - values: ArrayChunks<'_, P>, - filter: &Bitmap, - validity: &Bitmap, - target: &mut Vec, - dfn: D, -) where - T: NativeType, - P: ParquetNativeType, - i64: num_traits::AsPrimitive

, - D: DecoderFunction, -{ - /// # Safety - /// - /// All of the below need to hold: - /// - `values.len() >= validity.count_ones()` - /// - it needs to be safe to write to out[..filter.count_ones()] - unsafe fn decode_iter_allow_last( - values: ArrayChunks<'_, P>, - mut filter: u64, - mut validity: u64, - out: *mut T, - dfn: D, - ) -> (usize, usize) - where - T: NativeType, - P: ParquetNativeType, - i64: num_traits::AsPrimitive

, - D: DecoderFunction, - { - let mut num_read = 0; - let mut num_written = 0; - - while filter != 0 { - let offset = filter.trailing_zeros(); - - num_read += (validity & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; - validity >>= offset; - - let v = if validity & 1 != 0 { - dfn.decode(unsafe { values.get_unchecked(num_read) }) - } else { - T::zeroed() - }; - - *out.add(num_written) = v; - - num_written += 1; - num_read += (validity & 1) as usize; - - filter >>= offset + 1; - validity >>= 1; - } - - num_read += validity.count_ones() as usize; - - (num_read, num_written) - } - - let start_length = target.len(); - let num_rows = filter.set_bits(); - let num_values = validity.set_bits(); - - assert_eq!(filter.len(), validity.len()); - assert!(values.len() >= num_values); - - let mut filter_iter = filter.fast_iter_u56(); - let mut validity_iter = validity.fast_iter_u56(); - - let mut values = values; - - target.reserve(num_rows); - let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - - let mut num_rows_done = 0; - let mut num_values_done = 0; - - for (f, v) in filter_iter.by_ref().zip(validity_iter.by_ref()) { - let nv = v.count_ones() as usize; - let nr = f.count_ones() as usize; - - // @NOTE: - // If we cannot guarantee that there are more values than we need, we run the branching - // variant and know we can write the remaining values as dummy values because they will be - // None. - if num_values_done + nv == num_values { - // We don't really need to update `values` or `target_ptr` because we won't be - // using them after this. - _ = unsafe { decode_iter_allow_last::(values, f, v, target_ptr, dfn) }; - - num_rows_done += nr; - unsafe { target.set_len(start_length + num_rows_done) }; - target.resize(start_length + num_rows, T::zeroed()); - return; - } - - let mut v = v; - let mut f = f; - let mut num_read = 0; - let mut num_written = 0; - - while f != 0 { - let offset = f.trailing_zeros(); - - num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; - v >>= offset; - - unsafe { - *target_ptr.add(num_written) = dfn.decode(values.get_unchecked(num_read)); - } - - num_written += 1; - num_read += (v & 1) as usize; - - f >>= offset + 1; - v >>= 1; - } - - num_read += v.count_ones() as usize; - - values = ArrayChunks { - bytes: unsafe { values.bytes.get_unchecked(num_read..) }, - }; - target_ptr = unsafe { target_ptr.add(num_written) }; - - num_values_done += nv; - num_rows_done += nr; - } - - let (f, fl) = filter_iter.remainder(); - let (v, vl) = validity_iter.remainder(); - - assert_eq!(fl, vl); - - _ = unsafe { decode_iter_allow_last::(values, f, v, target_ptr, dfn) }; - - unsafe { target.set_len(start_length + num_rows) }; -} - impl utils::DictDecodable for IntDecoder where T: NativeType, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs index 54ed87c6b42b..d0099df10c09 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs @@ -1,15 +1,15 @@ use arrow::types::NativeType; use num_traits::AsPrimitive; -use crate::parquet::types::{decode, NativeType as ParquetNativeType}; +use crate::parquet::types::NativeType as ParquetNativeType; mod float; mod integer; +mod plain; pub(crate) use float::FloatDecoder; pub(crate) use integer::IntDecoder; -use super::utils::array_chunks::ArrayChunks; use super::utils::BatchableCollector; use super::ParquetResult; use crate::parquet::encoding::delta_bitpacked::{self, DeltaGatherer}; @@ -114,65 +114,6 @@ where } } -pub(crate) struct PlainDecoderFnCollector<'a, 'b, P, T, D> -where - T: NativeType, - P: ParquetNativeType, - D: DecoderFunction, -{ - pub(crate) chunks: &'b mut ArrayChunks<'a, P>, - pub(crate) decoder: D, - pub(crate) _pd: std::marker::PhantomData, -} - -impl<'a, 'b, P, T, D: DecoderFunction> BatchableCollector<(), Vec> - for PlainDecoderFnCollector<'a, 'b, P, T, D> -where - T: NativeType, - P: ParquetNativeType, - D: DecoderFunction, -{ - fn reserve(target: &mut Vec, n: usize) { - target.reserve(n); - } - - fn push_n(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - let n = usize::min(self.chunks.len(), n); - let (items, remainder) = self.chunks.bytes.split_at(n); - let decoder = self.decoder; - target.extend( - items - .iter() - .map(|chunk| decoder.decode(P::from_le_bytes(*chunk))), - ); - self.chunks.bytes = remainder; - Ok(()) - } - - fn push_n_nulls(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - target.resize(target.len() + n, T::default()); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.chunks.skip_in_place(n); - Ok(()) - } -} - -fn deserialize_plain(values: &[u8], decoder: D) -> Vec -where - T: NativeType, - P: ParquetNativeType, - D: DecoderFunction, -{ - values - .chunks_exact(size_of::

()) - .map(decode) - .map(|v| decoder.decode(v)) - .collect::>() -} - struct DeltaTranslator where T: NativeType, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs new file mode 100644 index 000000000000..8f5bdb41e7cb --- /dev/null +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -0,0 +1,376 @@ +use arrow::bitmap::{Bitmap, MutableBitmap}; +use arrow::types::NativeType; +use polars_compute::filter::filter_boolean_kernel; + +use super::DecoderFunction; +use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; +use crate::parquet::error::ParquetResult; +use crate::parquet::types::NativeType as ParquetNativeType; +use crate::read::deserialize::utils::array_chunks::ArrayChunks; +use crate::read::deserialize::utils::{decode_page_validity, filter_from_range}; +use crate::read::Filter; + +pub fn decode>( + values: ArrayChunks<'_, P>, + is_optional: bool, + page_validity: Option>, + filter: Option, + validity: &mut MutableBitmap, + target: &mut Vec, + dfn: D, +) -> ParquetResult<()> { + // @TODO: It would be really nice to reduce monomorphizations here. All decode kernels only + // dependent on the alignment and size of `T` so we could make it downcast here to types that + // are `Pod` and have the same alignment and size. + + decode_plain_dispatch( + values, + is_optional, + page_validity, + filter, + validity, + target, + dfn, + ) +} + + +#[inline(never)] +fn decode_plain_dispatch>( + values: ArrayChunks<'_, P>, + is_optional: bool, + page_validity: Option>, + filter: Option, + validity: &mut MutableBitmap, + target: &mut Vec, + dfn: D, +) -> ParquetResult<()> { + let page_validity = match (page_validity, filter.as_ref()) { + (None, _) => None, + (Some(page_validity), None) => Some(decode_page_validity(page_validity, None)?), + (Some(page_validity), Some(filter)) => Some(decode_page_validity( + page_validity, + Some(filter.max_offset()), + )?), + }; + + if is_optional { + match (page_validity.as_ref(), filter.as_ref()) { + (None, None) => validity.extend_constant(values.len(), true), + (None, Some(f)) => validity.extend_constant(f.num_rows(), true), + (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), + (Some(page_validity), Some(Filter::Range(rng))) => { + validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) + }, + (Some(page_validity), Some(Filter::Mask(mask))) => { + validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, &mask)) + }, + } + } + + match (filter, page_validity) { + (None, None) => decode_required(values, None, target, dfn), + (Some(Filter::Range(rng)), None) if rng.start == 0 => { + decode_required(values, Some(rng.end), target, dfn) + }, + (None, Some(page_validity)) => decode_optional(values, &page_validity, target, dfn), + (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { + decode_optional(values, &page_validity, target, dfn) + }, + (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target, dfn), + (Some(Filter::Mask(filter)), Some(page_validity)) => { + decode_masked_optional(values, &filter, &page_validity, target, dfn) + }, + (Some(Filter::Range(rng)), None) => { + decode_masked_required(values, &filter_from_range(rng.clone()), target, dfn) + }, + (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( + values, + &filter_from_range(rng.clone()), + &page_validity, + target, + dfn, + ), + } +} + +#[inline(never)] +fn decode_required>( + values: ArrayChunks<'_, P>, + limit: Option, + target: &mut Vec, + dfn: D, +) -> ParquetResult<()> { + let limit = limit.unwrap_or(values.len()); + assert!(limit <= values.len()); + + target.extend((0..limit).map(|i| { + let v = unsafe { values.get_unchecked(i) }; + dfn.decode(v) + })); + + Ok(()) +} + +#[inline(never)] +fn decode_optional>( + values: ArrayChunks<'_, P>, + validity: &Bitmap, + target: &mut Vec, + dfn: D, +) -> ParquetResult<()> { + let num_values = validity.set_bits(); + + if num_values == validity.len() { + return decode_required(values, Some(validity.len()), target, dfn); + } + + let mut limit = validity.len(); + + assert!(num_values <= values.len()); + + let start_length = target.len(); + let end_length = target.len() + limit; + target.reserve(limit); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut validity_iter = validity.fast_iter_u56(); + let mut num_values_remaining = num_values; + let mut value_offset = 0; + + let mut iter = |v: u64, len: usize| { + debug_assert!(len < 64); + + let num_chunk_values = v.count_ones() as usize; + + if num_values_remaining == num_chunk_values { + for i in 0..len { + let is_valid = v & 1 != 0; + let value = if is_valid { + let value = unsafe { values.get_unchecked(value_offset) }; + dfn.decode(value) + } else { + T::zeroed() + }; + unsafe { target_ptr.add(i).write(value) }; + + value_offset += (v & 1) as usize; + } + } else { + for i in 0..len { + let value = unsafe { values.get_unchecked(value_offset) }; + let value = dfn.decode(value); + unsafe { target_ptr.add(i).write(value) }; + + value_offset += (v & 1) as usize; + } + } + + num_values_remaining -= num_chunk_values; + unsafe { + target_ptr = target_ptr.add(len); + } + }; + + for v in validity_iter.by_ref() { + if limit < 56 { + iter(v, limit); + } else { + iter(v, 56); + } + limit -= 56; + } + + let (v, vl) = validity_iter.remainder(); + + iter(v, vl.min(limit)); + + unsafe { target.set_len(end_length) }; + + Ok(()) +} + +#[inline(never)] +fn decode_masked_required>( + values: ArrayChunks<'_, P>, + mask: &Bitmap, + target: &mut Vec, + dfn: D, +) -> ParquetResult<()> { + let num_rows = mask.set_bits(); + + if num_rows == mask.len() { + return decode_required(values, Some(num_rows), target, dfn); + } + + assert!(mask.len() <= values.len()); + + let start_length = target.len(); + target.reserve(num_rows); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut mask_iter = mask.fast_iter_u56(); + let mut num_rows_left = num_rows; + let mut value_offset = 0; + + let mut iter = |mut f: u64, len: usize| { + if num_rows_left == 0 { + return false; + } + + let mut num_read = 0; + let mut num_written = 0; + + while f != 0 { + let offset = f.trailing_zeros() as usize; + + num_read += offset; + + // SAFETY: + // 1. `values_buffer` starts out as only zeros, which we know is in the + // dictionary following the original `dict.is_empty` check. + // 2. Each time we write to `values_buffer`, it is followed by a + // `verify_dict_indices`. + let value = unsafe { values.get_unchecked(value_offset + num_read) }; + let value = dfn.decode(value); + unsafe { target_ptr.add(num_written).write(value) }; + + num_written += 1; + num_read += 1; + + f >>= offset + 1; // Clear least significant bit. + } + + unsafe { + target_ptr = target_ptr.add(num_written); + } + value_offset += len; + num_rows_left -= num_written; + + true + }; + + for f in mask_iter.by_ref() { + if !iter(f, 56) { + break; + } + } + + let (f, fl) = mask_iter.remainder(); + + iter(f, fl); + + unsafe { target.set_len(start_length + num_rows) }; + + Ok(()) +} + +#[inline(never)] +fn decode_masked_optional>( + values: ArrayChunks<'_, P>, + validity: &Bitmap, + mask: &Bitmap, + target: &mut Vec, + dfn: D, +) -> ParquetResult<()> { + let num_rows = mask.set_bits(); + let num_values = validity.set_bits(); + + if num_rows == mask.len() { + return decode_optional(values, validity, target, dfn); + } + + if num_values == validity.len() { + return decode_masked_required(values, mask, target, dfn); + } + + assert!(mask.len() <= values.len()); + + let start_length = target.len(); + target.reserve(num_rows); + let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; + + let mut validity_iter = validity.fast_iter_u56(); + let mut mask_iter = mask.fast_iter_u56(); + let mut num_values_left = num_values; + let mut num_rows_left = num_rows; + let mut value_offset = 0; + + let mut iter = |mut f: u64, mut v: u64, len: usize| { + if num_rows_left == 0 { + return false; + } + + let num_chunk_values = v.count_ones() as usize; + + let mut num_read = 0; + let mut num_written = 0; + + if num_chunk_values == num_values_left { + while f != 0 { + let offset = f.trailing_zeros() as usize; + + num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + v >>= offset; + + let is_valid = v & 1 != 0; + let value = if is_valid { + let value = unsafe { values.get_unchecked(value_offset + num_read) }; + dfn.decode(value) + } else { + T::zeroed() + }; + unsafe { target_ptr.add(num_written).write(value) }; + + num_written += 1; + num_read += (v & 1) as usize; + + f >>= offset + 1; // Clear least significant bit. + v >>= 1; + } + } else { + while f != 0 { + let offset = f.trailing_zeros() as usize; + + num_read += (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + v >>= offset; + + let value = unsafe { values.get_unchecked(value_offset + num_read) }; + let value = dfn.decode(value); + unsafe { target_ptr.add(num_written).write(value) }; + + num_written += 1; + num_read += (v & 1) as usize; + + f >>= offset + 1; // Clear least significant bit. + v >>= 1; + } + } + + unsafe { + target_ptr = target_ptr.add(num_written); + } + value_offset += len; + num_rows_left -= num_written; + num_values_left -= num_chunk_values; + + true + }; + + for (f, v) in mask_iter.by_ref().zip(validity_iter.by_ref()) { + if !iter(f, v, 56) { + break; + } + } + + let (f, fl) = mask_iter.remainder(); + let (v, vl) = validity_iter.remainder(); + + assert_eq!(fl, vl); + + iter(f, v, fl); + + unsafe { target.set_len(start_length + num_rows) }; + + Ok(()) +} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index b578f55892ca..270c0e4d78ca 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -1,5 +1,3 @@ -use std::ops::Range; - use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::types::NativeType; use bytemuck::Pod; @@ -9,43 +7,7 @@ use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::read::{Filter, ParquetError}; -fn filter_from_range(rng: Range) -> Bitmap { - let mut bm = MutableBitmap::with_capacity(rng.end); - - bm.extend_constant(rng.start, false); - bm.extend_constant(rng.len(), true); - - bm.freeze() -} - -fn decode_page_validity( - mut page_validity: HybridRleDecoder<'_>, - limit: Option, -) -> ParquetResult { - let mut limit = limit.unwrap_or(page_validity.len()); - let mut bm = MutableBitmap::with_capacity(limit); - - while let Some(chunk) = page_validity.next_chunk()? { - if limit == 0 { - break; - } - - match chunk { - HybridRleChunk::Rle(value, size) => { - let size = size.min(limit); - bm.extend_constant(size, value != 0); - limit -= size; - }, - HybridRleChunk::Bitpacked(decoder) => { - let len = decoder.len().min(limit); - bm.extend_from_slice(decoder.as_slice(), 0, len); - limit -= len; - }, - } - } - - Ok(bm.freeze()) -} +use super::{decode_page_validity, filter_from_range}; pub fn decode_dict( values: HybridRleDecoder<'_>, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 16e5b5cb6aa6..635102ddcaa2 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -2,6 +2,8 @@ pub(crate) mod array_chunks; pub(crate) mod dict_encoded; pub(crate) mod filter; +use std::ops::Range; + use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; @@ -12,7 +14,7 @@ use super::BasicDecompressor; use crate::parquet::encoding::hybrid_rle::gatherer::{ HybridRleGatherer, ZeroCount, ZeroCountGatherer, }; -use crate::parquet::encoding::hybrid_rle::{self, HybridRleDecoder}; +use crate::parquet::encoding::hybrid_rle::{self, HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::schema::Repetition; @@ -715,3 +717,41 @@ pub(crate) fn hybrid_rle_count_zeros( .gather_into(&mut count, &ZeroCountGatherer)?; Ok(count.num_zero) } + +pub(crate) fn filter_from_range(rng: Range) -> Bitmap { + let mut bm = MutableBitmap::with_capacity(rng.end); + + bm.extend_constant(rng.start, false); + bm.extend_constant(rng.len(), true); + + bm.freeze() +} + +pub(crate) fn decode_page_validity( + mut page_validity: HybridRleDecoder<'_>, + limit: Option, +) -> ParquetResult { + let mut limit = limit.unwrap_or(page_validity.len()); + let mut bm = MutableBitmap::with_capacity(limit); + + while let Some(chunk) = page_validity.next_chunk()? { + if limit == 0 { + break; + } + + match chunk { + HybridRleChunk::Rle(value, size) => { + let size = size.min(limit); + bm.extend_constant(size, value != 0); + limit -= size; + }, + HybridRleChunk::Bitpacked(decoder) => { + let len = decoder.len().min(limit); + bm.extend_from_slice(decoder.as_slice(), 0, len); + limit -= len; + }, + } + } + + Ok(bm.freeze()) +} From c544e675f0968a23bcf8cc61e6c445edaa341e4f Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Tue, 15 Oct 2024 11:59:31 +0200 Subject: [PATCH 06/32] working nested --- .../src/arrow/read/deserialize/binview.rs | 40 ++---- .../src/arrow/read/deserialize/boolean.rs | 31 +---- .../src/arrow/read/deserialize/dictionary.rs | 33 +---- .../read/deserialize/fixed_size_binary.rs | 36 ++--- .../arrow/read/deserialize/nested_utils.rs | 127 ++++++++++------- .../src/arrow/read/deserialize/null.rs | 27 +--- .../arrow/read/deserialize/primitive/float.rs | 49 ++----- .../read/deserialize/primitive/integer.rs | 78 ++--------- .../arrow/read/deserialize/primitive/plain.rs | 28 ++-- .../read/deserialize/utils/dict_encoded.rs | 27 ++-- .../src/arrow/read/deserialize/utils/mod.rs | 128 +++++++----------- .../src/parquet/encoding/hybrid_rle/mod.rs | 11 +- 12 files changed, 220 insertions(+), 395 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index d987419f8237..6eb1ad15051c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -5,7 +5,7 @@ use arrow::array::{ Array, BinaryViewArray, DictionaryArray, DictionaryKey, MutableBinaryViewArray, PrimitiveArray, Utf8ViewArray, View, }; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::buffer::Buffer; use arrow::datatypes::{ArrowDataType, PhysicalType}; @@ -14,7 +14,7 @@ use crate::parquet::encoding::delta_bitpacked::{lin_natural_sum, DeltaGatherer}; use crate::parquet::encoding::{delta_byte_array, delta_length_byte_array, hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; -use crate::read::deserialize::utils::{self, extend_from_decoder, Decoder, PageValidity}; +use crate::read::deserialize::utils::{self, extend_from_decoder, Decoder}; use crate::read::PrimitiveLogicalType; type DecodedStateTuple = (MutableBinaryViewArray<[u8]>, MutableBitmap); @@ -26,7 +26,7 @@ impl<'a> utils::StateTranslation<'a, BinViewDecoder> for StateTranslation<'a> { decoder: &BinViewDecoder, page: &'a DataPage, dict: Option<&'a ::Dict>, - _page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult { let is_string = matches!( page.descriptor.primitive_type.logical_type, @@ -41,7 +41,8 @@ impl<'a> utils::StateTranslation<'a, BinViewDecoder> for StateTranslation<'a> { Ok(Self::Plain(values)) }, (Encoding::PlainDictionary | Encoding::RleDictionary, Some(_)) => { - let values = dict_indices_decoder(page)?; + let values = + dict_indices_decoder(page, page_validity.map_or(0, |bm| bm.unset_bits()))?; Ok(Self::Dictionary(values)) }, (Encoding::DeltaLengthByteArray, _) => { @@ -90,7 +91,7 @@ impl<'a> utils::StateTranslation<'a, BinViewDecoder> for StateTranslation<'a> { decoder: &mut BinViewDecoder, decoded: &mut ::DecodedState, is_optional: bool, - page_validity: &mut Option>, + page_validity: &mut Option, dict: Option<&'a ::Dict>, additional: usize, ) -> ParquetResult<()> { @@ -609,7 +610,7 @@ impl utils::Decoder for BinViewDecoder { (values, validity): &mut Self::DecodedState, page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, + page_validity: Option<&mut Bitmap>, limit: usize, ) -> ParquetResult<()> { let views_offset = values.views().len(); @@ -702,7 +703,7 @@ impl utils::Decoder for BinViewDecoder { _decoded: &mut Self::DecodedState, _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { @@ -711,7 +712,7 @@ impl utils::Decoder for BinViewDecoder { fn extend_filtered_with_state<'a>( &mut self, - state: &mut utils::State<'a, Self>, + mut state: utils::State<'a, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -733,7 +734,7 @@ impl utils::Decoder for BinViewDecoder { indexes.clone(), &dict, state.is_optional, - state.page_validity.clone(), + state.page_validity.as_ref(), filter, &mut decoded.1, unsafe { decoded.0.views_mut() }, @@ -755,7 +756,7 @@ impl utils::Decoder for BinViewDecoder { // @NOTE: Needed for compatibility now. indexes.skip_in_place(max_offset)?; if let Some(ref mut page_validity) = state.page_validity { - page_validity.skip_in_place(max_offset)?; + page_validity.slice(max_offset, page_validity.len() - max_offset); } Ok(()) @@ -826,25 +827,6 @@ impl utils::DictDecodable for BinViewDecoder { } } -impl utils::NestedDecoder for BinViewDecoder { - fn validity_extend( - _: &mut utils::State<'_, Self>, - (_, validity): &mut Self::DecodedState, - value: bool, - n: usize, - ) { - validity.extend_constant(n, value); - } - - fn values_extend_nulls( - _: &mut utils::State<'_, Self>, - (values, _): &mut Self::DecodedState, - n: usize, - ) { - values.extend_constant(n, >::None); - } -} - #[derive(Debug)] pub struct BinaryIter<'a> { values: &'a [u8], diff --git a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs index af2e504d2646..361028d2bab2 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs @@ -1,6 +1,6 @@ use arrow::array::BooleanArray; use arrow::bitmap::utils::BitmapIter; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use super::utils::{self, extend_from_decoder, freeze_validity, Decoder, ExactSize}; @@ -9,7 +9,7 @@ use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; use crate::parquet::encoding::Encoding; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; -use crate::read::deserialize::utils::{BatchableCollector, PageValidity}; +use crate::read::deserialize::utils::BatchableCollector; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -25,7 +25,7 @@ impl<'a> utils::StateTranslation<'a, BooleanDecoder> for StateTranslation<'a> { _decoder: &BooleanDecoder, page: &'a DataPage, _dict: Option<&'a ::Dict>, - page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult { let values = split_buffer(page)?.values; @@ -90,7 +90,7 @@ impl<'a> utils::StateTranslation<'a, BooleanDecoder> for StateTranslation<'a> { decoder: &mut BooleanDecoder, decoded: &mut ::DecodedState, is_optional: bool, - page_validity: &mut Option>, + page_validity: &mut Option, _: Option<&'a ::Dict>, additional: usize, ) -> ParquetResult<()> { @@ -216,7 +216,7 @@ impl Decoder for BooleanDecoder { (values, validity): &mut Self::DecodedState, page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, + page_validity: Option<&mut Bitmap>, limit: usize, ) -> ParquetResult<()> { match page_validity { @@ -240,7 +240,7 @@ impl Decoder for BooleanDecoder { _decoded: &mut Self::DecodedState, _page_values: &mut HybridRleDecoder<'a>, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { @@ -257,22 +257,3 @@ impl Decoder for BooleanDecoder { Ok(BooleanArray::new(dtype, values.into(), validity)) } } - -impl utils::NestedDecoder for BooleanDecoder { - fn validity_extend( - _: &mut utils::State<'_, Self>, - (_, validity): &mut Self::DecodedState, - value: bool, - n: usize, - ) { - validity.extend_constant(n, value); - } - - fn values_extend_nulls( - _: &mut utils::State<'_, Self>, - (values, _): &mut Self::DecodedState, - n: usize, - ) { - values.extend_constant(n, false); - } -} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs index de2bfe2e47f3..06251691eb30 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs @@ -1,12 +1,12 @@ use std::sync::atomic::AtomicUsize; use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use super::utils::{ self, dict_indices_decoder, extend_from_decoder, freeze_validity, BatchableCollector, Decoder, - DictDecodable, ExactSize, PageValidity, StateTranslation, + ExactSize, StateTranslation, }; use super::ParquetError; use crate::parquet::encoding::hybrid_rle::{self, HybridRleDecoder, Translator}; @@ -23,7 +23,7 @@ impl<'a, K: DictionaryKey, D: utils::DictDecodable> StateTranslation<'a, Diction _decoder: &DictionaryDecoder, page: &'a DataPage, _dict: Option<&'a as Decoder>::Dict>, - _page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult { if !matches!( page.encoding(), @@ -32,7 +32,7 @@ impl<'a, K: DictionaryKey, D: utils::DictDecodable> StateTranslation<'a, Diction return Err(utils::not_implemented(page)); } - dict_indices_decoder(page) + dict_indices_decoder(page, page_validity.map_or(0, |bm| bm.unset_bits())) } fn len_when_not_nullable(&self) -> usize { @@ -48,7 +48,7 @@ impl<'a, K: DictionaryKey, D: utils::DictDecodable> StateTranslation<'a, Diction decoder: &mut DictionaryDecoder, decoded: &mut as Decoder>::DecodedState, is_optional: bool, - page_validity: &mut Option>, + page_validity: &mut Option, _: Option<&'a as Decoder>::Dict>, additional: usize, ) -> ParquetResult<()> { @@ -137,7 +137,7 @@ impl utils::Decoder for DictionaryDec _decoded: &mut Self::DecodedState, _page_values: &mut as StateTranslation<'a, Self>>::PlainDecoder, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _limit: usize, ) -> ParquetResult<()> { unreachable!() @@ -148,7 +148,7 @@ impl utils::Decoder for DictionaryDec _decoded: &mut Self::DecodedState, _page_values: &mut HybridRleDecoder<'a>, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { @@ -156,25 +156,6 @@ impl utils::Decoder for DictionaryDec } } -impl utils::NestedDecoder for DictionaryDecoder { - fn validity_extend( - _: &mut utils::State<'_, Self>, - (_, validity): &mut Self::DecodedState, - value: bool, - n: usize, - ) { - validity.extend_constant(n, value); - } - - fn values_extend_nulls( - _: &mut utils::State<'_, Self>, - (values, _): &mut Self::DecodedState, - n: usize, - ) { - values.resize(values.len() + n, K::default()); - } -} - pub(crate) struct DictArrayCollector<'a, 'b> { values: &'b mut hybrid_rle::HybridRleDecoder<'a>, dict_size: usize, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 3825d528c8f5..8f6aab926fa3 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -1,5 +1,5 @@ use arrow::array::{DictionaryArray, DictionaryKey, FixedSizeBinaryArray, PrimitiveArray}; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use super::utils::{dict_indices_decoder, extend_from_decoder, freeze_validity, Decoder}; @@ -7,7 +7,7 @@ use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; use crate::parquet::encoding::{hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; -use crate::read::deserialize::utils::{self, BatchableCollector, GatheredHybridRle, PageValidity}; +use crate::read::deserialize::utils::{self, BatchableCollector, GatheredHybridRle}; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -29,7 +29,7 @@ impl<'a> utils::StateTranslation<'a, BinaryDecoder> for StateTranslation<'a> { decoder: &BinaryDecoder, page: &'a DataPage, dict: Option<&'a ::Dict>, - _page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult { match (page.encoding(), dict) { (Encoding::Plain, _) => { @@ -44,7 +44,8 @@ impl<'a> utils::StateTranslation<'a, BinaryDecoder> for StateTranslation<'a> { Ok(Self::Plain(values, decoder.size)) }, (Encoding::PlainDictionary | Encoding::RleDictionary, Some(_)) => { - let values = dict_indices_decoder(page)?; + let values = + dict_indices_decoder(page, page_validity.map_or(0, |bm| bm.unset_bits()))?; Ok(Self::Dictionary(values)) }, _ => Err(utils::not_implemented(page)), @@ -76,7 +77,7 @@ impl<'a> utils::StateTranslation<'a, BinaryDecoder> for StateTranslation<'a> { decoder: &mut BinaryDecoder, decoded: &mut ::DecodedState, is_optional: bool, - page_validity: &mut Option>, + page_validity: &mut Option, dict: Option<&'a ::Dict>, additional: usize, ) -> ParquetResult<()> { @@ -146,7 +147,7 @@ impl Decoder for BinaryDecoder { (values, validity): &mut Self::DecodedState, page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, + page_validity: Option<&mut Bitmap>, limit: usize, ) -> ParquetResult<()> { struct FixedSizeBinaryCollector<'a, 'b> { @@ -208,7 +209,7 @@ impl Decoder for BinaryDecoder { (values, validity): &mut Self::DecodedState, page_values: &mut hybrid_rle::HybridRleDecoder<'a>, is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, + page_validity: Option<&mut Bitmap>, dict: &Self::Dict, limit: usize, ) -> ParquetResult<()> { @@ -333,24 +334,3 @@ impl utils::DictDecodable for BinaryDecoder { Ok(DictionaryArray::try_new(dtype, keys, Box::new(dict)).unwrap()) } } - -impl utils::NestedDecoder for BinaryDecoder { - fn validity_extend( - _: &mut utils::State<'_, Self>, - (_, validity): &mut Self::DecodedState, - value: bool, - n: usize, - ) { - validity.extend_constant(n, value); - } - - fn values_extend_nulls( - _: &mut utils::State<'_, Self>, - (values, _): &mut Self::DecodedState, - n: usize, - ) { - values - .values - .resize(values.values.len() + n * values.size, 0); - } -} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index a622153dfca8..5b407b6ce4f0 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -9,7 +9,7 @@ use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage}; use crate::parquet::read::levels::get_bit_width; -use crate::read::deserialize::utils::{hybrid_rle_count_zeros, BatchedCollector}; +use crate::read::deserialize::utils::{hybrid_rle_count_zeros, BatchedCollector, State}; #[derive(Debug)] pub struct Nested { @@ -210,30 +210,34 @@ impl Nested { } } -pub struct BatchedNestedDecoder<'a, 'b, 'c, D: utils::NestedDecoder> { - state: &'b mut utils::State<'a, D>, - decoder: &'c mut D, +pub struct BatchedNestedDecoder<'a> { + filter: &'a mut MutableBitmap, + validity: &'a mut MutableBitmap, } -impl<'a, 'b, 'c, D: utils::NestedDecoder> BatchableCollector<(), D::DecodedState> - for BatchedNestedDecoder<'a, 'b, 'c, D> -{ - fn reserve(_target: &mut D::DecodedState, _n: usize) { +impl<'a> BatchableCollector<(), ()> for BatchedNestedDecoder<'a> { + fn reserve(_target: &mut (), _n: usize) { unreachable!() } - fn push_n(&mut self, target: &mut D::DecodedState, n: usize) -> ParquetResult<()> { - self.decoder.push_n_valids(self.state, target, n)?; + fn push_n(&mut self, _target: &mut (), n: usize) -> ParquetResult<()> { + self.filter.extend_constant(n, true); + self.validity.extend_constant(n, true); Ok(()) } - fn push_n_nulls(&mut self, target: &mut D::DecodedState, n: usize) -> ParquetResult<()> { - self.decoder.push_n_nulls(self.state, target, n); + fn push_n_nulls(&mut self, _target: &mut (), n: usize) -> ParquetResult<()> { + self.filter.extend_constant(n, true); + self.validity.extend_constant(n, false); + Ok(()) } fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.state.skip_in_place(n) + self.filter.extend_constant(n, false); + self.validity.extend_constant(n, true); + + Ok(()) } } @@ -465,15 +469,10 @@ fn idx_to_limit(rep_iter: &HybridRleDecoder<'_>, idx: usize) -> ParquetResult( +fn extend_offsets2<'a>( mut def_iter: HybridRleDecoder<'a>, mut rep_iter: HybridRleDecoder<'a>, - batched_collector: &mut BatchedCollector< - '_, - (), - D::DecodedState, - BatchedNestedDecoder<'a, '_, '_, D>, - >, + batched_collector: &mut BatchedCollector<'_, (), (), BatchedNestedDecoder<'a>>, nested: &mut [Nested], filter: Option, @@ -570,15 +569,10 @@ fn extend_offsets2<'a, D: utils::NestedDecoder>( } } -fn extend_offsets_limited<'a, D: utils::NestedDecoder>( +fn extend_offsets_limited<'a>( def_iter: &mut HybridRleDecoder<'a>, rep_iter: &mut HybridRleDecoder<'a>, - batched_collector: &mut BatchedCollector< - '_, - (), - D::DecodedState, - BatchedNestedDecoder<'a, '_, '_, D>, - >, + batched_collector: &mut BatchedCollector<'_, (), (), BatchedNestedDecoder<'a>>, nested: &mut [Nested], mut limit: usize, // Amortized allocations @@ -740,7 +734,7 @@ fn extend_offsets_limited<'a, D: utils::NestedDecoder>( Ok(()) } -pub struct PageNestedDecoder { +pub struct PageNestedDecoder { pub iter: BasicDecompressor, pub dtype: ArrowDataType, pub dict: Option, @@ -763,7 +757,7 @@ fn level_iters(page: &DataPage) -> ParquetResult<(HybridRleDecoder, HybridRleDec Ok((def_iter, rep_iter)) } -impl PageNestedDecoder { +impl PageNestedDecoder { pub fn new( mut iter: BasicDecompressor, dtype: ArrowDataType, @@ -804,17 +798,19 @@ impl PageNestedDecoder { let page = page?; let page = page.decompress(&mut self.iter)?; - let mut state = - utils::State::new_nested(&self.decoder, &page, self.dict.as_ref())?; + let mut filter = MutableBitmap::new(); + let mut validity = MutableBitmap::new(); + let mut _tgt = (); + let (def_iter, rep_iter) = level_iters(&page)?; // @TODO: move this to outside the loop. let mut batched_collector = BatchedCollector::new( BatchedNestedDecoder { - state: &mut state, - decoder: &mut self.decoder, + filter: &mut filter, + validity: &mut validity, }, - &mut target, + &mut _tgt, ); extend_offsets2( @@ -829,7 +825,18 @@ impl PageNestedDecoder { batched_collector.finalize()?; - drop(state); + let state = utils::State::new_nested( + &self.decoder, + &page, + self.dict.as_ref(), + Some(validity.freeze()), + )?; + state.decode( + &mut self.decoder, + &mut target, + Some(Filter::Mask(filter.freeze())), + )?; + self.iter.reuse_page_buffer(page); } }, @@ -853,7 +860,7 @@ impl PageNestedDecoder { }; let page = page?; // We cannot lazily decompress because we don't have the number of row values - // at this point. We need repetition levels for that. *sign*. In general, lazy + // at this point. We need repetition levels for that. *sigh*. In general, lazy // decompression is quite difficult with nested values. // // @TODO @@ -864,7 +871,7 @@ impl PageNestedDecoder { let (mut def_iter, mut rep_iter) = level_iters(&page)?; - let mut state; + let state; let mut batched_collector; let start_length = nested_state.len(); @@ -875,6 +882,10 @@ impl PageNestedDecoder { let state_filter; (state_filter, filter) = Filter::split_at(&filter, num_row_values); + let mut leaf_filter = MutableBitmap::new(); + let mut leaf_validity = MutableBitmap::new(); + let mut _tgt = (); + match last_row_value_action { PageStartAction::Skip => { // Fast path: skip the whole page. @@ -889,14 +900,12 @@ impl PageNestedDecoder { // We just saw that we had at least one row value. debug_assert!(limit < rep_iter.len()); - state = - utils::State::new_nested(&self.decoder, &page, self.dict.as_ref())?; batched_collector = BatchedCollector::new( BatchedNestedDecoder { - state: &mut state, - decoder: &mut self.decoder, + filter: &mut leaf_filter, + validity: &mut leaf_validity, }, - &mut target, + &mut _tgt, ); let num_leaf_values = @@ -920,14 +929,12 @@ impl PageNestedDecoder { continue; } - state = - utils::State::new_nested(&self.decoder, &page, self.dict.as_ref())?; batched_collector = BatchedCollector::new( BatchedNestedDecoder { - state: &mut state, - decoder: &mut self.decoder, + filter: &mut leaf_filter, + validity: &mut leaf_validity, }, - &mut target, + &mut _tgt, ); extend_offsets_limited( @@ -944,12 +951,23 @@ impl PageNestedDecoder { if rep_iter.len() == 0 { batched_collector.finalize()?; + state = State::new_nested( + &self.decoder, + &page, + self.dict.as_ref(), + Some(leaf_validity.freeze()), + )?; + state.decode( + &mut self.decoder, + &mut target, + Some(Filter::Mask(leaf_filter.freeze())), + )?; + let num_done = nested_state.len() - start_length; debug_assert!(num_done <= num_rows_remaining); debug_assert!(num_done <= num_row_values); num_rows_remaining -= num_done; - drop(state); self.iter.reuse_page_buffer(page); continue; @@ -981,12 +999,23 @@ impl PageNestedDecoder { batched_collector.finalize()?; + state = State::new_nested( + &self.decoder, + &page, + self.dict.as_ref(), + Some(leaf_validity.freeze()), + )?; + state.decode( + &mut self.decoder, + &mut target, + Some(Filter::Mask(leaf_filter.freeze())), + )?; + let num_done = nested_state.len() - start_length; debug_assert!(num_done <= num_rows_remaining); debug_assert!(num_done <= num_row_values); num_rows_remaining -= num_done; - drop(state); self.iter.reuse_page_buffer(page); } }, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index e12757fe2e20..16a225d9a91c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -3,6 +3,7 @@ //! `DecodedState`. use arrow::array::{Array, NullArray}; +use arrow::bitmap::Bitmap; use arrow::datatypes::ArrowDataType; use super::utils; @@ -30,7 +31,7 @@ impl<'a> utils::StateTranslation<'a, NullDecoder> for () { _decoder: &NullDecoder, _page: &'a DataPage, _dict: Option<&'a ::Dict>, - _page_validity: Option<&utils::PageValidity<'a>>, + _page_validity: Option<&Bitmap>, ) -> ParquetResult { Ok(()) } @@ -48,7 +49,7 @@ impl<'a> utils::StateTranslation<'a, NullDecoder> for () { _decoder: &mut NullDecoder, decoded: &mut ::DecodedState, _is_optional: bool, - _page_validity: &mut Option>, + _page_validity: &mut Option, _: Option<&'a ::Dict>, additional: usize, ) -> ParquetResult<()> { @@ -77,7 +78,7 @@ impl utils::Decoder for NullDecoder { _decoded: &mut Self::DecodedState, _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, _is_optional: bool, - _page_validity: Option<&mut utils::PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _limit: usize, ) -> ParquetResult<()> { unimplemented!() @@ -88,7 +89,7 @@ impl utils::Decoder for NullDecoder { _decoded: &mut Self::DecodedState, _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, _is_optional: bool, - _page_validity: Option<&mut utils::PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { @@ -105,24 +106,6 @@ impl utils::Decoder for NullDecoder { } } -impl utils::NestedDecoder for NullDecoder { - fn validity_extend( - _: &mut utils::State<'_, Self>, - _: &mut Self::DecodedState, - _value: bool, - _n: usize, - ) { - } - - fn values_extend_nulls( - _state: &mut utils::State<'_, Self>, - decoded: &mut Self::DecodedState, - n: usize, - ) { - decoded.length += n; - } -} - use super::BasicDecompressor; /// Converts [`PagesIter`] to an [`ArrayIter`] diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index 73adb50a4d88..b1388feb1e85 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -1,5 +1,5 @@ use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::types::NativeType; @@ -14,7 +14,7 @@ use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{ - dict_indices_decoder, freeze_validity, Decoder, PageValidity, + dict_indices_decoder, freeze_validity, Decoder, }; use crate::read::{Filter, ParquetError}; @@ -38,11 +38,12 @@ where _decoder: &FloatDecoder, page: &'a DataPage, dict: Option<&'a as utils::Decoder>::Dict>, - _page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult { match (page.encoding(), dict) { (Encoding::PlainDictionary | Encoding::RleDictionary, Some(_)) => { - let values = dict_indices_decoder(page)?; + let values = + dict_indices_decoder(page, page_validity.map_or(0, |bm| bm.unset_bits()))?; Ok(Self::Dictionary(values)) }, (Encoding::Plain, _) => { @@ -88,7 +89,7 @@ where decoder: &mut FloatDecoder, decoded: &mut as utils::Decoder>::DecodedState, is_optional: bool, - page_validity: &mut Option>, + page_validity: &mut Option, dict: Option<&'a as utils::Decoder>::Dict>, additional: usize, ) -> ParquetResult<()> { @@ -238,7 +239,7 @@ where _decoded: &mut Self::DecodedState, _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _limit: usize, ) -> ParquetResult<()> { unreachable!() @@ -249,7 +250,7 @@ where _decoded: &mut Self::DecodedState, _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { @@ -258,7 +259,7 @@ where fn extend_filtered_with_state<'a>( &mut self, - state: &mut utils::State<'a, Self>, + mut state: utils::State<'a, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -275,7 +276,7 @@ where super::plain::decode( values.clone(), state.is_optional, - state.page_validity.clone(), + state.page_validity.as_ref(), filter, &mut decoded.1, &mut decoded.0, @@ -285,7 +286,7 @@ where // @NOTE: Needed for compatibility now. values.skip_in_place(max_offset); if let Some(ref mut page_validity) = state.page_validity { - page_validity.skip_in_place(max_offset)?; + page_validity.slice(max_offset, page_validity.len() - max_offset); } Ok(()) @@ -295,7 +296,7 @@ where indexes.clone(), state.dict.unwrap(), state.is_optional, - state.page_validity.clone(), + state.page_validity.as_ref(), filter, &mut decoded.1, &mut decoded.0, @@ -304,7 +305,7 @@ where // @NOTE: Needed for compatibility now. indexes.skip_in_place(max_offset)?; if let Some(ref mut page_validity) = state.page_validity { - page_validity.skip_in_place(max_offset)?; + page_validity.slice(max_offset, page_validity.len() - max_offset); } Ok(()) @@ -346,27 +347,3 @@ where Ok(DictionaryArray::try_new(dtype, keys, dict).unwrap()) } } - -impl utils::NestedDecoder for FloatDecoder -where - T: NativeType, - P: ParquetNativeType, - D: DecoderFunction, -{ - fn validity_extend( - _: &mut utils::State<'_, Self>, - (_, validity): &mut Self::DecodedState, - value: bool, - n: usize, - ) { - validity.extend_constant(n, value); - } - - fn values_extend_nulls( - _: &mut utils::State<'_, Self>, - (values, _): &mut Self::DecodedState, - n: usize, - ) { - values.resize(values.len() + n, T::default()); - } -} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 17d40c9c2998..f6b778bd2255 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -1,5 +1,5 @@ use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::types::NativeType; @@ -13,9 +13,7 @@ use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; -use crate::read::deserialize::utils::{ - dict_indices_decoder, freeze_validity, Decoder, PageValidity, -}; +use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity}; use crate::read::{Filter, ParquetError}; #[allow(clippy::large_enum_variant)] @@ -40,11 +38,12 @@ where _decoder: &IntDecoder, page: &'a DataPage, dict: Option<&'a as utils::Decoder>::Dict>, - _page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult { match (page.encoding(), dict) { (Encoding::PlainDictionary | Encoding::RleDictionary, Some(_)) => { - let values = dict_indices_decoder(page)?; + let values = + dict_indices_decoder(page, page_validity.map_or(0, |bm| bm.unset_bits()))?; Ok(Self::Dictionary(values)) }, (Encoding::Plain, _) => { @@ -98,26 +97,11 @@ where decoder: &mut IntDecoder, decoded: &mut as utils::Decoder>::DecodedState, is_optional: bool, - page_validity: &mut Option>, - dict: Option<&'a as utils::Decoder>::Dict>, + page_validity: &mut Option, + _dict: Option<&'a as utils::Decoder>::Dict>, additional: usize, ) -> ParquetResult<()> { match self { - Self::Plain(page_values) => decoder.decode_plain_encoded( - decoded, - page_values, - is_optional, - page_validity.as_mut(), - additional, - )?, - Self::Dictionary(ref mut page) => decoder.decode_dictionary_encoded( - decoded, - page, - is_optional, - page_validity.as_mut(), - dict.unwrap(), - additional, - )?, Self::ByteStreamSplit(page_values) => { let (values, validity) = decoded; @@ -173,6 +157,7 @@ where )?, } }, + _ => unreachable!(), } Ok(()) @@ -292,7 +277,7 @@ where _decoded: &mut Self::DecodedState, _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _limit: usize, ) -> ParquetResult<()> { unreachable!() @@ -303,7 +288,7 @@ where _decoded: &mut Self::DecodedState, _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, _is_optional: bool, - _page_validity: Option<&mut PageValidity<'a>>, + _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { @@ -322,7 +307,7 @@ where fn extend_filtered_with_state<'a>( &mut self, - state: &mut utils::State<'a, Self>, + mut state: utils::State<'a, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -339,19 +324,13 @@ where super::plain::decode( values.clone(), state.is_optional, - state.page_validity.clone(), + state.page_validity.as_ref(), filter, &mut decoded.1, &mut decoded.0, self.0.decoder, )?; - // @NOTE: Needed for compatibility now. - values.skip_in_place(max_offset); - if let Some(ref mut page_validity) = state.page_validity { - page_validity.skip_in_place(max_offset)?; - } - Ok(()) }, StateTranslation::Dictionary(ref mut indexes) => { @@ -359,18 +338,12 @@ where indexes.clone(), state.dict.unwrap(), state.is_optional, - state.page_validity.clone(), + state.page_validity.as_ref(), filter, &mut decoded.1, &mut decoded.0, )?; - // @NOTE: Needed for compatibility now. - indexes.skip_in_place(max_offset)?; - if let Some(ref mut page_validity) = state.page_validity { - page_validity.skip_in_place(max_offset)?; - } - Ok(()) }, _ => self.extend_filtered_with_state_default(state, decoded, filter), @@ -401,28 +374,3 @@ where Ok(DictionaryArray::try_new(dtype, keys, dict).unwrap()) } } - -impl utils::NestedDecoder for IntDecoder -where - T: NativeType, - P: ParquetNativeType, - i64: num_traits::AsPrimitive

, - D: DecoderFunction, -{ - fn validity_extend( - _: &mut utils::State<'_, Self>, - (_, validity): &mut Self::DecodedState, - value: bool, - n: usize, - ) { - validity.extend_constant(n, value); - } - - fn values_extend_nulls( - _: &mut utils::State<'_, Self>, - (values, _): &mut Self::DecodedState, - n: usize, - ) { - values.resize(values.len() + n, T::default()); - } -} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 8f5bdb41e7cb..009af6667e2f 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -3,17 +3,16 @@ use arrow::types::NativeType; use polars_compute::filter::filter_boolean_kernel; use super::DecoderFunction; -use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; use crate::parquet::error::ParquetResult; use crate::parquet::types::NativeType as ParquetNativeType; use crate::read::deserialize::utils::array_chunks::ArrayChunks; -use crate::read::deserialize::utils::{decode_page_validity, filter_from_range}; +use crate::read::deserialize::utils::filter_from_range; use crate::read::Filter; pub fn decode>( values: ArrayChunks<'_, P>, is_optional: bool, - page_validity: Option>, + page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, target: &mut Vec, @@ -34,33 +33,24 @@ pub fn decode>( ) } - #[inline(never)] fn decode_plain_dispatch>( values: ArrayChunks<'_, P>, is_optional: bool, - page_validity: Option>, + page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, target: &mut Vec, dfn: D, ) -> ParquetResult<()> { - let page_validity = match (page_validity, filter.as_ref()) { - (None, _) => None, - (Some(page_validity), None) => Some(decode_page_validity(page_validity, None)?), - (Some(page_validity), Some(filter)) => Some(decode_page_validity( - page_validity, - Some(filter.max_offset()), - )?), - }; - if is_optional { match (page_validity.as_ref(), filter.as_ref()) { (None, None) => validity.extend_constant(values.len(), true), (None, Some(f)) => validity.extend_constant(f.num_rows(), true), (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), (Some(page_validity), Some(Filter::Range(rng))) => { - validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) + let page_validity = (*page_validity).clone(); + validity.extend_from_bitmap(&page_validity.sliced(rng.start, rng.len())) }, (Some(page_validity), Some(Filter::Mask(mask))) => { validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, &mask)) @@ -79,15 +69,15 @@ fn decode_plain_dispatch decode_masked_required(values, &filter, target, dfn), (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_optional(values, &filter, &page_validity, target, dfn) + decode_masked_optional(values, &page_validity, &filter, target, dfn) }, (Some(Filter::Range(rng)), None) => { decode_masked_required(values, &filter_from_range(rng.clone()), target, dfn) }, (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( values, - &filter_from_range(rng.clone()), &page_validity, + &filter_from_range(rng.clone()), target, dfn, ), @@ -138,7 +128,7 @@ fn decode_optional let mut num_values_remaining = num_values; let mut value_offset = 0; - let mut iter = |v: u64, len: usize| { + let mut iter = |mut v: u64, len: usize| { debug_assert!(len < 64); let num_chunk_values = v.count_ones() as usize; @@ -155,6 +145,7 @@ fn decode_optional unsafe { target_ptr.add(i).write(value) }; value_offset += (v & 1) as usize; + v >>= 1; } } else { for i in 0..len { @@ -163,6 +154,7 @@ fn decode_optional unsafe { target_ptr.add(i).write(value) }; value_offset += (v & 1) as usize; + v >>= 1; } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index 270c0e4d78ca..d3c89064b593 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -3,17 +3,16 @@ use arrow::types::NativeType; use bytemuck::Pod; use polars_compute::filter::filter_boolean_kernel; +use super::filter_from_range; use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::read::{Filter, ParquetError}; -use super::{decode_page_validity, filter_from_range}; - -pub fn decode_dict( +pub fn decode_dict( values: HybridRleDecoder<'_>, dict: &[T], is_optional: bool, - page_validity: Option>, + page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, target: &mut Vec, @@ -34,31 +33,23 @@ pub fn decode_dict( } #[inline(never)] -fn decode_dict_dispatch( +fn decode_dict_dispatch( values: HybridRleDecoder<'_>, dict: &[T], is_optional: bool, - page_validity: Option>, + page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, target: &mut Vec, ) -> ParquetResult<()> { - let page_validity = match (page_validity, filter.as_ref()) { - (None, _) => None, - (Some(page_validity), None) => Some(decode_page_validity(page_validity, None)?), - (Some(page_validity), Some(filter)) => Some(decode_page_validity( - page_validity, - Some(filter.max_offset()), - )?), - }; - if is_optional { match (page_validity.as_ref(), filter.as_ref()) { (None, None) => validity.extend_constant(values.len(), true), (None, Some(f)) => validity.extend_constant(f.num_rows(), true), (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), (Some(page_validity), Some(Filter::Range(rng))) => { - validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) + let pv = (*page_validity).clone(); + validity.extend_from_bitmap(&pv.sliced(rng.start, rng.len())) }, (Some(page_validity), Some(Filter::Mask(mask))) => { validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, &mask)) @@ -219,7 +210,7 @@ pub fn decode_optional_dict( // Dispatch to the required kernel if all rows are valid anyway. if num_valid_values == validity.len() { - return decode_required_dict(values, dict, None, target); + return decode_required_dict(values, dict, Some(validity.len()), target); } if dict.is_empty() && num_valid_values > 0 { @@ -604,7 +595,7 @@ pub fn decode_masked_required_dict( // Dispatch to the non-filter kernel if all rows are needed anyway. if num_rows == filter.len() { - return decode_required_dict(values, dict, None, target); + return decode_required_dict(values, dict, Some(filter.len()), target); } if dict.is_empty() && values.len() > 0 { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 635102ddcaa2..076739a248f7 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -23,7 +23,7 @@ use crate::parquet::schema::Repetition; pub(crate) struct State<'a, D: Decoder> { pub(crate) dict: Option<&'a D::Dict>, pub(crate) is_optional: bool, - pub(crate) page_validity: Option>, + pub(crate) page_validity: Option, pub(crate) translation: D::Translation<'a>, } @@ -34,7 +34,7 @@ pub(crate) trait StateTranslation<'a, D: Decoder>: Sized { decoder: &D, page: &'a DataPage, dict: Option<&'a D::Dict>, - page_validity: Option<&PageValidity<'a>>, + page_validity: Option<&Bitmap>, ) -> ParquetResult; fn len_when_not_nullable(&self) -> usize; fn skip_in_place(&mut self, n: usize) -> ParquetResult<()>; @@ -46,7 +46,7 @@ pub(crate) trait StateTranslation<'a, D: Decoder>: Sized { decoder: &mut D, decoded: &mut D::DecodedState, is_optional: bool, - page_validity: &mut Option>, + page_validity: &mut Option, dict: Option<&'a D::Dict>, additional: usize, ) -> ParquetResult<()>; @@ -57,18 +57,16 @@ impl<'a, D: Decoder> State<'a, D> { let is_optional = page.descriptor.primitive_type.field_info.repetition == Repetition::Optional; - let mut page_validity = is_optional - .then(|| page_validity_decoder(page)) - .transpose()?; + let mut page_validity = None; // Make the page_validity None if there are no nulls in the page - let null_count = page - .null_count() - .map(Ok) - .or_else(|| page_validity.as_ref().map(hybrid_rle_count_zeros)) - .transpose()?; - if null_count == Some(0) { - page_validity = None; + if is_optional && !page.null_count().is_some_and(|nc| nc == 0) { + let pv = page_validity_decoder(page)?; + let pv = decode_page_validity(pv, None)?; + + if pv.unset_bits() > 0 { + page_validity = Some(pv); + } } let translation = D::Translation::new(decoder, page, dict, page_validity.as_ref())?; @@ -85,16 +83,22 @@ impl<'a, D: Decoder> State<'a, D> { decoder: &D, page: &'a DataPage, dict: Option<&'a D::Dict>, + mut page_validity: Option, ) -> ParquetResult { let translation = D::Translation::new(decoder, page, dict, None)?; + let is_optional = + page.descriptor.primitive_type.field_info.repetition == Repetition::Optional; + + if page_validity.as_ref().is_some_and(|bm| bm.unset_bits() == 0) { + page_validity = None; + } + Ok(Self { dict, translation, - - // Nested values may be optional, but all that is handled elsewhere. - is_optional: false, - page_validity: None, + is_optional, + page_validity, }) } @@ -110,20 +114,17 @@ impl<'a, D: Decoder> State<'a, D> { return Ok(()); } - let n = self - .page_validity - .as_mut() - .map_or(ParquetResult::Ok(n), |page_validity| { - let mut zc = ZeroCount::default(); - page_validity.gather_n_into(&mut zc, n, &ZeroCountGatherer)?; - Ok(zc.num_nonzero) - })?; + let n = self.page_validity.as_mut().map_or(n, |page_validity| { + let mut pv = page_validity.clone(); + pv.slice(0, n); + pv.unset_bits() + }); self.translation.skip_in_place(n) } - pub fn extend_from_state( - &mut self, + pub fn decode( + self, decoder: &mut D, decoded: &mut D::DecodedState, filter: Option, @@ -315,7 +316,7 @@ impl<'a, I, T, C: BatchableCollector> HybridRleGatherer for BatchGath /// Extends a [`Pushable`] from an iterator of non-null values and an hybrid-rle decoder pub(super) fn extend_from_decoder>( validity: &mut MutableBitmap, - page_validity: &mut PageValidity, + page_validity: &mut Bitmap, limit: Option, target: &mut T, collector: C, @@ -325,13 +326,20 @@ pub(super) fn extend_from_decoder>( validity.reserve(num_elements); C::reserve(target, num_elements); - let batched_collector = BatchedCollector::new(collector, target); - let mut target = (validity, batched_collector); - let gatherer = BatchGatherer(Default::default()); + let mut batched_collector = BatchedCollector::new(collector, target); + + let mut pv = page_validity.clone(); + pv.slice(0, num_elements); - page_validity.gather_n_into(&mut target, num_elements, &gatherer)?; + // @TODO: This is terribly slow now. + validity.extend_from_bitmap(&pv); + let mut iter = pv.iter(); + while iter.num_remaining() > 0 { + batched_collector.push_n_valids(iter.take_leading_ones())?; + batched_collector.push_n_invalids(iter.take_leading_zeros()); + } - target.1.finalize()?; + batched_collector.finalize()?; Ok(()) } @@ -443,7 +451,7 @@ pub(super) trait Decoder: Sized { fn extend_filtered_with_state<'a>( &mut self, - state: &mut State<'a, Self>, + state: State<'a, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -452,7 +460,7 @@ pub(super) trait Decoder: Sized { fn extend_filtered_with_state_default<'a>( &mut self, - state: &mut State<'a, Self>, + mut state: State<'a, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -546,7 +554,7 @@ pub(super) trait Decoder: Sized { decoded: &mut Self::DecodedState, page_values: &mut as StateTranslation<'a, Self>>::PlainDecoder, is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, + page_validity: Option<&mut Bitmap>, limit: usize, ) -> ParquetResult<()>; fn decode_dictionary_encoded<'a>( @@ -554,7 +562,7 @@ pub(super) trait Decoder: Sized { decoded: &mut Self::DecodedState, page_values: &mut HybridRleDecoder<'a>, is_optional: bool, - page_validity: Option<&mut PageValidity<'a>>, + page_validity: Option<&mut Bitmap>, dict: &Self::Dict, limit: usize, ) -> ParquetResult<()>; @@ -567,38 +575,6 @@ pub(super) trait Decoder: Sized { ) -> ParquetResult; } -pub(crate) trait NestedDecoder: Decoder { - fn validity_extend( - state: &mut State<'_, Self>, - decoded: &mut Self::DecodedState, - value: bool, - n: usize, - ); - fn values_extend_nulls(state: &mut State<'_, Self>, decoded: &mut Self::DecodedState, n: usize); - - fn push_n_valids( - &mut self, - state: &mut State<'_, Self>, - decoded: &mut Self::DecodedState, - n: usize, - ) -> ParquetResult<()> { - state.extend_from_state(self, decoded, Some(Filter::new_limited(n)))?; - Self::validity_extend(state, decoded, true, n); - - Ok(()) - } - - fn push_n_nulls( - &self, - state: &mut State<'_, Self>, - decoded: &mut Self::DecodedState, - n: usize, - ) { - Self::validity_extend(state, decoded, false, n); - Self::values_extend_nulls(state, decoded, n); - } -} - pub trait DictDecodable: Decoder { fn finalize_dict_array( &self, @@ -657,17 +633,14 @@ impl PageDecoder { let page = page.decompress(&mut self.iter)?; - let mut state = State::new(&self.decoder, &page, self.dict.as_ref())?; + let state = State::new(&self.decoder, &page, self.dict.as_ref())?; let start_length = target.len(); - state.extend_from_state(&mut self.decoder, &mut target, state_filter)?; + state.decode(&mut self.decoder, &mut target, state_filter)?; let end_length = target.len(); num_rows_remaining -= end_length - start_length; - debug_assert!(state.len() == 0 || num_rows_remaining == 0); - - drop(state); self.iter.reuse_page_buffer(page); } @@ -676,7 +649,10 @@ impl PageDecoder { } #[inline] -pub(super) fn dict_indices_decoder(page: &DataPage) -> ParquetResult { +pub(super) fn dict_indices_decoder( + page: &DataPage, + null_count: usize, +) -> ParquetResult { let indices_buffer = split_buffer(page)?.values; // SPEC: Data page format: the bit width used to encode the entry ids stored as 1 byte (max bit width = 32), @@ -687,7 +663,7 @@ pub(super) fn dict_indices_decoder(page: &DataPage) -> ParquetResult { decoder: HybridRleDecoder<'a> } +#[derive(Debug)] pub enum HybridRleChunk<'a> { Rle(u32, usize), Bitpacked(bitpacked::Decoder<'a, u32>), @@ -89,7 +90,7 @@ impl<'a> HybridRleDecoder<'a> { } pub fn next_chunk(&mut self) -> ParquetResult>> { - if self.len() == 0 { + if self.len() == 0 || self.data.is_empty() { return Ok(None); } @@ -106,7 +107,9 @@ impl<'a> HybridRleDecoder<'a> { // is bitpacking let bytes = (indicator as usize >> 1) * self.num_bits; let bytes = std::cmp::min(bytes, self.data.len()); - let (packed, remaining) = self.data.split_at(bytes); + let Some((packed, remaining)) = self.data.split_at_checked(bytes) else { + return Err(ParquetError::oos("Not enough bytes for bitpacked data")); + }; self.data = remaining; let length = std::cmp::min(packed.len() * 8 / self.num_bits, self.num_values); @@ -120,7 +123,9 @@ impl<'a> HybridRleDecoder<'a> { let run_length = indicator as usize >> 1; // repeated-value := value that is repeated, using a fixed-width of round-up-to-next-byte(bit-width) let rle_bytes = self.num_bits.div_ceil(8); - let (pack, remaining) = self.data.split_at(rle_bytes); + let Some((pack, remaining)) = self.data.split_at_checked(rle_bytes) else { + return Err(ParquetError::oos("Not enough bytes for RLE encoded data")); + }; self.data = remaining; let mut bytes = [0u8; std::mem::size_of::()]; From 88898a6e2439f8a090473b524b53cca2ad6369f6 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Tue, 15 Oct 2024 13:15:54 +0200 Subject: [PATCH 07/32] use BitMask --- crates/polars-arrow/src/array/binview/mod.rs | 8 +- crates/polars-arrow/src/bitmap/bitmask.rs | 17 ++- crates/polars-arrow/src/bitmap/iterator.rs | 79 ------------- .../src/arrow/read/deserialize/binview.rs | 14 --- .../arrow/read/deserialize/primitive/float.rs | 68 ++++-------- .../read/deserialize/primitive/integer.rs | 52 +++------ .../read/deserialize/utils/dict_encoded.rs | 104 ++++++++++++------ .../src/arrow/read/deserialize/utils/mod.rs | 2 - crates/polars-utils/src/bit_fns.rs | 30 ----- crates/polars-utils/src/lib.rs | 1 - 10 files changed, 132 insertions(+), 243 deletions(-) delete mode 100644 crates/polars-utils/src/bit_fns.rs diff --git a/crates/polars-arrow/src/array/binview/mod.rs b/crates/polars-arrow/src/array/binview/mod.rs index 2e0eee85d2ef..349384bb6635 100644 --- a/crates/polars-arrow/src/array/binview/mod.rs +++ b/crates/polars-arrow/src/array/binview/mod.rs @@ -170,7 +170,13 @@ impl BinaryViewArrayGeneric { // actual_total_buffer_len += buffer.len(); // } - for view in views.iter() { + for (i, view) in views.iter().enumerate() { + let is_valid = validity.as_ref().map_or(true, |v| v.get_bit(i)); + + if !is_valid { + continue; + } + // actual_total_bytes_len += view.length as usize; if view.length > View::MAX_INLINE_SIZE { assert!((view.buffer_idx as usize) < (buffers.len())); diff --git a/crates/polars-arrow/src/bitmap/bitmask.rs b/crates/polars-arrow/src/bitmap/bitmask.rs index 4d6457c07956..8a19d1653d4a 100644 --- a/crates/polars-arrow/src/bitmap/bitmask.rs +++ b/crates/polars-arrow/src/bitmap/bitmask.rs @@ -5,6 +5,9 @@ use polars_utils::slice::load_padded_le_u64; use crate::bitmap::Bitmap; +use super::iterator::FastU56BitmapIter; +use super::utils::count_zeros; + /// Returns the nth set bit in w, if n+1 bits are set. The indexing is /// zero-based, nth_set_bit_u32(w, 0) returns the least significant set bit in w. fn nth_set_bit_u32(w: u32, n: u32) -> Option { @@ -110,6 +113,18 @@ impl<'a> BitMask<'a> { (left, right) } + pub fn unset_bits(&self) -> usize { + count_zeros(self.bytes, self.offset, self.len) + } + + pub fn set_bits(&self) -> usize { + self.len - self.unset_bits() + } + + pub fn fast_iter_u56(&self) -> FastU56BitmapIter { + FastU56BitmapIter::new(self.bytes, self.offset, self.len) + } + #[cfg(feature = "simd")] #[inline] pub fn get_simd(&self, idx: usize) -> Mask @@ -162,7 +177,7 @@ impl<'a> BitMask<'a> { /// Computes the index of the nth set bit after start. /// - /// Both are zero-indexed, so nth_set_bit_idx(0, 0) finds the index of the + /// Both are zero-indexed, so `nth_set_bit_idx(0, 0)` finds the index of the /// first bit set (which can be 0 as well). The returned index is absolute, /// not relative to start. pub fn nth_set_bit_idx(&self, mut n: usize, mut start: usize) -> Option { diff --git a/crates/polars-arrow/src/bitmap/iterator.rs b/crates/polars-arrow/src/bitmap/iterator.rs index 66db8a6260c7..39d192419305 100644 --- a/crates/polars-arrow/src/bitmap/iterator.rs +++ b/crates/polars-arrow/src/bitmap/iterator.rs @@ -1,4 +1,3 @@ -use polars_utils::bit_fns::nth_setbit; use polars_utils::slice::load_padded_le_u64; use super::bitmask::BitMask; @@ -222,84 +221,6 @@ impl<'a> FastU56BitmapIter<'a> { let hi = self.next_remainder().unwrap_or(0); ((hi << 56) | lo, bits_left) } - - /// Return how many bits need to be consumed until we have seen `num_ones` 1s. - /// - /// It will return `self.bits_left()` if there are less 1s in the remaining values that - /// `num_ones`. - pub fn num_bits_before_nth_one(&self, mut num_ones: usize) -> usize { - let mut num_skipped = 0; - let mut slf = self.clone(); - - loop { - if num_ones < 56 { - break; - } - - let Some(v) = slf.next() else { - break; - }; - - num_ones -= v.count_ones() as usize; - num_skipped += 56; - } - - for v in slf.by_ref() { - let v_num_ones = v.count_ones() as usize; - - if num_ones < v_num_ones { - return num_skipped + nth_setbit(v, num_ones as u32) as usize; - } - - num_ones -= v_num_ones; - num_skipped += 56; - } - - let (v, _) = slf.remainder(); - let v_num_ones = v.count_ones() as usize; - - if num_ones < v_num_ones { - num_skipped + nth_setbit(v, num_ones as u32) as usize - } else { - self.bits_left - } - } - - /// Limit the iterator `num_bits`. - pub fn limit_to_bits(&self, num_bits: usize) -> Self { - let mut new = self.clone(); - new.bits_left = new.bits_left.min(num_bits); - new - } - - pub fn count_ones(&self) -> usize { - let mut slf = self.clone(); - - let base = slf.by_ref().map(|v| v.count_ones() as usize).sum::(); - let rem = slf.remainder().0.count_ones() as usize; - - base + rem - } - - /// Advance the iterator by `num_bits` returning the amount. - pub fn advance_by_bits(&mut self, num_bits: usize) { - if self.bits_left <= num_bits { - self.bytes = &[]; - self.bits_left = 0; - return; - } - - self.bits_left -= num_bits; - - self.shift += (num_bits % 8) as u32; - let extra = self.shift >= 8; - self.shift %= 8; - - self.bytes = unsafe { - self.bytes - .get_unchecked(num_bits / 8 + usize::from(extra)..) - }; - } } impl<'a> Iterator for FastU56BitmapIter<'a> { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index 6eb1ad15051c..e694a3ca8a6a 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -716,14 +716,6 @@ impl utils::Decoder for BinViewDecoder { decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { - let num_rows = state.len(); - let mut max_offset = num_rows; - - if let Some(ref filter) = filter { - max_offset = filter.max_offset(); - assert!(filter.max_offset() <= num_rows); - } - match state.translation { StateTranslation::Dictionary(ref mut indexes) => { let (dict, _) = state.dict.unwrap(); @@ -753,12 +745,6 @@ impl utils::Decoder for BinViewDecoder { .set_total_bytes_len(decoded.0.total_bytes_len() + total_length); } - // @NOTE: Needed for compatibility now. - indexes.skip_in_place(max_offset)?; - if let Some(ref mut page_validity) = state.page_validity { - page_validity.slice(max_offset, page_validity.len() - max_offset); - } - Ok(()) }, _ => self.extend_filtered_with_state_default(state, decoded, filter), diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index b1388feb1e85..58ac995093a4 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -13,9 +13,7 @@ use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; use crate::read::deserialize::utils::array_chunks::ArrayChunks; -use crate::read::deserialize::utils::{ - dict_indices_decoder, freeze_validity, Decoder, -}; +use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity, Decoder}; use crate::read::{Filter, ParquetError}; #[allow(clippy::large_enum_variant)] @@ -263,53 +261,25 @@ where decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { - let num_rows = state.len(); - let mut max_offset = num_rows; - - if let Some(ref filter) = filter { - max_offset = filter.max_offset(); - assert!(filter.max_offset() <= num_rows); - } - match state.translation { - StateTranslation::Plain(ref mut values) => { - super::plain::decode( - values.clone(), - state.is_optional, - state.page_validity.as_ref(), - filter, - &mut decoded.1, - &mut decoded.0, - self.0.decoder, - )?; - - // @NOTE: Needed for compatibility now. - values.skip_in_place(max_offset); - if let Some(ref mut page_validity) = state.page_validity { - page_validity.slice(max_offset, page_validity.len() - max_offset); - } - - Ok(()) - }, - StateTranslation::Dictionary(ref mut indexes) => { - utils::dict_encoded::decode_dict( - indexes.clone(), - state.dict.unwrap(), - state.is_optional, - state.page_validity.as_ref(), - filter, - &mut decoded.1, - &mut decoded.0, - )?; - - // @NOTE: Needed for compatibility now. - indexes.skip_in_place(max_offset)?; - if let Some(ref mut page_validity) = state.page_validity { - page_validity.slice(max_offset, page_validity.len() - max_offset); - } - - Ok(()) - }, + StateTranslation::Plain(ref mut values) => super::plain::decode( + values.clone(), + state.is_optional, + state.page_validity.as_ref(), + filter, + &mut decoded.1, + &mut decoded.0, + self.0.decoder, + ), + StateTranslation::Dictionary(ref mut indexes) => utils::dict_encoded::decode_dict( + indexes.clone(), + state.dict.unwrap(), + state.is_optional, + state.page_validity.as_ref(), + filter, + &mut decoded.1, + &mut decoded.0, + ), _ => self.extend_filtered_with_state_default(state, decoded, filter), } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index f6b778bd2255..9d2a36951350 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -311,41 +311,25 @@ where decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { - let num_rows = state.len(); - let mut max_offset = num_rows; - - if let Some(ref filter) = filter { - max_offset = filter.max_offset(); - assert!(filter.max_offset() <= num_rows); - } - match state.translation { - StateTranslation::Plain(ref mut values) => { - super::plain::decode( - values.clone(), - state.is_optional, - state.page_validity.as_ref(), - filter, - &mut decoded.1, - &mut decoded.0, - self.0.decoder, - )?; - - Ok(()) - }, - StateTranslation::Dictionary(ref mut indexes) => { - utils::dict_encoded::decode_dict( - indexes.clone(), - state.dict.unwrap(), - state.is_optional, - state.page_validity.as_ref(), - filter, - &mut decoded.1, - &mut decoded.0, - )?; - - Ok(()) - }, + StateTranslation::Plain(ref mut values) => super::plain::decode( + values.clone(), + state.is_optional, + state.page_validity.as_ref(), + filter, + &mut decoded.1, + &mut decoded.0, + self.0.decoder, + ), + StateTranslation::Dictionary(ref mut indexes) => utils::dict_encoded::decode_dict( + indexes.clone(), + state.dict.unwrap(), + state.is_optional, + state.page_validity.as_ref(), + filter, + &mut decoded.1, + &mut decoded.0, + ), _ => self.extend_filtered_with_state_default(state, decoded, filter), } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index d3c89064b593..aa39ddfa2422 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -1,3 +1,4 @@ +use arrow::bitmap::bitmask::BitMask; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::types::NativeType; use bytemuck::Pod; @@ -219,11 +220,12 @@ pub fn decode_optional_dict( assert!(num_valid_values <= values.len()); let start_length = target.len(); + let end_length = start_length + validity.len(); target.reserve(validity.len()); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - let mut validity_iter = validity.fast_iter_u56(); + let mut validity = BitMask::from_bitmap(validity); let mut values_buffer = [0u32; 128]; let values_buffer = &mut values_buffer; @@ -246,9 +248,9 @@ pub fn decode_optional_dict( // 2. Fill `num_rows` values into the target buffer. // 3. Advance the validity mask by `num_rows` values. - let num_chunk_rows = validity_iter.num_bits_before_nth_one(size); + let num_chunk_rows = validity.nth_set_bit_idx(size, 0).unwrap_or(validity.len()); - validity_iter.advance_by_bits(num_chunk_rows); + (_, validity) = unsafe { validity.split_at_unchecked(num_chunk_rows) }; let Some(&value) = dict.get(value as usize) else { return Err(oob_dict_idx()); @@ -276,10 +278,10 @@ pub fn decode_optional_dict( { let mut num_done = 0; - let mut cpy_validity_iter = validity_iter.clone(); + let mut validity_iter = validity.fast_iter_u56(); 'outer: while limit >= 64 { - let v = cpy_validity_iter.next().unwrap(); + let v = validity_iter.next().unwrap(); while num_buffered < v.count_ones() as usize { let buffer_part = <&mut [u32; 32]>::try_from( @@ -323,15 +325,19 @@ pub fn decode_optional_dict( limit -= 56; } - validity_iter.advance_by_bits(num_done); + (_, validity) = unsafe { validity.split_at_unchecked(num_done) }; } let num_decoder_remaining = num_buffered + chunked.decoder.len(); - let decoder_limit = validity_iter.num_bits_before_nth_one(num_decoder_remaining); + let decoder_limit = validity + .nth_set_bit_idx(num_decoder_remaining, 0) + .unwrap_or(validity.len()); let num_remaining = limit.min(decoder_limit); - let (v, _) = validity_iter.limit_to_bits(num_remaining).remainder(); - validity_iter.advance_by_bits(num_remaining); + let current_validity; + (current_validity, validity) = + unsafe { validity.split_at_unchecked(num_remaining) }; + let (v, _) = current_validity.fast_iter_u56().remainder(); while num_buffered < v.count_ones() as usize { let buffer_part = <&mut [u32; 32]>::try_from( @@ -365,8 +371,19 @@ pub fn decode_optional_dict( } } + if cfg!(debug_assertions) { + assert_eq!(validity.set_bits(), 0); + } + + let target_slice; + unsafe { + target_slice = std::slice::from_raw_parts_mut(target_ptr, limit); + } + + target_slice.fill(T::zeroed()); + unsafe { - target.set_len(start_length + validity.len()); + target.set_len(end_length); } Ok(()) @@ -404,8 +421,8 @@ pub fn decode_masked_optional_dict( target.reserve(num_rows); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - let mut filter_iter = filter.fast_iter_u56(); - let mut validity_iter = validity.fast_iter_u56(); + let mut filter = BitMask::from_bitmap(filter); + let mut validity = BitMask::from_bitmap(validity); let mut values_buffer = [0u32; 128]; let values_buffer = &mut values_buffer; @@ -436,11 +453,13 @@ pub fn decode_masked_optional_dict( // 2. Fill `num_rows` values into the target buffer. // 3. Advance the validity mask by `num_rows` values. - let num_chunk_values = validity_iter.num_bits_before_nth_one(size); - let num_chunk_rows = filter_iter.limit_to_bits(num_chunk_values).count_ones(); + let num_chunk_values = validity.nth_set_bit_idx(size, 0).unwrap_or(validity.len()); + + let current_filter; + (_, validity) = unsafe { validity.split_at_unchecked(num_chunk_values) }; + (current_filter, filter) = unsafe { validity.split_at_unchecked(num_chunk_values) }; - validity_iter.advance_by_bits(num_chunk_values); - filter_iter.advance_by_bits(num_chunk_values); + let num_chunk_rows = current_filter.set_bits(); if num_chunk_rows > 0 { // SAFETY: Bounds check done before. @@ -474,18 +493,19 @@ pub fn decode_masked_optional_dict( let size = decoder.len(); let mut chunked = decoder.chunked(); - let num_chunk_values = validity_iter.num_bits_before_nth_one(size); + let num_chunk_values = validity.nth_set_bit_idx(size, 0).unwrap_or(validity.len()); let mut buffer_part_idx = 0; let mut values_offset = 0; let mut num_buffered: usize = 0; let mut skip_values = 0; - let mut f_cpy = filter_iter.limit_to_bits(num_chunk_values); - let mut v_cpy = validity_iter.limit_to_bits(num_chunk_values); + let current_filter; + let current_validity; - filter_iter.advance_by_bits(num_chunk_values); - validity_iter.advance_by_bits(num_chunk_values); + (current_filter, filter) = unsafe { filter.split_at_unchecked(num_chunk_values) }; + (current_validity, validity) = + unsafe { validity.split_at_unchecked(num_chunk_values) }; let mut iter = |mut f: u64, mut v: u64| { // Skip chunk if we don't any values from here. @@ -563,12 +583,15 @@ pub fn decode_masked_optional_dict( ParquetResult::Ok(()) }; - for (f, v) in f_cpy.by_ref().zip(v_cpy.by_ref()) { + let mut f_iter = current_filter.fast_iter_u56(); + let mut v_iter = current_validity.fast_iter_u56(); + + for (f, v) in f_iter.by_ref().zip(v_iter.by_ref()) { iter(f, v)?; } - let (f, fl) = f_cpy.remainder(); - let (v, vl) = v_cpy.remainder(); + let (f, fl) = f_iter.remainder(); + let (v, vl) = v_iter.remainder(); assert_eq!(fl, vl); @@ -577,6 +600,18 @@ pub fn decode_masked_optional_dict( } } + if cfg!(debug_assertions) { + assert_eq!(validity.set_bits(), 0); + } + + let target_slice; + unsafe { + target_slice = std::slice::from_raw_parts_mut(target_ptr, num_rows_left); + } + + target_slice.fill(T::zeroed()); + + unsafe { target.set_len(start_length + num_rows); } @@ -607,7 +642,7 @@ pub fn decode_masked_required_dict( target.reserve(num_rows); let mut target_ptr = unsafe { target.as_mut_ptr().add(start_length) }; - let mut filter_iter = filter.fast_iter_u56(); + let mut filter = BitMask::from_bitmap(filter); let mut values_buffer = [0u32; 128]; let values_buffer = &mut values_buffer; @@ -629,6 +664,8 @@ pub fn decode_masked_required_dict( return Err(oob_dict_idx()); } + let size = size.min(filter.len()); + // If we know that we have `size` times `value` that we can append, but there might // be nulls in between those values. // @@ -637,9 +674,10 @@ pub fn decode_masked_required_dict( // 2. Fill `num_rows` values into the target buffer. // 3. Advance the validity mask by `num_rows` values. - let num_chunk_rows = filter_iter.limit_to_bits(size).count_ones(); + let current_filter; - filter_iter.advance_by_bits(size); + (current_filter, filter) = unsafe { filter.split_at_unchecked(size) }; + let num_chunk_rows = current_filter.set_bits(); if num_chunk_rows > 0 { // SAFETY: Bounds check done before. @@ -661,7 +699,7 @@ pub fn decode_masked_required_dict( } }, HybridRleChunk::Bitpacked(mut decoder) => { - let size = decoder.len(); + let size = decoder.len().min(filter.len()); let mut chunked = decoder.chunked(); let mut buffer_part_idx = 0; @@ -669,9 +707,9 @@ pub fn decode_masked_required_dict( let mut num_buffered: usize = 0; let mut skip_values = 0; - let mut f_cpy = filter_iter.limit_to_bits(size); + let current_filter; - filter_iter.advance_by_bits(size); + (current_filter, filter) = unsafe { filter.split_at_unchecked(size) }; let mut iter = |mut f: u64, len: usize| { debug_assert!(len <= 64); @@ -747,11 +785,13 @@ pub fn decode_masked_required_dict( ParquetResult::Ok(()) }; - for f in f_cpy.by_ref() { + let mut f_iter = current_filter.fast_iter_u56(); + + for f in f_iter.by_ref() { iter(f, 56)?; } - let (f, fl) = f_cpy.remainder(); + let (f, fl) = f_iter.remainder(); iter(f, fl)?; }, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 076739a248f7..23dbf1567c5c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -503,8 +503,6 @@ pub(super) trait Decoder: Sized { Ok(()) }, Filter::Mask(bitmap) => { - debug_assert!(bitmap.len() == state.len()); - let mut iter = bitmap.iter(); while iter.num_remaining() > 0 && state.len() > 0 { let prev_state_len = state.len(); diff --git a/crates/polars-utils/src/bit_fns.rs b/crates/polars-utils/src/bit_fns.rs deleted file mode 100644 index 50d683f347f8..000000000000 --- a/crates/polars-utils/src/bit_fns.rs +++ /dev/null @@ -1,30 +0,0 @@ -/// Get the index of the `n`th set bit in `x`. -/// -/// This is starting from `0`. -pub fn nth_setbit(x: u64, n: u32) -> u32 { - #[cfg(target_feature = "bmi2")] - { - use std::arch::x86_64::*; - - return unsafe { _tzcnt_u64(_pdep_u64(1u64 << n, x)) as u32 }; - } - - - #[cfg(not(target_feature = "bmi2"))] - { - let mut x = x; - let mut n = n; - - assert!(x.count_ones() > n); - - let mut idx = 0; - - while x & 1 == 0 || n > 0 { - n -= (x & 1) as u32; - x >>= 1; - idx += 1; - } - - return idx; - } -} diff --git a/crates/polars-utils/src/lib.rs b/crates/polars-utils/src/lib.rs index 109414da657a..38475e411ef2 100644 --- a/crates/polars-utils/src/lib.rs +++ b/crates/polars-utils/src/lib.rs @@ -7,7 +7,6 @@ pub mod abs_diff; pub mod arena; pub mod atomic; pub mod binary_search; -pub mod bit_fns; pub mod cache; pub mod cell; pub mod chunks; From 99eea4e0b759fe8ddd4c274448643cbcb624b6ab Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Tue, 15 Oct 2024 14:38:34 +0200 Subject: [PATCH 08/32] clippy and format --- crates/polars-arrow/src/array/binview/mod.rs | 2 +- crates/polars-arrow/src/bitmap/bitmask.rs | 3 +- .../src/arrow/read/deserialize/binview.rs | 10 ++-- .../src/arrow/read/deserialize/boolean.rs | 4 +- .../src/arrow/read/deserialize/dictionary.rs | 4 +- .../read/deserialize/fixed_size_binary.rs | 4 +- .../src/arrow/read/deserialize/null.rs | 4 +- .../arrow/read/deserialize/primitive/float.rs | 10 ++-- .../read/deserialize/primitive/integer.rs | 10 ++-- .../arrow/read/deserialize/primitive/plain.rs | 16 +++--- .../read/deserialize/utils/dict_encoded.rs | 55 ++++++++++--------- .../src/arrow/read/deserialize/utils/mod.rs | 19 ++++--- .../src/parquet/encoding/bitpacked/mod.rs | 2 +- .../src/parquet/encoding/hybrid_rle/mod.rs | 4 +- crates/polars-utils/src/chunks.rs | 6 +- 15 files changed, 77 insertions(+), 76 deletions(-) diff --git a/crates/polars-arrow/src/array/binview/mod.rs b/crates/polars-arrow/src/array/binview/mod.rs index 349384bb6635..e9c73a51df22 100644 --- a/crates/polars-arrow/src/array/binview/mod.rs +++ b/crates/polars-arrow/src/array/binview/mod.rs @@ -172,7 +172,7 @@ impl BinaryViewArrayGeneric { for (i, view) in views.iter().enumerate() { let is_valid = validity.as_ref().map_or(true, |v| v.get_bit(i)); - + if !is_valid { continue; } diff --git a/crates/polars-arrow/src/bitmap/bitmask.rs b/crates/polars-arrow/src/bitmap/bitmask.rs index 8a19d1653d4a..7e8a3164156d 100644 --- a/crates/polars-arrow/src/bitmap/bitmask.rs +++ b/crates/polars-arrow/src/bitmap/bitmask.rs @@ -3,10 +3,9 @@ use std::simd::{LaneCount, Mask, MaskElement, SupportedLaneCount}; use polars_utils::slice::load_padded_le_u64; -use crate::bitmap::Bitmap; - use super::iterator::FastU56BitmapIter; use super::utils::count_zeros; +use crate::bitmap::Bitmap; /// Returns the nth set bit in w, if n+1 bits are set. The indexing is /// zero-based, nth_set_bit_u32(w, 0) returns the least significant set bit in w. diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index e694a3ca8a6a..b7d6961e3c7c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -698,10 +698,10 @@ impl utils::Decoder for BinViewDecoder { Ok(()) } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, _is_optional: bool, _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, @@ -710,9 +710,9 @@ impl utils::Decoder for BinViewDecoder { unreachable!() } - fn extend_filtered_with_state<'a>( + fn extend_filtered_with_state( &mut self, - mut state: utils::State<'a, Self>, + mut state: utils::State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -724,7 +724,7 @@ impl utils::Decoder for BinViewDecoder { utils::dict_encoded::decode_dict( indexes.clone(), - &dict, + dict, state.is_optional, state.page_validity.as_ref(), filter, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs index 361028d2bab2..92c98a6bda09 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs @@ -235,10 +235,10 @@ impl Decoder for BooleanDecoder { Ok(()) } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut HybridRleDecoder<'a>, + _page_values: &mut HybridRleDecoder<'_>, _is_optional: bool, _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs index 06251691eb30..0aa37a151b94 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs @@ -143,10 +143,10 @@ impl utils::Decoder for DictionaryDec unreachable!() } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut HybridRleDecoder<'a>, + _page_values: &mut HybridRleDecoder<'_>, _is_optional: bool, _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 8f6aab926fa3..2d9a3da5a98c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -204,10 +204,10 @@ impl Decoder for BinaryDecoder { Ok(()) } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, (values, validity): &mut Self::DecodedState, - page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + page_values: &mut hybrid_rle::HybridRleDecoder<'_>, is_optional: bool, page_validity: Option<&mut Bitmap>, dict: &Self::Dict, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index 16a225d9a91c..5826f8fc3171 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -84,10 +84,10 @@ impl utils::Decoder for NullDecoder { unimplemented!() } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, _is_optional: bool, _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index 58ac995093a4..7683091e2727 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -243,10 +243,10 @@ where unreachable!() } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, _is_optional: bool, _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, @@ -255,15 +255,15 @@ where unreachable!() } - fn extend_filtered_with_state<'a>( + fn extend_filtered_with_state( &mut self, - mut state: utils::State<'a, Self>, + mut state: utils::State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { match state.translation { StateTranslation::Plain(ref mut values) => super::plain::decode( - values.clone(), + *values, state.is_optional, state.page_validity.as_ref(), filter, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 9d2a36951350..74b6aae2a419 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -283,10 +283,10 @@ where unreachable!() } - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'a>, + _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, _is_optional: bool, _page_validity: Option<&mut Bitmap>, _dict: &Self::Dict, @@ -305,15 +305,15 @@ where Ok(PrimitiveArray::try_new(dtype, values.into(), validity).unwrap()) } - fn extend_filtered_with_state<'a>( + fn extend_filtered_with_state( &mut self, - mut state: utils::State<'a, Self>, + mut state: utils::State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { match state.translation { StateTranslation::Plain(ref mut values) => super::plain::decode( - values.clone(), + *values, state.is_optional, state.page_validity.as_ref(), filter, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 009af6667e2f..07735283876e 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -44,16 +44,16 @@ fn decode_plain_dispatch ParquetResult<()> { if is_optional { - match (page_validity.as_ref(), filter.as_ref()) { + match (page_validity, filter.as_ref()) { (None, None) => validity.extend_constant(values.len(), true), (None, Some(f)) => validity.extend_constant(f.num_rows(), true), (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), (Some(page_validity), Some(Filter::Range(rng))) => { - let page_validity = (*page_validity).clone(); - validity.extend_from_bitmap(&page_validity.sliced(rng.start, rng.len())) + let page_validity = page_validity.clone(); + validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) }, (Some(page_validity), Some(Filter::Mask(mask))) => { - validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, &mask)) + validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, mask)) }, } } @@ -63,20 +63,20 @@ fn decode_plain_dispatch { decode_required(values, Some(rng.end), target, dfn) }, - (None, Some(page_validity)) => decode_optional(values, &page_validity, target, dfn), + (None, Some(page_validity)) => decode_optional(values, page_validity, target, dfn), (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - decode_optional(values, &page_validity, target, dfn) + decode_optional(values, page_validity, target, dfn) }, (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target, dfn), (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_optional(values, &page_validity, &filter, target, dfn) + decode_masked_optional(values, page_validity, &filter, target, dfn) }, (Some(Filter::Range(rng)), None) => { decode_masked_required(values, &filter_from_range(rng.clone()), target, dfn) }, (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( values, - &page_validity, + page_validity, &filter_from_range(rng.clone()), target, dfn, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index aa39ddfa2422..9af624e3a78e 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -44,16 +44,15 @@ fn decode_dict_dispatch( target: &mut Vec, ) -> ParquetResult<()> { if is_optional { - match (page_validity.as_ref(), filter.as_ref()) { + match (page_validity, filter.as_ref()) { (None, None) => validity.extend_constant(values.len(), true), (None, Some(f)) => validity.extend_constant(f.num_rows(), true), (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), (Some(page_validity), Some(Filter::Range(rng))) => { - let pv = (*page_validity).clone(); - validity.extend_from_bitmap(&pv.sliced(rng.start, rng.len())) + validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) }, (Some(page_validity), Some(Filter::Mask(mask))) => { - validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, &mask)) + validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, mask)) }, } } @@ -63,15 +62,15 @@ fn decode_dict_dispatch( (Some(Filter::Range(rng)), None) if rng.start == 0 => { decode_required_dict(values, dict, Some(rng.end), target) }, - (None, Some(page_validity)) => decode_optional_dict(values, dict, &page_validity, target), + (None, Some(page_validity)) => decode_optional_dict(values, dict, page_validity, target), (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - decode_optional_dict(values, dict, &page_validity, target) + decode_optional_dict(values, dict, page_validity, target) }, (Some(Filter::Mask(filter)), None) => { decode_masked_required_dict(values, dict, &filter, target) }, (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_optional_dict(values, dict, &filter, &page_validity, target) + decode_masked_optional_dict(values, dict, &filter, page_validity, target) }, (Some(Filter::Range(rng)), None) => { decode_masked_required_dict(values, dict, &filter_from_range(rng.clone()), target) @@ -80,7 +79,7 @@ fn decode_dict_dispatch( values, dict, &filter_from_range(rng.clone()), - &page_validity, + page_validity, target, ), } @@ -160,9 +159,10 @@ pub fn decode_required_dict( verify_dict_indices(&chunk, dict.len())?; - for i in 0..32 { - let v = unsafe { dict.get_unchecked(chunk[i] as usize) }; - unsafe { target_ptr.add(i).write(*v) }; + for (i, &idx) in chunk.iter().enumerate() { + let value = unsafe { dict.get_unchecked(idx as usize) }; + let value = *value; + unsafe { target_ptr.add(i).write(value) }; } unsafe { @@ -177,9 +177,10 @@ pub fn decode_required_dict( let highest_idx = chunk[..chunk_size].iter().copied().max().unwrap(); assert!((highest_idx as usize) < dict.len()); - for i in 0..chunk_size { - let v = unsafe { dict.get_unchecked(chunk[i] as usize) }; - unsafe { target_ptr.add(i).write(*v) }; + for (i, &idx) in chunk[..chunk_size].iter().enumerate() { + let value = unsafe { dict.get_unchecked(idx as usize) }; + let value = *value; + unsafe { target_ptr.add(i).write(value) }; } unsafe { @@ -292,7 +293,7 @@ pub fn decode_optional_dict( break 'outer; }; - verify_dict_indices(&buffer_part, dict.len())?; + verify_dict_indices(buffer_part, dict.len())?; num_buffered += num_added; @@ -311,7 +312,8 @@ pub fn decode_optional_dict( // 2. Each time we write to `values_buffer`, it is followed by a // `verify_dict_indices`. let value = unsafe { dict.get_unchecked(idx as usize) }; - unsafe { target_ptr.add(i).write(*value) }; + let value = *value; + unsafe { target_ptr.add(i).write(value) }; num_read += ((v >> i) & 1) as usize; } @@ -346,7 +348,7 @@ pub fn decode_optional_dict( .unwrap(); let num_added = chunked.next_into(buffer_part).unwrap(); - verify_dict_indices(&buffer_part, dict.len())?; + verify_dict_indices(buffer_part, dict.len())?; num_buffered += num_added; @@ -358,7 +360,8 @@ pub fn decode_optional_dict( for i in 0..num_remaining { let idx = values_buffer[(values_offset + num_read) % 128]; - let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + let value = unsafe { dict.get_unchecked(idx as usize) }; + let value = *value; unsafe { *target_ptr.add(i) = value }; num_read += ((v >> i) & 1) as usize; } @@ -522,7 +525,7 @@ pub fn decode_masked_optional_dict( // If we skipped plenty already, just skip decoding those chunks instead of // decoding them and throwing them away. - chunked.decoder.skip_chunks((skip_values / 32) as usize); + chunked.decoder.skip_chunks(skip_values / 32); // The leftovers we have to decode but we can also just skip. skip_values %= 32; @@ -533,9 +536,9 @@ pub fn decode_masked_optional_dict( .unwrap(); let num_added = chunked.next_into(buffer_part).unwrap(); - verify_dict_indices(&buffer_part, dict.len())?; + verify_dict_indices(buffer_part, dict.len())?; - let skip_chunk_values = (skip_values as usize).min(num_added); + let skip_chunk_values = skip_values.min(num_added); values_offset += skip_chunk_values; num_buffered += num_added - skip_chunk_values; @@ -560,7 +563,8 @@ pub fn decode_masked_optional_dict( // dictionary following the original `dict.is_empty` check. // 2. Each time we write to `values_buffer`, it is followed by a // `verify_dict_indices`. - let value = unsafe { dict.get_unchecked(idx as usize) }.clone(); + let value = unsafe { dict.get_unchecked(idx as usize) }; + let value = *value; unsafe { target_ptr.add(num_written).write(value) }; num_written += 1; @@ -611,7 +615,6 @@ pub fn decode_masked_optional_dict( target_slice.fill(T::zeroed()); - unsafe { target.set_len(start_length + num_rows); } @@ -728,7 +731,7 @@ pub fn decode_masked_required_dict( // If we skipped plenty already, just skip decoding those chunks instead of // decoding them and throwing them away. - chunked.decoder.skip_chunks((skip_values / 32) as usize); + chunked.decoder.skip_chunks(skip_values / 32); // The leftovers we have to decode but we can also just skip. skip_values %= 32; @@ -739,9 +742,9 @@ pub fn decode_masked_required_dict( .unwrap(); let num_added = chunked.next_into(buffer_part).unwrap(); - verify_dict_indices(&buffer_part, dict.len())?; + verify_dict_indices(buffer_part, dict.len())?; - let skip_chunk_values = (skip_values as usize).min(num_added); + let skip_chunk_values = skip_values.min(num_added); values_offset += skip_chunk_values; num_buffered += num_added - skip_chunk_values; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 23dbf1567c5c..52a6f9de58d7 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -90,7 +90,10 @@ impl<'a, D: Decoder> State<'a, D> { let is_optional = page.descriptor.primitive_type.field_info.repetition == Repetition::Optional; - if page_validity.as_ref().is_some_and(|bm| bm.unset_bits() == 0) { + if page_validity + .as_ref() + .is_some_and(|bm| bm.unset_bits() == 0) + { page_validity = None; } @@ -327,7 +330,7 @@ pub(super) fn extend_from_decoder>( C::reserve(target, num_elements); let mut batched_collector = BatchedCollector::new(collector, target); - + let mut pv = page_validity.clone(); pv.slice(0, num_elements); @@ -449,18 +452,18 @@ pub(super) trait Decoder: Sized { /// Deserializes a [`DictPage`] into [`Self::Dict`]. fn deserialize_dict(&self, page: DictPage) -> ParquetResult; - fn extend_filtered_with_state<'a>( + fn extend_filtered_with_state( &mut self, - state: State<'a, Self>, + state: State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { self.extend_filtered_with_state_default(state, decoded, filter) } - fn extend_filtered_with_state_default<'a>( + fn extend_filtered_with_state_default( &mut self, - mut state: State<'a, Self>, + mut state: State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { @@ -555,10 +558,10 @@ pub(super) trait Decoder: Sized { page_validity: Option<&mut Bitmap>, limit: usize, ) -> ParquetResult<()>; - fn decode_dictionary_encoded<'a>( + fn decode_dictionary_encoded( &mut self, decoded: &mut Self::DecodedState, - page_values: &mut HybridRleDecoder<'a>, + page_values: &mut HybridRleDecoder<'_>, is_optional: bool, page_validity: Option<&mut Bitmap>, dict: &Self::Dict, diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs index 0a9045ef40d5..057bb5d6b7fe 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs @@ -57,7 +57,7 @@ mod encode; mod pack; mod unpack; -pub use decode::{Decoder, ChunkedDecoder}; +pub use decode::{ChunkedDecoder, Decoder}; pub use encode::{encode, encode_pack}; /// A byte slice (e.g. `[u8; 8]`) denoting types that represent complete packs. diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs index 45c851834ba2..3f3444ef28e1 100644 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs @@ -52,7 +52,7 @@ pub struct HybridRleDecoder<'a> { } pub struct HybridRleChunkIter<'a> { - decoder: HybridRleDecoder<'a> + decoder: HybridRleDecoder<'a>, } #[derive(Debug)] @@ -114,7 +114,7 @@ impl<'a> HybridRleDecoder<'a> { let length = std::cmp::min(packed.len() * 8 / self.num_bits, self.num_values); let decoder = bitpacked::Decoder::::try_new(packed, self.num_bits, length)?; - + self.num_values -= length; HybridRleChunk::Bitpacked(decoder) diff --git a/crates/polars-utils/src/chunks.rs b/crates/polars-utils/src/chunks.rs index 1c5d0e8c4f8e..21fc37e99aff 100644 --- a/crates/polars-utils/src/chunks.rs +++ b/crates/polars-utils/src/chunks.rs @@ -32,11 +32,7 @@ impl<'a, T> DoubleEndedIterator for Chunks<'a, T> { } let rem = self.slice.len() % self.chunk_size; - let offset = if rem == 0 { - self.chunk_size - } else { - rem - }; + let offset = if rem == 0 { self.chunk_size } else { rem }; let item; (self.slice, item) = self.slice.split_at(self.slice.len() - offset); From 0e89ea3903670c2f953a4861d37b4b7ab4c4106a Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Tue, 15 Oct 2024 15:26:02 +0200 Subject: [PATCH 09/32] fix failing tests --- crates/polars-arrow/src/array/binview/mod.rs | 4 +++ crates/polars-arrow/src/bitmap/bitmask.rs | 17 ++++++++++ .../src/arrow/read/deserialize/binview.rs | 2 +- .../arrow/read/deserialize/primitive/plain.rs | 31 ++++++++++++++++--- .../read/deserialize/utils/dict_encoded.rs | 29 ++++++++++++++--- .../src/arrow/read/deserialize/utils/mod.rs | 8 +++++ .../polars/tests/it/io/parquet/arrow/mod.rs | 1 + 7 files changed, 83 insertions(+), 9 deletions(-) diff --git a/crates/polars-arrow/src/array/binview/mod.rs b/crates/polars-arrow/src/array/binview/mod.rs index e9c73a51df22..ebc03af76dc7 100644 --- a/crates/polars-arrow/src/array/binview/mod.rs +++ b/crates/polars-arrow/src/array/binview/mod.rs @@ -162,6 +162,10 @@ impl BinaryViewArrayGeneric { // Verify the invariants #[cfg(debug_assertions)] { + if let Some(validity) = validity.as_ref() { + assert_eq!(validity.len(), views.len()); + } + // @TODO: Enable this. This is currently bugged with concatenate. // let mut actual_total_buffer_len = 0; // let mut actual_total_bytes_len = 0; diff --git a/crates/polars-arrow/src/bitmap/bitmask.rs b/crates/polars-arrow/src/bitmap/bitmask.rs index 7e8a3164156d..2ccf2a6b5266 100644 --- a/crates/polars-arrow/src/bitmap/bitmask.rs +++ b/crates/polars-arrow/src/bitmap/bitmask.rs @@ -112,6 +112,23 @@ impl<'a> BitMask<'a> { (left, right) } + #[inline] + pub fn sliced(&self, offset: usize, length: usize) -> Self { + assert!(offset.checked_add(length).unwrap() <= self.len); + unsafe { self.sliced_unchecked(offset, length) } + } + + /// # Safety + /// The index must be in-bounds. + #[inline] + pub unsafe fn sliced_unchecked(&self, offset: usize, length: usize) -> Self { + if cfg!(debug_assertions) { + assert!(offset.checked_add(length).unwrap() <= self.len); + } + + Self { bytes: self.bytes, offset: self.offset + offset, len: length } + } + pub fn unset_bits(&self) -> usize { count_zeros(self.bytes, self.offset, self.len) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index b7d6961e3c7c..c3f31836a8f5 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -385,7 +385,7 @@ impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for &mut Delta n: usize, ) -> ParquetResult<()> { self.flush(target); - target.extend_constant(n, >::None); + target.extend_constant(n, Some(&[])); Ok(()) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 07735283876e..554bc3284113 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -58,25 +58,46 @@ fn decode_plain_dispatch values.len(), + (Some(f), v) => { + if cfg!(debug_assertions) { + if let Some(v) = v { + assert!(v.len() >= f.max_offset()); + } + } + + f.max_offset() + }, + }; + + let page_validity = page_validity.map(|pv| { + if pv.len() > num_unfiltered_rows { + pv.clone().sliced(0, num_unfiltered_rows) + } else { + pv.clone() + } + }); + match (filter, page_validity) { (None, None) => decode_required(values, None, target, dfn), (Some(Filter::Range(rng)), None) if rng.start == 0 => { decode_required(values, Some(rng.end), target, dfn) }, - (None, Some(page_validity)) => decode_optional(values, page_validity, target, dfn), + (None, Some(page_validity)) => decode_optional(values, &page_validity, target, dfn), (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - decode_optional(values, page_validity, target, dfn) + decode_optional(values, &page_validity, target, dfn) }, (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target, dfn), (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_optional(values, page_validity, &filter, target, dfn) + decode_masked_optional(values, &page_validity, &filter, target, dfn) }, (Some(Filter::Range(rng)), None) => { decode_masked_required(values, &filter_from_range(rng.clone()), target, dfn) }, (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( values, - page_validity, + &page_validity, &filter_from_range(rng.clone()), target, dfn, @@ -117,6 +138,8 @@ fn decode_optional let mut limit = validity.len(); + dbg!(limit); + assert!(num_values <= values.len()); let start_length = target.len(); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index 9af624e3a78e..bb7ae690ebd2 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -57,20 +57,41 @@ fn decode_dict_dispatch( } } + let num_unfiltered_rows = match (filter.as_ref(), page_validity) { + (None, _) => values.len(), + (Some(f), v) => { + if cfg!(debug_assertions) { + if let Some(v) = v { + assert!(v.len() >= f.max_offset()); + } + } + + f.max_offset() + }, + }; + + let page_validity = page_validity.map(|pv| { + if pv.len() > num_unfiltered_rows { + pv.clone().sliced(0, num_unfiltered_rows) + } else { + pv.clone() + } + }); + match (filter, page_validity) { (None, None) => decode_required_dict(values, dict, None, target), (Some(Filter::Range(rng)), None) if rng.start == 0 => { decode_required_dict(values, dict, Some(rng.end), target) }, - (None, Some(page_validity)) => decode_optional_dict(values, dict, page_validity, target), + (None, Some(page_validity)) => decode_optional_dict(values, dict, &page_validity, target), (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - decode_optional_dict(values, dict, page_validity, target) + decode_optional_dict(values, dict, &page_validity, target) }, (Some(Filter::Mask(filter)), None) => { decode_masked_required_dict(values, dict, &filter, target) }, (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_optional_dict(values, dict, &filter, page_validity, target) + decode_masked_optional_dict(values, dict, &filter, &page_validity, target) }, (Some(Filter::Range(rng)), None) => { decode_masked_required_dict(values, dict, &filter_from_range(rng.clone()), target) @@ -79,7 +100,7 @@ fn decode_dict_dispatch( values, dict, &filter_from_range(rng.clone()), - page_validity, + &page_validity, target, ), } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 52a6f9de58d7..27581fc15e24 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -631,15 +631,23 @@ impl PageDecoder { if state_filter.as_ref().is_some_and(|f| f.num_rows() == 0) { continue; } + + dbg!(&state_filter); let page = page.decompress(&mut self.iter)?; + dbg!(page.encoding()); + let state = State::new(&self.decoder, &page, self.dict.as_ref())?; let start_length = target.len(); state.decode(&mut self.decoder, &mut target, state_filter)?; let end_length = target.len(); + dbg!(num_rows_remaining); + dbg!(start_length); + dbg!(end_length); + num_rows_remaining -= end_length - start_length; self.iter.reuse_page_buffer(page); diff --git a/crates/polars/tests/it/io/parquet/arrow/mod.rs b/crates/polars/tests/it/io/parquet/arrow/mod.rs index 0a573eb4a186..3e9944365fc0 100644 --- a/crates/polars/tests/it/io/parquet/arrow/mod.rs +++ b/crates/polars/tests/it/io/parquet/arrow/mod.rs @@ -1617,6 +1617,7 @@ fn list_int_nullable() -> PolarsResult<()> { #[test] fn limit() -> PolarsResult<()> { let (schema, chunk) = generic_data()?; + dbg!(&chunk); assert_roundtrip(schema, chunk, Some(2)) } From 64a938f84ea6066081df5d91cb31a7129dadff40 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Tue, 15 Oct 2024 15:27:45 +0200 Subject: [PATCH 10/32] clippy --- crates/polars-arrow/src/bitmap/bitmask.rs | 8 ++++++-- .../src/arrow/read/deserialize/primitive/plain.rs | 2 -- .../src/arrow/read/deserialize/utils/mod.rs | 8 -------- crates/polars/tests/it/io/parquet/arrow/mod.rs | 1 - 4 files changed, 6 insertions(+), 13 deletions(-) diff --git a/crates/polars-arrow/src/bitmap/bitmask.rs b/crates/polars-arrow/src/bitmap/bitmask.rs index 2ccf2a6b5266..b1bc5cca3d86 100644 --- a/crates/polars-arrow/src/bitmap/bitmask.rs +++ b/crates/polars-arrow/src/bitmap/bitmask.rs @@ -123,10 +123,14 @@ impl<'a> BitMask<'a> { #[inline] pub unsafe fn sliced_unchecked(&self, offset: usize, length: usize) -> Self { if cfg!(debug_assertions) { - assert!(offset.checked_add(length).unwrap() <= self.len); + assert!(offset.checked_add(length).unwrap() <= self.len); } - Self { bytes: self.bytes, offset: self.offset + offset, len: length } + Self { + bytes: self.bytes, + offset: self.offset + offset, + len: length, + } } pub fn unset_bits(&self) -> usize { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 554bc3284113..f4c7be6397ed 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -138,8 +138,6 @@ fn decode_optional let mut limit = validity.len(); - dbg!(limit); - assert!(num_values <= values.len()); let start_length = target.len(); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 27581fc15e24..52a6f9de58d7 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -631,23 +631,15 @@ impl PageDecoder { if state_filter.as_ref().is_some_and(|f| f.num_rows() == 0) { continue; } - - dbg!(&state_filter); let page = page.decompress(&mut self.iter)?; - dbg!(page.encoding()); - let state = State::new(&self.decoder, &page, self.dict.as_ref())?; let start_length = target.len(); state.decode(&mut self.decoder, &mut target, state_filter)?; let end_length = target.len(); - dbg!(num_rows_remaining); - dbg!(start_length); - dbg!(end_length); - num_rows_remaining -= end_length - start_length; self.iter.reuse_page_buffer(page); diff --git a/crates/polars/tests/it/io/parquet/arrow/mod.rs b/crates/polars/tests/it/io/parquet/arrow/mod.rs index 3e9944365fc0..0a573eb4a186 100644 --- a/crates/polars/tests/it/io/parquet/arrow/mod.rs +++ b/crates/polars/tests/it/io/parquet/arrow/mod.rs @@ -1617,7 +1617,6 @@ fn list_int_nullable() -> PolarsResult<()> { #[test] fn limit() -> PolarsResult<()> { let (schema, chunk) = generic_data()?; - dbg!(&chunk); assert_roundtrip(schema, chunk, Some(2)) } From 913bac79e10ab1f189a2dd71077e14896fe6f43a Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 17 Oct 2024 17:39:02 +0200 Subject: [PATCH 11/32] massively reduce monomorphisations for kernels --- .../polars-arrow/src/types/aligned_bytes.rs | 2 +- crates/polars-arrow/src/types/mod.rs | 2 + .../src/arrow/read/deserialize/binview.rs | 2 +- .../src/arrow/read/deserialize/boolean.rs | 2 +- .../src/arrow/read/deserialize/dictionary.rs | 2 +- .../read/deserialize/fixed_size_binary.rs | 2 +- .../arrow/read/deserialize/nested_utils.rs | 2 +- .../src/arrow/read/deserialize/null.rs | 2 +- .../arrow/read/deserialize/primitive/float.rs | 32 ++-- .../read/deserialize/primitive/integer.rs | 32 ++-- .../arrow/read/deserialize/primitive/mod.rs | 22 ++- .../arrow/read/deserialize/primitive/plain.rs | 180 ++++++++++-------- .../read/deserialize/utils/array_chunks.rs | 29 ++- .../read/deserialize/utils/dict_encoded.rs | 106 ++++++----- .../src/arrow/read/deserialize/utils/mod.rs | 4 +- crates/polars-parquet/src/parquet/types.rs | 19 +- py-polars/tests/unit/io/test_parquet.py | 1 + 17 files changed, 252 insertions(+), 189 deletions(-) diff --git a/crates/polars-arrow/src/types/aligned_bytes.rs b/crates/polars-arrow/src/types/aligned_bytes.rs index 2c9bf9aed977..48ff2dfabade 100644 --- a/crates/polars-arrow/src/types/aligned_bytes.rs +++ b/crates/polars-arrow/src/types/aligned_bytes.rs @@ -49,7 +49,7 @@ macro_rules! impl_aligned_bytes { ) => { $( /// Bytes with a size and alignment. - /// + /// /// This is used to reduce the monomorphizations for routines that solely rely on the size /// and alignment of types. #[derive(Debug, Copy, Clone, PartialEq, Eq, Default, Pod, Zeroable)] diff --git a/crates/polars-arrow/src/types/mod.rs b/crates/polars-arrow/src/types/mod.rs index c6f653a32311..e720254294c6 100644 --- a/crates/polars-arrow/src/types/mod.rs +++ b/crates/polars-arrow/src/types/mod.rs @@ -25,6 +25,8 @@ pub use aligned_bytes::*; mod bit_chunk; pub use bit_chunk::{BitChunk, BitChunkIter, BitChunkOnes}; mod index; +mod aligned_bytes; +pub use aligned_bytes::*; pub mod simd; pub use index::*; mod native; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index c3f31836a8f5..a1ee405b1cdf 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -548,7 +548,7 @@ impl utils::Decoder for BinViewDecoder { Ok(()) } - fn deserialize_dict(&self, page: DictPage) -> ParquetResult { + fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { let values = &page.buffer; let num_values = page.num_values; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs index 92c98a6bda09..c2deee921334 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs @@ -207,7 +207,7 @@ impl Decoder for BooleanDecoder { ) } - fn deserialize_dict(&self, _: DictPage) -> ParquetResult { + fn deserialize_dict(&mut self, _: DictPage) -> ParquetResult { Ok(()) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs index 0aa37a151b94..6fc00f64993f 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs @@ -112,7 +112,7 @@ impl utils::Decoder for DictionaryDec ) } - fn deserialize_dict(&self, page: DictPage) -> ParquetResult { + fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { let dict = self.decoder.deserialize_dict(page)?; self.dict_size .store(dict.len(), std::sync::atomic::Ordering::Relaxed); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 2d9a3da5a98c..dde3aed07e3b 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -138,7 +138,7 @@ impl Decoder for BinaryDecoder { ) } - fn deserialize_dict(&self, page: DictPage) -> ParquetResult { + fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { Ok(page.buffer.into_vec()) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index 5b407b6ce4f0..038d54f2481a 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -761,7 +761,7 @@ impl PageNestedDecoder { pub fn new( mut iter: BasicDecompressor, dtype: ArrowDataType, - decoder: D, + mut decoder: D, init: Vec, ) -> ParquetResult { let dict_page = iter.read_dict_page()?; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index 5826f8fc3171..a0e1f0963e0f 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -69,7 +69,7 @@ impl utils::Decoder for NullDecoder { NullArrayLength { length: 0 } } - fn deserialize_dict(&self, _: DictPage) -> ParquetResult { + fn deserialize_dict(&mut self, _: DictPage) -> ParquetResult { Ok(()) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index 7683091e2727..e19aac3b500c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -12,25 +12,24 @@ use crate::parquet::encoding::{byte_stream_split, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; -use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity, Decoder}; -use crate::read::{Filter, ParquetError}; +use crate::read::Filter; #[allow(clippy::large_enum_variant)] #[derive(Debug)] -pub(crate) enum StateTranslation<'a, P: ParquetNativeType> { - Plain(ArrayChunks<'a, P>), +pub(crate) enum StateTranslation<'a> { + Plain(&'a [u8]), Dictionary(hybrid_rle::HybridRleDecoder<'a>), ByteStreamSplit(byte_stream_split::Decoder<'a>), } -impl<'a, P, T, D> utils::StateTranslation<'a, FloatDecoder> for StateTranslation<'a, P> +impl<'a, P, T, D> utils::StateTranslation<'a, FloatDecoder> for StateTranslation<'a> where T: NativeType, P: ParquetNativeType, D: DecoderFunction, { - type PlainDecoder = ArrayChunks<'a, P>; + type PlainDecoder = &'a [u8]; fn new( _decoder: &FloatDecoder, @@ -46,8 +45,7 @@ where }, (Encoding::Plain, _) => { let values = split_buffer(page)?.values; - let chunks = ArrayChunks::new(values).unwrap(); - Ok(Self::Plain(chunks)) + Ok(Self::Plain(values)) }, (Encoding::ByteStreamSplit, _) => { let values = split_buffer(page)?.values; @@ -62,7 +60,7 @@ where fn len_when_not_nullable(&self) -> usize { match self { - Self::Plain(n) => n.len(), + Self::Plain(n) => n.len() / size_of::

(), Self::Dictionary(n) => n.len(), Self::ByteStreamSplit(n) => n.len(), } @@ -74,7 +72,7 @@ where } match self { - Self::Plain(t) => t.skip_in_place(n), + Self::Plain(t) => *t = &t[usize::min(t.len(), size_of::

() * n)..], Self::Dictionary(t) => t.skip_in_place(n)?, Self::ByteStreamSplit(t) => _ = t.iter_converted(|_| ()).nth(n - 1), } @@ -200,7 +198,7 @@ where P: ParquetNativeType, D: DecoderFunction, { - type Translation<'a> = StateTranslation<'a, P>; + type Translation<'a> = StateTranslation<'a>; type Dict = Vec; type DecodedState = (Vec, MutableBitmap); type Output = PrimitiveArray; @@ -212,20 +210,17 @@ where ) } - fn deserialize_dict(&self, page: DictPage) -> ParquetResult { - let Some(values) = ArrayChunks::

::new(page.buffer.as_ref()) else { - return Err(ParquetError::oos( - "Primitive dictionary page size is not a multiple of primitive size", - )); - }; + fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { + let values = page.buffer.as_ref(); - let mut target = Vec::new(); + let mut target = Vec::with_capacity(page.num_values); super::plain::decode( values, false, None, None, &mut MutableBitmap::new(), + &mut self.0.intermediate, &mut target, self.0.decoder, )?; @@ -268,6 +263,7 @@ where state.page_validity.as_ref(), filter, &mut decoded.1, + &mut self.0.intermediate, &mut decoded.0, self.0.decoder, ), diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 74b6aae2a419..085389c2e2e1 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -12,27 +12,26 @@ use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, hybrid_rle, E use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; -use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity}; -use crate::read::{Filter, ParquetError}; +use crate::read::Filter; #[allow(clippy::large_enum_variant)] #[derive(Debug)] -pub(crate) enum StateTranslation<'a, P: ParquetNativeType> { - Plain(ArrayChunks<'a, P>), +pub(crate) enum StateTranslation<'a> { + Plain(&'a [u8]), Dictionary(hybrid_rle::HybridRleDecoder<'a>), ByteStreamSplit(byte_stream_split::Decoder<'a>), DeltaBinaryPacked(delta_bitpacked::Decoder<'a>), } -impl<'a, P, T, D> utils::StateTranslation<'a, IntDecoder> for StateTranslation<'a, P> +impl<'a, P, T, D> utils::StateTranslation<'a, IntDecoder> for StateTranslation<'a> where T: NativeType, P: ParquetNativeType, i64: num_traits::AsPrimitive

, D: DecoderFunction, { - type PlainDecoder = ArrayChunks<'a, P>; + type PlainDecoder = &'a [u8]; fn new( _decoder: &IntDecoder, @@ -48,8 +47,7 @@ where }, (Encoding::Plain, _) => { let values = split_buffer(page)?.values; - let chunks = ArrayChunks::new(values).unwrap(); - Ok(Self::Plain(chunks)) + Ok(Self::Plain(values)) }, (Encoding::ByteStreamSplit, _) => { let values = split_buffer(page)?.values; @@ -70,7 +68,7 @@ where fn len_when_not_nullable(&self) -> usize { match self { - Self::Plain(v) => v.len(), + Self::Plain(v) => v.len() / size_of::

(), Self::Dictionary(v) => v.len(), Self::ByteStreamSplit(v) => v.len(), Self::DeltaBinaryPacked(v) => v.len(), @@ -83,7 +81,7 @@ where } match self { - Self::Plain(v) => v.skip_in_place(n), + Self::Plain(v) => *v = &v[usize::min(v.len(), size_of::

() * n)..], Self::Dictionary(v) => v.skip_in_place(n)?, Self::ByteStreamSplit(v) => _ = v.iter_converted(|_| ()).nth(n - 1), Self::DeltaBinaryPacked(v) => v.skip_in_place(n)?, @@ -240,7 +238,7 @@ where i64: num_traits::AsPrimitive

, D: DecoderFunction, { - type Translation<'a> = StateTranslation<'a, P>; + type Translation<'a> = StateTranslation<'a>; type Dict = Vec; type DecodedState = (Vec, MutableBitmap); type Output = PrimitiveArray; @@ -252,20 +250,17 @@ where ) } - fn deserialize_dict(&self, page: DictPage) -> ParquetResult { - let Some(values) = ArrayChunks::

::new(page.buffer.as_ref()) else { - return Err(ParquetError::oos( - "Primitive dictionary page size is not a multiple of primitive size", - )); - }; + fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { + let values = page.buffer.as_ref(); - let mut target = Vec::new(); + let mut target = Vec::with_capacity(page.num_values); super::plain::decode( values, false, None, None, &mut MutableBitmap::new(), + &mut self.0.intermediate, &mut target, self.0.decoder, )?; @@ -318,6 +313,7 @@ where state.page_validity.as_ref(), filter, &mut decoded.1, + &mut self.0.intermediate, &mut decoded.0, self.0.decoder, ), diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs index d0099df10c09..f6287fdb7066 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs @@ -22,6 +22,7 @@ where D: DecoderFunction, { pub(crate) decoder: D, + pub(crate) intermediate: Vec

, _pd: std::marker::PhantomData<(P, T)>, } @@ -35,6 +36,7 @@ where pub(crate) fn new(decoder: D) -> Self { Self { decoder, + intermediate: Vec::new(), _pd: std::marker::PhantomData, } } @@ -49,12 +51,22 @@ where T: NativeType, P: ParquetNativeType, { + const NEED_TO_DECODE: bool; + const CAN_TRANSMUTE: bool = { + let has_same_size = size_of::

() == size_of::(); + let has_same_alignment = align_of::

() == align_of::(); + + has_same_size && has_same_alignment + }; + fn decode(self, x: P) -> T; } #[derive(Default, Clone, Copy)] pub(crate) struct UnitDecoderFunction(std::marker::PhantomData); impl DecoderFunction for UnitDecoderFunction { + const NEED_TO_DECODE: bool = false; + #[inline(always)] fn decode(self, x: T) -> T { x @@ -62,11 +74,15 @@ impl DecoderFunction for UnitDecoderFun } #[derive(Default, Clone, Copy)] -pub(crate) struct AsDecoderFunction(std::marker::PhantomData<(P, T)>); +pub(crate) struct AsDecoderFunction( + std::marker::PhantomData<(P, T)>, +); macro_rules! as_decoder_impl { ($($p:ty => $t:ty,)+) => { $( impl DecoderFunction<$p, $t> for AsDecoderFunction<$p, $t> { + const NEED_TO_DECODE: bool = Self::CAN_TRANSMUTE; + #[inline(always)] fn decode(self, x : $p) -> $t { x as $t @@ -94,6 +110,8 @@ where P: ParquetNativeType + Into, T: NativeType, { + const NEED_TO_DECODE: bool = true; + #[inline(always)] fn decode(self, x: P) -> T { x.into() @@ -108,6 +126,8 @@ where T: NativeType, F: Copy + Fn(P) -> T, { + const NEED_TO_DECODE: bool = true; + #[inline(always)] fn decode(self, x: P) -> T { (self.0)(x) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index f4c7be6397ed..47621b0e68f2 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -1,65 +1,94 @@ use arrow::bitmap::{Bitmap, MutableBitmap}; -use arrow::types::NativeType; -use polars_compute::filter::filter_boolean_kernel; +use arrow::types::{AlignedBytes, NativeType}; use super::DecoderFunction; use crate::parquet::error::ParquetResult; use crate::parquet::types::NativeType as ParquetNativeType; use crate::read::deserialize::utils::array_chunks::ArrayChunks; +use crate::read::deserialize::utils::dict_encoded::append_validity; use crate::read::deserialize::utils::filter_from_range; -use crate::read::Filter; +use crate::read::{Filter, ParquetError}; pub fn decode>( - values: ArrayChunks<'_, P>, + values: &[u8], is_optional: bool, page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, + intermediate: &mut Vec

, target: &mut Vec, dfn: D, ) -> ParquetResult<()> { - // @TODO: It would be really nice to reduce monomorphizations here. All decode kernels only - // dependent on the alignment and size of `T` so we could make it downcast here to types that - // are `Pod` and have the same alignment and size. - - decode_plain_dispatch( - values, - is_optional, - page_validity, - filter, - validity, - target, - dfn, - ) + if cfg!(debug_assertions) && is_optional { + assert_eq!(target.len(), validity.len()); + } + + if D::CAN_TRANSMUTE { + let values = ArrayChunks::<'_, T::AlignedBytes>::new(values).ok_or_else(|| { + ParquetError::oos("Page content does not align with expected element size") + })?; + + let start_length = target.len(); + decode_aligned_bytes_dispatch( + values, + is_optional, + page_validity, + filter, + validity, + ::cast_vec_ref_mut(target), + )?; + + if D::NEED_TO_DECODE { + let to_decode: &mut [P] = bytemuck::cast_slice_mut(&mut target[start_length..]); + + for v in to_decode { + *v = bytemuck::cast(dfn.decode(*v)); + } + } + } else { + let values = ArrayChunks::<'_, P::AlignedBytes>::new(values).ok_or_else(|| { + ParquetError::oos("Page content does not align with expected element size") + })?; + + intermediate.clear(); + decode_aligned_bytes_dispatch( + values, + is_optional, + page_validity, + filter, + validity, + ::cast_vec_ref_mut(intermediate), + )?; + + target.extend(intermediate.iter().copied().map(|v| dfn.decode(v))); + } + + if cfg!(debug_assertions) && is_optional { + assert_eq!(target.len(), validity.len()); + } + + Ok(()) } #[inline(never)] -fn decode_plain_dispatch>( - values: ArrayChunks<'_, P>, +pub fn decode_aligned_bytes_dispatch( + values: ArrayChunks<'_, B>, is_optional: bool, page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, - target: &mut Vec, - dfn: D, + target: &mut Vec, ) -> ParquetResult<()> { if is_optional { - match (page_validity, filter.as_ref()) { - (None, None) => validity.extend_constant(values.len(), true), - (None, Some(f)) => validity.extend_constant(f.num_rows(), true), - (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), - (Some(page_validity), Some(Filter::Range(rng))) => { - let page_validity = page_validity.clone(); - validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) - }, - (Some(page_validity), Some(Filter::Mask(mask))) => { - validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, mask)) - }, - } + append_validity(page_validity, filter.as_ref(), validity, values.len()); } let num_unfiltered_rows = match (filter.as_ref(), page_validity) { - (None, _) => values.len(), + (None, None) => values.len(), + (None, Some(pv)) => { + debug_assert!(pv.len() >= values.len()); + pv.len() + }, (Some(f), v) => { if cfg!(debug_assertions) { if let Some(v) = v { @@ -79,61 +108,61 @@ fn decode_plain_dispatch decode_required(values, None, target, dfn), + (None, None) => decode_required(values, None, target), (Some(Filter::Range(rng)), None) if rng.start == 0 => { - decode_required(values, Some(rng.end), target, dfn) + decode_required(values, Some(rng.end), target) }, - (None, Some(page_validity)) => decode_optional(values, &page_validity, target, dfn), + (None, Some(page_validity)) => decode_optional(values, &page_validity, target), (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - decode_optional(values, &page_validity, target, dfn) + decode_optional(values, &page_validity, target) }, - (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target, dfn), + (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target), (Some(Filter::Mask(filter)), Some(page_validity)) => { - decode_masked_optional(values, &page_validity, &filter, target, dfn) + decode_masked_optional(values, &page_validity, &filter, target) }, + // @TODO: Use values.skip_in_place(rng.start) (Some(Filter::Range(rng)), None) => { - decode_masked_required(values, &filter_from_range(rng.clone()), target, dfn) + decode_masked_required(values, dbg!(&filter_from_range(rng.clone())), target) }, + // @TODO: Use values.skip_in_place(page_validity.sliced(0, rng.start).set_bits()) (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( values, - &page_validity, - &filter_from_range(rng.clone()), + dbg!(&page_validity), + dbg!(&filter_from_range(rng.clone())), target, - dfn, ), - } + }?; + + Ok(()) } #[inline(never)] -fn decode_required>( - values: ArrayChunks<'_, P>, +fn decode_required( + values: ArrayChunks<'_, B>, limit: Option, - target: &mut Vec, - dfn: D, + target: &mut Vec, ) -> ParquetResult<()> { let limit = limit.unwrap_or(values.len()); assert!(limit <= values.len()); - target.extend((0..limit).map(|i| { - let v = unsafe { values.get_unchecked(i) }; - dfn.decode(v) - })); + target.extend(values.take(limit).map(|v| B::from_unaligned(*v))); Ok(()) } #[inline(never)] -fn decode_optional>( - values: ArrayChunks<'_, P>, +fn decode_optional( + values: ArrayChunks<'_, B>, validity: &Bitmap, - target: &mut Vec, - dfn: D, + target: &mut Vec, ) -> ParquetResult<()> { let num_values = validity.set_bits(); if num_values == validity.len() { - return decode_required(values, Some(validity.len()), target, dfn); + return decode_required(values, Some(validity.len()), target); } let mut limit = validity.len(); @@ -158,10 +187,9 @@ fn decode_optional for i in 0..len { let is_valid = v & 1 != 0; let value = if is_valid { - let value = unsafe { values.get_unchecked(value_offset) }; - dfn.decode(value) + unsafe { values.get_unchecked(value_offset) } } else { - T::zeroed() + B::zeroed() }; unsafe { target_ptr.add(i).write(value) }; @@ -171,7 +199,6 @@ fn decode_optional } else { for i in 0..len { let value = unsafe { values.get_unchecked(value_offset) }; - let value = dfn.decode(value); unsafe { target_ptr.add(i).write(value) }; value_offset += (v & 1) as usize; @@ -204,16 +231,15 @@ fn decode_optional } #[inline(never)] -fn decode_masked_required>( - values: ArrayChunks<'_, P>, +fn decode_masked_required( + values: ArrayChunks<'_, B>, mask: &Bitmap, - target: &mut Vec, - dfn: D, + target: &mut Vec, ) -> ParquetResult<()> { let num_rows = mask.set_bits(); if num_rows == mask.len() { - return decode_required(values, Some(num_rows), target, dfn); + return decode_required(values, Some(num_rows), target); } assert!(mask.len() <= values.len()); @@ -245,7 +271,6 @@ fn decode_masked_required>( - values: ArrayChunks<'_, P>, +fn decode_masked_optional( + values: ArrayChunks<'_, B>, validity: &Bitmap, mask: &Bitmap, - target: &mut Vec, - dfn: D, + target: &mut Vec, ) -> ParquetResult<()> { let num_rows = mask.set_bits(); let num_values = validity.set_bits(); if num_rows == mask.len() { - return decode_optional(values, validity, target, dfn); + return decode_optional(values, validity, target); } if num_values == validity.len() { - return decode_masked_required(values, mask, target, dfn); + return decode_masked_required(values, mask, target); } - assert!(mask.len() <= values.len()); + assert!(num_values <= values.len()); let start_length = target.len(); target.reserve(num_rows); @@ -328,10 +352,9 @@ fn decode_masked_optional>= offset; let value = unsafe { values.get_unchecked(value_offset + num_read) }; - let value = dfn.decode(value); unsafe { target_ptr.add(num_written).write(value) }; num_written += 1; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs index eede2f910ffa..ecd33c5850f6 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs @@ -1,22 +1,26 @@ -use crate::parquet::types::NativeType as ParquetNativeType; +use arrow::types::AlignedBytes; -/// A slice of chunks that fit the `P` type. +/// A slice of chunks that fit an [`AlignedBytes`] type. /// /// This is essentially the equivalent of [`ChunksExact`][std::slice::ChunksExact], but with a size /// and type known at compile-time. This makes the compiler able to reason much more about the /// code. Especially, since the chunk-sizes for this type are almost always powers of 2 and /// bitshifts or special instructions would be much better to use. #[derive(Debug, Clone, Copy)] -pub(crate) struct ArrayChunks<'a, P: ParquetNativeType> { - pub(crate) bytes: &'a [P::Bytes], +pub(crate) struct ArrayChunks<'a, B: AlignedBytes> { + pub(crate) bytes: &'a [B::Unaligned], } -impl<'a, P: ParquetNativeType> ArrayChunks<'a, P> { +impl<'a, B: AlignedBytes> ArrayChunks<'a, B> { /// Create a new [`ArrayChunks`] /// /// This returns null if the `bytes` slice's length is not a multiple of the size of `P::Bytes`. pub(crate) fn new(bytes: &'a [u8]) -> Option { +<<<<<<< HEAD if bytes.len() % size_of::() != 0 { +======= + if bytes.len() % B::SIZE != 0 { +>>>>>>> 26e3481c1d (massively reduce monomorphisations for kernels) return None; } @@ -25,18 +29,13 @@ impl<'a, P: ParquetNativeType> ArrayChunks<'a, P> { Some(Self { bytes }) } - pub(crate) unsafe fn get_unchecked(&self, at: usize) -> P { - P::from_le_bytes(*unsafe { self.bytes.get_unchecked(at) }) - } - - pub(crate) fn skip_in_place(&mut self, n: usize) { - let n = usize::min(self.bytes.len(), n); - self.bytes = &self.bytes[n..]; + pub(crate) unsafe fn get_unchecked(&self, at: usize) -> B { + B::from_unaligned(*unsafe { self.bytes.get_unchecked(at) }) } } -impl<'a, P: ParquetNativeType> Iterator for ArrayChunks<'a, P> { - type Item = &'a P::Bytes; +impl<'a, B: AlignedBytes> Iterator for ArrayChunks<'a, B> { + type Item = &'a B::Unaligned; #[inline(always)] fn next(&mut self) -> Option { @@ -51,4 +50,4 @@ impl<'a, P: ParquetNativeType> Iterator for ArrayChunks<'a, P> { } } -impl<'a, P: ParquetNativeType> ExactSizeIterator for ArrayChunks<'a, P> {} +impl<'a, B: AlignedBytes> ExactSizeIterator for ArrayChunks<'a, B> {} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index bb7ae690ebd2..327093cbec0b 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -1,7 +1,6 @@ use arrow::bitmap::bitmask::BitMask; use arrow::bitmap::{Bitmap, MutableBitmap}; -use arrow::types::NativeType; -use bytemuck::Pod; +use arrow::types::{AlignedBytes, NativeType}; use polars_compute::filter::filter_boolean_kernel; use super::filter_from_range; @@ -9,7 +8,7 @@ use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::read::{Filter, ParquetError}; -pub fn decode_dict( +pub fn decode_dict( values: HybridRleDecoder<'_>, dict: &[T], is_optional: bool, @@ -18,47 +17,61 @@ pub fn decode_dict( validity: &mut MutableBitmap, target: &mut Vec, ) -> ParquetResult<()> { - // @TODO: It would be really nice to reduce monomorphizations here. All decode kernels only - // dependent on the alignment and size of `T` so we could make it downcast here to types that - // are `Pod` and have the same alignment and size. - decode_dict_dispatch( values, - dict, + bytemuck::cast_slice(dict), is_optional, page_validity, filter, validity, - target, + ::cast_vec_ref_mut(target), ) } +pub(crate) fn append_validity( + page_validity: Option<&Bitmap>, + filter: Option<&Filter>, + validity: &mut MutableBitmap, + values_len: usize, +) { + match (page_validity, filter) { + (None, None) => validity.extend_constant(values_len, true), + (None, Some(f)) => validity.extend_constant(f.num_rows(), true), + (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), + (Some(page_validity), Some(Filter::Range(rng))) => { + let page_validity = page_validity.clone(); + validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) + }, + (Some(page_validity), Some(Filter::Mask(mask))) => { + validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, mask)) + }, + } +} + #[inline(never)] -fn decode_dict_dispatch( +pub fn decode_dict_dispatch( values: HybridRleDecoder<'_>, - dict: &[T], + dict: &[B], is_optional: bool, page_validity: Option<&Bitmap>, filter: Option, validity: &mut MutableBitmap, - target: &mut Vec, + target: &mut Vec, ) -> ParquetResult<()> { + if cfg!(debug_assertions) && is_optional { + assert_eq!(target.len(), validity.len()); + } + if is_optional { - match (page_validity, filter.as_ref()) { - (None, None) => validity.extend_constant(values.len(), true), - (None, Some(f)) => validity.extend_constant(f.num_rows(), true), - (Some(page_validity), None) => validity.extend_from_bitmap(page_validity), - (Some(page_validity), Some(Filter::Range(rng))) => { - validity.extend_from_bitmap(&page_validity.clone().sliced(rng.start, rng.len())) - }, - (Some(page_validity), Some(Filter::Mask(mask))) => { - validity.extend_from_bitmap(&filter_boolean_kernel(page_validity, mask)) - }, - } + append_validity(page_validity, filter.as_ref(), validity, values.len()); } let num_unfiltered_rows = match (filter.as_ref(), page_validity) { - (None, _) => values.len(), + (None, None) => values.len(), + (None, Some(pv)) => { + debug_assert!(pv.len() >= values.len()); + pv.len() + }, (Some(f), v) => { if cfg!(debug_assertions) { if let Some(v) = v { @@ -78,6 +91,8 @@ fn decode_dict_dispatch( } }); + dbg!(&filter); + match (filter, page_validity) { (None, None) => decode_required_dict(values, dict, None, target), (Some(Filter::Range(rng)), None) if rng.start == 0 => { @@ -94,20 +109,25 @@ fn decode_dict_dispatch( decode_masked_optional_dict(values, dict, &filter, &page_validity, target) }, (Some(Filter::Range(rng)), None) => { - decode_masked_required_dict(values, dict, &filter_from_range(rng.clone()), target) + decode_masked_required_dict(values, dict, dbg!(&filter_from_range(rng.clone())), target) }, (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional_dict( values, dict, - &filter_from_range(rng.clone()), - &page_validity, + dbg!(&filter_from_range(rng.clone())), + dbg!(&page_validity), target, ), + }?; + + if cfg!(debug_assertions) && is_optional { + assert_eq!(target.len(), validity.len()); } + + Ok(()) } #[cold] -#[inline(always)] fn oob_dict_idx() -> ParquetError { ParquetError::oos("Dictionary Index is out-of-bounds") } @@ -127,11 +147,11 @@ fn verify_dict_indices(indices: &[u32; 32], dict_size: usize) -> ParquetResult<( } #[inline(never)] -pub fn decode_required_dict( +pub fn decode_required_dict( mut values: HybridRleDecoder<'_>, - dict: &[T], + dict: &[B], limit: Option, - target: &mut Vec, + target: &mut Vec, ) -> ParquetResult<()> { if dict.is_empty() && values.len() > 0 { return Err(oob_dict_idx()); @@ -222,11 +242,11 @@ pub fn decode_required_dict( } #[inline(never)] -pub fn decode_optional_dict( +pub fn decode_optional_dict( mut values: HybridRleDecoder<'_>, - dict: &[T], + dict: &[B], validity: &Bitmap, - target: &mut Vec, + target: &mut Vec, ) -> ParquetResult<()> { let mut limit = validity.len(); let num_valid_values = validity.set_bits(); @@ -404,7 +424,7 @@ pub fn decode_optional_dict( target_slice = std::slice::from_raw_parts_mut(target_ptr, limit); } - target_slice.fill(T::zeroed()); + target_slice.fill(B::zeroed()); unsafe { target.set_len(end_length); @@ -414,12 +434,12 @@ pub fn decode_optional_dict( } #[inline(never)] -pub fn decode_masked_optional_dict( +pub fn decode_masked_optional_dict( mut values: HybridRleDecoder<'_>, - dict: &[T], + dict: &[B], filter: &Bitmap, validity: &Bitmap, - target: &mut Vec, + target: &mut Vec, ) -> ParquetResult<()> { let num_rows = filter.set_bits(); let num_valid_values = validity.set_bits(); @@ -481,7 +501,7 @@ pub fn decode_masked_optional_dict( let current_filter; (_, validity) = unsafe { validity.split_at_unchecked(num_chunk_values) }; - (current_filter, filter) = unsafe { validity.split_at_unchecked(num_chunk_values) }; + (current_filter, filter) = unsafe { filter.split_at_unchecked(num_chunk_values) }; let num_chunk_rows = current_filter.set_bits(); @@ -634,7 +654,7 @@ pub fn decode_masked_optional_dict( target_slice = std::slice::from_raw_parts_mut(target_ptr, num_rows_left); } - target_slice.fill(T::zeroed()); + target_slice.fill(B::zeroed()); unsafe { target.set_len(start_length + num_rows); @@ -644,11 +664,11 @@ pub fn decode_masked_optional_dict( } #[inline(never)] -pub fn decode_masked_required_dict( +pub fn decode_masked_required_dict( mut values: HybridRleDecoder<'_>, - dict: &[T], + dict: &[B], filter: &Bitmap, - target: &mut Vec, + target: &mut Vec, ) -> ParquetResult<()> { let num_rows = filter.set_bits(); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 52a6f9de58d7..9054fe93e02d 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -450,7 +450,7 @@ pub(super) trait Decoder: Sized { fn with_capacity(&self, capacity: usize) -> Self::DecodedState; /// Deserializes a [`DictPage`] into [`Self::Dict`]. - fn deserialize_dict(&self, page: DictPage) -> ParquetResult; + fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult; fn extend_filtered_with_state( &mut self, @@ -596,7 +596,7 @@ impl PageDecoder { pub fn new( mut iter: BasicDecompressor, dtype: ArrowDataType, - decoder: D, + mut decoder: D, ) -> ParquetResult { let dict_page = iter.read_dict_page()?; let dict = dict_page.map(|d| decoder.deserialize_dict(d)).transpose()?; diff --git a/crates/polars-parquet/src/parquet/types.rs b/crates/polars-parquet/src/parquet/types.rs index 1dd65d0ff622..3d74ba7e0563 100644 --- a/crates/polars-parquet/src/parquet/types.rs +++ b/crates/polars-parquet/src/parquet/types.rs @@ -1,7 +1,9 @@ +use arrow::types::{AlignedBytes, Bytes12Alignment4, Bytes4Alignment4, Bytes8Alignment8, SameBytes}; + use crate::parquet::schema::types::PhysicalType; /// A physical native representation of a Parquet fixed-sized type. -pub trait NativeType: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { +pub trait NativeType: std::fmt::Debug + Send + Sync + 'static + Copy + Clone + SameBytes { type Bytes: AsRef<[u8]> + bytemuck::Pod + IntoIterator @@ -9,6 +11,7 @@ pub trait NativeType: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { + std::fmt::Debug + Clone + Copy; + type AlignedBytes: AlignedBytes + From + Into; fn to_le_bytes(&self) -> Self::Bytes; @@ -20,9 +23,11 @@ pub trait NativeType: std::fmt::Debug + Send + Sync + 'static + Copy + Clone { } macro_rules! native { - ($type:ty, $physical_type:expr) => { + ($type:ty, $unaligned:ty, $physical_type:expr) => { impl NativeType for $type { type Bytes = [u8; size_of::()]; + type AlignedBytes = $unaligned; + #[inline] fn to_le_bytes(&self) -> Self::Bytes { Self::to_le_bytes(*self) @@ -43,15 +48,17 @@ macro_rules! native { }; } -native!(i32, PhysicalType::Int32); -native!(i64, PhysicalType::Int64); -native!(f32, PhysicalType::Float); -native!(f64, PhysicalType::Double); +native!(i32, Bytes4Alignment4, PhysicalType::Int32); +native!(i64, Bytes8Alignment8, PhysicalType::Int64); +native!(f32, Bytes4Alignment4, PhysicalType::Float); +native!(f64, Bytes8Alignment8, PhysicalType::Double); impl NativeType for [u32; 3] { const TYPE: PhysicalType = PhysicalType::Int96; type Bytes = [u8; size_of::()]; + type AlignedBytes = Bytes12Alignment4; + #[inline] fn to_le_bytes(&self) -> Self::Bytes { let mut bytes = [0; 12]; diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py index c818fd77d838..2ff36017da18 100644 --- a/py-polars/tests/unit/io/test_parquet.py +++ b/py-polars/tests/unit/io/test_parquet.py @@ -1578,6 +1578,7 @@ def test_slice_roundtrip(df: pl.DataFrame, offset: int, length: int) -> None: df.write_parquet(f) f.seek(0) + print((offset, length)) scanned = pl.scan_parquet(f).slice(offset, length).collect() assert_frame_equal(scanned, df.slice(offset, length)) From 295a7a7c90f82094aa2ca47c99c50c4b63044d30 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 17 Oct 2024 18:01:14 +0200 Subject: [PATCH 12/32] move fixed size binary to branchless --- .../read/deserialize/fixed_size_binary.rs | 173 +++++++++++------- .../arrow/read/deserialize/primitive/mod.rs | 2 +- .../arrow/read/deserialize/primitive/plain.rs | 27 +-- .../read/deserialize/utils/dict_encoded.rs | 44 +++-- 4 files changed, 133 insertions(+), 113 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index dde3aed07e3b..3f03099f3cc4 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -1,12 +1,20 @@ use arrow::array::{DictionaryArray, DictionaryKey, FixedSizeBinaryArray, PrimitiveArray}; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; +use arrow::types::{ + Bytes12Alignment4, Bytes16Alignment16, Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, + Bytes4Alignment4, Bytes8Alignment8, +}; +use super::utils::array_chunks::ArrayChunks; +use super::utils::dict_encoded::append_validity; use super::utils::{dict_indices_decoder, extend_from_decoder, freeze_validity, Decoder}; +use super::Filter; use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; use crate::parquet::encoding::{hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; +use crate::read::deserialize::utils::dict_encoded::constrain_page_validity; use crate::read::deserialize::utils::{self, BatchableCollector, GatheredHybridRle}; #[allow(clippy::large_enum_variant)] @@ -108,34 +116,116 @@ pub(crate) struct BinaryDecoder { pub(crate) size: usize, } +enum FSBTarget { + Size1(Vec), + Size2(Vec), + Size4(Vec), + Size8(Vec), + Size12(Vec), + Size16(Vec), + Size32(Vec), + Other(Vec, usize), +} + impl utils::ExactSize for Vec { fn len(&self) -> usize { Vec::len(self) } } -impl utils::ExactSize for (FixedSizeBinary, MutableBitmap) { +impl utils::ExactSize for FSBTarget { + fn len(&self) -> usize { + match self { + FSBTarget::Size1(vec) => vec.len(), + FSBTarget::Size2(vec) => vec.len(), + FSBTarget::Size4(vec) => vec.len(), + FSBTarget::Size8(vec) => vec.len(), + FSBTarget::Size16(vec) => vec.len(), + FSBTarget::Size32(vec) => vec.len(), + FSBTarget::Other(vec, size) => vec.len() / size, + } + } +} + +impl utils::ExactSize for (FSBTarget, MutableBitmap) { fn len(&self) -> usize { - self.0.values.len() / self.0.size + self.0.len() + } +} + +fn decode_fsb_plain( + size: usize, + values: &[u8], + target: &mut FSBTarget, + validity: &mut MutableBitmap, + is_optional: bool, + filter: Option, + page_validity: Option<&Bitmap>, +) -> ParquetResult<()> { + assert_ne!(size, 0); + assert_eq!(values.len() % size, 0); + + if is_optional { + append_validity( + page_validity, + filter.as_ref(), + validity, + values.len() / size, + ); + } + + let page_validity = constrain_page_validity(values.len() / size, page_validity, filter.as_ref()); + + macro_rules! decode_static_size { + ($target:ident) => {{ + let values = ArrayChunks::new(values).ok_or_else(|| { + ParquetError::oos("Page content does not align with expected element size") + })?; + super::primitive::plain::decode_aligned_bytes_dispatch( + values, + is_optional, + page_validity.as_ref(), + filter, + validity, + $target, + ) + }}; + } + + use FSBTarget as T; + match target { + T::Size1(target) => decode_static_size!(target), + T::Size2(target) => decode_static_size!(target), + T::Size4(target) => decode_static_size!(target), + T::Size8(target) => decode_static_size!(target), + T::Size12(target) => decode_static_size!(target), + T::Size16(target) => decode_static_size!(target), + T::Size32(target) => decode_static_size!(target), + T::Other(_target, _) => todo!(), } } impl Decoder for BinaryDecoder { type Translation<'a> = StateTranslation<'a>; type Dict = Vec; - type DecodedState = (FixedSizeBinary, MutableBitmap); + type DecodedState = (FSBTarget, MutableBitmap); type Output = FixedSizeBinaryArray; fn with_capacity(&self, capacity: usize) -> Self::DecodedState { let size = self.size; - ( - FixedSizeBinary { - values: Vec::with_capacity(capacity * size), - size, - }, - MutableBitmap::with_capacity(capacity), - ) + let values = match size { + 1 => FSBTarget::Size1(Vec::with_capacity(capacity)), + 2 => FSBTarget::Size2(Vec::with_capacity(capacity)), + 4 => FSBTarget::Size4(Vec::with_capacity(capacity)), + 8 => FSBTarget::Size8(Vec::with_capacity(capacity)), + 12 => FSBTarget::Size12(Vec::with_capacity(capacity)), + 16 => FSBTarget::Size16(Vec::with_capacity(capacity)), + 32 => FSBTarget::Size32(Vec::with_capacity(capacity)), + _ => FSBTarget::Other(Vec::with_capacity(capacity * size), size), + }; + + (values, MutableBitmap::with_capacity(capacity)) } fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { @@ -144,64 +234,13 @@ impl Decoder for BinaryDecoder { fn decode_plain_encoded<'a>( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - is_optional: bool, - page_validity: Option<&mut Bitmap>, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, + _is_optional: bool, + _page_validity: Option<&mut Bitmap>, + _limit: usize, ) -> ParquetResult<()> { - struct FixedSizeBinaryCollector<'a, 'b> { - slice: &'b mut &'a [u8], - size: usize, - } - - impl<'a, 'b> BatchableCollector<(), Vec> for FixedSizeBinaryCollector<'a, 'b> { - fn reserve(target: &mut Vec, n: usize) { - target.reserve(n); - } - - fn push_n(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - let n = usize::min(n, self.slice.len() / self.size); - target.extend_from_slice(&self.slice[..n * self.size]); - *self.slice = &self.slice[n * self.size..]; - Ok(()) - } - - fn push_n_nulls(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - target.resize(target.len() + n * self.size, 0); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - let n = usize::min(n, self.slice.len() / self.size); - *self.slice = &self.slice[n * self.size..]; - Ok(()) - } - } - - let mut collector = FixedSizeBinaryCollector { - slice: page_values, - size: self.size, - }; - - match page_validity { - None => { - collector.push_n(&mut values.values, limit)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => extend_from_decoder( - validity, - page_validity, - Some(limit), - &mut values.values, - collector, - )?, - } - - Ok(()) + unreachable!() } fn decode_dictionary_encoded( diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs index f6287fdb7066..18ba498171fe 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs @@ -5,7 +5,7 @@ use crate::parquet::types::NativeType as ParquetNativeType; mod float; mod integer; -mod plain; +pub(crate) mod plain; pub(crate) use float::FloatDecoder; pub(crate) use integer::IntDecoder; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 47621b0e68f2..96b7a89dcc0a 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -5,7 +5,7 @@ use super::DecoderFunction; use crate::parquet::error::ParquetResult; use crate::parquet::types::NativeType as ParquetNativeType; use crate::read::deserialize::utils::array_chunks::ArrayChunks; -use crate::read::deserialize::utils::dict_encoded::append_validity; +use crate::read::deserialize::utils::dict_encoded::{append_validity, constrain_page_validity}; use crate::read::deserialize::utils::filter_from_range; use crate::read::{Filter, ParquetError}; @@ -83,30 +83,7 @@ pub fn decode_aligned_bytes_dispatch( append_validity(page_validity, filter.as_ref(), validity, values.len()); } - let num_unfiltered_rows = match (filter.as_ref(), page_validity) { - (None, None) => values.len(), - (None, Some(pv)) => { - debug_assert!(pv.len() >= values.len()); - pv.len() - }, - (Some(f), v) => { - if cfg!(debug_assertions) { - if let Some(v) = v { - assert!(v.len() >= f.max_offset()); - } - } - - f.max_offset() - }, - }; - - let page_validity = page_validity.map(|pv| { - if pv.len() > num_unfiltered_rows { - pv.clone().sliced(0, num_unfiltered_rows) - } else { - pv.clone() - } - }); + let page_validity = constrain_page_validity(values.len(), page_validity, filter.as_ref()); dbg!(&filter); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index 327093cbec0b..e6803c8bd624 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -48,24 +48,7 @@ pub(crate) fn append_validity( } } -#[inline(never)] -pub fn decode_dict_dispatch( - values: HybridRleDecoder<'_>, - dict: &[B], - is_optional: bool, - page_validity: Option<&Bitmap>, - filter: Option, - validity: &mut MutableBitmap, - target: &mut Vec, -) -> ParquetResult<()> { - if cfg!(debug_assertions) && is_optional { - assert_eq!(target.len(), validity.len()); - } - - if is_optional { - append_validity(page_validity, filter.as_ref(), validity, values.len()); - } - +pub(crate) fn constrain_page_validity(values_len: usize, page_validity: Option<&Bitmap>, filter: Option<&Filter>) -> Option { let num_unfiltered_rows = match (filter.as_ref(), page_validity) { (None, None) => values.len(), (None, Some(pv)) => { @@ -83,13 +66,34 @@ pub fn decode_dict_dispatch( }, }; - let page_validity = page_validity.map(|pv| { + page_validity.map(|pv| { if pv.len() > num_unfiltered_rows { pv.clone().sliced(0, num_unfiltered_rows) } else { pv.clone() } - }); + }) +} + +#[inline(never)] +pub fn decode_dict_dispatch( + values: HybridRleDecoder<'_>, + dict: &[B], + is_optional: bool, + page_validity: Option<&Bitmap>, + filter: Option, + validity: &mut MutableBitmap, + target: &mut Vec, +) -> ParquetResult<()> { + if cfg!(debug_assertions) && is_optional { + assert_eq!(target.len(), validity.len()); + } + + if is_optional { + append_validity(page_validity, filter.as_ref(), validity, values.len()); + } + + let page_validity = constrain_page_validity(values.len(), page_validity, filter.as_ref()); dbg!(&filter); From 5867824863d7130f348635e7895346570eaef787 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 10:19:17 +0200 Subject: [PATCH 13/32] impl fsb kernels --- .../polars-arrow/src/types/aligned_bytes.rs | 6 + .../read/deserialize/fixed_size_binary.rs | 333 ++++++++++-------- .../read/deserialize/utils/dict_encoded.rs | 4 +- .../src/arrow/read/deserialize/utils/mod.rs | 55 --- 4 files changed, 198 insertions(+), 200 deletions(-) diff --git a/crates/polars-arrow/src/types/aligned_bytes.rs b/crates/polars-arrow/src/types/aligned_bytes.rs index 48ff2dfabade..5d766c75ceeb 100644 --- a/crates/polars-arrow/src/types/aligned_bytes.rs +++ b/crates/polars-arrow/src/types/aligned_bytes.rs @@ -100,13 +100,19 @@ macro_rules! impl_aligned_bytes { impl_aligned_bytes! { (Bytes1Alignment1, 1, 1, [u8, i8]), + (Bytes2Alignment1, 2, 1, []), (Bytes2Alignment2, 2, 2, [u16, i16, f16]), + (Bytes4Alignment1, 4, 1, []), (Bytes4Alignment4, 4, 4, [u32, i32, f32]), + (Bytes8Alignment1, 8, 1, []), (Bytes8Alignment8, 8, 8, [u64, i64, f64]), (Bytes8Alignment4, 8, 4, [days_ms]), + (Bytes12Alignment1, 8, 1, []), (Bytes12Alignment4, 12, 4, [[u32; 3]]), + (Bytes16Alignment1, 16, 1, []), (Bytes16Alignment4, 16, 4, [View]), (Bytes16Alignment8, 16, 8, [months_days_ns]), (Bytes16Alignment16, 16, 16, [u128, i128]), + (Bytes32Alignment1, 32, 1, []), (Bytes32Alignment16, 32, 16, [i256]), } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 3f03099f3cc4..53875f92b99d 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -2,20 +2,20 @@ use arrow::array::{DictionaryArray, DictionaryKey, FixedSizeBinaryArray, Primiti use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::types::{ - Bytes12Alignment4, Bytes16Alignment16, Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, - Bytes4Alignment4, Bytes8Alignment8, + Bytes12Alignment1, Bytes16Alignment1, Bytes1Alignment1, Bytes2Alignment1, Bytes32Alignment1, + Bytes4Alignment1, Bytes8Alignment1, }; use super::utils::array_chunks::ArrayChunks; use super::utils::dict_encoded::append_validity; -use super::utils::{dict_indices_decoder, extend_from_decoder, freeze_validity, Decoder}; +use super::utils::{dict_indices_decoder, freeze_validity, Decoder}; use super::Filter; -use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; +use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; use crate::parquet::encoding::{hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; +use crate::read::deserialize::utils; use crate::read::deserialize::utils::dict_encoded::constrain_page_validity; -use crate::read::deserialize::utils::{self, BatchableCollector, GatheredHybridRle}; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -24,12 +24,6 @@ pub(crate) enum StateTranslation<'a> { Dictionary(hybrid_rle::HybridRleDecoder<'a>), } -#[derive(Debug)] -pub struct FixedSizeBinary { - pub values: Vec, - pub size: usize, -} - impl<'a> utils::StateTranslation<'a, BinaryDecoder> for StateTranslation<'a> { type PlainDecoder = &'a [u8]; @@ -116,38 +110,67 @@ pub(crate) struct BinaryDecoder { pub(crate) size: usize, } -enum FSBTarget { +pub(crate) enum FSBVec { Size1(Vec), - Size2(Vec), - Size4(Vec), - Size8(Vec), - Size12(Vec), - Size16(Vec), - Size32(Vec), + Size2(Vec), + Size4(Vec), + Size8(Vec), + Size12(Vec), + Size16(Vec), + Size32(Vec), Other(Vec, usize), } +impl FSBVec { + pub fn new(size: usize) -> FSBVec { + match size { + 1 => Self::Size1(Vec::new()), + 2 => Self::Size2(Vec::new()), + 4 => Self::Size4(Vec::new()), + 8 => Self::Size8(Vec::new()), + 12 => Self::Size12(Vec::new()), + 16 => Self::Size16(Vec::new()), + 32 => Self::Size32(Vec::new()), + _ => Self::Other(Vec::new(), size), + } + } + + pub fn to_bytes_vec(self) -> Vec { + match self { + FSBVec::Size1(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Size2(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Size4(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Size8(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Size12(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Size16(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Size32(vec) => bytemuck::allocation::cast_vec(vec), + FSBVec::Other(vec, _) => vec, + } + } +} + impl utils::ExactSize for Vec { fn len(&self) -> usize { Vec::len(self) } } -impl utils::ExactSize for FSBTarget { +impl utils::ExactSize for FSBVec { fn len(&self) -> usize { match self { - FSBTarget::Size1(vec) => vec.len(), - FSBTarget::Size2(vec) => vec.len(), - FSBTarget::Size4(vec) => vec.len(), - FSBTarget::Size8(vec) => vec.len(), - FSBTarget::Size16(vec) => vec.len(), - FSBTarget::Size32(vec) => vec.len(), - FSBTarget::Other(vec, size) => vec.len() / size, + FSBVec::Size1(vec) => vec.len(), + FSBVec::Size2(vec) => vec.len(), + FSBVec::Size4(vec) => vec.len(), + FSBVec::Size8(vec) => vec.len(), + FSBVec::Size12(vec) => vec.len(), + FSBVec::Size16(vec) => vec.len(), + FSBVec::Size32(vec) => vec.len(), + FSBVec::Other(vec, size) => vec.len() / size, } } } -impl utils::ExactSize for (FSBTarget, MutableBitmap) { +impl utils::ExactSize for (FSBVec, MutableBitmap) { fn len(&self) -> usize { self.0.len() } @@ -156,7 +179,7 @@ impl utils::ExactSize for (FSBTarget, MutableBitmap) { fn decode_fsb_plain( size: usize, values: &[u8], - target: &mut FSBTarget, + target: &mut FSBVec, validity: &mut MutableBitmap, is_optional: bool, filter: Option, @@ -174,7 +197,8 @@ fn decode_fsb_plain( ); } - let page_validity = constrain_page_validity(values.len() / size, page_validity, filter.as_ref()); + let page_validity = + constrain_page_validity(values.len() / size, page_validity, filter.as_ref()); macro_rules! decode_static_size { ($target:ident) => {{ @@ -192,7 +216,7 @@ fn decode_fsb_plain( }}; } - use FSBTarget as T; + use FSBVec as T; match target { T::Size1(target) => decode_static_size!(target), T::Size2(target) => decode_static_size!(target), @@ -201,35 +225,116 @@ fn decode_fsb_plain( T::Size12(target) => decode_static_size!(target), T::Size16(target) => decode_static_size!(target), T::Size32(target) => decode_static_size!(target), - T::Other(_target, _) => todo!(), + T::Other(_target, _) => { + todo!() + // match (page_validity, filter.as_ref()) { + // (None, None) => target.extend_from_slice(values), + // (None, Some(filter)) => match filter { + // Filter::Range(range) => target.extend_from_slice(&values[range.start * size..range.end * size]), + // Filter::Mask(bitmap) => { + // for v in + // }, + // }, + // (Some(_), None) => todo!(), + // (Some(_), Some(_)) => todo!(), + // } + // + // Ok(()) + }, + } +} + +fn decode_fsb_dict( + size: usize, + values: HybridRleDecoder<'_>, + dict: &FSBVec, + target: &mut FSBVec, + validity: &mut MutableBitmap, + is_optional: bool, + filter: Option, + page_validity: Option<&Bitmap>, +) -> ParquetResult<()> { + assert_ne!(size, 0); + assert_eq!(values.len() % size, 0); + + macro_rules! decode_static_size { + ($dict:ident, $target:ident) => {{ + super::utils::dict_encoded::decode_dict_dispatch( + values, + $dict, + is_optional, + page_validity, + filter, + validity, + $target, + ) + }}; + } + + use FSBVec as T; + match (dict, target) { + (T::Size1(dict), T::Size1(target)) => decode_static_size!(dict, target), + (T::Size2(dict), T::Size2(target)) => decode_static_size!(dict, target), + (T::Size4(dict), T::Size4(target)) => decode_static_size!(dict, target), + (T::Size8(dict), T::Size8(target)) => decode_static_size!(dict, target), + (T::Size12(dict), T::Size12(target)) => decode_static_size!(dict, target), + (T::Size16(dict), T::Size16(target)) => decode_static_size!(dict, target), + (T::Size32(dict), T::Size32(target)) => decode_static_size!(dict, target), + (T::Other(_dict, _), T::Other(_target, _)) => { + todo!() + // match (page_validity, filter.as_ref()) { + // (None, None) => target.extend_from_slice(values), + // (None, Some(filter)) => match filter { + // Filter::Range(range) => target.extend_from_slice(&values[range.start * size..range.end * size]), + // Filter::Mask(bitmap) => { + // for v in + // }, + // }, + // (Some(_), None) => todo!(), + // (Some(_), Some(_)) => todo!(), + // } + // + // Ok(()) + }, + _ => unreachable!(), } } impl Decoder for BinaryDecoder { type Translation<'a> = StateTranslation<'a>; - type Dict = Vec; - type DecodedState = (FSBTarget, MutableBitmap); + type Dict = FSBVec; + type DecodedState = (FSBVec, MutableBitmap); type Output = FixedSizeBinaryArray; fn with_capacity(&self, capacity: usize) -> Self::DecodedState { let size = self.size; let values = match size { - 1 => FSBTarget::Size1(Vec::with_capacity(capacity)), - 2 => FSBTarget::Size2(Vec::with_capacity(capacity)), - 4 => FSBTarget::Size4(Vec::with_capacity(capacity)), - 8 => FSBTarget::Size8(Vec::with_capacity(capacity)), - 12 => FSBTarget::Size12(Vec::with_capacity(capacity)), - 16 => FSBTarget::Size16(Vec::with_capacity(capacity)), - 32 => FSBTarget::Size32(Vec::with_capacity(capacity)), - _ => FSBTarget::Other(Vec::with_capacity(capacity * size), size), + 1 => FSBVec::Size1(Vec::with_capacity(capacity)), + 2 => FSBVec::Size2(Vec::with_capacity(capacity)), + 4 => FSBVec::Size4(Vec::with_capacity(capacity)), + 8 => FSBVec::Size8(Vec::with_capacity(capacity)), + 12 => FSBVec::Size12(Vec::with_capacity(capacity)), + 16 => FSBVec::Size16(Vec::with_capacity(capacity)), + 32 => FSBVec::Size32(Vec::with_capacity(capacity)), + _ => FSBVec::Other(Vec::with_capacity(capacity * size), size), }; (values, MutableBitmap::with_capacity(capacity)) } fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { - Ok(page.buffer.into_vec()) + let mut target = FSBVec::new(self.size); + decode_fsb_plain( + self.size, + page.buffer.as_ref(), + &mut target, + &mut MutableBitmap::new(), + false, + None, + None, + )?; + Ok(target) } fn decode_plain_encoded<'a>( @@ -245,105 +350,14 @@ impl Decoder for BinaryDecoder { fn decode_dictionary_encoded( &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut hybrid_rle::HybridRleDecoder<'_>, - is_optional: bool, - page_validity: Option<&mut Bitmap>, - dict: &Self::Dict, - limit: usize, + _decoded: &mut Self::DecodedState, + _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, + _is_optional: bool, + _page_validity: Option<&mut Bitmap>, + _dict: &Self::Dict, + _limit: usize, ) -> ParquetResult<()> { - struct FixedSizeBinaryGatherer<'a> { - dict: &'a [u8], - size: usize, - } - - impl<'a> HybridRleGatherer<&'a [u8]> for FixedSizeBinaryGatherer<'a> { - type Target = Vec; - - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n * self.size); - } - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.len() / self.size - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult<&'a [u8]> { - let value = value as usize; - - if value * self.size >= self.dict.len() { - return Err(ParquetError::oos( - "Fixed size binary dictionary index out-of-range", - )); - } - - Ok(&self.dict[value * self.size..(value + 1) * self.size]) - } - - fn gather_one(&self, target: &mut Self::Target, value: &'a [u8]) -> ParquetResult<()> { - // We make the null value length 0, which allows us to do this. - if value.is_empty() { - target.resize(target.len() + self.size, 0); - return Ok(()); - } - - target.extend_from_slice(value); - Ok(()) - } - - fn gather_repeated( - &self, - target: &mut Self::Target, - value: &'a [u8], - n: usize, - ) -> ParquetResult<()> { - // We make the null value length 0, which allows us to do this. - if value.is_empty() { - target.resize(target.len() + n * self.size, 0); - return Ok(()); - } - - debug_assert_eq!(value.len(), self.size); - for _ in 0..n { - target.extend(value); - } - - Ok(()) - } - } - - let gatherer = FixedSizeBinaryGatherer { - dict, - size: self.size, - }; - - // @NOTE: - // This is a special case in our gatherer. If the length of the value is 0, then we just - // resize with the appropriate size. Important is that this also works for FSL with size=0. - let null_value = &[]; - - match page_validity { - None => { - page_values.gather_n_into(&mut values.values, limit, &gatherer)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - let collector = GatheredHybridRle::new(page_values, &gatherer, null_value); - - extend_from_decoder( - validity, - page_validity, - Some(limit), - &mut values.values, - collector, - )?; - }, - } - - Ok(()) + unreachable!() } fn finalize( @@ -353,12 +367,42 @@ impl Decoder for BinaryDecoder { (values, validity): Self::DecodedState, ) -> ParquetResult { let validity = freeze_validity(validity); + Ok(FixedSizeBinaryArray::new( dtype, - values.values.into(), + values.to_bytes_vec().into(), validity, )) } + + fn extend_filtered_with_state( + &mut self, + state: utils::State<'_, Self>, + decoded: &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + match state.translation { + StateTranslation::Plain(values, size) => decode_fsb_plain( + size, + values, + &mut decoded.0, + &mut decoded.1, + state.is_optional, + filter, + state.page_validity.as_ref(), + ), + StateTranslation::Dictionary(values) => decode_fsb_dict( + self.size, + values, + state.dict.unwrap(), + &mut decoded.0, + &mut decoded.1, + state.is_optional, + filter, + state.page_validity.as_ref(), + ), + } + } } impl utils::DictDecodable for BinaryDecoder { @@ -368,8 +412,11 @@ impl utils::DictDecodable for BinaryDecoder { dict: Self::Dict, keys: PrimitiveArray, ) -> ParquetResult> { - let dict = - FixedSizeBinaryArray::new(ArrowDataType::FixedSizeBinary(self.size), dict.into(), None); + let dict = FixedSizeBinaryArray::new( + ArrowDataType::FixedSizeBinary(self.size), + dict.to_bytes_vec().into(), + None, + ); Ok(DictionaryArray::try_new(dtype, keys, Box::new(dict)).unwrap()) } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index e6803c8bd624..e89b8525ea84 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -50,9 +50,9 @@ pub(crate) fn append_validity( pub(crate) fn constrain_page_validity(values_len: usize, page_validity: Option<&Bitmap>, filter: Option<&Filter>) -> Option { let num_unfiltered_rows = match (filter.as_ref(), page_validity) { - (None, None) => values.len(), + (None, None) => values_len, (None, Some(pv)) => { - debug_assert!(pv.len() >= values.len()); + debug_assert!(pv.len() >= values_len); pv.len() }, (Some(f), v) => { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 9054fe93e02d..e4b62b990030 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -347,61 +347,6 @@ pub(super) fn extend_from_decoder>( Ok(()) } -pub struct GatheredHybridRle<'a, 'b, 'c, O, G> -where - O: Clone, - G: HybridRleGatherer, -{ - decoder: &'a mut HybridRleDecoder<'b>, - gatherer: &'c G, - null_value: O, - _pd: std::marker::PhantomData, -} - -impl<'a, 'b, 'c, O, G> GatheredHybridRle<'a, 'b, 'c, O, G> -where - O: Clone, - G: HybridRleGatherer, -{ - pub fn new(decoder: &'a mut HybridRleDecoder<'b>, gatherer: &'c G, null_value: O) -> Self { - Self { - decoder, - gatherer, - null_value, - _pd: Default::default(), - } - } -} - -impl<'a, 'b, 'c, O, G> BatchableCollector> for GatheredHybridRle<'a, 'b, 'c, O, G> -where - O: Clone, - G: HybridRleGatherer>, -{ - #[inline] - fn reserve(target: &mut Vec, n: usize) { - target.reserve(n); - } - - #[inline] - fn push_n(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - self.decoder.gather_n_into(target, n, self.gatherer)?; - Ok(()) - } - - #[inline] - fn push_n_nulls(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - self.gatherer - .gather_repeated(target, self.null_value.clone(), n)?; - Ok(()) - } - - #[inline] - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } -} - impl, I: Iterator> BatchableCollector for I { #[inline] fn reserve(target: &mut P, n: usize) { From 45d5159fd789182b1c3cf904e78029b72b11cc9f Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 14:26:24 +0200 Subject: [PATCH 14/32] properly support and check fixed size binary --- crates/polars-arrow/src/bitmap/bitmask.rs | 6 +- .../polars-arrow/src/types/aligned_bytes.rs | 2 +- .../read/deserialize/fixed_size_binary.rs | 258 ++++++++++++++---- .../arrow/read/deserialize/primitive/plain.rs | 10 +- .../read/deserialize/utils/dict_encoded.rs | 8 +- .../src/parquet/encoding/hybrid_rle/mod.rs | 4 + py-polars/tests/unit/io/test_parquet.py | 74 ++++- 7 files changed, 301 insertions(+), 61 deletions(-) diff --git a/crates/polars-arrow/src/bitmap/bitmask.rs b/crates/polars-arrow/src/bitmap/bitmask.rs index b1bc5cca3d86..1421984be207 100644 --- a/crates/polars-arrow/src/bitmap/bitmask.rs +++ b/crates/polars-arrow/src/bitmap/bitmask.rs @@ -4,7 +4,7 @@ use std::simd::{LaneCount, Mask, MaskElement, SupportedLaneCount}; use polars_utils::slice::load_padded_le_u64; use super::iterator::FastU56BitmapIter; -use super::utils::count_zeros; +use super::utils::{count_zeros, BitmapIter}; use crate::bitmap::Bitmap; /// Returns the nth set bit in w, if n+1 bits are set. The indexing is @@ -280,6 +280,10 @@ impl<'a> BitMask<'a> { false } } + + pub fn iter(&self) -> BitmapIter { + BitmapIter::new(&self.bytes, self.offset, self.len) + } } #[cfg(test)] diff --git a/crates/polars-arrow/src/types/aligned_bytes.rs b/crates/polars-arrow/src/types/aligned_bytes.rs index 5d766c75ceeb..738cb0174fc1 100644 --- a/crates/polars-arrow/src/types/aligned_bytes.rs +++ b/crates/polars-arrow/src/types/aligned_bytes.rs @@ -107,7 +107,7 @@ impl_aligned_bytes! { (Bytes8Alignment1, 8, 1, []), (Bytes8Alignment8, 8, 8, [u64, i64, f64]), (Bytes8Alignment4, 8, 4, [days_ms]), - (Bytes12Alignment1, 8, 1, []), + (Bytes12Alignment1, 12, 1, []), (Bytes12Alignment4, 12, 4, [[u32; 3]]), (Bytes16Alignment1, 16, 1, []), (Bytes16Alignment4, 16, 4, [View]), diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 53875f92b99d..5f711c1d6d26 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -1,4 +1,6 @@ -use arrow::array::{DictionaryArray, DictionaryKey, FixedSizeBinaryArray, PrimitiveArray}; +use arrow::array::{ + DictionaryArray, DictionaryKey, FixedSizeBinaryArray, PrimitiveArray, Splitable, +}; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use arrow::types::{ @@ -10,7 +12,7 @@ use super::utils::array_chunks::ArrayChunks; use super::utils::dict_encoded::append_validity; use super::utils::{dict_indices_decoder, freeze_validity, Decoder}; use super::Filter; -use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; +use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::encoding::{hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; @@ -188,18 +190,6 @@ fn decode_fsb_plain( assert_ne!(size, 0); assert_eq!(values.len() % size, 0); - if is_optional { - append_validity( - page_validity, - filter.as_ref(), - validity, - values.len() / size, - ); - } - - let page_validity = - constrain_page_validity(values.len() / size, page_validity, filter.as_ref()); - macro_rules! decode_static_size { ($target:ident) => {{ let values = ArrayChunks::new(values).ok_or_else(|| { @@ -208,7 +198,7 @@ fn decode_fsb_plain( super::primitive::plain::decode_aligned_bytes_dispatch( values, is_optional, - page_validity.as_ref(), + page_validity, filter, validity, $target, @@ -225,21 +215,92 @@ fn decode_fsb_plain( T::Size12(target) => decode_static_size!(target), T::Size16(target) => decode_static_size!(target), T::Size32(target) => decode_static_size!(target), - T::Other(_target, _) => { - todo!() - // match (page_validity, filter.as_ref()) { - // (None, None) => target.extend_from_slice(values), - // (None, Some(filter)) => match filter { - // Filter::Range(range) => target.extend_from_slice(&values[range.start * size..range.end * size]), - // Filter::Mask(bitmap) => { - // for v in - // }, - // }, - // (Some(_), None) => todo!(), - // (Some(_), Some(_)) => todo!(), - // } - // - // Ok(()) + T::Other(target, _) => { + // @NOTE: All these kernels are quite slow, but they should be very uncommon and the + // general case requires arbitrary length memcopies anyway. + + if is_optional { + append_validity( + page_validity, + filter.as_ref(), + validity, + values.len() / size, + ); + } + + let page_validity = + constrain_page_validity(values.len() / size, page_validity, filter.as_ref()); + + match (page_validity, filter.as_ref()) { + (None, None) => target.extend_from_slice(values), + (None, Some(filter)) => match filter { + Filter::Range(range) => { + target.extend_from_slice(&values[range.start * size..range.end * size]) + }, + Filter::Mask(bitmap) => { + let mut iter = bitmap.iter(); + let mut offset = 0; + + while iter.num_remaining() > 0 { + let num_selected = iter.take_leading_ones(); + target + .extend_from_slice(&values[offset * size..][..num_selected * size]); + offset += num_selected; + + let num_filtered = iter.take_leading_zeros(); + offset += num_filtered; + } + }, + }, + (Some(validity), None) => { + let mut iter = validity.iter(); + let mut offset = 0; + + while iter.num_remaining() > 0 { + let num_valid = iter.take_leading_ones(); + target.extend_from_slice(&values[offset * size..][..num_valid * size]); + offset += num_valid; + + let num_filtered = iter.take_leading_zeros(); + target.resize(target.len() + num_filtered * size, 0); + } + }, + (Some(validity), Some(filter)) => match filter { + Filter::Range(range) => { + let (skipped, active) = validity.split_at(range.start); + + let active = active.sliced(0, range.len()); + + let mut iter = active.iter(); + let mut offset = skipped.set_bits(); + + while iter.num_remaining() > 0 { + let num_valid = iter.take_leading_ones(); + target.extend_from_slice(&values[offset * size..][..num_valid * size]); + offset += num_valid; + + let num_filtered = iter.take_leading_zeros(); + target.resize(target.len() + num_filtered * size, 0); + } + }, + Filter::Mask(filter) => { + let mut offset = 0; + for (is_selected, is_valid) in filter.iter().zip(validity.iter()) { + if is_selected { + if is_valid { + target.extend_from_slice(&values[offset * size..][..size]); + } else { + target.resize(target.len() + size, 0); + } + } + + offset += usize::from(is_valid); + } + }, + }, + } + + Ok(()) }, } } @@ -255,7 +316,6 @@ fn decode_fsb_dict( page_validity: Option<&Bitmap>, ) -> ParquetResult<()> { assert_ne!(size, 0); - assert_eq!(values.len() % size, 0); macro_rules! decode_static_size { ($dict:ident, $target:ident) => {{ @@ -280,21 +340,127 @@ fn decode_fsb_dict( (T::Size12(dict), T::Size12(target)) => decode_static_size!(dict, target), (T::Size16(dict), T::Size16(target)) => decode_static_size!(dict, target), (T::Size32(dict), T::Size32(target)) => decode_static_size!(dict, target), - (T::Other(_dict, _), T::Other(_target, _)) => { - todo!() - // match (page_validity, filter.as_ref()) { - // (None, None) => target.extend_from_slice(values), - // (None, Some(filter)) => match filter { - // Filter::Range(range) => target.extend_from_slice(&values[range.start * size..range.end * size]), - // Filter::Mask(bitmap) => { - // for v in - // }, - // }, - // (Some(_), None) => todo!(), - // (Some(_), Some(_)) => todo!(), - // } - // - // Ok(()) + (T::Other(dict, _), T::Other(target, _)) => { + // @NOTE: All these kernels are quite slow, but they should be very uncommon and the + // general case requires arbitrary length memcopies anyway. + + if is_optional { + append_validity( + page_validity, + filter.as_ref(), + validity, + values.len() / size, + ); + } + + let page_validity = + constrain_page_validity(values.len() / size, page_validity, filter.as_ref()); + + let mut indexes = Vec::with_capacity(values.len()); + + for chunk in values.into_chunk_iter() { + match chunk? { + HybridRleChunk::Rle(value, repeats) => { + indexes.resize(indexes.len() + repeats, value) + }, + HybridRleChunk::Bitpacked(decoder) => decoder.collect_into(&mut indexes), + } + } + + match (page_validity, filter.as_ref()) { + (None, None) => target.extend( + indexes + .into_iter() + .map(|v| &dict[(v as usize) * size..][..size]) + .flatten(), + ), + (None, Some(filter)) => match filter { + Filter::Range(range) => target.extend( + indexes[range.start..range.end] + .into_iter() + .map(|v| &dict[(*v as usize) * size..][..size]) + .flatten(), + ), + Filter::Mask(bitmap) => { + let mut iter = bitmap.iter(); + let mut offset = 0; + + while iter.num_remaining() > 0 { + let num_selected = iter.take_leading_ones(); + target.extend( + indexes[offset..][..num_selected] + .into_iter() + .map(|v| &dict[(*v as usize) * size..][..size]) + .flatten(), + ); + offset += num_selected; + + let num_filtered = iter.take_leading_zeros(); + offset += num_filtered; + } + }, + }, + (Some(validity), None) => { + let mut iter = validity.iter(); + let mut offset = 0; + + while iter.num_remaining() > 0 { + let num_valid = iter.take_leading_ones(); + target.extend( + indexes[offset..][..num_valid] + .into_iter() + .map(|v| &dict[(*v as usize) * size..][..size]) + .flatten(), + ); + offset += num_valid; + + let num_filtered = iter.take_leading_zeros(); + target.resize(target.len() + num_filtered * size, 0); + } + }, + (Some(validity), Some(filter)) => match filter { + Filter::Range(range) => { + let (skipped, active) = validity.split_at(range.start); + + let active = active.sliced(0, range.len()); + + let mut iter = active.iter(); + let mut offset = skipped.set_bits(); + + while iter.num_remaining() > 0 { + let num_valid = iter.take_leading_ones(); + target.extend( + indexes[offset..][..num_valid] + .into_iter() + .map(|v| &dict[(*v as usize) * size..][..size]) + .flatten(), + ); + offset += num_valid; + + let num_filtered = iter.take_leading_zeros(); + target.resize(target.len() + num_filtered * size, 0); + } + }, + Filter::Mask(filter) => { + let mut offset = 0; + for (is_selected, is_valid) in filter.iter().zip(validity.iter()) { + if is_selected { + if is_valid { + target.extend_from_slice( + &dict[(indexes[offset] as usize) * size..][..size], + ); + } else { + target.resize(target.len() + size, 0); + } + } + + offset += usize::from(is_valid); + } + }, + }, + } + + Ok(()) }, _ => unreachable!(), } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 96b7a89dcc0a..2861e29e377b 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -66,7 +66,7 @@ pub fn decode>( if cfg!(debug_assertions) && is_optional { assert_eq!(target.len(), validity.len()); } - + Ok(()) } @@ -85,8 +85,6 @@ pub fn decode_aligned_bytes_dispatch( let page_validity = constrain_page_validity(values.len(), page_validity, filter.as_ref()); - dbg!(&filter); - match (filter, page_validity) { (None, None) => decode_required(values, None, target), (Some(Filter::Range(rng)), None) if rng.start == 0 => { @@ -102,13 +100,13 @@ pub fn decode_aligned_bytes_dispatch( }, // @TODO: Use values.skip_in_place(rng.start) (Some(Filter::Range(rng)), None) => { - decode_masked_required(values, dbg!(&filter_from_range(rng.clone())), target) + decode_masked_required(values, &filter_from_range(rng.clone()), target) }, // @TODO: Use values.skip_in_place(page_validity.sliced(0, rng.start).set_bits()) (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( values, - dbg!(&page_validity), - dbg!(&filter_from_range(rng.clone())), + &page_validity, + &filter_from_range(rng.clone()), target, ), }?; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index e89b8525ea84..36e65e82ede4 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -95,8 +95,6 @@ pub fn decode_dict_dispatch( let page_validity = constrain_page_validity(values.len(), page_validity, filter.as_ref()); - dbg!(&filter); - match (filter, page_validity) { (None, None) => decode_required_dict(values, dict, None, target), (Some(Filter::Range(rng)), None) if rng.start == 0 => { @@ -113,13 +111,13 @@ pub fn decode_dict_dispatch( decode_masked_optional_dict(values, dict, &filter, &page_validity, target) }, (Some(Filter::Range(rng)), None) => { - decode_masked_required_dict(values, dict, dbg!(&filter_from_range(rng.clone())), target) + decode_masked_required_dict(values, dict, &filter_from_range(rng.clone()), target) }, (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional_dict( values, dict, - dbg!(&filter_from_range(rng.clone())), - dbg!(&page_validity), + &filter_from_range(rng.clone()), + &page_validity, target, ), }?; diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs index 3f3444ef28e1..7196ab576b52 100644 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs @@ -142,6 +142,10 @@ impl<'a> HybridRleDecoder<'a> { })) } + pub fn limit_to(&mut self, length: usize) { + self.num_values = self.num_values.min(length); + } + fn gather_limited_once>( &mut self, target: &mut G::Target, diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py index 2ff36017da18..e2059bdb4d61 100644 --- a/py-polars/tests/unit/io/test_parquet.py +++ b/py-polars/tests/unit/io/test_parquet.py @@ -3,7 +3,7 @@ import io from datetime import datetime, time, timezone from decimal import Decimal -from typing import IO, TYPE_CHECKING, Any, Literal, cast +from typing import IO, TYPE_CHECKING, Any, Callable, Literal, cast import fsspec import numpy as np @@ -2042,7 +2042,77 @@ def test_conserve_sortedness( ) -def test_f16() -> None: +@pytest.mark.parametrize("use_dictionary", [True, False]) +@pytest.mark.parametrize( + "values", + [ + (size, x) + for size in [1, 2, 3, 4, 8, 12, 15, 16, 32] + for x in [ + [list(range(size)), list(range(7, 7 + size))], + [list(range(size)), None], + [list(range(i, i + size)) for i in range(13)], + [list(range(i, i + size)) if i % 3 < 2 else None for i in range(13)], + ] + ] +) +@pytest.mark.parametrize( + "filt", + [ + lambda _: None, + lambda _: pl.col.f > 0, + lambda _: pl.col.f > 1, + lambda _: pl.col.f < 5, + lambda _: pl.col.f % 2 == 0, + lambda _: pl.col.f % 5 < 4, + lambda values: (0, min(1, len(values))), + lambda _: (1, 1), + lambda _: (-2, 1), + lambda values: (1, len(values) - 2), + ], +) +def test_fixed_size_binary( + use_dictionary: bool, + values: tuple[int, list[None | list[int]]], + filt: Callable[[list[None | list[int]]], None | pl.Expr | tuple[int, int]], +) -> None: + size, elems = values + bs = [bytes(v) if v is not None else None for v in elems] + + tbl = pa.table( + { + "a": bs, + "f": range(len(bs)), + }, + schema=pa.schema( + [ + pa.field("a", pa.binary(length=size), nullable=True), + pa.field("f", pa.int32(), nullable=True), + ] + ), + ) + + df = pl.DataFrame(tbl) + + f = io.BytesIO() + pq.write_table(tbl, f, use_dictionary=use_dictionary) + + f.seek(0) + + loaded: pl.DataFrame + if isinstance(filt, pl.Expr): + loaded = pl.scan_parquet(f).filter(filt).collect() + df = df.filter(filt) + elif isinstance(filt, tuple): + loaded = pl.scan_parquet(f).slice(filt[0], filt[1]).collect() + df = df.slice(filt[0], filt[1]) + else: + loaded = pl.read_parquet(f) + + assert_frame_equal(loaded, df) + + +def test_decode_f16() -> None: values = [float("nan"), 0.0, 0.5, 1.0, 1.5] table = pa.Table.from_pydict( From 279e94a61b7de559084dc5fd5b155639e9533f86 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 15:10:21 +0200 Subject: [PATCH 15/32] add size and alignment casting to sharedstorage --- crates/polars-arrow/src/storage.rs | 3 + .../polars-arrow/src/types/aligned_bytes.rs | 68 ++++++++++++++----- crates/polars-arrow/src/types/mod.rs | 2 - 3 files changed, 54 insertions(+), 19 deletions(-) diff --git a/crates/polars-arrow/src/storage.rs b/crates/polars-arrow/src/storage.rs index 76116f917316..d5dd983f944c 100644 --- a/crates/polars-arrow/src/storage.rs +++ b/crates/polars-arrow/src/storage.rs @@ -33,6 +33,9 @@ impl VecVTable { } use crate::ffi::InternalArrowArray; +use crate::types::{ + AlignedBytes, Bytes12Alignment4, Bytes16Alignment16, Bytes16Alignment4, Bytes16Alignment8, Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, Bytes4Alignment4, Bytes8Alignment4, Bytes8Alignment8, PrimitiveSizeAlignmentPair +}; enum BackingStorage { Vec { diff --git a/crates/polars-arrow/src/types/aligned_bytes.rs b/crates/polars-arrow/src/types/aligned_bytes.rs index 738cb0174fc1..3e0351089788 100644 --- a/crates/polars-arrow/src/types/aligned_bytes.rs +++ b/crates/polars-arrow/src/types/aligned_bytes.rs @@ -14,6 +14,7 @@ pub unsafe trait AlignedBytesCast: Pod {} pub trait AlignedBytes: Pod + Zeroable + Copy + Default + Eq { const ALIGNMENT: usize; const SIZE: usize; + const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair; type Unaligned: AsRef<[u8]> + AsMut<[u8]> @@ -45,7 +46,7 @@ pub trait AlignedBytes: Pod + Zeroable + Copy + Default + Eq { macro_rules! impl_aligned_bytes { ( - $(($name:ident, $size:literal, $alignment:literal, [$($eq_type:ty),*]),)+ + $(($name:ident, $size:literal, $alignment:literal, $sap:ident, [$($eq_type:ty),*]),)+ ) => { $( /// Bytes with a size and alignment. @@ -59,6 +60,7 @@ macro_rules! impl_aligned_bytes { impl AlignedBytes for $name { const ALIGNMENT: usize = $alignment; const SIZE: usize = $size; + const SIZE_ALIGNMENT_PAIR: PrimitiveSizeAlignmentPair = PrimitiveSizeAlignmentPair::$sap; type Unaligned = [u8; $size]; @@ -98,21 +100,53 @@ macro_rules! impl_aligned_bytes { } } +#[derive(Clone, Copy)] +pub enum PrimitiveSizeAlignmentPair { + S1A1, + S2A2, + S4A4, + S8A4, + S8A8, + S12A4, + S16A4, + S16A8, + S16A16, + S32A16, +} + +impl PrimitiveSizeAlignmentPair { + pub const fn size(self) -> usize { + match self { + Self::S1A1 => 1, + Self::S2A2 => 2, + Self::S4A4 => 4, + Self::S8A4 | Self::S8A8 => 8, + Self::S12A4 => 12, + Self::S16A4 | Self::S16A8 | Self::S16A16 => 16, + Self::S32A16 => 32, + } + } + + pub const fn alignment(self) -> usize { + match self { + Self::S1A1 => 1, + Self::S2A2 => 2, + Self::S4A4 | Self::S8A4 | Self::S12A4 | Self::S16A4 => 4, + Self::S8A8 | Self::S16A8 => 8, + Self::S16A16 | Self::S32A16 => 16, + } + } +} + impl_aligned_bytes! { - (Bytes1Alignment1, 1, 1, [u8, i8]), - (Bytes2Alignment1, 2, 1, []), - (Bytes2Alignment2, 2, 2, [u16, i16, f16]), - (Bytes4Alignment1, 4, 1, []), - (Bytes4Alignment4, 4, 4, [u32, i32, f32]), - (Bytes8Alignment1, 8, 1, []), - (Bytes8Alignment8, 8, 8, [u64, i64, f64]), - (Bytes8Alignment4, 8, 4, [days_ms]), - (Bytes12Alignment1, 12, 1, []), - (Bytes12Alignment4, 12, 4, [[u32; 3]]), - (Bytes16Alignment1, 16, 1, []), - (Bytes16Alignment4, 16, 4, [View]), - (Bytes16Alignment8, 16, 8, [months_days_ns]), - (Bytes16Alignment16, 16, 16, [u128, i128]), - (Bytes32Alignment1, 32, 1, []), - (Bytes32Alignment16, 32, 16, [i256]), + (Bytes1Alignment1, 1, 1, S1A1, [u8, i8]), + (Bytes2Alignment2, 2, 2, S2A2, [u16, i16, f16]), + (Bytes4Alignment4, 4, 4, S4A4, [u32, i32, f32]), + (Bytes8Alignment8, 8, 8, S8A8, [u64, i64, f64]), + (Bytes8Alignment4, 8, 4, S8A4, [days_ms]), + (Bytes12Alignment4, 12, 4, S12A4, [[u32; 3]]), + (Bytes16Alignment4, 16, 4, S16A4, [View]), + (Bytes16Alignment8, 16, 8, S16A8, [months_days_ns]), + (Bytes16Alignment16, 16, 16, S16A16, [u128, i128]), + (Bytes32Alignment16, 32, 16, S32A16, [i256]), } diff --git a/crates/polars-arrow/src/types/mod.rs b/crates/polars-arrow/src/types/mod.rs index e720254294c6..c6f653a32311 100644 --- a/crates/polars-arrow/src/types/mod.rs +++ b/crates/polars-arrow/src/types/mod.rs @@ -25,8 +25,6 @@ pub use aligned_bytes::*; mod bit_chunk; pub use bit_chunk::{BitChunk, BitChunkIter, BitChunkOnes}; mod index; -mod aligned_bytes; -pub use aligned_bytes::*; pub mod simd; pub use index::*; mod native; From 335b656a02b7b67def069ba218f2e0ff0beb552d Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 15:25:35 +0200 Subject: [PATCH 16/32] fixedsizebinary with existing kernels --- crates/polars-arrow/src/buffer/immutable.rs | 2 +- crates/polars-arrow/src/storage.rs | 4 +- .../read/deserialize/fixed_size_binary.rs | 44 ++++++++++--------- 3 files changed, 27 insertions(+), 23 deletions(-) diff --git a/crates/polars-arrow/src/buffer/immutable.rs b/crates/polars-arrow/src/buffer/immutable.rs index 1c6e5b5aa4ff..eb27a62341b9 100644 --- a/crates/polars-arrow/src/buffer/immutable.rs +++ b/crates/polars-arrow/src/buffer/immutable.rs @@ -79,7 +79,7 @@ impl Buffer { } /// Auxiliary method to create a new Buffer - pub(crate) fn from_storage(storage: SharedStorage) -> Self { + pub fn from_storage(storage: SharedStorage) -> Self { let ptr = storage.as_ptr(); let length = storage.len(); Buffer { diff --git a/crates/polars-arrow/src/storage.rs b/crates/polars-arrow/src/storage.rs index d5dd983f944c..a1b2050bc14b 100644 --- a/crates/polars-arrow/src/storage.rs +++ b/crates/polars-arrow/src/storage.rs @@ -34,7 +34,9 @@ impl VecVTable { use crate::ffi::InternalArrowArray; use crate::types::{ - AlignedBytes, Bytes12Alignment4, Bytes16Alignment16, Bytes16Alignment4, Bytes16Alignment8, Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, Bytes4Alignment4, Bytes8Alignment4, Bytes8Alignment8, PrimitiveSizeAlignmentPair + AlignedBytes, Bytes12Alignment4, Bytes16Alignment16, Bytes16Alignment4, Bytes16Alignment8, + Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, Bytes4Alignment4, Bytes8Alignment4, + Bytes8Alignment8, PrimitiveSizeAlignmentPair, }; enum BackingStorage { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 5f711c1d6d26..5e36e7ec3aa8 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -2,10 +2,12 @@ use arrow::array::{ DictionaryArray, DictionaryKey, FixedSizeBinaryArray, PrimitiveArray, Splitable, }; use arrow::bitmap::{Bitmap, MutableBitmap}; +use arrow::buffer::Buffer; use arrow::datatypes::ArrowDataType; +use arrow::storage::SharedStorage; use arrow::types::{ - Bytes12Alignment1, Bytes16Alignment1, Bytes1Alignment1, Bytes2Alignment1, Bytes32Alignment1, - Bytes4Alignment1, Bytes8Alignment1, + Bytes12Alignment4, Bytes16Alignment16, Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, + Bytes4Alignment4, Bytes8Alignment8, }; use super::utils::array_chunks::ArrayChunks; @@ -114,12 +116,12 @@ pub(crate) struct BinaryDecoder { pub(crate) enum FSBVec { Size1(Vec), - Size2(Vec), - Size4(Vec), - Size8(Vec), - Size12(Vec), - Size16(Vec), - Size32(Vec), + Size2(Vec), + Size4(Vec), + Size8(Vec), + Size12(Vec), + Size16(Vec), + Size32(Vec), Other(Vec, usize), } @@ -137,17 +139,17 @@ impl FSBVec { } } - pub fn to_bytes_vec(self) -> Vec { - match self { - FSBVec::Size1(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Size2(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Size4(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Size8(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Size12(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Size16(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Size32(vec) => bytemuck::allocation::cast_vec(vec), - FSBVec::Other(vec, _) => vec, - } + pub fn to_bytes_buffer(self) -> Buffer { + Buffer::from_storage(match self { + FSBVec::Size1(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size2(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size4(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size8(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size12(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size16(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size32(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Other(vec, _) => SharedStorage::from_vec(vec), + }) } } @@ -536,7 +538,7 @@ impl Decoder for BinaryDecoder { Ok(FixedSizeBinaryArray::new( dtype, - values.to_bytes_vec().into(), + values.to_bytes_buffer(), validity, )) } @@ -580,7 +582,7 @@ impl utils::DictDecodable for BinaryDecoder { ) -> ParquetResult> { let dict = FixedSizeBinaryArray::new( ArrowDataType::FixedSizeBinary(self.size), - dict.to_bytes_vec().into(), + dict.to_bytes_buffer(), None, ); Ok(DictionaryArray::try_new(dtype, keys, Box::new(dict)).unwrap()) From 69f856b2b96e0b40096184e75b5ec9924f803d5b Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 15:44:32 +0200 Subject: [PATCH 17/32] properly drop backingstorage with original_type --- crates/polars-arrow/src/storage.rs | 5 --- .../read/deserialize/fixed_size_binary.rs | 33 ++++--------------- .../arrow/read/deserialize/primitive/float.rs | 16 +-------- 3 files changed, 8 insertions(+), 46 deletions(-) diff --git a/crates/polars-arrow/src/storage.rs b/crates/polars-arrow/src/storage.rs index a1b2050bc14b..76116f917316 100644 --- a/crates/polars-arrow/src/storage.rs +++ b/crates/polars-arrow/src/storage.rs @@ -33,11 +33,6 @@ impl VecVTable { } use crate::ffi::InternalArrowArray; -use crate::types::{ - AlignedBytes, Bytes12Alignment4, Bytes16Alignment16, Bytes16Alignment4, Bytes16Alignment8, - Bytes1Alignment1, Bytes2Alignment2, Bytes32Alignment16, Bytes4Alignment4, Bytes8Alignment4, - Bytes8Alignment8, PrimitiveSizeAlignmentPair, -}; enum BackingStorage { Vec { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 5e36e7ec3aa8..414b724a0ac8 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -80,33 +80,14 @@ impl<'a> utils::StateTranslation<'a, BinaryDecoder> for StateTranslation<'a> { fn extend_from_state( &mut self, - decoder: &mut BinaryDecoder, - decoded: &mut ::DecodedState, - is_optional: bool, - page_validity: &mut Option, - dict: Option<&'a ::Dict>, - additional: usize, + _decoder: &mut BinaryDecoder, + _decoded: &mut ::DecodedState, + _is_optional: bool, + _page_validity: &mut Option, + _dict: Option<&'a ::Dict>, + _additional: usize, ) -> ParquetResult<()> { - use StateTranslation as T; - match self { - T::Plain(page_values, _) => decoder.decode_plain_encoded( - decoded, - page_values, - is_optional, - page_validity.as_mut(), - additional, - )?, - T::Dictionary(page_values) => decoder.decode_dictionary_encoded( - decoded, - page_values, - is_optional, - page_validity.as_mut(), - dict.unwrap(), - additional, - )?, - } - - Ok(()) + unreachable!() } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index e19aac3b500c..85783a4ea3f1 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -90,21 +90,6 @@ where additional: usize, ) -> ParquetResult<()> { match self { - Self::Plain(page_values) => decoder.decode_plain_encoded( - decoded, - page_values, - is_optional, - page_validity.as_mut(), - additional, - )?, - Self::Dictionary(ref mut page) => decoder.decode_dictionary_encoded( - decoded, - page, - is_optional, - page_validity.as_mut(), - dict.unwrap(), - additional, - )?, Self::ByteStreamSplit(page_values) => { let (values, validity) = decoded; @@ -129,6 +114,7 @@ where )?, } }, + _ => unreachable!(), } Ok(()) From aeb56a23fffd200957c34ef1e49b4ab5b33d04d1 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 15:57:37 +0200 Subject: [PATCH 18/32] some work --- .../src/arrow/read/deserialize/binview.rs | 22 ++++++------------- .../src/arrow/read/deserialize/null.rs | 4 ++-- .../arrow/read/deserialize/primitive/float.rs | 2 +- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index a1ee405b1cdf..f88b6cb49287 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -113,21 +113,7 @@ impl<'a> utils::StateTranslation<'a, BinViewDecoder> for StateTranslation<'a> { // Already done in decode_plain_encoded validate_utf8 = false; }, - Self::Dictionary(ref mut page) => { - let dict = dict.unwrap(); - - decoder.decode_dictionary_encoded( - decoded, - page, - is_optional, - page_validity.as_mut(), - dict, - additional, - )?; - - // Already done in decode_plain_encoded - validate_utf8 = false; - }, + Self::Dictionary(_) => unreachable!(), Self::DeltaLengthByteArray(ref mut page_values, ref mut lengths) => { let (values, validity) = decoded; @@ -519,6 +505,12 @@ impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for DeltaBytes } } +// fn decode_plain( +// values: BinaryIter<'_>, +// ) -> ParquetResult<()> { +// +// } +// impl utils::Decoder for BinViewDecoder { type Translation<'a> = StateTranslation<'a>; type Dict = (Vec, Vec>); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index a0e1f0963e0f..6a74f8b70ffc 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -134,12 +134,12 @@ pub fn iter_to_arrays( continue; } - let num_rows = match state_filter { + let num_page_rows = match state_filter { None => page.num_values(), Some(filter) => filter.num_rows(), }; - len = (len + num_rows).min(num_rows); + len = (len + num_page_rows).min(num_rows); } Ok(Box::new(NullArray::new(dtype, len))) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index 85783a4ea3f1..da0491046b01 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -86,7 +86,7 @@ where decoded: &mut as utils::Decoder>::DecodedState, is_optional: bool, page_validity: &mut Option, - dict: Option<&'a as utils::Decoder>::Dict>, + _dict: Option<&'a as utils::Decoder>::Dict>, additional: usize, ) -> ParquetResult<()> { match self { From a30a924ec526ab4e522e77711e91dfa7c0895306 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Fri, 18 Oct 2024 17:41:48 +0200 Subject: [PATCH 19/32] get started on string filters --- .../src/arrow/read/deserialize/binview.rs | 114 +++++++++++++++++- 1 file changed, 108 insertions(+), 6 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index f88b6cb49287..8f252493c330 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -9,7 +9,9 @@ use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::buffer::Buffer; use arrow::datatypes::{ArrowDataType, PhysicalType}; +use super::utils::dict_encoded::{append_validity, constrain_page_validity}; use super::utils::{dict_indices_decoder, freeze_validity, BatchableCollector}; +use super::Filter; use crate::parquet::encoding::delta_bitpacked::{lin_natural_sum, DeltaGatherer}; use crate::parquet::encoding::{delta_byte_array, delta_length_byte_array, hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; @@ -505,12 +507,112 @@ impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for DeltaBytes } } -// fn decode_plain( -// values: BinaryIter<'_>, -// ) -> ParquetResult<()> { -// -// } -// +fn decode_plain( + mut values: BinaryIter<'_>, + is_optional: bool, + page_validity: Option<&Bitmap>, + filter: Option, + validity: &mut MutableBitmap, + target: &mut Vec, + buffers: &mut Vec>, +) -> ParquetResult<()> { + let mut has_valid_last_bytes = true; + + // @TODO: Scratch and reduce capacity + let mut has_size_below_128 = MutableBitmap::with_capacity(values.max_num_values); + let mut transition_offsets = Vec::with_capacity(values.max_num_values); + + if is_optional { + append_validity( + page_validity, + filter.as_ref(), + validity, + values.max_num_values, + ); + } + + let page_validity = + constrain_page_validity(values.max_num_values, page_validity, filter.as_ref()); + + todo!(); + // match (filter, page_validity) { + // (None, None) => decode_required(values, None, target), + // (Some(Filter::Range(rng)), None) if rng.start == 0 => { + // decode_required(values, Some(rng.end), target) + // }, + // (None, Some(page_validity)) => decode_optional(values, &page_validity, target), + // (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { + // decode_optional(values, &page_validity, target) + // }, + // (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target), + // (Some(Filter::Mask(filter)), Some(page_validity)) => { + // decode_masked_optional(values, &page_validity, &filter, target) + // }, + // // @TODO: Use values.skip_in_place(rng.start) + // (Some(Filter::Range(rng)), None) => { + // decode_masked_required(values, &filter_from_range(rng.clone()), target) + // }, + // // @TODO: Use values.skip_in_place(page_validity.sliced(0, rng.start).set_bits()) + // (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( + // values, + // &page_validity, + // &filter_from_range(rng.clone()), + // target, + // ), + // }?; + + fn decode_required( + mut values: BinaryIter<'_>, + limit: Option, + target: &mut Vec, + buffers: &mut Vec>, + ) -> ParquetResult<()> { + let limit = limit.unwrap_or(values.max_num_values); + + assert!(limit <= values.max_num_values); + + if limit == 0 { + return Ok(()); + } + + let proportional_memory_consumption = (values.values.len() as f64) / (values.max_num_values as f64) * (limit as f64); + let margin = (values.values.len() as f64) * 0.05; + + let estimated_memory_consumption = ((proportional_memory_consumption + margin).ceil() as usize).min(u32::MIN as usize).min(values.values.len()); + + let mut current_buffer = match buffers.pop() { + None => Vec::with_capacity(estimated_memory_consumption), + Some(b) if b.capacity() - b.len() < estimated_memory_consumption => { + buffers.push(b); + Vec::with_capacity(estimated_memory_consumption) + } + Some(b) => b, + }; + + target.reserve(limit); + + let buffer_idx = buffers.len() as u32; + let mut offset = 0; + for value in values.take(limit) { + let idx = target.len(); + let view = unsafe { target.get_unchecked_mut(idx) }; + *view = View::new_from_bytes(value, buffer_idx, offset as u32); + unsafe { target.set_len(target.len() + 1) }; + + if value.len() > View::MAX_INLINE_SIZE as usize { + current_buffer.extend_from_slice(value); + offset += value.len(); + } + } + + buffer.push(current_buffer); + + Ok(()) + } + + Ok(()) +} + impl utils::Decoder for BinViewDecoder { type Translation<'a> = StateTranslation<'a>; type Dict = (Vec, Vec>); From 8420ac553ed5bd57307ad47511e0aca57377c33e Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Tue, 22 Oct 2024 16:28:49 +0200 Subject: [PATCH 20/32] add binview and booleans --- .../polars-arrow/src/array/binview/mutable.rs | 6 +- crates/polars-arrow/src/array/binview/view.rs | 13 + crates/polars-arrow/src/bitmap/bitmask.rs | 2 +- .../src/arrow/read/deserialize/binview.rs | 554 +++++++++++------- .../src/arrow/read/deserialize/boolean.rs | 390 +++++++++--- .../src/arrow/read/deserialize/dictionary.rs | 12 - .../read/deserialize/fixed_size_binary.rs | 42 +- .../src/arrow/read/deserialize/null.rs | 13 - .../arrow/read/deserialize/primitive/float.rs | 16 +- .../read/deserialize/primitive/integer.rs | 14 +- .../arrow/read/deserialize/primitive/plain.rs | 1 + .../src/arrow/read/deserialize/utils/mod.rs | 31 +- .../src/parquet/encoding/hybrid_rle/mod.rs | 4 + py-polars/tests/unit/io/test_parquet.py | 55 ++ 14 files changed, 734 insertions(+), 419 deletions(-) diff --git a/crates/polars-arrow/src/array/binview/mutable.rs b/crates/polars-arrow/src/array/binview/mutable.rs index 0d7dcac94b5a..e2516c18a4aa 100644 --- a/crates/polars-arrow/src/array/binview/mutable.rs +++ b/crates/polars-arrow/src/array/binview/mutable.rs @@ -509,7 +509,7 @@ impl MutableBinaryViewArray { Self::from_iterator(slice.as_ref().iter().map(|opt_v| opt_v.as_ref())) } - fn finish_in_progress(&mut self) -> bool { + pub fn finish_in_progress(&mut self) -> bool { if !self.in_progress_buffer.is_empty() { self.completed_buffers .push(std::mem::take(&mut self.in_progress_buffer).into()); @@ -531,6 +531,10 @@ impl MutableBinaryViewArray { arr } + pub fn take(self) -> (Vec, Vec>) { + (self.views, self.completed_buffers) + } + #[inline] pub fn value(&self, i: usize) -> &T { assert!(i < self.len()); diff --git a/crates/polars-arrow/src/array/binview/view.rs b/crates/polars-arrow/src/array/binview/view.rs index 15a744f804e9..f9737ac5eb56 100644 --- a/crates/polars-arrow/src/array/binview/view.rs +++ b/crates/polars-arrow/src/array/binview/view.rs @@ -154,6 +154,19 @@ impl View { } } + /// Construct a byte slice from an inline view. + /// + /// # Safety + /// + /// Assumes that this view is inlinable. + pub unsafe fn get_inlined_slice_unchecked(&self) -> &[u8] { + debug_assert!(self.length <= View::MAX_INLINE_SIZE); + + let ptr = self as *const View as *const u8; + // SAFETY: Invariant of function + unsafe { std::slice::from_raw_parts(ptr.add(4), self.length as usize) } + } + /// Extend a `Vec` with inline views slices of `src` with `width`. /// /// This tries to use SIMD to optimize the copying and can be massively faster than doing a diff --git a/crates/polars-arrow/src/bitmap/bitmask.rs b/crates/polars-arrow/src/bitmap/bitmask.rs index 1421984be207..2e4e45195266 100644 --- a/crates/polars-arrow/src/bitmap/bitmask.rs +++ b/crates/polars-arrow/src/bitmap/bitmask.rs @@ -282,7 +282,7 @@ impl<'a> BitMask<'a> { } pub fn iter(&self) -> BitmapIter { - BitmapIter::new(&self.bytes, self.offset, self.len) + BitmapIter::new(self.bytes, self.offset, self.len) } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index 8f252493c330..b0ec85a0a718 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -10,7 +10,7 @@ use arrow::buffer::Buffer; use arrow::datatypes::{ArrowDataType, PhysicalType}; use super::utils::dict_encoded::{append_validity, constrain_page_validity}; -use super::utils::{dict_indices_decoder, freeze_validity, BatchableCollector}; +use super::utils::{dict_indices_decoder, filter_from_range, freeze_validity, BatchableCollector}; use super::Filter; use crate::parquet::encoding::delta_bitpacked::{lin_natural_sum, DeltaGatherer}; use crate::parquet::encoding::{delta_byte_array, delta_length_byte_array, hybrid_rle, Encoding}; @@ -94,7 +94,7 @@ impl<'a> utils::StateTranslation<'a, BinViewDecoder> for StateTranslation<'a> { decoded: &mut ::DecodedState, is_optional: bool, page_validity: &mut Option, - dict: Option<&'a ::Dict>, + _dict: Option<&'a ::Dict>, additional: usize, ) -> ParquetResult<()> { let views_offset = decoded.0.views().len(); @@ -507,109 +507,336 @@ impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for DeltaBytes } } -fn decode_plain( - mut values: BinaryIter<'_>, +#[allow(clippy::too_many_arguments)] +pub fn decode_plain( + values: &[u8], + max_num_values: usize, + target: &mut MutableBinaryViewArray<[u8]>, + is_optional: bool, + validity: &mut MutableBitmap, + page_validity: Option<&Bitmap>, filter: Option, - validity: &mut MutableBitmap, - target: &mut Vec, - buffers: &mut Vec>, + + verify_utf8: bool, ) -> ParquetResult<()> { - let mut has_valid_last_bytes = true; + if is_optional { + append_validity(page_validity, filter.as_ref(), validity, max_num_values); + } + let page_validity = constrain_page_validity(max_num_values, page_validity, filter.as_ref()); + + match (filter, page_validity) { + (None, None) => decode_required_plain(max_num_values, values, None, target, verify_utf8), + (Some(Filter::Range(rng)), None) if rng.start == 0 => { + decode_required_plain(max_num_values, values, Some(rng.end), target, verify_utf8) + }, + (None, Some(page_validity)) => decode_optional_plain( + page_validity.set_bits(), + values, + target, + &page_validity, + verify_utf8, + ), + (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => decode_optional_plain( + page_validity.set_bits(), + values, + target, + &page_validity, + verify_utf8, + ), + (Some(Filter::Mask(mask)), None) => { + decode_masked_required_plain(max_num_values, values, target, &mask, verify_utf8) + }, + (Some(Filter::Mask(mask)), Some(page_validity)) => decode_masked_optional_plain( + page_validity.set_bits(), + values, + target, + &page_validity, + &mask, + verify_utf8, + ), + (Some(Filter::Range(rng)), None) => decode_masked_required_plain( + max_num_values, + values, + target, + &filter_from_range(rng.clone()), + verify_utf8, + ), + (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional_plain( + page_validity.set_bits(), + values, + target, + &page_validity, + &filter_from_range(rng.clone()), + verify_utf8, + ), + }?; - // @TODO: Scratch and reduce capacity - let mut has_size_below_128 = MutableBitmap::with_capacity(values.max_num_values); - let mut transition_offsets = Vec::with_capacity(values.max_num_values); + Ok(()) +} - if is_optional { - append_validity( +#[cold] +fn invalid_input_err() -> ParquetError { + ParquetError::oos("String data does not match given length") +} + +#[cold] +fn invalid_utf8_err() -> ParquetError { + ParquetError::oos("String data contained invalid UTF-8") +} + +fn decode_required_plain( + num_expected_values: usize, + values: &[u8], + limit: Option, + target: &mut MutableBinaryViewArray<[u8]>, + + verify_utf8: bool, +) -> ParquetResult<()> { + let limit = limit.unwrap_or(num_expected_values); + + let mut idx = 0; + decode_plain_generic( + values, + target, + limit, + || { + if idx >= limit { + return None; + } + + idx += 1; + + Some((true, true)) + }, + verify_utf8, + ) +} + +fn decode_optional_plain( + num_expected_values: usize, + values: &[u8], + target: &mut MutableBinaryViewArray<[u8]>, + page_validity: &Bitmap, + + verify_utf8: bool, +) -> ParquetResult<()> { + if page_validity.unset_bits() == 0 { + return decode_required_plain( + num_expected_values, + values, + Some(page_validity.len()), + target, + verify_utf8, + ); + } + + let mut validity_iter = page_validity.iter(); + decode_plain_generic( + values, + target, + page_validity.len(), + || Some((validity_iter.next()?, true)), + verify_utf8, + ) +} + +fn decode_masked_required_plain( + num_expected_values: usize, + values: &[u8], + target: &mut MutableBinaryViewArray<[u8]>, + + mask: &Bitmap, + + verify_utf8: bool, +) -> ParquetResult<()> { + if mask.unset_bits() == 0 { + return decode_required_plain( + num_expected_values, + values, + Some(mask.len()), + target, + verify_utf8, + ); + } + + let mut mask_iter = mask.iter(); + decode_plain_generic( + values, + target, + mask.set_bits(), + || Some((true, mask_iter.next()?)), + verify_utf8, + ) +} + +fn decode_masked_optional_plain( + num_expected_values: usize, + values: &[u8], + target: &mut MutableBinaryViewArray<[u8]>, + + page_validity: &Bitmap, + mask: &Bitmap, + + verify_utf8: bool, +) -> ParquetResult<()> { + assert_eq!(page_validity.len(), mask.len()); + + if mask.unset_bits() == 0 { + return decode_optional_plain( + num_expected_values, + values, + target, page_validity, - filter.as_ref(), - validity, - values.max_num_values, + verify_utf8, ); } - let page_validity = - constrain_page_validity(values.max_num_values, page_validity, filter.as_ref()); - - todo!(); - // match (filter, page_validity) { - // (None, None) => decode_required(values, None, target), - // (Some(Filter::Range(rng)), None) if rng.start == 0 => { - // decode_required(values, Some(rng.end), target) - // }, - // (None, Some(page_validity)) => decode_optional(values, &page_validity, target), - // (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - // decode_optional(values, &page_validity, target) - // }, - // (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target), - // (Some(Filter::Mask(filter)), Some(page_validity)) => { - // decode_masked_optional(values, &page_validity, &filter, target) - // }, - // // @TODO: Use values.skip_in_place(rng.start) - // (Some(Filter::Range(rng)), None) => { - // decode_masked_required(values, &filter_from_range(rng.clone()), target) - // }, - // // @TODO: Use values.skip_in_place(page_validity.sliced(0, rng.start).set_bits()) - // (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( - // values, - // &page_validity, - // &filter_from_range(rng.clone()), - // target, - // ), - // }?; - - fn decode_required( - mut values: BinaryIter<'_>, - limit: Option, - target: &mut Vec, - buffers: &mut Vec>, - ) -> ParquetResult<()> { - let limit = limit.unwrap_or(values.max_num_values); - - assert!(limit <= values.max_num_values); + if page_validity.unset_bits() == 0 { + return decode_masked_required_plain( + num_expected_values, + values, + target, + page_validity, + verify_utf8, + ); + } - if limit == 0 { - return Ok(()); + let mut validity_iter = page_validity.iter(); + let mut mask_iter = mask.iter(); + decode_plain_generic( + values, + target, + mask.set_bits(), + || Some((validity_iter.next()?, mask_iter.next()?)), + verify_utf8, + ) +} + +pub fn decode_plain_generic( + values: &[u8], + target: &mut MutableBinaryViewArray<[u8]>, + + num_rows: usize, + mut next: impl FnMut() -> Option<(bool, bool)>, + + verify_utf8: bool, +) -> ParquetResult<()> { + target.finish_in_progress(); + unsafe { target.views_mut() }.reserve(num_rows); + + let buffer_idx = target.completed_buffers().len() as u32; + let mut buffer = Vec::with_capacity(values.len() + 1); + let mut none_starting_with_continuation_byte = true; // Whether the transition from between strings is valid + // UTF-8 + let mut all_len_below_128 = true; // Whether all the lengths of the values are below 128, this + // allows us to make UTF-8 verification a lot faster. + + let mut total_bytes_len = 0; + let mut num_seen = 0; + let mut num_inlined = 0; + + let mut mvalues = values; + while let Some((is_valid, is_selected)) = next() { + if !is_valid { + if is_selected { + unsafe { target.views_mut() }.push(unsafe { View::new_inline_unchecked(&[]) }); + } + continue; + } + + if mvalues.len() < 4 { + return Err(invalid_input_err()); } - let proportional_memory_consumption = (values.values.len() as f64) / (values.max_num_values as f64) * (limit as f64); - let margin = (values.values.len() as f64) * 0.05; + let length; + (length, mvalues) = mvalues.split_at(4); + let length: &[u8; 4] = unsafe { length.try_into().unwrap_unchecked() }; + let length = u32::from_le_bytes(*length); - let estimated_memory_consumption = ((proportional_memory_consumption + margin).ceil() as usize).min(u32::MIN as usize).min(values.values.len()); + if mvalues.len() < length as usize { + return Err(invalid_input_err()); + } - let mut current_buffer = match buffers.pop() { - None => Vec::with_capacity(estimated_memory_consumption), - Some(b) if b.capacity() - b.len() < estimated_memory_consumption => { - buffers.push(b); - Vec::with_capacity(estimated_memory_consumption) + let value; + (value, mvalues) = mvalues.split_at(length as usize); + + num_seen += 1; + all_len_below_128 &= value.len() < 128; + // Everything starting with 10.. .... is a continuation byte. + none_starting_with_continuation_byte &= + value.is_empty() || value[0] & 0b1100_0000 != 0b1000_0000; + + if !is_selected { + continue; + } + + let offset = buffer.len() as u32; + + if value.len() <= View::MAX_INLINE_SIZE as usize { + unsafe { target.views_mut() }.push(unsafe { View::new_inline_unchecked(value) }); + num_inlined += 1; + } else { + buffer.extend_from_slice(value); + unsafe { target.views_mut() } + .push(unsafe { View::new_noninline_unchecked(value, buffer_idx, offset) }); + } + + total_bytes_len += value.len(); + } + + unsafe { + target.set_total_bytes_len(target.total_bytes_len() + total_bytes_len); + } + + if verify_utf8 { + if num_inlined == 0 { + if !none_starting_with_continuation_byte || simdutf8::basic::from_utf8(&buffer).is_err() + { + return Err(invalid_utf8_err()); + } + } else if all_len_below_128 { + if simdutf8::basic::from_utf8(values).is_err() { + return Err(invalid_utf8_err()); + } + } else { + if !none_starting_with_continuation_byte || simdutf8::basic::from_utf8(&buffer).is_err() + { + return Err(invalid_utf8_err()); } - Some(b) => b, - }; - target.reserve(limit); + let mut all_inlined_are_ascii = true; - let buffer_idx = buffers.len() as u32; - let mut offset = 0; - for value in values.take(limit) { - let idx = target.len(); - let view = unsafe { target.get_unchecked_mut(idx) }; - *view = View::new_from_bytes(value, buffer_idx, offset as u32); - unsafe { target.set_len(target.len() + 1) }; + const MASK: [u128; 2] = [ + 0x0000_0000_8080_8080_8080_8080_8080_8080, + 0x0000_0000_0000_0000_0000_0000_0000_0000, + ]; - if value.len() > View::MAX_INLINE_SIZE as usize { - current_buffer.extend_from_slice(value); - offset += value.len(); + for view in &target.views()[target.len() - num_seen..] { + all_inlined_are_ascii &= + view.as_u128() & MASK[usize::from(view.length > View::MAX_INLINE_SIZE)] == 0; } - } - buffer.push(current_buffer); + if !all_inlined_are_ascii { + let mut is_valid = true; + for view in &target.views()[target.len() - num_seen..] { + if view.length <= View::MAX_INLINE_SIZE { + is_valid &= + std::str::from_utf8(unsafe { view.get_inlined_slice_unchecked() }) + .is_ok(); + } + } - Ok(()) + if !is_valid { + return Err(invalid_utf8_err()); + } + } + } } - + + target.push_buffer(buffer.into()); + Ok(()) } @@ -646,159 +873,26 @@ impl utils::Decoder for BinViewDecoder { let values = &page.buffer; let num_values = page.num_values; - // Each value is prepended by the length which is 4 bytes. - let num_bytes = values.len() - 4 * num_values; - - let mut views = Vec::with_capacity(num_values); - let mut buffer = Vec::with_capacity(num_bytes); - - let mut buffers = Vec::with_capacity(1); - - let mut offset = 0; - let mut max_length = 0; - views.extend(BinaryIter::new(values, num_values).map(|v| { - let length = v.len(); - max_length = usize::max(length, max_length); - if length <= View::MAX_INLINE_SIZE as usize { - View::new_inline(v) - } else { - if offset >= u32::MAX as usize { - let full_buffer = std::mem::take(&mut buffer); - let num_bytes = full_buffer.capacity() - full_buffer.len(); - buffers.push(Buffer::from(full_buffer)); - buffer.reserve(num_bytes); - offset = 0; - } - - buffer.extend_from_slice(v); - let view = View::new_from_bytes(v, buffers.len() as u32, offset as u32); - offset += v.len(); - view - } - })); + let mut arr = MutableBinaryViewArray::new(); + decode_required_plain( + num_values, + values, + None, + &mut arr, + self.check_utf8.load(Ordering::Relaxed), + )?; - buffers.push(Buffer::from(buffer)); - - if self.check_utf8.load(Ordering::Relaxed) { - // This is a small trick that allows us to check the Parquet buffer instead of the view - // buffer. Batching the UTF-8 verification is more performant. For this to be allowed, - // all the interleaved lengths need to be valid UTF-8. - // - // Every strings prepended by 4 bytes (L, 0, 0, 0), since we check here L < 128. L is - // only a valid first byte of a UTF-8 code-point and (L, 0, 0, 0) is valid UTF-8. - // Consequently, it is valid to just check the whole buffer. - if max_length < 128 { - simdutf8::basic::from_utf8(values) - .map_err(|_| ParquetError::oos("String data contained invalid UTF-8"))?; - } else { - arrow::array::validate_utf8_view(&views, &buffers) - .map_err(|_| ParquetError::oos("String data contained invalid UTF-8"))?; - } - } + let (views, buffers) = arr.take(); Ok((views, buffers)) } fn decode_plain_encoded<'a>( - &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - is_optional: bool, - page_validity: Option<&mut Bitmap>, - limit: usize, - ) -> ParquetResult<()> { - let views_offset = values.views().len(); - let buffer_offset = values.completed_buffers().len(); - - struct Collector<'a, 'b> { - iter: &'b mut BinaryIter<'a>, - max_length: &'b mut usize, - } - - impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for Collector<'a, 'b> { - fn reserve(target: &mut MutableBinaryViewArray<[u8]>, n: usize) { - target.reserve(n); - } - - fn push_n( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - for x in self.iter.take(n) { - *self.max_length = usize::max(*self.max_length, x.len()); - target.push_value(x); - } - Ok(()) - } - - fn push_n_nulls( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - target.extend_constant(n, >::None); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n > 0 { - _ = self.iter.nth(n - 1); - } - Ok(()) - } - } - - let mut max_length = 0; - let buffer = page_values.values; - let mut collector = Collector { - iter: page_values, - max_length: &mut max_length, - }; - - match page_validity { - None => { - collector.push_n(values, limit)?; - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - extend_from_decoder(validity, page_validity, Some(limit), values, collector)? - }, - } - - let buffer = &buffer[..buffer.len() - page_values.values.len()]; - - if self.check_utf8.load(Ordering::Relaxed) { - // This is a small trick that allows us to check the Parquet buffer instead of the view - // buffer. Batching the UTF-8 verification is more performant. For this to be allowed, - // all the interleaved lengths need to be valid UTF-8. - // - // Every strings prepended by 4 bytes (L, 0, 0, 0), since we check here L < 128. L is - // only a valid first byte of a UTF-8 code-point and (L, 0, 0, 0) is valid UTF-8. - // Consequently, it is valid to just check the whole buffer. - if max_length < 128 { - simdutf8::basic::from_utf8(buffer) - .map_err(|_| ParquetError::oos("String data contained invalid UTF-8"))?; - } else { - values - .validate_utf8(buffer_offset, views_offset) - .map_err(|_| ParquetError::oos("String data contained invalid UTF-8"))? - } - } - - Ok(()) - } - - fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, + _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, _is_optional: bool, _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { unreachable!() @@ -811,6 +905,16 @@ impl utils::Decoder for BinViewDecoder { filter: Option, ) -> ParquetResult<()> { match state.translation { + StateTranslation::Plain(iter) => decode_plain( + iter.values, + iter.max_num_values, + &mut decoded.0, + state.is_optional, + &mut decoded.1, + state.page_validity.as_ref(), + filter, + self.check_utf8.load(Ordering::Relaxed), + ), StateTranslation::Dictionary(ref mut indexes) => { let (dict, _) = state.dict.unwrap(); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs index c2deee921334..0830404a96ad 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs @@ -1,15 +1,20 @@ use arrow::array::BooleanArray; +use arrow::bitmap::bitmask::BitMask; use arrow::bitmap::utils::BitmapIter; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; - -use super::utils::{self, extend_from_decoder, freeze_validity, Decoder, ExactSize}; -use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; -use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; +use polars_compute::filter::filter_boolean_kernel; + +use super::utils::dict_encoded::{append_validity, constrain_page_validity}; +use super::utils::{ + self, decode_hybrid_rle_into_bitmap, filter_from_range, freeze_validity, + Decoder, ExactSize, +}; +use super::Filter; +use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::encoding::Encoding; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; -use crate::read::deserialize::utils::BatchableCollector; #[allow(clippy::large_enum_variant)] #[derive(Debug)] @@ -87,97 +92,207 @@ impl<'a> utils::StateTranslation<'a, BooleanDecoder> for StateTranslation<'a> { fn extend_from_state( &mut self, - decoder: &mut BooleanDecoder, - decoded: &mut ::DecodedState, - is_optional: bool, - page_validity: &mut Option, + _decoder: &mut BooleanDecoder, + _decoded: &mut ::DecodedState, + _is_optional: bool, + _page_validity: &mut Option, _: Option<&'a ::Dict>, - additional: usize, + _additional: usize, ) -> ParquetResult<()> { - match self { - Self::Plain(page_values) => decoder.decode_plain_encoded( - decoded, - page_values, - is_optional, - page_validity.as_mut(), - additional, - )?, - Self::Rle(page_values) => { - let (values, validity) = decoded; - match page_validity { - None => { - page_values.gather_n_into(values, additional, &BitmapGatherer)?; - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => utils::extend_from_decoder( - validity, - page_validity, - Some(additional), - values, - BitmapCollector(page_values), - )?, + unreachable!() + } +} + +fn decode_required_rle( + values: HybridRleDecoder<'_>, + limit: Option, + target: &mut MutableBitmap, +) -> ParquetResult<()> { + decode_hybrid_rle_into_bitmap(values, limit, target)?; + Ok(()) +} + +fn decode_optional_rle( + values: HybridRleDecoder<'_>, + target: &mut MutableBitmap, + page_validity: &Bitmap, +) -> ParquetResult<()> { + debug_assert!(page_validity.set_bits() <= values.len()); + + if page_validity.unset_bits() == 0 { + return decode_required_rle(values, Some(page_validity.len()), target); + } + + target.reserve(page_validity.len()); + + let mut validity_mask = BitMask::from_bitmap(page_validity); + + for chunk in values.into_chunk_iter() { + let chunk = chunk?; + + match chunk { + HybridRleChunk::Rle(value, size) => { + let offset = validity_mask + .nth_set_bit_idx(size, 0) + .unwrap_or(validity_mask.len()); + + let t; + (t, validity_mask) = validity_mask.split_at(offset); + + target.extend_constant(t.len(), value != 0); + }, + HybridRleChunk::Bitpacked(decoder) => { + let decoder_slice = decoder.as_slice(); + let offset = validity_mask + .nth_set_bit_idx(decoder.len(), 0) + .unwrap_or(validity_mask.len()); + + let decoder_validity; + (decoder_validity, validity_mask) = validity_mask.split_at(offset); + + let mut offset = 0; + let mut validity_iter = decoder_validity.iter(); + while validity_iter.num_remaining() > 0 { + let num_valid = validity_iter.take_leading_ones(); + target.extend_from_slice(decoder_slice, offset, num_valid); + offset += num_valid; + + let num_invalid = validity_iter.take_leading_zeros(); + target.extend_constant(num_invalid, false); } }, } + } - Ok(()) + if cfg!(debug_assertions) { + assert_eq!(validity_mask.set_bits(), 0); } + target.extend_constant(validity_mask.len(), false); + + Ok(()) } -struct BitmapGatherer; -impl HybridRleGatherer for BitmapGatherer { - type Target = MutableBitmap; +fn decode_masked_required_rle( + values: HybridRleDecoder<'_>, + target: &mut MutableBitmap, + mask: &Bitmap, +) -> ParquetResult<()> { + debug_assert!(mask.len() <= values.len()); - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); + if mask.unset_bits() == 0 { + return decode_required_rle(values, Some(mask.len()), target); } - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.len() + let mut im_target = MutableBitmap::new(); + decode_required_rle(values, Some(mask.len()), &mut im_target)?; + + target.extend_from_bitmap(&filter_boolean_kernel(&im_target.freeze(), mask)); + + Ok(()) +} + +fn decode_masked_optional_rle( + values: HybridRleDecoder<'_>, + target: &mut MutableBitmap, + page_validity: &Bitmap, + mask: &Bitmap, +) -> ParquetResult<()> { + debug_assert_eq!(page_validity.len(), mask.len()); + debug_assert!(mask.len() <= values.len()); + + if mask.unset_bits() == 0 { + return decode_optional_rle(values, target, page_validity); } - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(value) + if page_validity.unset_bits() == 0 { + return decode_masked_required_rle(values, target, mask); } - fn gather_one(&self, target: &mut Self::Target, value: u32) -> ParquetResult<()> { - target.push(value != 0); - Ok(()) + let mut im_target = MutableBitmap::new(); + decode_optional_rle(values, &mut im_target, page_validity)?; + + target.extend_from_bitmap(&filter_boolean_kernel(&im_target.freeze(), mask)); + + Ok(()) +} + +fn decode_required_plain( + mut values: BitmapIter<'_>, + limit: Option, + target: &mut MutableBitmap, +) -> ParquetResult<()> { + let limit = limit.unwrap_or(values.len()); + values.collect_n_into(target, limit); + Ok(()) +} + +fn decode_optional_plain( + mut values: BitmapIter<'_>, + target: &mut MutableBitmap, + page_validity: &Bitmap, +) -> ParquetResult<()> { + debug_assert!(page_validity.set_bits() <= values.len()); + + if page_validity.unset_bits() == 0 { + return decode_required_plain(values, Some(page_validity.len()), target); } - fn gather_repeated( - &self, - target: &mut Self::Target, - value: u32, - n: usize, - ) -> ParquetResult<()> { - target.extend_constant(n, value != 0); - Ok(()) + target.reserve(page_validity.len()); + + let mut validity_iter = page_validity.iter(); + while validity_iter.num_remaining() > 0 { + let num_valid = validity_iter.take_leading_ones(); + values.collect_n_into(target, num_valid); + + let num_invalid = validity_iter.take_leading_zeros(); + target.extend_constant(num_invalid, false); } - // @TODO: The slice impl here can speed some stuff up + Ok(()) } -struct BitmapCollector<'a, 'b>(&'b mut HybridRleDecoder<'a>); -impl<'a, 'b> BatchableCollector for BitmapCollector<'a, 'b> { - fn reserve(target: &mut MutableBitmap, n: usize) { - target.reserve(n); - } - fn push_n(&mut self, target: &mut MutableBitmap, n: usize) -> ParquetResult<()> { - self.0.gather_n_into(target, n, &BitmapGatherer) +fn decode_masked_required_plain( + values: BitmapIter<'_>, + target: &mut MutableBitmap, + mask: &Bitmap, +) -> ParquetResult<()> { + debug_assert!(mask.len() <= values.len()); + + if mask.unset_bits() == 0 { + return decode_required_plain(values, Some(mask.len()), target); } - fn push_n_nulls(&mut self, target: &mut MutableBitmap, n: usize) -> ParquetResult<()> { - target.extend_constant(n, false); - Ok(()) + let mut im_target = MutableBitmap::new(); + decode_required_plain(values, Some(mask.len()), &mut im_target)?; + + target.extend_from_bitmap(&filter_boolean_kernel(&im_target.freeze(), mask)); + + Ok(()) +} + +fn decode_masked_optional_plain( + values: BitmapIter<'_>, + target: &mut MutableBitmap, + page_validity: &Bitmap, + mask: &Bitmap, +) -> ParquetResult<()> { + debug_assert_eq!(page_validity.len(), mask.len()); + debug_assert!(page_validity.set_bits() <= values.len()); + + if mask.unset_bits() == 0 { + return decode_optional_plain(values, target, page_validity); } - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.0.skip_in_place(n) + if page_validity.unset_bits() == 0 { + return decode_masked_required_plain(values, target, mask); } + + let mut im_target = MutableBitmap::new(); + decode_optional_plain(values, &mut im_target, page_validity)?; + + target.extend_from_bitmap(&filter_boolean_kernel(&im_target.freeze(), mask)); + + Ok(()) } impl ExactSize for (MutableBitmap, MutableBitmap) { @@ -212,39 +327,14 @@ impl Decoder for BooleanDecoder { } fn decode_plain_encoded<'a>( - &mut self, - (values, validity): &mut Self::DecodedState, - page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - is_optional: bool, - page_validity: Option<&mut Bitmap>, - limit: usize, - ) -> ParquetResult<()> { - match page_validity { - None => { - page_values.collect_n_into(values, limit); - - if is_optional { - validity.extend_constant(limit, true); - } - }, - Some(page_validity) => { - extend_from_decoder(validity, page_validity, Some(limit), values, page_values)? - }, - } - - Ok(()) - } - - fn decode_dictionary_encoded( &mut self, _decoded: &mut Self::DecodedState, - _page_values: &mut HybridRleDecoder<'_>, + _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, _is_optional: bool, _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, _limit: usize, ) -> ParquetResult<()> { - unimplemented!() + unreachable!() } fn finalize( @@ -254,6 +344,112 @@ impl Decoder for BooleanDecoder { (values, validity): Self::DecodedState, ) -> ParquetResult { let validity = freeze_validity(validity); - Ok(BooleanArray::new(dtype, values.into(), validity)) + Ok(BooleanArray::new(dtype, values.freeze(), validity)) + } + + fn extend_filtered_with_state( + &mut self, + state: utils::State<'_, Self>, + (target, validity): &mut Self::DecodedState, + filter: Option, + ) -> ParquetResult<()> { + match state.translation { + StateTranslation::Plain(values) => { + if state.is_optional { + append_validity( + state.page_validity.as_ref(), + filter.as_ref(), + validity, + values.len(), + ); + } + + let page_validity = constrain_page_validity( + values.len(), + state.page_validity.as_ref(), + filter.as_ref(), + ); + + match (filter, page_validity) { + (None, None) => decode_required_plain(values, None, target), + (Some(Filter::Range(rng)), None) if rng.start == 0 => { + decode_required_plain(values, Some(rng.end), target) + }, + (None, Some(page_validity)) => { + decode_optional_plain(values, target, &page_validity) + }, + (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { + decode_optional_plain(values, target, &page_validity) + }, + (Some(Filter::Mask(mask)), None) => { + decode_masked_required_plain(values, target, &mask) + }, + (Some(Filter::Mask(mask)), Some(page_validity)) => { + decode_masked_optional_plain(values, target, &page_validity, &mask) + }, + (Some(Filter::Range(rng)), None) => decode_masked_required_plain( + values, + target, + &filter_from_range(rng.clone()), + ), + (Some(Filter::Range(rng)), Some(page_validity)) => { + decode_masked_optional_plain( + values, + target, + &page_validity, + &filter_from_range(rng.clone()), + ) + }, + }?; + + Ok(()) + }, + StateTranslation::Rle(values) => { + if state.is_optional { + append_validity( + state.page_validity.as_ref(), + filter.as_ref(), + validity, + values.len(), + ); + } + + let page_validity = constrain_page_validity( + values.len(), + state.page_validity.as_ref(), + filter.as_ref(), + ); + + match (filter, page_validity) { + (None, None) => decode_required_rle(values, None, target), + (Some(Filter::Range(rng)), None) if rng.start == 0 => { + decode_required_rle(values, Some(rng.end), target) + }, + (None, Some(page_validity)) => { + decode_optional_rle(values, target, &page_validity) + }, + (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { + decode_optional_rle(values, target, &page_validity) + }, + (Some(Filter::Mask(filter)), None) => { + decode_masked_required_rle(values, target, &filter) + }, + (Some(Filter::Mask(filter)), Some(page_validity)) => { + decode_masked_optional_rle(values, target, &page_validity, &filter) + }, + (Some(Filter::Range(rng)), None) => { + decode_masked_required_rle(values, target, &filter_from_range(rng.clone())) + }, + (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional_rle( + values, + target, + &page_validity, + &filter_from_range(rng.clone()), + ), + }?; + + Ok(()) + }, + } } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs index 6fc00f64993f..7bfafee9a149 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs @@ -142,18 +142,6 @@ impl utils::Decoder for DictionaryDec ) -> ParquetResult<()> { unreachable!() } - - fn decode_dictionary_encoded( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut HybridRleDecoder<'_>, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } } pub(crate) struct DictArrayCollector<'a, 'b> { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 414b724a0ac8..1be147402f9c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -120,7 +120,7 @@ impl FSBVec { } } - pub fn to_bytes_buffer(self) -> Buffer { + pub fn into_bytes_buffer(self) -> Buffer { Buffer::from_storage(match self { FSBVec::Size1(vec) => SharedStorage::bytes_from_aligned_bytes(vec), FSBVec::Size2(vec) => SharedStorage::bytes_from_aligned_bytes(vec), @@ -288,6 +288,7 @@ fn decode_fsb_plain( } } +#[allow(clippy::too_many_arguments)] fn decode_fsb_dict( size: usize, values: HybridRleDecoder<'_>, @@ -354,15 +355,13 @@ fn decode_fsb_dict( (None, None) => target.extend( indexes .into_iter() - .map(|v| &dict[(v as usize) * size..][..size]) - .flatten(), + .flat_map(|v| &dict[(v as usize) * size..][..size]), ), (None, Some(filter)) => match filter { Filter::Range(range) => target.extend( indexes[range.start..range.end] - .into_iter() - .map(|v| &dict[(*v as usize) * size..][..size]) - .flatten(), + .iter() + .flat_map(|v| &dict[(*v as usize) * size..][..size]), ), Filter::Mask(bitmap) => { let mut iter = bitmap.iter(); @@ -372,9 +371,8 @@ fn decode_fsb_dict( let num_selected = iter.take_leading_ones(); target.extend( indexes[offset..][..num_selected] - .into_iter() - .map(|v| &dict[(*v as usize) * size..][..size]) - .flatten(), + .iter() + .flat_map(|v| &dict[(*v as usize) * size..][..size]), ); offset += num_selected; @@ -391,9 +389,8 @@ fn decode_fsb_dict( let num_valid = iter.take_leading_ones(); target.extend( indexes[offset..][..num_valid] - .into_iter() - .map(|v| &dict[(*v as usize) * size..][..size]) - .flatten(), + .iter() + .flat_map(|v| &dict[(*v as usize) * size..][..size]), ); offset += num_valid; @@ -414,9 +411,8 @@ fn decode_fsb_dict( let num_valid = iter.take_leading_ones(); target.extend( indexes[offset..][..num_valid] - .into_iter() - .map(|v| &dict[(*v as usize) * size..][..size]) - .flatten(), + .iter() + .flat_map(|v| &dict[(*v as usize) * size..][..size]), ); offset += num_valid; @@ -497,18 +493,6 @@ impl Decoder for BinaryDecoder { unreachable!() } - fn decode_dictionary_encoded( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn finalize( &self, dtype: ArrowDataType, @@ -519,7 +503,7 @@ impl Decoder for BinaryDecoder { Ok(FixedSizeBinaryArray::new( dtype, - values.to_bytes_buffer(), + values.into_bytes_buffer(), validity, )) } @@ -563,7 +547,7 @@ impl utils::DictDecodable for BinaryDecoder { ) -> ParquetResult> { let dict = FixedSizeBinaryArray::new( ArrowDataType::FixedSizeBinary(self.size), - dict.to_bytes_buffer(), + dict.into_bytes_buffer(), None, ); Ok(DictionaryArray::try_new(dtype, keys, Box::new(dict)).unwrap()) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index 6a74f8b70ffc..01a39baacd30 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -8,7 +8,6 @@ use arrow::datatypes::ArrowDataType; use super::utils; use super::utils::filter::Filter; -use crate::parquet::encoding::hybrid_rle; use crate::parquet::error::ParquetResult; use crate::parquet::page::{DataPage, DictPage}; @@ -84,18 +83,6 @@ impl utils::Decoder for NullDecoder { unimplemented!() } - fn decode_dictionary_encoded( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, - _limit: usize, - ) -> ParquetResult<()> { - unimplemented!() - } - fn finalize( &self, dtype: ArrowDataType, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index da0491046b01..5b5c22bcd23f 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -12,7 +12,7 @@ use crate::parquet::encoding::{byte_stream_split, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; -use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity, Decoder}; +use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity}; use crate::read::Filter; #[allow(clippy::large_enum_variant)] @@ -224,18 +224,6 @@ where unreachable!() } - fn decode_dictionary_encoded( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn extend_filtered_with_state( &mut self, mut state: utils::State<'_, Self>, @@ -244,7 +232,7 @@ where ) -> ParquetResult<()> { match state.translation { StateTranslation::Plain(ref mut values) => super::plain::decode( - *values, + values, state.is_optional, state.page_validity.as_ref(), filter, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 085389c2e2e1..877b03391948 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -278,18 +278,6 @@ where unreachable!() } - fn decode_dictionary_encoded( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut hybrid_rle::HybridRleDecoder<'_>, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _dict: &Self::Dict, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn finalize( &self, dtype: ArrowDataType, @@ -308,7 +296,7 @@ where ) -> ParquetResult<()> { match state.translation { StateTranslation::Plain(ref mut values) => super::plain::decode( - *values, + values, state.is_optional, state.page_validity.as_ref(), filter, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 2861e29e377b..af02a836273f 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -9,6 +9,7 @@ use crate::read::deserialize::utils::dict_encoded::{append_validity, constrain_p use crate::read::deserialize::utils::filter_from_range; use crate::read::{Filter, ParquetError}; +#[allow(clippy::too_many_arguments)] pub fn decode>( values: &[u8], is_optional: bool, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index e4b62b990030..88f94b02ad4d 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -503,15 +503,6 @@ pub(super) trait Decoder: Sized { page_validity: Option<&mut Bitmap>, limit: usize, ) -> ParquetResult<()>; - fn decode_dictionary_encoded( - &mut self, - decoded: &mut Self::DecodedState, - page_values: &mut HybridRleDecoder<'_>, - is_optional: bool, - page_validity: Option<&mut Bitmap>, - dict: &Self::Dict, - limit: usize, - ) -> ParquetResult<()>; fn finalize( &self, @@ -649,12 +640,15 @@ pub(crate) fn filter_from_range(rng: Range) -> Bitmap { bm.freeze() } -pub(crate) fn decode_page_validity( +pub(crate) fn decode_hybrid_rle_into_bitmap( mut page_validity: HybridRleDecoder<'_>, limit: Option, -) -> ParquetResult { + bitmap: &mut MutableBitmap, +) -> ParquetResult<()> { + assert!(page_validity.num_bits() <= 1); + let mut limit = limit.unwrap_or(page_validity.len()); - let mut bm = MutableBitmap::with_capacity(limit); + bitmap.reserve(limit); while let Some(chunk) = page_validity.next_chunk()? { if limit == 0 { @@ -664,16 +658,25 @@ pub(crate) fn decode_page_validity( match chunk { HybridRleChunk::Rle(value, size) => { let size = size.min(limit); - bm.extend_constant(size, value != 0); + bitmap.extend_constant(size, value != 0); limit -= size; }, HybridRleChunk::Bitpacked(decoder) => { let len = decoder.len().min(limit); - bm.extend_from_slice(decoder.as_slice(), 0, len); + bitmap.extend_from_slice(decoder.as_slice(), 0, len); limit -= len; }, } } + Ok(()) +} + +pub(crate) fn decode_page_validity( + page_validity: HybridRleDecoder<'_>, + limit: Option, +) -> ParquetResult { + let mut bm = MutableBitmap::new(); + decode_hybrid_rle_into_bitmap(page_validity, limit, &mut bm)?; Ok(bm.freeze()) } diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs index 7196ab576b52..30fd4798ccb2 100644 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs @@ -85,6 +85,10 @@ impl<'a> HybridRleDecoder<'a> { self.num_values } + pub fn num_bits(&self) -> usize { + self.num_bits + } + pub fn into_chunk_iter(self) -> HybridRleChunkIter<'a> { HybridRleChunkIter { decoder: self } } diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py index e2059bdb4d61..f55dcb8f0ec6 100644 --- a/py-polars/tests/unit/io/test_parquet.py +++ b/py-polars/tests/unit/io/test_parquet.py @@ -2140,3 +2140,58 @@ def test_decode_f16() -> None: pl.scan_parquet(f).slice(1, 3).collect(), df.slice(1, 3), ) + + +def test_invalid_utf8_binary() -> None: + a = pl.Series('a', [b"\x80"], pl.Binary).to_frame() + f = io.BytesIO() + + a.write_parquet(f) + f.seek(0) + out = pl.read_parquet(f) + + assert_frame_equal(a, out) + + +@pytest.mark.parametrize( + "dtype", + [ + pl.Null, + pl.Int8, + pl.Int32, + pl.Datetime(), + pl.String, + pl.Binary, + pl.Boolean, + pl.Struct({ 'x': pl.Int32 }), + pl.List(pl.Int32), + pl.Array(pl.Int32, 0), + pl.Array(pl.Int32, 2), + ] +) +@pytest.mark.parametrize( + "filt", + [ + pl.col.f == 0, + pl.col.f != 0, + pl.col.f == 1, + pl.col.f != 1, + pl.col.f == 2, + pl.col.f != 2, + pl.col.f == 3, + pl.col.f != 3, + ] +) +def test_filter_only_invalid(dtype: pl.DataType, filt: pl.Expr) -> None: + df = pl.DataFrame([ + pl.Series('a', [None, None, None], dtype), + pl.Series('f', range(3), pl.Int32), + ]) + + f = io.BytesIO() + + df.write_parquet(f) + f.seek(0) + out = pl.scan_parquet(f, parallel='prefiltered').filter(filt).collect() + + assert_frame_equal(df.filter(filt), out) From 56448d213b19a559c628c35dce0559d99fea8496 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Wed, 23 Oct 2024 13:05:56 +0200 Subject: [PATCH 21/32] move plain required to memcpy --- .../arrow/read/deserialize/primitive/plain.rs | 58 +++++++++++-------- .../read/deserialize/utils/array_chunks.rs | 25 ++++++++ 2 files changed, 58 insertions(+), 25 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index af02a836273f..4885ca674c55 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -6,7 +6,6 @@ use crate::parquet::error::ParquetResult; use crate::parquet::types::NativeType as ParquetNativeType; use crate::read::deserialize::utils::array_chunks::ArrayChunks; use crate::read::deserialize::utils::dict_encoded::{append_validity, constrain_page_validity}; -use crate::read::deserialize::utils::filter_from_range; use crate::read::{Filter, ParquetError}; #[allow(clippy::too_many_arguments)] @@ -87,29 +86,23 @@ pub fn decode_aligned_bytes_dispatch( let page_validity = constrain_page_validity(values.len(), page_validity, filter.as_ref()); match (filter, page_validity) { - (None, None) => decode_required(values, None, target), - (Some(Filter::Range(rng)), None) if rng.start == 0 => { - decode_required(values, Some(rng.end), target) - }, + (None, None) => decode_required(values, target), (None, Some(page_validity)) => decode_optional(values, &page_validity, target), - (Some(Filter::Range(rng)), Some(page_validity)) if rng.start == 0 => { - decode_optional(values, &page_validity, target) - }, + + (Some(Filter::Range(rng)), None) => decode_required( + unsafe { values.slice_unchecked(rng.start, rng.end) }, + target, + ), + (Some(Filter::Range(rng)), Some(page_validity)) => decode_optional( + unsafe { values.slice_unchecked(rng.start, rng.end) }, + &page_validity.sliced(rng.start, rng.len()), + target, + ), + (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target), (Some(Filter::Mask(filter)), Some(page_validity)) => { decode_masked_optional(values, &page_validity, &filter, target) }, - // @TODO: Use values.skip_in_place(rng.start) - (Some(Filter::Range(rng)), None) => { - decode_masked_required(values, &filter_from_range(rng.clone()), target) - }, - // @TODO: Use values.skip_in_place(page_validity.sliced(0, rng.start).set_bits()) - (Some(Filter::Range(rng)), Some(page_validity)) => decode_masked_optional( - values, - &page_validity, - &filter_from_range(rng.clone()), - target, - ), }?; Ok(()) @@ -118,13 +111,28 @@ pub fn decode_aligned_bytes_dispatch( #[inline(never)] fn decode_required( values: ArrayChunks<'_, B>, - limit: Option, target: &mut Vec, ) -> ParquetResult<()> { - let limit = limit.unwrap_or(values.len()); - assert!(limit <= values.len()); + if values.is_empty() { + return Ok(()); + } - target.extend(values.take(limit).map(|v| B::from_unaligned(*v))); + target.reserve(values.len()); + + // SAFETY: Vec guarantees if the `capacity != 0` the pointer to valid since we just reserve + // that pointer. + let dst = unsafe { target.as_mut_ptr().add(target.len()) }; + let src = values.as_ptr(); + + // SAFETY: + // - `src` is valid for read of values.len() elements. + // - `dst` is valid for writes of values.len() elements, it was just reserved. + // - B::Unaligned is always aligned, since it has an alignment of 1 + // - The ranges for src and dst do not overlap + unsafe { + std::ptr::copy_nonoverlapping::(src.cast(), dst.cast(), values.len()); + target.set_len(target.len() + values.len()); + }; Ok(()) } @@ -138,7 +146,7 @@ fn decode_optional( let num_values = validity.set_bits(); if num_values == validity.len() { - return decode_required(values, Some(validity.len()), target); + return decode_required(values.truncate(validity.len()), target); } let mut limit = validity.len(); @@ -215,7 +223,7 @@ fn decode_masked_required( let num_rows = mask.set_bits(); if num_rows == mask.len() { - return decode_required(values, Some(num_rows), target); + return decode_required(values.truncate(num_rows), target); } assert!(mask.len() <= values.len()); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs index ecd33c5850f6..2df4371460f2 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs @@ -32,6 +32,31 @@ impl<'a, B: AlignedBytes> ArrayChunks<'a, B> { pub(crate) unsafe fn get_unchecked(&self, at: usize) -> B { B::from_unaligned(*unsafe { self.bytes.get_unchecked(at) }) } + + pub fn truncate(&self, length: usize) -> ArrayChunks<'a, B> { + let length = length.min(self.bytes.len()); + + Self { + bytes: unsafe { self.bytes.get_unchecked(..length) }, + } + } + + pub unsafe fn slice_unchecked(&self, start: usize, end: usize) -> ArrayChunks<'a, B> { + debug_assert!(start <= self.bytes.len()); + debug_assert!(end <= self.bytes.len()); + + Self { + bytes: unsafe { self.bytes.get_unchecked(start..end) }, + } + } + + pub fn as_ptr(&self) -> *const B::Unaligned { + self.bytes.as_ptr() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } } impl<'a, B: AlignedBytes> Iterator for ArrayChunks<'a, B> { From 9f6aea944d18cc9e65d966c26ffaff8c57f83890 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Wed, 23 Oct 2024 14:10:23 +0200 Subject: [PATCH 22/32] remove a whole load of unused code --- .../src/arrow/read/deserialize/binview.rs | 244 +------------- .../src/arrow/read/deserialize/boolean.rs | 44 --- .../src/arrow/read/deserialize/dictionary.rs | 145 ++------- .../read/deserialize/fixed_size_binary.rs | 43 --- .../arrow/read/deserialize/nested_utils.rs | 4 - .../src/arrow/read/deserialize/null.rs | 41 +-- .../arrow/read/deserialize/primitive/float.rs | 92 +----- .../read/deserialize/primitive/integer.rs | 140 ++------ .../arrow/read/deserialize/primitive/mod.rs | 106 ------ .../src/arrow/read/deserialize/utils/mod.rs | 308 ++++++++---------- .../encoding/delta_bitpacked/decoder.rs | 38 --- .../delta_length_byte_array/decoder.rs | 4 - 12 files changed, 238 insertions(+), 971 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index b0ec85a0a718..395c36fc5506 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -1,4 +1,3 @@ -use std::mem::MaybeUninit; use std::sync::atomic::{AtomicBool, Ordering}; use arrow::array::{ @@ -16,7 +15,7 @@ use crate::parquet::encoding::delta_bitpacked::{lin_natural_sum, DeltaGatherer}; use crate::parquet::encoding::{delta_byte_array, delta_length_byte_array, hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; -use crate::read::deserialize::utils::{self, extend_from_decoder, Decoder}; +use crate::read::deserialize::utils::{self}; use crate::read::PrimitiveLogicalType; type DecodedStateTuple = (MutableBinaryViewArray<[u8]>, MutableBitmap); @@ -63,122 +62,6 @@ impl<'a> utils::StateTranslation<'a, BinViewDecoder> for StateTranslation<'a> { _ => Err(utils::not_implemented(page)), } } - - fn len_when_not_nullable(&self) -> usize { - match self { - Self::Plain(v) => v.len_when_not_nullable(), - Self::Dictionary(v) => v.len(), - Self::DeltaLengthByteArray(v, _) => v.len(), - Self::DeltaBytes(v) => v.len(), - } - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n == 0 { - return Ok(()); - } - - match self { - Self::Plain(t) => _ = t.by_ref().nth(n - 1), - Self::Dictionary(t) => t.skip_in_place(n)?, - Self::DeltaLengthByteArray(t, _) => t.skip_in_place(n)?, - Self::DeltaBytes(t) => t.skip_in_place(n)?, - } - - Ok(()) - } - - fn extend_from_state( - &mut self, - decoder: &mut BinViewDecoder, - decoded: &mut ::DecodedState, - is_optional: bool, - page_validity: &mut Option, - _dict: Option<&'a ::Dict>, - additional: usize, - ) -> ParquetResult<()> { - let views_offset = decoded.0.views().len(); - let buffer_offset = decoded.0.completed_buffers().len(); - - let mut validate_utf8 = decoder.check_utf8.load(Ordering::Relaxed); - - match self { - Self::Plain(page_values) => { - decoder.decode_plain_encoded( - decoded, - page_values, - is_optional, - page_validity.as_mut(), - additional, - )?; - - // Already done in decode_plain_encoded - validate_utf8 = false; - }, - Self::Dictionary(_) => unreachable!(), - Self::DeltaLengthByteArray(ref mut page_values, ref mut lengths) => { - let (values, validity) = decoded; - - let mut collector = DeltaCollector { - gatherer: &mut StatGatherer::default(), - pushed_lengths: lengths, - decoder: page_values, - }; - - match page_validity { - None => { - (&mut collector).push_n(values, additional)?; - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => extend_from_decoder( - validity, - page_validity, - Some(additional), - values, - &mut collector, - )?, - } - - collector.flush(values); - }, - Self::DeltaBytes(ref mut page_values) => { - let (values, validity) = decoded; - - let mut collector = DeltaBytesCollector { - decoder: page_values, - }; - - match page_validity { - None => { - collector.push_n(values, additional)?; - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => extend_from_decoder( - validity, - page_validity, - Some(additional), - values, - collector, - )?, - } - }, - } - - if validate_utf8 { - decoded - .0 - .validate_utf8(buffer_offset, views_offset) - .map_err(|_| ParquetError::oos("Binary view contained invalid UTF-8"))? - } - - Ok(()) - } } #[derive(Default)] @@ -217,10 +100,6 @@ pub(crate) struct DeltaCollector<'a, 'b> { pub(crate) decoder: &'b mut delta_length_byte_array::Decoder<'a>, } -pub(crate) struct DeltaBytesCollector<'a, 'b> { - pub(crate) decoder: &'b mut delta_byte_array::Decoder<'a>, -} - /// A [`DeltaGatherer`] that gathers the minimum, maximum and summation of the values as `usize`s. pub(crate) struct StatGatherer { min: usize, @@ -351,10 +230,6 @@ impl DeltaGatherer for StatGatherer { } impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for &mut DeltaCollector<'a, 'b> { - fn reserve(target: &mut MutableBinaryViewArray<[u8]>, n: usize) { - target.reserve(n); - } - fn push_n( &mut self, _target: &mut MutableBinaryViewArray<[u8]>, @@ -416,97 +291,6 @@ impl<'a, 'b> DeltaCollector<'a, 'b> { } } -impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for DeltaBytesCollector<'a, 'b> { - fn reserve(target: &mut MutableBinaryViewArray<[u8]>, n: usize) { - target.reserve(n); - } - - fn push_n(&mut self, target: &mut MutableBinaryViewArray<[u8]>, n: usize) -> ParquetResult<()> { - struct MaybeUninitCollector(usize); - - impl DeltaGatherer for MaybeUninitCollector { - type Target = [MaybeUninit; BATCH_SIZE]; - - fn target_len(&self, _target: &Self::Target) -> usize { - self.0 - } - - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - fn gather_one(&mut self, target: &mut Self::Target, v: i64) -> ParquetResult<()> { - target[self.0] = MaybeUninit::new(v as usize); - self.0 += 1; - Ok(()) - } - } - - let decoder_len = self.decoder.len(); - let mut n = usize::min(n, decoder_len); - - if n == 0 { - return Ok(()); - } - - let mut buffer = Vec::new(); - target.reserve(n); - - const BATCH_SIZE: usize = 4096; - - let mut prefix_lengths = [const { MaybeUninit::::uninit() }; BATCH_SIZE]; - let mut suffix_lengths = [const { MaybeUninit::::uninit() }; BATCH_SIZE]; - - while n > 0 { - let num_elems = usize::min(n, BATCH_SIZE); - n -= num_elems; - - self.decoder.prefix_lengths.gather_n_into( - &mut prefix_lengths, - num_elems, - &mut MaybeUninitCollector(0), - )?; - self.decoder.suffix_lengths.gather_n_into( - &mut suffix_lengths, - num_elems, - &mut MaybeUninitCollector(0), - )?; - - for i in 0..num_elems { - let prefix_length = unsafe { prefix_lengths[i].assume_init() }; - let suffix_length = unsafe { suffix_lengths[i].assume_init() }; - - buffer.clear(); - - buffer.extend_from_slice(&self.decoder.last[..prefix_length]); - buffer.extend_from_slice( - &self.decoder.values[self.decoder.offset..self.decoder.offset + suffix_length], - ); - - target.push_value(&buffer); - - self.decoder.last.clear(); - std::mem::swap(&mut self.decoder.last, &mut buffer); - - self.decoder.offset += suffix_length; - } - } - - Ok(()) - } - - fn push_n_nulls( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - target.extend_constant(n, >::None); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } -} - #[allow(clippy::too_many_arguments)] pub fn decode_plain( values: &[u8], @@ -887,17 +671,6 @@ impl utils::Decoder for BinViewDecoder { Ok((views, buffers)) } - fn decode_plain_encoded<'a>( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn extend_filtered_with_state( &mut self, mut state: utils::State<'_, Self>, @@ -945,7 +718,15 @@ impl utils::Decoder for BinViewDecoder { Ok(()) }, - _ => self.extend_filtered_with_state_default(state, decoded, filter), + StateTranslation::DeltaLengthByteArray(_decoder, _vec) => { + dbg!("TODO!"); + todo!() + }, + StateTranslation::DeltaBytes(_decoder) => { + dbg!("TODO!"); + todo!() + }, + // _ => self.extend_filtered_with_state_default(state, decoded, filter), } } @@ -1029,11 +810,6 @@ impl<'a> BinaryIter<'a> { max_num_values, } } - - /// Return the length of the iterator when the data is not nullable. - pub fn len_when_not_nullable(&self) -> usize { - self.max_num_values - } } impl<'a> Iterator for BinaryIter<'a> { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs index 0830404a96ad..c567a2e46a01 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs @@ -68,39 +68,6 @@ impl<'a> utils::StateTranslation<'a, BooleanDecoder> for StateTranslation<'a> { _ => Err(utils::not_implemented(page)), } } - - fn len_when_not_nullable(&self) -> usize { - match self { - Self::Plain(v) => v.len(), - Self::Rle(v) => v.len(), - } - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n == 0 { - return Ok(()); - } - - // @TODO: Add a skip_in_place on BitmapIter - match self { - Self::Plain(t) => _ = t.nth(n - 1), - Self::Rle(t) => t.skip_in_place(n)?, - } - - Ok(()) - } - - fn extend_from_state( - &mut self, - _decoder: &mut BooleanDecoder, - _decoded: &mut ::DecodedState, - _is_optional: bool, - _page_validity: &mut Option, - _: Option<&'a ::Dict>, - _additional: usize, - ) -> ParquetResult<()> { - unreachable!() - } } fn decode_required_rle( @@ -326,17 +293,6 @@ impl Decoder for BooleanDecoder { Ok(()) } - fn decode_plain_encoded<'a>( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn finalize( &self, dtype: ArrowDataType, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs index 7bfafee9a149..1f81951d959c 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs @@ -5,11 +5,11 @@ use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; use super::utils::{ - self, dict_indices_decoder, extend_from_decoder, freeze_validity, BatchableCollector, Decoder, - ExactSize, StateTranslation, + self, dict_indices_decoder, freeze_validity, unspecialized_decode, Decoder, ExactSize, + StateTranslation, }; use super::ParquetError; -use crate::parquet::encoding::hybrid_rle::{self, HybridRleDecoder, Translator}; +use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; use crate::parquet::encoding::Encoding; use crate::parquet::error::ParquetResult; use crate::parquet::page::{DataPage, DictPage}; @@ -34,52 +34,6 @@ impl<'a, K: DictionaryKey, D: utils::DictDecodable> StateTranslation<'a, Diction dict_indices_decoder(page, page_validity.map_or(0, |bm| bm.unset_bits())) } - - fn len_when_not_nullable(&self) -> usize { - self.len() - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - HybridRleDecoder::skip_in_place(self, n) - } - - fn extend_from_state( - &mut self, - decoder: &mut DictionaryDecoder, - decoded: &mut as Decoder>::DecodedState, - is_optional: bool, - page_validity: &mut Option, - _: Option<&'a as Decoder>::Dict>, - additional: usize, - ) -> ParquetResult<()> { - let (values, validity) = decoded; - - let dict_size = decoder.dict_size.load(std::sync::atomic::Ordering::Relaxed); - - if dict_size == usize::MAX { - panic!("Dictionary not set for dictionary array"); - } - - let mut collector = DictArrayCollector { - values: self, - dict_size, - }; - - match page_validity { - None => { - collector.push_n(&mut decoded.0, additional)?; - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => { - extend_from_decoder(validity, page_validity, Some(additional), values, collector)? - }, - } - - Ok(()) - } } #[derive(Debug)] @@ -132,80 +86,37 @@ impl utils::Decoder for DictionaryDec self.decoder.finalize_dict_array(dtype, dict, keys) } - fn decode_plain_encoded<'a>( + fn extend_filtered_with_state( &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, + state: utils::State<'_, Self>, + decoded: &mut Self::DecodedState, + filter: Option, ) -> ParquetResult<()> { - unreachable!() - } -} + let keys = state.translation.collect()?; + let num_rows = keys.len(); + let mut iter = keys.into_iter(); -pub(crate) struct DictArrayCollector<'a, 'b> { - values: &'b mut hybrid_rle::HybridRleDecoder<'a>, - dict_size: usize, -} + let dict_size = self.dict_size.load(std::sync::atomic::Ordering::Relaxed); -pub(crate) struct DictArrayTranslator { - dict_size: usize, -} - -impl<'a, 'b, K: DictionaryKey> BatchableCollector<(), Vec> for DictArrayCollector<'a, 'b> { - fn reserve(target: &mut Vec, n: usize) { - target.reserve(n); - } + unspecialized_decode( + num_rows, + || { + let value = iter.next().unwrap(); - fn push_n(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - let translator = DictArrayTranslator { - dict_size: self.dict_size, - }; - self.values - .translate_and_collect_n_into(target, n, &translator) - } - - fn push_n_nulls(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - target.resize(target.len() + n, K::default()); - Ok(()) - } + let value = value as usize; - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.values.skip_in_place(n) - } -} - -impl Translator for DictArrayTranslator { - fn translate(&self, value: u32) -> ParquetResult { - let value = value as usize; - - if value >= self.dict_size || value > K::MAX_USIZE_VALUE { - return Err(ParquetError::oos("Dictionary index out-of-range")); - } - - // SAFETY: value for sure fits in K - Ok(unsafe { K::from_usize_unchecked(value) }) - } - - fn translate_slice(&self, target: &mut Vec, source: &[u32]) -> ParquetResult<()> { - let Some(max) = source.iter().max() else { - return Ok(()); - }; - - let max = *max as usize; - - if max >= self.dict_size || max > K::MAX_USIZE_VALUE { - return Err(ParquetError::oos("Dictionary index out-of-range")); - } - - // SAFETY: value for sure fits in K - target.extend( - source - .iter() - .map(|v| unsafe { K::from_usize_unchecked(*v as usize) }), - ); + if value >= dict_size || value > K::MAX_USIZE_VALUE { + return Err(ParquetError::oos("Dictionary index out-of-range")); + } - Ok(()) + // SAFETY: value for sure fits in K + Ok(unsafe { K::from_usize_unchecked(value) }) + }, + filter, + state.page_validity, + state.is_optional, + &mut decoded.1, + &mut decoded.0, + ) } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 1be147402f9c..13a2cf390820 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -57,38 +57,6 @@ impl<'a> utils::StateTranslation<'a, BinaryDecoder> for StateTranslation<'a> { _ => Err(utils::not_implemented(page)), } } - - fn len_when_not_nullable(&self) -> usize { - match self { - Self::Plain(v, size) => v.len() / size, - Self::Dictionary(v) => v.len(), - } - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n == 0 { - return Ok(()); - } - - match self { - Self::Plain(v, size) => *v = &v[usize::min(v.len(), n * *size)..], - Self::Dictionary(v) => v.skip_in_place(n)?, - } - - Ok(()) - } - - fn extend_from_state( - &mut self, - _decoder: &mut BinaryDecoder, - _decoded: &mut ::DecodedState, - _is_optional: bool, - _page_validity: &mut Option, - _dict: Option<&'a ::Dict>, - _additional: usize, - ) -> ParquetResult<()> { - unreachable!() - } } pub(crate) struct BinaryDecoder { @@ -482,17 +450,6 @@ impl Decoder for BinaryDecoder { Ok(target) } - fn decode_plain_encoded<'a>( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn finalize( &self, dtype: ArrowDataType, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index 038d54f2481a..5881a94c3bc1 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -216,10 +216,6 @@ pub struct BatchedNestedDecoder<'a> { } impl<'a> BatchableCollector<(), ()> for BatchedNestedDecoder<'a> { - fn reserve(_target: &mut (), _n: usize) { - unreachable!() - } - fn push_n(&mut self, _target: &mut (), n: usize) -> ParquetResult<()> { self.filter.extend_constant(n, true); self.validity.extend_constant(n, true); diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index 01a39baacd30..68a5b9a25049 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -34,27 +34,6 @@ impl<'a> utils::StateTranslation<'a, NullDecoder> for () { ) -> ParquetResult { Ok(()) } - - fn len_when_not_nullable(&self) -> usize { - usize::MAX - } - - fn skip_in_place(&mut self, _: usize) -> ParquetResult<()> { - Ok(()) - } - - fn extend_from_state( - &mut self, - _decoder: &mut NullDecoder, - decoded: &mut ::DecodedState, - _is_optional: bool, - _page_validity: &mut Option, - _: Option<&'a ::Dict>, - additional: usize, - ) -> ParquetResult<()> { - decoded.length += additional; - Ok(()) - } } impl utils::Decoder for NullDecoder { @@ -72,17 +51,6 @@ impl utils::Decoder for NullDecoder { Ok(()) } - fn decode_plain_encoded<'a>( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, - ) -> ParquetResult<()> { - unimplemented!() - } - fn finalize( &self, dtype: ArrowDataType, @@ -91,6 +59,15 @@ impl utils::Decoder for NullDecoder { ) -> ParquetResult { Ok(NullArray::new(dtype, decoded.length)) } + + fn extend_filtered_with_state( + &mut self, + _state: utils::State<'_, Self>, + _decoded: &mut Self::DecodedState, + _filter: Option, + ) -> ParquetResult<()> { + unreachable!() + } } use super::BasicDecompressor; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs index 5b5c22bcd23f..225738b0c1fd 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/float.rs @@ -12,7 +12,9 @@ use crate::parquet::encoding::{byte_stream_split, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; -use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity}; +use crate::read::deserialize::utils::{ + dict_indices_decoder, freeze_validity, unspecialized_decode, +}; use crate::read::Filter; #[allow(clippy::large_enum_variant)] @@ -57,68 +59,6 @@ where _ => Err(utils::not_implemented(page)), } } - - fn len_when_not_nullable(&self) -> usize { - match self { - Self::Plain(n) => n.len() / size_of::

(), - Self::Dictionary(n) => n.len(), - Self::ByteStreamSplit(n) => n.len(), - } - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n == 0 { - return Ok(()); - } - - match self { - Self::Plain(t) => *t = &t[usize::min(t.len(), size_of::

() * n)..], - Self::Dictionary(t) => t.skip_in_place(n)?, - Self::ByteStreamSplit(t) => _ = t.iter_converted(|_| ()).nth(n - 1), - } - - Ok(()) - } - - fn extend_from_state( - &mut self, - decoder: &mut FloatDecoder, - decoded: &mut as utils::Decoder>::DecodedState, - is_optional: bool, - page_validity: &mut Option, - _dict: Option<&'a as utils::Decoder>::Dict>, - additional: usize, - ) -> ParquetResult<()> { - match self { - Self::ByteStreamSplit(page_values) => { - let (values, validity) = decoded; - - match page_validity { - None => { - values.extend( - page_values - .iter_converted(|v| decoder.0.decoder.decode(decode(v))) - .take(additional), - ); - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => utils::extend_from_decoder( - validity, - page_validity, - Some(additional), - values, - &mut page_values.iter_converted(|v| decoder.0.decoder.decode(decode(v))), - )?, - } - }, - _ => unreachable!(), - } - - Ok(()) - } } #[derive(Debug)] @@ -213,17 +153,6 @@ where Ok(target) } - fn decode_plain_encoded<'a>( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn extend_filtered_with_state( &mut self, mut state: utils::State<'_, Self>, @@ -250,7 +179,20 @@ where &mut decoded.1, &mut decoded.0, ), - _ => self.extend_filtered_with_state_default(state, decoded, filter), + StateTranslation::ByteStreamSplit(mut decoder) => { + let num_rows = decoder.len(); + let mut iter = decoder.iter_converted(|v| self.0.decoder.decode(decode(v))); + + unspecialized_decode( + num_rows, + || Ok(iter.next().unwrap()), + filter, + state.page_validity, + state.is_optional, + &mut decoded.1, + &mut decoded.0, + ) + }, } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 877b03391948..73d03b7cb753 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -5,14 +5,14 @@ use arrow::types::NativeType; use super::super::utils; use super::{ - AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, DeltaCollector, DeltaTranslator, + AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, IntoDecoderFunction, PrimitiveDecoder, UnitDecoderFunction, }; use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; -use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity}; +use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity, unspecialized_decode}; use crate::read::Filter; #[allow(clippy::large_enum_variant)] @@ -65,101 +65,6 @@ where _ => Err(utils::not_implemented(page)), } } - - fn len_when_not_nullable(&self) -> usize { - match self { - Self::Plain(v) => v.len() / size_of::

(), - Self::Dictionary(v) => v.len(), - Self::ByteStreamSplit(v) => v.len(), - Self::DeltaBinaryPacked(v) => v.len(), - } - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n == 0 { - return Ok(()); - } - - match self { - Self::Plain(v) => *v = &v[usize::min(v.len(), size_of::

() * n)..], - Self::Dictionary(v) => v.skip_in_place(n)?, - Self::ByteStreamSplit(v) => _ = v.iter_converted(|_| ()).nth(n - 1), - Self::DeltaBinaryPacked(v) => v.skip_in_place(n)?, - } - - Ok(()) - } - - fn extend_from_state( - &mut self, - decoder: &mut IntDecoder, - decoded: &mut as utils::Decoder>::DecodedState, - is_optional: bool, - page_validity: &mut Option, - _dict: Option<&'a as utils::Decoder>::Dict>, - additional: usize, - ) -> ParquetResult<()> { - match self { - Self::ByteStreamSplit(page_values) => { - let (values, validity) = decoded; - - match page_validity { - None => { - values.extend( - page_values - .iter_converted(|v| decoder.0.decoder.decode(decode(v))) - .take(additional), - ); - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => { - utils::extend_from_decoder( - validity, - page_validity, - Some(additional), - values, - &mut page_values - .iter_converted(|v| decoder.0.decoder.decode(decode(v))), - )?; - }, - } - }, - Self::DeltaBinaryPacked(page_values) => { - let (values, validity) = decoded; - - let mut gatherer = DeltaTranslator { - dfn: decoder.0.decoder, - _pd: std::marker::PhantomData, - }; - - match page_validity { - None => { - page_values.gather_n_into(values, additional, &mut gatherer)?; - - if is_optional { - validity.extend_constant(additional, true); - } - }, - Some(page_validity) => utils::extend_from_decoder( - validity, - page_validity, - Some(additional), - values, - DeltaCollector { - decoder: page_values, - gatherer, - }, - )?, - } - }, - _ => unreachable!(), - } - - Ok(()) - } } /// Decoder of integer parquet type @@ -267,17 +172,6 @@ where Ok(target) } - fn decode_plain_encoded<'a>( - &mut self, - _decoded: &mut Self::DecodedState, - _page_values: &mut as utils::StateTranslation<'a, Self>>::PlainDecoder, - _is_optional: bool, - _page_validity: Option<&mut Bitmap>, - _limit: usize, - ) -> ParquetResult<()> { - unreachable!() - } - fn finalize( &self, dtype: ArrowDataType, @@ -314,7 +208,35 @@ where &mut decoded.1, &mut decoded.0, ), - _ => self.extend_filtered_with_state_default(state, decoded, filter), + StateTranslation::ByteStreamSplit(mut decoder) => { + let num_rows = decoder.len(); + let mut iter = decoder.iter_converted(|v| self.0.decoder.decode(decode(v))); + + unspecialized_decode( + num_rows, + || Ok(iter.next().unwrap()), + filter, + state.page_validity, + state.is_optional, + &mut decoded.1, + &mut decoded.0, + ) + }, + StateTranslation::DeltaBinaryPacked(mut _decoder) => { + // let num_rows = decoder.len(); + + dbg!("TODO!"); + todo!() + // unspecialized_decode( + // num_rows, + // || Ok(decoder.next().unwrap()), + // filter, + // state.page_validity, + // state.is_optional, + // &mut decoded.1, + // &mut decoded.0, + // ) + }, } } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs index 18ba498171fe..88b8a55932a7 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/mod.rs @@ -1,5 +1,4 @@ use arrow::types::NativeType; -use num_traits::AsPrimitive; use crate::parquet::types::NativeType as ParquetNativeType; @@ -10,10 +9,6 @@ pub(crate) mod plain; pub(crate) use float::FloatDecoder; pub(crate) use integer::IntDecoder; -use super::utils::BatchableCollector; -use super::ParquetResult; -use crate::parquet::encoding::delta_bitpacked::{self, DeltaGatherer}; - #[derive(Debug)] pub(crate) struct PrimitiveDecoder where @@ -133,104 +128,3 @@ where (self.0)(x) } } - -struct DeltaTranslator -where - T: NativeType, - P: ParquetNativeType, - i64: AsPrimitive

, - D: DecoderFunction, -{ - dfn: D, - _pd: std::marker::PhantomData<(P, T)>, -} - -struct DeltaCollector<'a, 'b, P, T, D> -where - T: NativeType, - P: ParquetNativeType, - i64: AsPrimitive

, - D: DecoderFunction, -{ - decoder: &'b mut delta_bitpacked::Decoder<'a>, - gatherer: DeltaTranslator, -} - -impl DeltaGatherer for DeltaTranslator -where - T: NativeType, - P: ParquetNativeType, - i64: AsPrimitive

, - D: DecoderFunction, -{ - type Target = Vec; - - fn target_len(&self, target: &Self::Target) -> usize { - target.len() - } - - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } - - fn gather_one(&mut self, target: &mut Self::Target, v: i64) -> ParquetResult<()> { - target.push(self.dfn.decode(v.as_())); - Ok(()) - } - - fn gather_constant( - &mut self, - target: &mut Self::Target, - v: i64, - delta: i64, - num_repeats: usize, - ) -> ParquetResult<()> { - target.extend((0..num_repeats).map(|i| self.dfn.decode((v + (i as i64) * delta).as_()))); - Ok(()) - } - - fn gather_slice(&mut self, target: &mut Self::Target, slice: &[i64]) -> ParquetResult<()> { - target.extend(slice.iter().copied().map(|v| self.dfn.decode(v.as_()))); - Ok(()) - } - - fn gather_chunk(&mut self, target: &mut Self::Target, chunk: &[i64; 64]) -> ParquetResult<()> { - target.extend(chunk.iter().copied().map(|v| self.dfn.decode(v.as_()))); - Ok(()) - } -} - -impl<'a, 'b, P, T, D> BatchableCollector<(), Vec> for DeltaCollector<'a, 'b, P, T, D> -where - T: NativeType, - P: ParquetNativeType, - i64: AsPrimitive

, - D: DecoderFunction, -{ - fn reserve(target: &mut Vec, n: usize) { - target.reserve(n); - } - - fn push_n(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - let start_length = target.len(); - let start_num_elems = self.decoder.len(); - - self.decoder.gather_n_into(target, n, &mut self.gatherer)?; - - let consumed_elements = usize::min(n, start_num_elems); - - debug_assert_eq!(self.decoder.len(), start_num_elems - consumed_elements); - debug_assert_eq!(target.len(), start_length + consumed_elements); - - Ok(()) - } - - fn push_n_nulls(&mut self, target: &mut Vec, n: usize) -> ParquetResult<()> { - target.resize(target.len() + n, T::default()); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } -} diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 88f94b02ad4d..d608338982aa 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -36,20 +36,6 @@ pub(crate) trait StateTranslation<'a, D: Decoder>: Sized { dict: Option<&'a D::Dict>, page_validity: Option<&Bitmap>, ) -> ParquetResult; - fn len_when_not_nullable(&self) -> usize; - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()>; - - /// extends [`Self::DecodedState`] by deserializing items in [`Self::State`]. - /// It guarantees that the length of `decoded` is at most `decoded.len() + additional`. - fn extend_from_state( - &mut self, - decoder: &mut D, - decoded: &mut D::DecodedState, - is_optional: bool, - page_validity: &mut Option, - dict: Option<&'a D::Dict>, - additional: usize, - ) -> ParquetResult<()>; } impl<'a, D: Decoder> State<'a, D> { @@ -105,27 +91,6 @@ impl<'a, D: Decoder> State<'a, D> { }) } - pub fn len(&self) -> usize { - match &self.page_validity { - Some(v) => v.len(), - None => self.translation.len_when_not_nullable(), - } - } - - pub fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n == 0 { - return Ok(()); - } - - let n = self.page_validity.as_mut().map_or(n, |page_validity| { - let mut pv = page_validity.clone(); - pv.slice(0, n); - pv.unset_bits() - }); - - self.translation.skip_in_place(n) - } - pub fn decode( self, decoder: &mut D, @@ -147,7 +112,6 @@ pub fn not_implemented(page: &DataPage) -> ParquetError { } pub trait BatchableCollector { - fn reserve(target: &mut T, n: usize); fn push_n(&mut self, target: &mut T, n: usize) -> ParquetResult<()>; fn push_n_nulls(&mut self, target: &mut T, n: usize) -> ParquetResult<()>; fn skip_in_place(&mut self, n: usize) -> ParquetResult<()>; @@ -244,6 +208,148 @@ pub(crate) fn page_validity_decoder(page: &DataPage) -> ParquetResult( + mut num_rows: usize, + + mut decode_one: impl FnMut() -> ParquetResult, + + mut filter: Option, + page_validity: Option, + + is_optional: bool, + + validity: &mut MutableBitmap, + target: &mut Vec, +) -> ParquetResult<()> { + match &filter { + None => {}, + Some(Filter::Range(range)) => { + match page_validity.as_ref() { + None => { + for _ in 0..range.start { + decode_one()?; + } + }, + Some(pv) => { + for _ in 0..pv.clone().sliced(0, range.start).set_bits() { + decode_one()?; + } + }, + } + + num_rows = range.len(); + filter = None; + }, + Some(Filter::Mask(mask)) => { + if mask.unset_bits() == 0 { + num_rows = mask.len(); + filter = None; + } + }, + }; + + match (filter, page_validity) { + (None, None) => { + target.reserve(num_rows); + for _ in 0..num_rows { + target.push(decode_one()?); + } + + if is_optional { + validity.extend_constant(num_rows, true); + } + }, + (None, Some(page_validity)) => { + target.reserve(page_validity.len()); + for is_valid in page_validity.iter() { + let v = if is_valid { + decode_one()? + } else { + T::default() + }; + target.push(v); + } + + validity.extend_from_bitmap(&page_validity); + }, + (Some(Filter::Range(_)), _) => unreachable!(), + (Some(Filter::Mask(mask)), None) => { + let num_rows = mask.set_bits(); + target.reserve(num_rows); + + let mut iter = mask.iter(); + while iter.num_remaining() > 0 { + let num_ones = iter.take_leading_ones(); + + if num_ones > 0 { + for _ in 0..num_rows { + target.push(decode_one()?); + } + } + + let num_zeros = iter.take_leading_zeros(); + for _ in 0..num_zeros { + decode_one()?; + } + } + + if is_optional { + validity.extend_constant(num_rows, true); + } + }, + (Some(Filter::Mask(mask)), Some(page_validity)) => { + assert_eq!(mask.len(), page_validity.len()); + + let num_rows = mask.set_bits(); + target.reserve(num_rows); + + let mut mask_iter = mask.fast_iter_u56(); + let mut validity_iter = page_validity.fast_iter_u56(); + + let mut iter = |mut f: u64, mut v: u64| { + while f != 0 { + let offset = f.trailing_ones(); + + if (v >> offset) & 1 != 0 { + target.push(decode_one()?); + } else { + target.push(T::default()); + } + + let skip = (v & (1u64 << offset).wrapping_sub(1)).count_ones() as usize; + for _ in 0..skip { + decode_one()?; + } + + v >>= offset + 1; + f >>= offset + 1; + } + + for _ in 0..v.count_ones() as usize { + decode_one()?; + } + + ParquetResult::Ok(()) + }; + + for (f, v) in mask_iter.by_ref().zip(validity_iter.by_ref()) { + iter(f, v)?; + } + + let (f, fl) = mask_iter.remainder(); + let (v, vl) = validity_iter.remainder(); + + assert_eq!(fl, vl); + + iter(f, v)?; + + validity.extend_from_bitmap(&page_validity); + }, + } + + Ok(()) +} + #[derive(Default)] pub(crate) struct BatchGatherer<'a, I, T, C: BatchableCollector>( std::marker::PhantomData<&'a (I, T, C)>, @@ -316,43 +422,7 @@ impl<'a, I, T, C: BatchableCollector> HybridRleGatherer for BatchGath } } -/// Extends a [`Pushable`] from an iterator of non-null values and an hybrid-rle decoder -pub(super) fn extend_from_decoder>( - validity: &mut MutableBitmap, - page_validity: &mut Bitmap, - limit: Option, - target: &mut T, - collector: C, -) -> ParquetResult<()> { - let num_elements = limit.map_or(page_validity.len(), |limit| limit.min(page_validity.len())); - - validity.reserve(num_elements); - C::reserve(target, num_elements); - - let mut batched_collector = BatchedCollector::new(collector, target); - - let mut pv = page_validity.clone(); - pv.slice(0, num_elements); - - // @TODO: This is terribly slow now. - validity.extend_from_bitmap(&pv); - let mut iter = pv.iter(); - while iter.num_remaining() > 0 { - batched_collector.push_n_valids(iter.take_leading_ones())?; - batched_collector.push_n_invalids(iter.take_leading_zeros()); - } - - batched_collector.finalize()?; - - Ok(()) -} - impl, I: Iterator> BatchableCollector for I { - #[inline] - fn reserve(target: &mut P, n: usize) { - target.reserve(n); - } - #[inline] fn push_n(&mut self, target: &mut P, n: usize) -> ParquetResult<()> { target.extend_n(n, self); @@ -402,90 +472,7 @@ pub(super) trait Decoder: Sized { state: State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, - ) -> ParquetResult<()> { - self.extend_filtered_with_state_default(state, decoded, filter) - } - - fn extend_filtered_with_state_default( - &mut self, - mut state: State<'_, Self>, - decoded: &mut Self::DecodedState, - filter: Option, - ) -> ParquetResult<()> { - match filter { - None => { - let num_rows = state.len(); - - if num_rows == 0 { - return Ok(()); - } - - state.translation.extend_from_state( - self, - decoded, - state.is_optional, - &mut state.page_validity, - state.dict, - num_rows, - ) - }, - Some(filter) => match filter { - Filter::Range(range) => { - let start = range.start; - let end = range.end; - - state.skip_in_place(start)?; - debug_assert!(end - start <= state.len()); - - if end - start > 0 { - state.translation.extend_from_state( - self, - decoded, - state.is_optional, - &mut state.page_validity, - state.dict, - end - start, - )?; - } - - Ok(()) - }, - Filter::Mask(bitmap) => { - let mut iter = bitmap.iter(); - while iter.num_remaining() > 0 && state.len() > 0 { - let prev_state_len = state.len(); - - let num_ones = iter.take_leading_ones(); - - if num_ones > 0 { - state.translation.extend_from_state( - self, - decoded, - state.is_optional, - &mut state.page_validity, - state.dict, - num_ones, - )?; - } - - if iter.num_remaining() == 0 || state.len() == 0 { - break; - } - - let num_zeros = iter.take_leading_zeros(); - state.skip_in_place(num_zeros)?; - - assert!( - prev_state_len != state.len(), - "No forward progress was booked in a filtered parquet file." - ); - } - - Ok(()) - }, - }, - } - } + ) -> ParquetResult<()>; fn apply_dictionary( &mut self, @@ -495,15 +482,6 @@ pub(super) trait Decoder: Sized { Ok(()) } - fn decode_plain_encoded<'a>( - &mut self, - decoded: &mut Self::DecodedState, - page_values: &mut as StateTranslation<'a, Self>>::PlainDecoder, - is_optional: bool, - page_validity: Option<&mut Bitmap>, - limit: usize, - ) -> ParquetResult<()>; - fn finalize( &self, dtype: ArrowDataType, diff --git a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs index 261e84ce2e23..6766722ffdca 100644 --- a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs +++ b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs @@ -69,7 +69,6 @@ struct MiniBlock<'a> { unpacked_end: usize, } -struct SkipGatherer; pub(crate) struct SumGatherer(pub(crate) usize); pub trait DeltaGatherer { @@ -109,38 +108,6 @@ pub trait DeltaGatherer { } } -impl DeltaGatherer for SkipGatherer { - type Target = usize; - - fn target_len(&self, target: &Self::Target) -> usize { - *target - } - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - fn gather_one(&mut self, target: &mut Self::Target, _v: i64) -> ParquetResult<()> { - *target += 1; - Ok(()) - } - fn gather_constant( - &mut self, - target: &mut Self::Target, - _v: i64, - _delta: i64, - num_repeats: usize, - ) -> ParquetResult<()> { - *target += num_repeats; - Ok(()) - } - fn gather_chunk(&mut self, target: &mut Self::Target, chunk: &[i64; 64]) -> ParquetResult<()> { - *target += chunk.len(); - Ok(()) - } - fn gather_slice(&mut self, target: &mut Self::Target, slice: &[i64]) -> ParquetResult<()> { - *target += slice.len(); - Ok(()) - } -} - impl DeltaGatherer for SumGatherer { type Target = usize; @@ -749,11 +716,6 @@ impl<'a> Decoder<'a> { Ok(()) } - pub fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - let mut gatherer = SkipGatherer; - self.gather_n_into(&mut 0usize, n, &mut gatherer) - } - #[cfg(test)] pub(crate) fn collect_n>( &mut self, diff --git a/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs b/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs index b3191e0a51ff..a57dcc201f51 100644 --- a/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs +++ b/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs @@ -30,10 +30,6 @@ impl<'a> Decoder<'a> { self.offset += sum; Ok(()) } - - pub fn len(&self) -> usize { - self.lengths.len() - } } #[cfg(test)] From 8bbb577afe79573591e3a66708829d7ad4413953 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Wed, 23 Oct 2024 16:27:48 +0200 Subject: [PATCH 23/32] remove a whole load more of unneeded code --- crates/polars-arrow/src/pushable.rs | 1 + .../src/arrow/read/deserialize/binview.rs | 275 +++++------------- .../src/arrow/read/deserialize/dictionary.rs | 11 +- .../src/arrow/read/deserialize/null.rs | 22 +- .../read/deserialize/primitive/integer.rs | 31 +- .../arrow/read/deserialize/primitive/plain.rs | 20 +- .../src/arrow/read/deserialize/utils/mod.rs | 2 +- .../encoding/delta_bitpacked/decoder.rs | 2 - .../parquet/encoding/delta_bitpacked/fuzz.rs | 76 ----- .../parquet/encoding/delta_bitpacked/mod.rs | 3 +- .../encoding/delta_byte_array/decoder.rs | 64 ++-- .../delta_length_byte_array/decoder.rs | 11 +- 12 files changed, 156 insertions(+), 362 deletions(-) delete mode 100644 crates/polars-parquet/src/parquet/encoding/delta_bitpacked/fuzz.rs diff --git a/crates/polars-arrow/src/pushable.rs b/crates/polars-arrow/src/pushable.rs index 7b8857ab3a15..29464b8df679 100644 --- a/crates/polars-arrow/src/pushable.rs +++ b/crates/polars-arrow/src/pushable.rs @@ -181,6 +181,7 @@ impl Pushable> for MutablePrimitiveArray { pub trait NoOption {} impl NoOption for &str {} impl NoOption for &[u8] {} +impl NoOption for Vec {} impl Pushable for MutableBinaryViewArray where diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index 395c36fc5506..6bc9831cd240 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -9,9 +9,10 @@ use arrow::buffer::Buffer; use arrow::datatypes::{ArrowDataType, PhysicalType}; use super::utils::dict_encoded::{append_validity, constrain_page_validity}; -use super::utils::{dict_indices_decoder, filter_from_range, freeze_validity, BatchableCollector}; +use super::utils::{ + dict_indices_decoder, filter_from_range, freeze_validity, unspecialized_decode, +}; use super::Filter; -use crate::parquet::encoding::delta_bitpacked::{lin_natural_sum, DeltaGatherer}; use crate::parquet::encoding::{delta_byte_array, delta_length_byte_array, hybrid_rle, Encoding}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; @@ -90,207 +91,6 @@ impl utils::ExactSize for (Vec, Vec>) { } } -pub(crate) struct DeltaCollector<'a, 'b> { - // We gatherer the decoded lengths into `pushed_lengths`. Then, we `flush` those to the - // `BinView` This allows us to group many memcopies into one and take better potential fast - // paths for inlineable views and such. - pub(crate) gatherer: &'b mut StatGatherer, - pub(crate) pushed_lengths: &'b mut Vec, - - pub(crate) decoder: &'b mut delta_length_byte_array::Decoder<'a>, -} - -/// A [`DeltaGatherer`] that gathers the minimum, maximum and summation of the values as `usize`s. -pub(crate) struct StatGatherer { - min: usize, - max: usize, - sum: usize, -} - -impl Default for StatGatherer { - fn default() -> Self { - Self { - min: usize::MAX, - max: usize::MIN, - sum: 0, - } - } -} - -impl DeltaGatherer for StatGatherer { - type Target = Vec; - - fn target_len(&self, target: &Self::Target) -> usize { - target.len() - } - - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } - - fn gather_one(&mut self, target: &mut Self::Target, v: i64) -> ParquetResult<()> { - if v < 0 { - return Err(ParquetError::oos("DELTA_LENGTH_BYTE_ARRAY length < 0")); - } - - if v > i64::from(u32::MAX) { - return Err(ParquetError::not_supported( - "DELTA_LENGTH_BYTE_ARRAY length > u32::MAX", - )); - } - - let v = v as usize; - - self.min = self.min.min(v); - self.max = self.max.max(v); - self.sum += v; - - target.push(v as u32); - - Ok(()) - } - - fn gather_slice(&mut self, target: &mut Self::Target, slice: &[i64]) -> ParquetResult<()> { - let mut is_invalid = false; - let mut is_too_large = false; - - target.extend(slice.iter().map(|&v| { - is_invalid |= v < 0; - is_too_large |= v > i64::from(u32::MAX); - - let v = v as usize; - - self.min = self.min.min(v); - self.max = self.max.max(v); - self.sum += v; - - v as u32 - })); - - if is_invalid { - target.truncate(target.len() - slice.len()); - return Err(ParquetError::oos("DELTA_LENGTH_BYTE_ARRAY length < 0")); - } - - if is_too_large { - return Err(ParquetError::not_supported( - "DELTA_LENGTH_BYTE_ARRAY length > u32::MAX", - )); - } - - Ok(()) - } - - fn gather_constant( - &mut self, - target: &mut Self::Target, - v: i64, - delta: i64, - num_repeats: usize, - ) -> ParquetResult<()> { - if v < 0 || (delta < 0 && num_repeats > 0 && (num_repeats - 1) as i64 * delta + v < 0) { - return Err(ParquetError::oos("DELTA_LENGTH_BYTE_ARRAY length < 0")); - } - - if v > i64::from(u32::MAX) || v + ((num_repeats - 1) as i64) * delta > i64::from(u32::MAX) { - return Err(ParquetError::not_supported( - "DELTA_LENGTH_BYTE_ARRAY length > u32::MAX", - )); - } - - target.extend((0..num_repeats).map(|i| (v + (i as i64) * delta) as u32)); - - let vstart = v; - let vend = v + (num_repeats - 1) as i64 * delta; - - let (min, max) = if delta < 0 { - (vend, vstart) - } else { - (vstart, vend) - }; - - let sum = lin_natural_sum(v, delta, num_repeats) as usize; - - #[cfg(debug_assertions)] - { - assert_eq!( - (0..num_repeats) - .map(|i| (v + (i as i64) * delta) as usize) - .sum::(), - sum - ); - } - - self.min = self.min.min(min as usize); - self.max = self.max.max(max as usize); - self.sum += sum; - - Ok(()) - } -} - -impl<'a, 'b> BatchableCollector<(), MutableBinaryViewArray<[u8]>> for &mut DeltaCollector<'a, 'b> { - fn push_n( - &mut self, - _target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - self.decoder - .lengths - .gather_n_into(self.pushed_lengths, n, self.gatherer)?; - - Ok(()) - } - - fn push_n_nulls( - &mut self, - target: &mut MutableBinaryViewArray<[u8]>, - n: usize, - ) -> ParquetResult<()> { - self.flush(target); - target.extend_constant(n, Some(&[])); - Ok(()) - } - - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - self.decoder.skip_in_place(n) - } -} - -impl<'a, 'b> DeltaCollector<'a, 'b> { - pub fn flush(&mut self, target: &mut MutableBinaryViewArray<[u8]>) { - if !self.pushed_lengths.is_empty() { - let start_bytes_len = target.total_bytes_len(); - let start_buffer_len = target.total_buffer_len(); - unsafe { - target.extend_from_lengths_with_stats( - &self.decoder.values[self.decoder.offset..], - self.pushed_lengths.iter().map(|&v| v as usize), - self.gatherer.min, - self.gatherer.max, - self.gatherer.sum, - ) - }; - debug_assert_eq!( - target.total_bytes_len() - start_bytes_len, - self.gatherer.sum, - ); - debug_assert_eq!( - target.total_buffer_len() - start_buffer_len, - self.pushed_lengths - .iter() - .map(|&v| v as usize) - .filter(|&v| v > View::MAX_INLINE_SIZE as usize) - .sum::(), - ); - - self.decoder.offset += self.gatherer.sum; - self.pushed_lengths.clear(); - *self.gatherer = StatGatherer::default(); - } - } -} - #[allow(clippy::too_many_arguments)] pub fn decode_plain( values: &[u8], @@ -718,15 +518,70 @@ impl utils::Decoder for BinViewDecoder { Ok(()) }, - StateTranslation::DeltaLengthByteArray(_decoder, _vec) => { - dbg!("TODO!"); - todo!() + StateTranslation::DeltaLengthByteArray(decoder, _vec) => { + let values = decoder.values; + let lengths = decoder.lengths.collect::>()?; + + if self.check_utf8.load(Ordering::Relaxed) { + let mut none_starting_with_continuation_byte = true; + let mut offset = 0; + for length in &lengths { + none_starting_with_continuation_byte &= + *length == 0 || values[offset] & 0xC0 != 0x80; + offset += *length as usize; + } + + if !none_starting_with_continuation_byte { + return Err(invalid_utf8_err()); + } + + if simdutf8::basic::from_utf8(&values[..offset]).is_err() { + return Err(invalid_utf8_err()); + } + } + + let mut i = 0; + let mut offset = 0; + unspecialized_decode( + lengths.len(), + || { + let length = lengths[i] as usize; + + let value = &values[offset..offset + length]; + + i += 1; + offset += length; + + Ok(value) + }, + filter, + state.page_validity, + state.is_optional, + &mut decoded.1, + &mut decoded.0, + ) }, - StateTranslation::DeltaBytes(_decoder) => { - dbg!("TODO!"); - todo!() + StateTranslation::DeltaBytes(mut decoder) => { + let check_utf8 = self.check_utf8.load(Ordering::Relaxed); + + unspecialized_decode( + decoder.len(), + || { + let value = decoder.next().unwrap()?; + + if check_utf8 && simdutf8::basic::from_utf8(&value[..]).is_err() { + return Err(invalid_utf8_err()); + } + + Ok(value) + }, + filter, + state.page_validity, + state.is_optional, + &mut decoded.1, + &mut decoded.0, + ) }, - // _ => self.extend_filtered_with_state_default(state, decoded, filter), } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs index 1f81951d959c..6c149803a887 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/dictionary.rs @@ -1,5 +1,3 @@ -use std::sync::atomic::AtomicUsize; - use arrow::array::{DictionaryArray, DictionaryKey, PrimitiveArray}; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; @@ -38,7 +36,7 @@ impl<'a, K: DictionaryKey, D: utils::DictDecodable> StateTranslation<'a, Diction #[derive(Debug)] pub struct DictionaryDecoder { - dict_size: AtomicUsize, + dict_size: usize, decoder: D, _pd: std::marker::PhantomData, } @@ -46,7 +44,7 @@ pub struct DictionaryDecoder { impl DictionaryDecoder { pub fn new(decoder: D) -> Self { Self { - dict_size: AtomicUsize::new(usize::MAX), + dict_size: usize::MAX, decoder, _pd: std::marker::PhantomData, } @@ -68,8 +66,7 @@ impl utils::Decoder for DictionaryDec fn deserialize_dict(&mut self, page: DictPage) -> ParquetResult { let dict = self.decoder.deserialize_dict(page)?; - self.dict_size - .store(dict.len(), std::sync::atomic::Ordering::Relaxed); + self.dict_size = dict.len(); Ok(dict) } @@ -96,7 +93,7 @@ impl utils::Decoder for DictionaryDec let num_rows = keys.len(); let mut iter = keys.into_iter(); - let dict_size = self.dict_size.load(std::sync::atomic::Ordering::Relaxed); + let dict_size = self.dict_size; unspecialized_decode( num_rows, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index 68a5b9a25049..22000905a54b 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -11,6 +11,9 @@ use super::utils::filter::Filter; use crate::parquet::error::ParquetResult; use crate::parquet::page::{DataPage, DictPage}; +pub(crate) struct NullTranslation { + length: usize, +} pub(crate) struct NullDecoder; #[derive(Debug)] pub(crate) struct NullArrayLength { @@ -23,21 +26,23 @@ impl utils::ExactSize for NullArrayLength { } } -impl<'a> utils::StateTranslation<'a, NullDecoder> for () { +impl<'a> utils::StateTranslation<'a, NullDecoder> for NullTranslation { type PlainDecoder = (); fn new( _decoder: &NullDecoder, - _page: &'a DataPage, + page: &'a DataPage, _dict: Option<&'a ::Dict>, _page_validity: Option<&Bitmap>, ) -> ParquetResult { - Ok(()) + Ok(Self { + length: page.num_values() + }) } } impl utils::Decoder for NullDecoder { - type Translation<'a> = (); + type Translation<'a> = NullTranslation; type Dict = (); type DecodedState = NullArrayLength; type Output = NullArray; @@ -62,11 +67,12 @@ impl utils::Decoder for NullDecoder { fn extend_filtered_with_state( &mut self, - _state: utils::State<'_, Self>, - _decoded: &mut Self::DecodedState, - _filter: Option, + state: utils::State<'_, Self>, + decoded: &mut Self::DecodedState, + filter: Option, ) -> ParquetResult<()> { - unreachable!() + decoded.length += Filter::opt_num_rows(&filter, state.translation.length); + Ok(()) } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index 73d03b7cb753..ba3c249ae772 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -222,20 +222,25 @@ where &mut decoded.0, ) }, - StateTranslation::DeltaBinaryPacked(mut _decoder) => { - // let num_rows = decoder.len(); + StateTranslation::DeltaBinaryPacked(decoder) => { + let num_rows = decoder.len(); + let values = decoder.collect::>()?; - dbg!("TODO!"); - todo!() - // unspecialized_decode( - // num_rows, - // || Ok(decoder.next().unwrap()), - // filter, - // state.page_validity, - // state.is_optional, - // &mut decoded.1, - // &mut decoded.0, - // ) + let mut i = 0; + unspecialized_decode( + num_rows, + || { + use num_traits::AsPrimitive; + let value = values[i]; + i += 1; + Ok(self.0.decoder.decode(value.as_())) + }, + filter, + state.page_validity, + state.is_optional, + &mut decoded.1, + &mut decoded.0, + ) }, } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs index 4885ca674c55..9e843f673072 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/plain.rs @@ -1,3 +1,4 @@ +use arrow::array::Splitable; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::types::{AlignedBytes, NativeType}; @@ -93,11 +94,20 @@ pub fn decode_aligned_bytes_dispatch( unsafe { values.slice_unchecked(rng.start, rng.end) }, target, ), - (Some(Filter::Range(rng)), Some(page_validity)) => decode_optional( - unsafe { values.slice_unchecked(rng.start, rng.end) }, - &page_validity.sliced(rng.start, rng.len()), - target, - ), + (Some(Filter::Range(rng)), Some(mut page_validity)) => { + let prevalidity; + (prevalidity, page_validity) = page_validity.split_at(rng.start); + + (page_validity, _) = page_validity.split_at(rng.len()); + + let values_start = prevalidity.set_bits(); + + decode_optional( + unsafe { values.slice_unchecked(values_start, values.len()) }, + &page_validity, + target, + ) + }, (Some(Filter::Mask(filter)), None) => decode_masked_required(values, &filter, target), (Some(Filter::Mask(filter)), Some(page_validity)) => { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index d608338982aa..428f1fae06fe 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -219,7 +219,7 @@ pub(crate) fn unspecialized_decode( is_optional: bool, validity: &mut MutableBitmap, - target: &mut Vec, + target: &mut impl Pushable, ) -> ParquetResult<()> { match &filter { None => {}, diff --git a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs index 6766722ffdca..f176ef9862d4 100644 --- a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs +++ b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/decoder.rs @@ -716,7 +716,6 @@ impl<'a> Decoder<'a> { Ok(()) } - #[cfg(test)] pub(crate) fn collect_n>( &mut self, e: &mut E, @@ -748,7 +747,6 @@ impl<'a> Decoder<'a> { self.gather_n_into(&mut target, n, &mut gatherer) } - #[cfg(test)] pub(crate) fn collect + Default>( mut self, ) -> ParquetResult { diff --git a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/fuzz.rs b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/fuzz.rs deleted file mode 100644 index dc16bc8353fd..000000000000 --- a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/fuzz.rs +++ /dev/null @@ -1,76 +0,0 @@ -#[ignore = "Fuzz test. Takes too long"] -#[test] -fn fuzz_test_delta_encoding() -> Result<(), Box> { - use rand::Rng; - - use super::DeltaGatherer; - use crate::parquet::error::ParquetResult; - - struct SimpleGatherer; - - impl DeltaGatherer for SimpleGatherer { - type Target = Vec; - - fn target_len(&self, target: &Self::Target) -> usize { - target.len() - } - - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } - - fn gather_one(&mut self, target: &mut Self::Target, v: i64) -> ParquetResult<()> { - target.push(v); - Ok(()) - } - } - - const MIN_VALUES: usize = 1; - const MAX_VALUES: usize = 515; - - const MIN: i64 = i64::MIN; - const MAX: i64 = i64::MAX; - - const NUM_ITERATIONS: usize = 1_000_000; - - let mut values = Vec::with_capacity(MAX_VALUES); - let mut rng = rand::thread_rng(); - - let mut encoded = Vec::with_capacity(MAX_VALUES); - let mut decoded = Vec::with_capacity(MAX_VALUES); - let mut gatherer = SimpleGatherer; - - for i in 0..NUM_ITERATIONS { - values.clear(); - - let num_values = rng.gen_range(MIN_VALUES..=MAX_VALUES); - values.extend(std::iter::from_fn(|| Some(rng.gen_range(MIN..=MAX))).take(num_values)); - - encoded.clear(); - decoded.clear(); - - super::encode( - values.iter().copied(), - &mut encoded, - 1 << rng.gen_range(0..=2), - ); - let (mut decoder, rem) = super::Decoder::try_new(&encoded)?; - - assert!(rem.is_empty()); - - let mut num_remaining = num_values; - while num_remaining > 0 { - let n = rng.gen_range(1usize..=num_remaining); - decoder.gather_n_into(&mut decoded, n, &mut gatherer)?; - num_remaining -= n; - } - - assert_eq!(values, decoded); - - if i % 1000 == 999 { - eprintln!("[INFO]: {} iterations done.", i + 1); - } - } - - Ok(()) -} diff --git a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/mod.rs b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/mod.rs index 4a32610a302e..040909a336bb 100644 --- a/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/delta_bitpacked/mod.rs @@ -1,8 +1,7 @@ mod decoder; mod encoder; -mod fuzz; -pub(crate) use decoder::{Decoder, DeltaGatherer, SumGatherer}; +pub(crate) use decoder::{Decoder, SumGatherer}; pub(crate) use encoder::encode; /// The sum of `start, start + delta, start + 2 * delta, ... len times`. diff --git a/crates/polars-parquet/src/parquet/encoding/delta_byte_array/decoder.rs b/crates/polars-parquet/src/parquet/encoding/delta_byte_array/decoder.rs index 03889e0aa5d3..1d1b82e7cd49 100644 --- a/crates/polars-parquet/src/parquet/encoding/delta_byte_array/decoder.rs +++ b/crates/polars-parquet/src/parquet/encoding/delta_byte_array/decoder.rs @@ -52,45 +52,51 @@ impl<'a> Decoder<'a> { } } -#[cfg(test)] -mod tests { - use super::*; +impl<'a> Iterator for Decoder<'a> { + type Item = ParquetResult>; - impl<'a> Iterator for Decoder<'a> { - type Item = ParquetResult>; + fn next(&mut self) -> Option { + if self.len() == 0 { + return None; + } - fn next(&mut self) -> Option { - if self.len() == 0 { - return None; - } + let mut prefix_length = vec![]; + let mut suffix_length = vec![]; + if let Err(e) = self.prefix_lengths.collect_n(&mut prefix_length, 1) { + return Some(Err(e)); + } + if let Err(e) = self.suffix_lengths.collect_n(&mut suffix_length, 1) { + return Some(Err(e)); + } + let prefix_length = prefix_length[0]; + let suffix_length = suffix_length[0]; - let mut prefix_length = vec![]; - let mut suffix_length = vec![]; - if let Err(e) = self.prefix_lengths.collect_n(&mut prefix_length, 1) { - return Some(Err(e)); - } - if let Err(e) = self.suffix_lengths.collect_n(&mut suffix_length, 1) { - return Some(Err(e)); - } - let prefix_length = prefix_length[0]; - let suffix_length = suffix_length[0]; + let prefix_length = prefix_length as usize; + let suffix_length = suffix_length as usize; - let prefix_length = prefix_length as usize; - let suffix_length = suffix_length as usize; + let mut value = Vec::with_capacity(prefix_length + suffix_length); - let mut value = Vec::with_capacity(prefix_length + suffix_length); + value.extend_from_slice(&self.last[..prefix_length]); + value.extend_from_slice(&self.values[self.offset..self.offset + suffix_length]); - value.extend_from_slice(&self.last[..prefix_length]); - value.extend_from_slice(&self.values[self.offset..self.offset + suffix_length]); + self.last.clear(); + self.last.extend_from_slice(&value); - self.last.clear(); - self.last.extend_from_slice(&value); + self.offset += suffix_length; - self.offset += suffix_length; + Some(Ok(value)) + } - Some(Ok(value)) - } + fn size_hint(&self) -> (usize, Option) { + (self.prefix_lengths.len(), Some(self.prefix_lengths.len())) } +} + +impl<'a> ExactSizeIterator for Decoder<'a> {} + +#[cfg(test)] +mod tests { + use super::*; #[test] fn test_bla() -> ParquetResult<()> { diff --git a/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs b/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs index a57dcc201f51..3bd46add609e 100644 --- a/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs +++ b/crates/polars-parquet/src/parquet/encoding/delta_length_byte_array/decoder.rs @@ -1,5 +1,4 @@ use super::super::delta_bitpacked; -use crate::parquet::encoding::delta_bitpacked::SumGatherer; use crate::parquet::error::ParquetResult; /// Decodes [Delta-length byte array](https://github.com/apache/parquet-format/blob/master/Encodings.md#delta-length-byte-array-delta_length_byte_array--6) @@ -10,6 +9,7 @@ use crate::parquet::error::ParquetResult; pub(crate) struct Decoder<'a> { pub(crate) lengths: delta_bitpacked::Decoder<'a>, pub(crate) values: &'a [u8], + #[cfg(test)] pub(crate) offset: usize, } @@ -19,17 +19,10 @@ impl<'a> Decoder<'a> { Ok(Self { lengths, values, + #[cfg(test)] offset: 0, }) } - - pub(crate) fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - let mut sum = 0usize; - self.lengths - .gather_n_into(&mut sum, n, &mut SumGatherer(0))?; - self.offset += sum; - Ok(()) - } } #[cfg(test)] From 99aa500707177603c48f7f1641999b3cea789f03 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 10:45:40 +0200 Subject: [PATCH 24/32] remove even more code and simplify nested by a lot --- .../arrow/read/deserialize/nested_utils.rs | 802 +++++------------- .../src/arrow/read/deserialize/null.rs | 19 +- .../src/arrow/read/deserialize/utils/mod.rs | 85 -- .../src/parquet/encoding/bitpacked/decode.rs | 5 + .../src/parquet/encoding/bitpacked/mod.rs | 32 + .../src/parquet/encoding/bitpacked/pack.rs | 1 + .../src/parquet/encoding/bitpacked/unpack.rs | 1 + .../parquet/encoding/hybrid_rle/buffered.rs | 280 ------ .../src/parquet/encoding/hybrid_rle/fuzz.rs | 390 --------- .../parquet/encoding/hybrid_rle/gatherer.rs | 545 ------------ .../src/parquet/encoding/hybrid_rle/mod.rs | 382 +-------- .../polars/tests/it/io/parquet/read/binary.rs | 9 +- crates/polars/tests/it/io/parquet/read/mod.rs | 30 +- .../tests/it/io/parquet/read/primitive.rs | 6 +- py-polars/tests/unit/io/test_parquet.py | 21 + 15 files changed, 338 insertions(+), 2270 deletions(-) delete mode 100644 crates/polars-parquet/src/parquet/encoding/hybrid_rle/buffered.rs delete mode 100644 crates/polars-parquet/src/parquet/encoding/hybrid_rle/fuzz.rs delete mode 100644 crates/polars-parquet/src/parquet/encoding/hybrid_rle/gatherer.rs diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index 5881a94c3bc1..c1a6a5b75507 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -1,15 +1,14 @@ -use arrow::bitmap::MutableBitmap; +use arrow::bitmap::utils::BitmapIter; +use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; -use polars_error::PolarsResult; use super::utils::{self, BatchableCollector}; use super::{BasicDecompressor, Filter}; -use crate::parquet::encoding::hybrid_rle::gatherer::HybridRleGatherer; -use crate::parquet::encoding::hybrid_rle::HybridRleDecoder; +use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage}; use crate::parquet::read::levels::get_bit_width; -use crate::read::deserialize::utils::{hybrid_rle_count_zeros, BatchedCollector, State}; +use crate::read::deserialize::utils::BatchedCollector; #[derive(Debug)] pub struct Nested { @@ -321,410 +320,183 @@ impl NestedState { } } -/// Calculate the number of leaf values that are covered by the first `limit` definition level -/// values. -fn limit_to_num_values( - def_iter: &HybridRleDecoder<'_>, - def_levels: &[u16], - limit: usize, -) -> ParquetResult { - struct NumValuesGatherer { - leaf_def_level: u16, - } - struct NumValuesState { - num_values: usize, - length: usize, - } - - impl HybridRleGatherer for NumValuesGatherer { - type Target = NumValuesState; - - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.length - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(value) - } +fn collect_level_values( + target: &mut Vec, + hybrid_rle: HybridRleDecoder<'_>, +) -> ParquetResult<()> { + target.reserve(hybrid_rle.len()); - fn gather_one(&self, target: &mut Self::Target, value: u32) -> ParquetResult<()> { - target.num_values += usize::from(value == self.leaf_def_level as u32); - target.length += 1; - Ok(()) - } + for chunk in hybrid_rle.into_chunk_iter() { + let chunk = chunk?; - fn gather_repeated( - &self, - target: &mut Self::Target, - value: u32, - n: usize, - ) -> ParquetResult<()> { - target.num_values += n * usize::from(value == self.leaf_def_level as u32); - target.length += n; - Ok(()) + match chunk { + HybridRleChunk::Rle(value, size) => { + target.resize(target.len() + size, value as u16); + }, + HybridRleChunk::Bitpacked(decoder) => { + decoder.lower_element::()?.collect_into(target); + }, } } - let mut state = NumValuesState { - num_values: 0, - length: 0, - }; - def_iter.clone().gather_n_into( - &mut state, - limit, - &NumValuesGatherer { - leaf_def_level: *def_levels.last().unwrap(), - }, - )?; - - Ok(state.num_values) + Ok(()) } -fn idx_to_limit(rep_iter: &HybridRleDecoder<'_>, idx: usize) -> ParquetResult { - struct RowIdxOffsetGatherer; - struct RowIdxOffsetState { - num_elements_seen: usize, - top_level_limit: usize, - found: Option, - } - - impl HybridRleGatherer for RowIdxOffsetGatherer { - type Target = RowIdxOffsetState; - - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.num_elements_seen - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(value == 0) - } - - fn gather_one(&self, target: &mut Self::Target, value: bool) -> ParquetResult<()> { - let idx = target.num_elements_seen; - target.num_elements_seen += 1; - - if !value || target.found.is_some() { - return Ok(()); - } - - if target.top_level_limit > 0 { - target.top_level_limit -= 1; - return Ok(()); - } - - target.found = Some(idx); - - Ok(()) - } - - fn gather_repeated( - &self, - target: &mut Self::Target, - value: bool, - n: usize, - ) -> ParquetResult<()> { - let idx = target.num_elements_seen; - target.num_elements_seen += n; - - if !value || target.found.is_some() { - return Ok(()); - } - - if target.top_level_limit >= n { - target.top_level_limit -= n; - return Ok(()); - } - - target.found = Some(idx + target.top_level_limit); - target.top_level_limit = 0; - - Ok(()) - } - - // @TODO: Add specialization for other methods - } - - let mut state = RowIdxOffsetState { - num_elements_seen: 0, - top_level_limit: idx, - found: None, - }; - - const ROW_IDX_BATCH_SIZE: usize = 1024; - - let mut row_idx_iter = rep_iter.clone(); - while row_idx_iter.len() > 0 && state.found.is_none() { - row_idx_iter.gather_n_into(&mut state, ROW_IDX_BATCH_SIZE, &RowIdxOffsetGatherer)?; - } - - Ok(state.found.unwrap_or(rep_iter.len())) +struct DecodingState { + num_skips: Option, + num_collects: Option, } #[allow(clippy::too_many_arguments)] -fn extend_offsets2<'a>( - mut def_iter: HybridRleDecoder<'a>, - mut rep_iter: HybridRleDecoder<'a>, - batched_collector: &mut BatchedCollector<'_, (), (), BatchedNestedDecoder<'a>>, - nested: &mut [Nested], - filter: Option, - - def_levels: &[u16], - rep_levels: &[u16], -) -> PolarsResult<()> { - debug_assert_eq!(def_iter.len(), rep_iter.len()); - - match filter { - None => { - let limit = def_iter.len(); - - extend_offsets_limited( - &mut def_iter, - &mut rep_iter, - batched_collector, - nested, - limit, - def_levels, - rep_levels, - )?; - - debug_assert_eq!(def_iter.len(), rep_iter.len()); - debug_assert_eq!(def_iter.len(), 0); +fn decode_nested( + mut current_def_levels: &[u16], + mut current_rep_levels: &[u16], - Ok(()) - }, - Some(Filter::Range(range)) => { - let start = range.start; - let end = range.end; - - if start > 0 { - let start_cell = idx_to_limit(&rep_iter, start)?; - - let num_skipped_values = limit_to_num_values(&def_iter, def_levels, start_cell)?; - batched_collector.skip_in_place(num_skipped_values)?; - - rep_iter.skip_in_place(start_cell)?; - def_iter.skip_in_place(start_cell)?; - } - - if end - start > 0 { - let limit = idx_to_limit(&rep_iter, end - start)?; - - extend_offsets_limited( - &mut def_iter, - &mut rep_iter, - batched_collector, - nested, - limit, - def_levels, - rep_levels, - )?; - } - - // @NOTE: This is kind of unused - let last_skip = def_iter.len(); - let num_skipped_values = limit_to_num_values(&def_iter, def_levels, last_skip)?; - batched_collector.skip_in_place(num_skipped_values)?; - rep_iter.skip_in_place(last_skip)?; - def_iter.skip_in_place(last_skip)?; - - Ok(()) - }, - Some(Filter::Mask(bitmap)) => { - let mut iter = bitmap.iter(); - while iter.num_remaining() > 0 { - let num_zeros = iter.take_leading_zeros(); - if num_zeros > 0 { - let offset = idx_to_limit(&rep_iter, num_zeros)?; - let num_skipped_values = limit_to_num_values(&def_iter, def_levels, offset)?; - batched_collector.skip_in_place(num_skipped_values)?; - rep_iter.skip_in_place(offset)?; - def_iter.skip_in_place(offset)?; - } - - let num_ones = iter.take_leading_ones(); - if num_ones > 0 { - let limit = idx_to_limit(&rep_iter, num_ones)?; - extend_offsets_limited( - &mut def_iter, - &mut rep_iter, - batched_collector, - nested, - limit, - def_levels, - rep_levels, - )?; - } - } + batched_collector: &mut BatchedCollector<'_, (), (), BatchedNestedDecoder<'_>>, + nested: &mut [Nested], - Ok(()) - }, - } -} + state: &mut DecodingState, + top_level_filter: &mut BitmapIter<'_>, -fn extend_offsets_limited<'a>( - def_iter: &mut HybridRleDecoder<'a>, - rep_iter: &mut HybridRleDecoder<'a>, - batched_collector: &mut BatchedCollector<'_, (), (), BatchedNestedDecoder<'a>>, - nested: &mut [Nested], - mut limit: usize, // Amortized allocations def_levels: &[u16], rep_levels: &[u16], -) -> PolarsResult<()> { - #[derive(Default)] - struct LevelGatherer<'a>(std::marker::PhantomData<&'a ()>); - struct LevelGathererState<'a> { - offset: usize, - slice: &'a mut [u16], - } - - impl<'a> HybridRleGatherer for LevelGatherer<'a> { - type Target = LevelGathererState<'a>; - - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.offset - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - debug_assert!(value <= u16::MAX as u32); - Ok(value as u16) - } - - fn gather_one(&self, target: &mut Self::Target, value: u16) -> ParquetResult<()> { - debug_assert!(target.offset < target.slice.len()); - - target.slice[target.offset] = value; - target.offset += 1; - - Ok(()) - } +) -> ParquetResult<()> { + let max_depth = nested.len(); + let leaf_def_level = *def_levels.last().unwrap(); + + while !current_def_levels.is_empty() { + debug_assert_eq!(current_def_levels.len(), current_rep_levels.len()); + + // Handle skips + if let Some(ref mut num_skips) = state.num_skips { + let mut i = 0; + let mut num_skipped_values = 0; + while i < current_def_levels.len() && (*num_skips > 0 || current_rep_levels[i] != 0) { + let def = current_def_levels[i]; + let rep = current_rep_levels[i]; + + *num_skips -= usize::from(rep == 0); + i += 1; + + // @NOTE: + // We don't need to account for higher def-levels that imply extra values, since we + // don't have those higher levels either. + num_skipped_values += usize::from(def == leaf_def_level); + } + batched_collector.skip_in_place(num_skipped_values)?; - fn gather_repeated( - &self, - target: &mut Self::Target, - value: u16, - n: usize, - ) -> ParquetResult<()> { - debug_assert!(target.offset + n <= target.slice.len()); + current_def_levels = ¤t_def_levels[i..]; + current_rep_levels = ¤t_rep_levels[i..]; - for i in 0..n { - target.slice[target.offset + i] = value; + if current_def_levels.is_empty() { + break; + } else { + state.num_skips = None; } - target.offset += n; - - Ok(()) } - // @TODO: Add specialization for other methods - } - - let mut def_values = [0u16; DECODE_BATCH_SIZE]; - let mut rep_values = [0u16; DECODE_BATCH_SIZE]; - - let max_depth = nested.len(); - - const DECODE_BATCH_SIZE: usize = 1024; - while def_iter.len() > 0 && limit > 0 { - let additional = usize::min(limit, DECODE_BATCH_SIZE); - - let mut def_state = LevelGathererState { - offset: 0, - slice: &mut def_values, - }; - let mut rep_state = LevelGathererState { - offset: 0, - slice: &mut rep_values, - }; + // Handle collects + if let Some(ref mut num_collects) = state.num_collects { + let mut i = 0; + while i < current_def_levels.len() && (*num_collects > 0 || current_rep_levels[i] != 0) + { + let def = current_def_levels[i]; + let rep = current_rep_levels[i]; - def_iter.gather_n_into(&mut def_state, additional, &LevelGatherer::default())?; - rep_iter.gather_n_into(&mut rep_state, additional, &LevelGatherer::default())?; + *num_collects -= usize::from(rep == 0); + i += 1; - debug_assert_eq!(def_state.offset, rep_state.offset); - debug_assert_eq!(def_state.offset, additional); + let mut is_required = false; - for i in 0..additional { - let def = def_values[i]; - let rep = rep_values[i]; + for depth in 0..max_depth { + // Defines whether this element is defined at `depth` + // + // e.g. [ [ [ 1 ] ] ] is defined at [ ... ], [ [ ... ] ], [ [ [ ... ] ] ] and + // [ [ [ 1 ] ] ]. + let is_defined_at_this_depth = + rep <= rep_levels[depth] && def >= def_levels[depth]; - let mut is_required = false; + let length = nested + .get(depth + 1) + .map(|x| x.len() as i64) + // the last depth is the leaf, which is always increased by 1 + .unwrap_or(1); - for depth in 0..max_depth { - // Defines whether this element is defined at `depth` - // - // e.g. [ [ [ 1 ] ] ] is defined at [ ... ], [ [ ... ] ], [ [ [ ... ] ] ] and - // [ [ [ 1 ] ] ]. - let is_defined_at_this_depth = rep <= rep_levels[depth] && def >= def_levels[depth]; + let nest = &mut nested[depth]; - let length = nested - .get(depth + 1) - .map(|x| x.len() as i64) - // the last depth is the leaf, which is always increased by 1 - .unwrap_or(1); + let is_valid = !nest.is_nullable() || def > def_levels[depth]; - let nest = &mut nested[depth]; + if is_defined_at_this_depth && !is_valid { + let mut num_elements = 1; - let is_valid = !nest.is_nullable() || def > def_levels[depth]; + nest.push(length, is_valid); - if is_defined_at_this_depth && !is_valid { - let mut num_elements = 1; + for embed_depth in depth..max_depth { + let embed_length = nested + .get(embed_depth + 1) + .map(|x| x.len() as i64) + // the last depth is the leaf, which is always increased by 1 + .unwrap_or(1); - nest.push(length, is_valid); + let embed_nest = &mut nested[embed_depth]; - for embed_depth in depth..max_depth { - let embed_length = nested - .get(embed_depth + 1) - .map(|x| x.len() as i64) - // the last depth is the leaf, which is always increased by 1 - .unwrap_or(1); + if embed_depth > depth { + for _ in 0..num_elements { + embed_nest.push_default(embed_length); + } + } - let embed_nest = &mut nested[embed_depth]; + let embed_num_values = embed_nest.invalid_num_values(); + num_elements *= embed_num_values; - if embed_depth > depth { - for _ in 0..num_elements { - embed_nest.push_default(embed_length); + if embed_num_values == 0 { + break; } } - let embed_num_values = embed_nest.invalid_num_values(); - num_elements *= embed_num_values; + batched_collector.push_n_invalids(num_elements); - if embed_num_values == 0 { - break; - } + break; } - batched_collector.push_n_invalids(num_elements); + if is_required || is_defined_at_this_depth { + nest.push(length, is_valid); - break; - } - - if is_required || is_defined_at_this_depth { - nest.push(length, is_valid); + if depth == max_depth - 1 { + // the leaf / primitive + let is_valid = (def != def_levels[depth]) || !nest.is_nullable(); - if depth == max_depth - 1 { - // the leaf / primitive - let is_valid = (def != def_levels[depth]) || !nest.is_nullable(); - - if is_valid { - batched_collector.push_valid()?; - } else { - batched_collector.push_invalid(); + if is_valid { + batched_collector.push_valid()?; + } else { + batched_collector.push_invalid(); + } } } + + is_required = (is_required || is_defined_at_this_depth) + && nest.is_required() + && !is_valid; } + } + + current_def_levels = ¤t_def_levels[i..]; + current_rep_levels = ¤t_rep_levels[i..]; - is_required = - (is_required || is_defined_at_this_depth) && nest.is_required() && !is_valid; + if current_def_levels.is_empty() { + break; + } else { + state.num_collects = None; } } - limit -= additional; + if top_level_filter.num_remaining() == 0 { + break; + } + + state.num_skips = Some(top_level_filter.take_leading_zeros()).filter(|v| *v != 0); + state.num_collects = Some(top_level_filter.take_leading_ones()).filter(|v| *v != 0); } Ok(()) @@ -785,236 +557,94 @@ impl PageNestedDecoder { // Amortize the allocations. let (def_levels, rep_levels) = nested_state.levels(); - match filter { - None => { - loop { - let Some(page) = self.iter.next() else { - break; - }; - let page = page?; - let page = page.decompress(&mut self.iter)?; - - let mut filter = MutableBitmap::new(); - let mut validity = MutableBitmap::new(); - let mut _tgt = (); - - let (def_iter, rep_iter) = level_iters(&page)?; - - // @TODO: move this to outside the loop. - let mut batched_collector = BatchedCollector::new( - BatchedNestedDecoder { - filter: &mut filter, - validity: &mut validity, - }, - &mut _tgt, - ); - - extend_offsets2( - def_iter, - rep_iter, - &mut batched_collector, - &mut nested_state.nested, - None, - &def_levels, - &rep_levels, - )?; - - batched_collector.finalize()?; - - let state = utils::State::new_nested( - &self.decoder, - &page, - self.dict.as_ref(), - Some(validity.freeze()), - )?; - state.decode( - &mut self.decoder, - &mut target, - Some(Filter::Mask(filter.freeze())), - )?; - - self.iter.reuse_page_buffer(page); - } - }, - Some(mut filter) => { - enum PageStartAction { - Skip, - Collect, - } - - // We may have an action (skip / collect) for one row value left over from the - // previous page. Every page may state what the next page needs to do until the - // first of its own row values (rep_lvl = 0). - let mut last_row_value_action = PageStartAction::Skip; - let mut num_rows_remaining = filter.num_rows(); - - while num_rows_remaining > 0 - || matches!(last_row_value_action, PageStartAction::Collect) - { - let Some(page) = self.iter.next() else { - break; - }; - let page = page?; - // We cannot lazily decompress because we don't have the number of row values - // at this point. We need repetition levels for that. *sigh*. In general, lazy - // decompression is quite difficult with nested values. - // - // @TODO - // Lazy decompression is quite doable in the V2 specification since that does - // not compress the repetition and definition levels. However, not a lot of - // people use the V2 specification. So let us ignore that for now. - let page = page.decompress(&mut self.iter)?; - - let (mut def_iter, mut rep_iter) = level_iters(&page)?; - - let state; - let mut batched_collector; - - let start_length = nested_state.len(); - - // rep lvl == 0 ==> row value - let num_row_values = hybrid_rle_count_zeros(&rep_iter)?; - - let state_filter; - (state_filter, filter) = Filter::split_at(&filter, num_row_values); + let mut current_def_levels = Vec::::new(); + let mut current_rep_levels = Vec::::new(); + + let (mut decode_state, top_level_filter) = match filter { + None => ( + DecodingState { + num_skips: None, + num_collects: Some(usize::MAX), + }, + Bitmap::new(), + ), + Some(Filter::Range(range)) => ( + DecodingState { + num_skips: Some(range.start), + num_collects: Some(range.len()), + }, + Bitmap::new(), + ), + Some(Filter::Mask(mask)) => ( + DecodingState { + num_skips: None, + num_collects: None, + }, + mask, + ), + }; - let mut leaf_filter = MutableBitmap::new(); - let mut leaf_validity = MutableBitmap::new(); - let mut _tgt = (); + let mut top_level_filter = top_level_filter.iter(); + + loop { + let Some(page) = self.iter.next() else { + break; + }; + let page = page?; + let page = page.decompress(&mut self.iter)?; + + let (mut def_iter, mut rep_iter) = level_iters(&page)?; + + let num_levels = def_iter.len().min(rep_iter.len()); + def_iter.limit_to(num_levels); + rep_iter.limit_to(num_levels); + + current_def_levels.clear(); + current_rep_levels.clear(); + + collect_level_values(&mut current_def_levels, def_iter)?; + collect_level_values(&mut current_rep_levels, rep_iter)?; + + let mut leaf_filter = MutableBitmap::new(); + let mut leaf_validity = MutableBitmap::new(); + + let mut _tgt = (); + + // @TODO: move this to outside the loop. + let mut batched_collector = BatchedCollector::new( + BatchedNestedDecoder { + filter: &mut leaf_filter, + validity: &mut leaf_validity, + }, + &mut _tgt, + ); + + decode_nested( + ¤t_def_levels, + ¤t_rep_levels, + &mut batched_collector, + &mut nested_state.nested, + &mut decode_state, + &mut top_level_filter, + &def_levels, + &rep_levels, + )?; - match last_row_value_action { - PageStartAction::Skip => { - // Fast path: skip the whole page. - // No new row values or we don't care about any of the row values. - if num_row_values == 0 && state_filter.num_rows() == 0 { - self.iter.reuse_page_buffer(page); - continue; - } + batched_collector.finalize()?; - let limit = idx_to_limit(&rep_iter, 0)?; - - // We just saw that we had at least one row value. - debug_assert!(limit < rep_iter.len()); - - batched_collector = BatchedCollector::new( - BatchedNestedDecoder { - filter: &mut leaf_filter, - validity: &mut leaf_validity, - }, - &mut _tgt, - ); - - let num_leaf_values = - limit_to_num_values(&def_iter, &def_levels, limit)?; - batched_collector.skip_in_place(num_leaf_values)?; - rep_iter.skip_in_place(limit)?; - def_iter.skip_in_place(limit)?; - }, - PageStartAction::Collect => { - let limit = if num_row_values == 0 { - rep_iter.len() - } else { - idx_to_limit(&rep_iter, 0)? - }; - - // Fast path: we are not interested in any of the row values in this - // page. - if limit == 0 && state_filter.num_rows() == 0 { - self.iter.reuse_page_buffer(page); - last_row_value_action = PageStartAction::Skip; - continue; - } - - batched_collector = BatchedCollector::new( - BatchedNestedDecoder { - filter: &mut leaf_filter, - validity: &mut leaf_validity, - }, - &mut _tgt, - ); - - extend_offsets_limited( - &mut def_iter, - &mut rep_iter, - &mut batched_collector, - &mut nested_state.nested, - limit, - &def_levels, - &rep_levels, - )?; - - // No new row values. Keep collecting. - if rep_iter.len() == 0 { - batched_collector.finalize()?; - - state = State::new_nested( - &self.decoder, - &page, - self.dict.as_ref(), - Some(leaf_validity.freeze()), - )?; - state.decode( - &mut self.decoder, - &mut target, - Some(Filter::Mask(leaf_filter.freeze())), - )?; - - let num_done = nested_state.len() - start_length; - debug_assert!(num_done <= num_rows_remaining); - debug_assert!(num_done <= num_row_values); - num_rows_remaining -= num_done; - - self.iter.reuse_page_buffer(page); - - continue; - } - }, - } + let state = utils::State::new_nested( + &self.decoder, + &page, + self.dict.as_ref(), + Some(leaf_validity.freeze()), + )?; + state.decode( + &mut self.decoder, + &mut target, + Some(Filter::Mask(leaf_filter.freeze())), + )?; - // Two cases: - // 1. First page: Must always start with a row value. - // 2. Other pages: If they did not have a row value, they would have been - // handled by the last_row_value_action. - debug_assert!(num_row_values > 0); - - last_row_value_action = if state_filter.do_include_at(num_row_values - 1) { - PageStartAction::Collect - } else { - PageStartAction::Skip - }; - - extend_offsets2( - def_iter, - rep_iter, - &mut batched_collector, - &mut nested_state.nested, - Some(state_filter), - &def_levels, - &rep_levels, - )?; - - batched_collector.finalize()?; - - state = State::new_nested( - &self.decoder, - &page, - self.dict.as_ref(), - Some(leaf_validity.freeze()), - )?; - state.decode( - &mut self.decoder, - &mut target, - Some(Filter::Mask(leaf_filter.freeze())), - )?; - - let num_done = nested_state.len() - start_length; - debug_assert!(num_done <= num_rows_remaining); - debug_assert!(num_done <= num_row_values); - num_rows_remaining -= num_done; - - self.iter.reuse_page_buffer(page); - } - }, + self.iter.reuse_page_buffer(page); } // we pop the primitive off here. diff --git a/crates/polars-parquet/src/arrow/read/deserialize/null.rs b/crates/polars-parquet/src/arrow/read/deserialize/null.rs index 22000905a54b..7ec884715318 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/null.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/null.rs @@ -11,9 +11,6 @@ use super::utils::filter::Filter; use crate::parquet::error::ParquetResult; use crate::parquet::page::{DataPage, DictPage}; -pub(crate) struct NullTranslation { - length: usize, -} pub(crate) struct NullDecoder; #[derive(Debug)] pub(crate) struct NullArrayLength { @@ -26,23 +23,21 @@ impl utils::ExactSize for NullArrayLength { } } -impl<'a> utils::StateTranslation<'a, NullDecoder> for NullTranslation { +impl<'a> utils::StateTranslation<'a, NullDecoder> for () { type PlainDecoder = (); fn new( _decoder: &NullDecoder, - page: &'a DataPage, + _page: &'a DataPage, _dict: Option<&'a ::Dict>, _page_validity: Option<&Bitmap>, ) -> ParquetResult { - Ok(Self { - length: page.num_values() - }) + Ok(()) } } impl utils::Decoder for NullDecoder { - type Translation<'a> = NullTranslation; + type Translation<'a> = (); type Dict = (); type DecodedState = NullArrayLength; type Output = NullArray; @@ -67,11 +62,13 @@ impl utils::Decoder for NullDecoder { fn extend_filtered_with_state( &mut self, - state: utils::State<'_, Self>, + _state: utils::State<'_, Self>, decoded: &mut Self::DecodedState, filter: Option, ) -> ParquetResult<()> { - decoded.length += Filter::opt_num_rows(&filter, state.translation.length); + // @NOTE: This is only used by nested decoders. Those will always supply a mask. + let filter = filter.unwrap(); + decoded.length += filter.num_rows(); Ok(()) } } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index 428f1fae06fe..c067a5f257b1 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -11,9 +11,6 @@ use arrow::pushable::Pushable; use self::filter::Filter; use super::BasicDecompressor; -use crate::parquet::encoding::hybrid_rle::gatherer::{ - HybridRleGatherer, ZeroCount, ZeroCountGatherer, -}; use crate::parquet::encoding::hybrid_rle::{self, HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::{ParquetError, ParquetResult}; use crate::parquet::page::{split_buffer, DataPage, DictPage}; @@ -350,78 +347,6 @@ pub(crate) fn unspecialized_decode( Ok(()) } -#[derive(Default)] -pub(crate) struct BatchGatherer<'a, I, T, C: BatchableCollector>( - std::marker::PhantomData<&'a (I, T, C)>, -); -impl<'a, I, T, C: BatchableCollector> HybridRleGatherer for BatchGatherer<'a, I, T, C> { - type Target = (&'a mut MutableBitmap, BatchedCollector<'a, I, T, C>); - - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.0.len() - } - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(value) - } - - fn gather_one(&self, (validity, values): &mut Self::Target, value: u32) -> ParquetResult<()> { - if value == 0 { - values.push_invalid(); - validity.extend_constant(1, false); - } else { - values.push_valid()?; - validity.extend_constant(1, true); - } - - Ok(()) - } - - fn gather_repeated( - &self, - (validity, values): &mut Self::Target, - value: u32, - n: usize, - ) -> ParquetResult<()> { - if value == 0 { - values.push_n_invalids(n); - validity.extend_constant(n, false); - } else { - values.push_n_valids(n)?; - validity.extend_constant(n, true); - } - - Ok(()) - } - - fn gather_slice(&self, target: &mut Self::Target, source: &[u32]) -> ParquetResult<()> { - let mut prev = 0u32; - let mut len = 0usize; - - for v in source { - let v = *v; - - if v == prev { - len += 1; - } else { - if len != 0 { - self.gather_repeated(target, prev, len)?; - } - prev = v; - len = 1; - } - } - - if len != 0 { - self.gather_repeated(target, prev, len)?; - } - - Ok(()) - } -} - impl, I: Iterator> BatchableCollector for I { #[inline] fn push_n(&mut self, target: &mut P, n: usize) -> ParquetResult<()> { @@ -599,16 +524,6 @@ pub fn freeze_validity(validity: MutableBitmap) -> Option { Some(validity) } -pub(crate) fn hybrid_rle_count_zeros( - decoder: &hybrid_rle::HybridRleDecoder<'_>, -) -> ParquetResult { - let mut count = ZeroCount::default(); - decoder - .clone() - .gather_into(&mut count, &ZeroCountGatherer)?; - Ok(count.num_zero) -} - pub(crate) fn filter_from_range(rng: Range) -> Bitmap { let mut bm = MutableBitmap::with_capacity(rng.end); diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs index f860989689e6..b8c24e229705 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/decode.rs @@ -111,6 +111,11 @@ impl<'a, T: Unpackable> Decoder<'a, T> { pub fn as_slice(&self) -> &[u8] { self.packed.as_slice() } + + pub fn lower_element(self) -> ParquetResult> { + let packed = self.packed.as_slice(); + Decoder::try_new(packed, self.num_bits, self.length) + } } /// A iterator over the exact chunks in a [`Decoder`]. diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs index 057bb5d6b7fe..10a1f61af963 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/mod.rs @@ -1,4 +1,20 @@ macro_rules! seq_macro { + ($i:ident in 1..15 $block:block) => { + seq_macro!($i in [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, + ] $block) + }; + ($i:ident in 0..16 $block:block) => { + seq_macro!($i in [ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, + ] $block) + }; + ($i:ident in 0..=16 $block:block) => { + seq_macro!($i in [ + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, + 16, + ] $block) + }; ($i:ident in 1..31 $block:block) => { seq_macro!($i in [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, @@ -155,10 +171,26 @@ impl Unpacked for [u64; 64] { pub trait Unpackable: Copy + Sized + Default { type Packed: Packed; type Unpacked: Unpacked; + fn unpack(packed: &[u8], num_bits: usize, unpacked: &mut Self::Unpacked); fn pack(unpacked: &Self::Unpacked, num_bits: usize, packed: &mut [u8]); } +impl Unpackable for u16 { + type Packed = [u8; 16 * 2]; + type Unpacked = [u16; 16]; + + #[inline] + fn unpack(packed: &[u8], num_bits: usize, unpacked: &mut Self::Unpacked) { + unpack::unpack16(packed, unpacked, num_bits) + } + + #[inline] + fn pack(packed: &Self::Unpacked, num_bits: usize, unpacked: &mut [u8]) { + pack::pack16(packed, unpacked, num_bits) + } +} + impl Unpackable for u32 { type Packed = [u8; 32 * 4]; type Unpacked = [u32; 32]; diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/pack.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/pack.rs index c318f42649d3..349d0f34ee87 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/pack.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/pack.rs @@ -81,6 +81,7 @@ macro_rules! pack { }; } +pack!(pack16, u16, 2, 16, 15); pack!(pack32, u32, 4, 32, 31); pack!(pack64, u64, 8, 64, 63); diff --git a/crates/polars-parquet/src/parquet/encoding/bitpacked/unpack.rs b/crates/polars-parquet/src/parquet/encoding/bitpacked/unpack.rs index 61d59925a39e..c52c17d21681 100644 --- a/crates/polars-parquet/src/parquet/encoding/bitpacked/unpack.rs +++ b/crates/polars-parquet/src/parquet/encoding/bitpacked/unpack.rs @@ -116,6 +116,7 @@ macro_rules! unpack { }; } +unpack!(unpack16, u16, 2, 16); unpack!(unpack32, u32, 4, 32); unpack!(unpack64, u64, 8, 64); diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/buffered.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/buffered.rs deleted file mode 100644 index 824638d253ad..000000000000 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/buffered.rs +++ /dev/null @@ -1,280 +0,0 @@ -use super::gatherer::HybridRleGatherer; -use crate::parquet::encoding::bitpacked::{self, Unpackable, Unpacked}; -use crate::parquet::error::ParquetResult; - -#[derive(Debug, Clone)] -pub struct BufferedBitpacked<'a> { - pub unpacked: [u32; 32], - pub unpacked_start: usize, - pub unpacked_end: usize, - - pub decoder: bitpacked::Decoder<'a, u32>, -} - -#[derive(Debug, Clone)] -pub struct BufferedRle { - pub value: u32, - pub length: usize, -} - -/// A buffered set of items for the [`HybridRleDecoder`]. This can be iterated over and stopped at -/// any time. -#[derive(Debug, Clone)] -pub enum HybridRleBuffered<'a> { - Bitpacked(BufferedBitpacked<'a>), - Rle(BufferedRle), -} - -impl Iterator for BufferedRle { - type Item = u32; - - fn next(&mut self) -> Option { - if self.length > 0 { - self.length -= 1; - Some(self.value) - } else { - None - } - } - - fn size_hint(&self) -> (usize, Option) { - (self.length, Some(self.length)) - } -} - -impl ExactSizeIterator for BufferedRle {} - -impl<'a> Iterator for BufferedBitpacked<'a> { - type Item = u32; - - fn next(&mut self) -> Option { - if self.unpacked_start < self.unpacked_end { - let value = self.unpacked[self.unpacked_start]; - self.unpacked_start += 1; - return Some(value); - } - - self.decoder - .chunked() - .next_inexact() - .map(|(unpacked, unpacked_length)| { - debug_assert!(unpacked_length > 0); - let value = unpacked[0]; - self.unpacked = unpacked; - self.unpacked_end = unpacked_length; - self.unpacked_start = 1; - value - }) - } - - fn size_hint(&self) -> (usize, Option) { - let unpacked_num_elements = self.unpacked_end - self.unpacked_start; - let exact = unpacked_num_elements + self.decoder.len(); - (exact, Some(exact)) - } -} - -impl<'a> ExactSizeIterator for BufferedBitpacked<'a> {} - -impl<'a> Iterator for HybridRleBuffered<'a> { - type Item = u32; - - fn next(&mut self) -> Option { - match self { - HybridRleBuffered::Bitpacked(b) => b.next(), - HybridRleBuffered::Rle(b) => b.next(), - } - } - - fn size_hint(&self) -> (usize, Option) { - match self { - HybridRleBuffered::Bitpacked(b) => b.size_hint(), - HybridRleBuffered::Rle(b) => b.size_hint(), - } - } -} - -impl<'a> ExactSizeIterator for HybridRleBuffered<'a> {} - -impl<'a> BufferedBitpacked<'a> { - fn gather_limited_into>( - &mut self, - target: &mut G::Target, - limit: usize, - gatherer: &G, - ) -> ParquetResult { - let unpacked_num_elements = self.unpacked_end - self.unpacked_start; - if limit <= unpacked_num_elements { - gatherer.gather_slice( - target, - &self.unpacked[self.unpacked_start..self.unpacked_start + limit], - )?; - self.unpacked_start += limit; - return Ok(limit); - } - - gatherer.gather_slice( - target, - &self.unpacked[self.unpacked_start..self.unpacked_end], - )?; - self.unpacked_end = 0; - self.unpacked_start = 0; - let limit = limit - unpacked_num_elements; - - let decoder = self.decoder.take(); - let decoder_len = decoder.len(); - if limit >= decoder_len { - gatherer.gather_bitpacked_all(target, decoder)?; - Ok(unpacked_num_elements + decoder_len) - } else { - let buffered = gatherer.gather_bitpacked_limited(target, decoder, limit)?; - *self = buffered; - Ok(unpacked_num_elements + limit) - } - } - - pub fn gather_into>( - self, - target: &mut G::Target, - gatherer: &G, - ) -> ParquetResult { - let unpacked_num_elements = self.unpacked_end - self.unpacked_start; - gatherer.gather_slice( - target, - &self.unpacked[self.unpacked_start..self.unpacked_end], - )?; - let decoder_len = self.decoder.len(); - gatherer.gather_bitpacked_all(target, self.decoder)?; - Ok(unpacked_num_elements + decoder_len) - } - - pub fn skip_in_place(&mut self, n: usize) -> usize { - let unpacked_num_elements = self.unpacked_end - self.unpacked_start; - - if n < unpacked_num_elements { - self.unpacked_start += n; - return n; - } - - let n = n - unpacked_num_elements; - - if self.decoder.len() > n { - let num_chunks = n / ::Unpacked::LENGTH; - let unpacked_offset = n % ::Unpacked::LENGTH; - self.decoder.skip_chunks(num_chunks); - let (unpacked, unpacked_length) = self.decoder.chunked().next_inexact().unwrap(); - debug_assert!(unpacked_offset < unpacked_length); - - self.unpacked = unpacked; - self.unpacked_start = unpacked_offset; - self.unpacked_end = unpacked_length; - - return unpacked_num_elements + n; - } - - // We skip the entire decoder. Essentially, just zero it out. - let decoder = self.decoder.take(); - self.unpacked_start = 0; - self.unpacked_end = 0; - - decoder.len() + unpacked_num_elements - } -} - -impl BufferedRle { - pub fn gather_limited_into>( - &mut self, - target: &mut G::Target, - limit: usize, - gatherer: &G, - ) -> ParquetResult { - let value = gatherer.hybridrle_to_target(self.value)?; - let num_elements = usize::min(self.length, limit); - self.length -= num_elements; - gatherer.gather_repeated(target, value, num_elements)?; - Ok(num_elements) - } - - pub fn gather_into>( - self, - target: &mut A::Target, - applicator: &A, - ) -> ParquetResult { - let value = applicator.hybridrle_to_target(self.value)?; - applicator.gather_repeated(target, value, self.length)?; - Ok(self.length) - } - - pub fn skip_in_place(&mut self, n: usize) -> usize { - let num_elements = usize::min(self.length, n); - self.length -= num_elements; - num_elements - } -} - -impl<'a> HybridRleBuffered<'a> { - pub fn gather_limited_into>( - &mut self, - target: &mut G::Target, - limit: usize, - gatherer: &G, - ) -> ParquetResult { - let start_target_length = gatherer.target_num_elements(target); - let start_length = self.len(); - - let num_processed = match self { - HybridRleBuffered::Bitpacked(b) => b.gather_limited_into(target, limit, gatherer), - HybridRleBuffered::Rle(b) => b.gather_limited_into(target, limit, gatherer), - }?; - - debug_assert!(num_processed <= limit); - debug_assert_eq!( - num_processed, - gatherer.target_num_elements(target) - start_target_length - ); - debug_assert_eq!(num_processed, start_length - self.len()); - - Ok(num_processed) - } - - pub fn gather_into>( - self, - target: &mut G::Target, - gatherer: &G, - ) -> ParquetResult { - let start_target_length = gatherer.target_num_elements(target); - let start_length = self.len(); - - let num_processed = match self { - HybridRleBuffered::Bitpacked(b) => b.gather_into(target, gatherer), - HybridRleBuffered::Rle(b) => b.gather_into(target, gatherer), - }?; - - debug_assert_eq!( - num_processed, - gatherer.target_num_elements(target) - start_target_length - ); - debug_assert_eq!(num_processed, start_length); - - Ok(num_processed) - } - - pub fn skip_in_place(&mut self, n: usize) -> usize { - let start_length = self.len(); - - let num_skipped = match self { - HybridRleBuffered::Bitpacked(b) => b.skip_in_place(n), - HybridRleBuffered::Rle(b) => b.skip_in_place(n), - }; - - debug_assert!(num_skipped <= n); - debug_assert_eq!( - num_skipped, - start_length - self.len(), - "{self:?}: {num_skipped} != {start_length} - {}", - self.len() - ); - - num_skipped - } -} diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/fuzz.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/fuzz.rs deleted file mode 100644 index f4a980fb5062..000000000000 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/fuzz.rs +++ /dev/null @@ -1,390 +0,0 @@ -/// Since the HybridRle decoder is very widely used within the Parquet reader and the code is quite -/// complex to facilitate performance. We create this small fuzzer -use std::collections::VecDeque; - -use rand::Rng; - -use super::*; - -fn run_iteration( - bs: &[u32], - collects: impl Iterator, - encoded: &mut Vec, - decoded: &mut Vec, - num_bits: u32, -) -> ParquetResult<()> { - encoded.clear(); - decoded.clear(); - - encoder::encode(encoded, bs.iter().copied(), num_bits).unwrap(); - - let mut decoder = HybridRleDecoder::new(&encoded[..], num_bits, bs.len()); - - for c in collects { - decoder.collect_n_into(decoded, c)?; - } - - Ok(()) -} - -/// Minimizes a failing case -fn minimize_failing_case( - bs: &mut Vec, - collects: &mut VecDeque, - encoded: &mut Vec, - decoded: &mut Vec, - num_bits: u32, -) -> ParquetResult<()> { - loop { - let initial_bs_len = bs.len(); - let initial_collects_len = collects.len(); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - assert_ne!(&bs, &decoded); - - while collects.len() > 2 { - let last = collects.pop_back().unwrap(); - - *collects.back_mut().unwrap() += last; - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - *collects.back_mut().unwrap() -= last; - collects.push_back(last); - break; - } - } - - while collects.len() > 2 { - let first = collects.pop_front().unwrap(); - - *collects.front_mut().unwrap() += first; - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - *collects.front_mut().unwrap() -= first; - collects.push_front(first); - break; - } - } - - while bs.len() > 1 { - let last = bs.pop().unwrap(); - *collects.back_mut().unwrap() -= 1; - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - bs.push(last); - *collects.back_mut().unwrap() += 1; - break; - } - - if *collects.back().unwrap() == 0 { - collects.pop_back().unwrap(); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - collects.push_back(0); - break; - } - } - } - - while bs.len() > 1 { - let last = bs.pop().unwrap(); - *collects.front_mut().unwrap() -= 1; - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - bs.push(last); - *collects.front_mut().unwrap() += 1; - break; - } - - if *collects.front().unwrap() == 0 { - collects.pop_front().unwrap(); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - collects.push_front(0); - break; - } - } - } - - while bs.len() > 1 { - let first = bs.remove(0); - *collects.back_mut().unwrap() -= 1; - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - bs.insert(0, first); - *collects.back_mut().unwrap() += 1; - break; - } - - if *collects.back().unwrap() == 0 { - collects.pop_back().unwrap(); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - collects.push_back(0); - break; - } - } - } - - while bs.len() > 1 { - let first = bs.remove(0); - *collects.front_mut().unwrap() -= 1; - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - bs.insert(0, first); - *collects.front_mut().unwrap() += 1; - break; - } - - if *collects.front().unwrap() == 0 { - collects.pop_front().unwrap(); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - collects.push_front(0); - break; - } - } - } - - let mut start_offset = collects[0]; - for i in 1..collects.len() - 1 { - loop { - let start_length = collects[i]; - - while collects[i] > 0 { - collects[i] -= 1; - let item = bs.remove(start_offset); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - bs.insert(start_offset, item); - collects[i] += 1; - break; - } - - if collects[i] == 0 { - collects.remove(i); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - collects.insert(i, 0); - break; - } - } - } - - while collects[i] > 0 { - collects[i] -= 1; - let end_offset = start_offset + collects[i] - 1; - let item = bs.remove(end_offset); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - bs.insert(end_offset, item); - collects[i] += 1; - break; - } - - if collects[i] == 0 { - collects.remove(i); - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - if bs == decoded { - collects.insert(i, 0); - break; - } - } - } - - if collects[i] == start_length { - break; - } - } - - start_offset += collects[i]; - } - - let now_bs_len = bs.len(); - let now_collects_len = collects.len(); - - if initial_bs_len == now_bs_len && initial_collects_len == now_collects_len { - break; - } - } - - run_iteration(bs, collects.iter().copied(), encoded, decoded, num_bits)?; - - Ok(()) -} - -fn fuzz_loops(num_loops: usize) -> ParquetResult<()> { - let mut rng = rand::thread_rng(); - - const MAX_LENGTH: usize = 10_000; - - let mut encoded = Vec::with_capacity(1024); - let mut decoded = Vec::with_capacity(1024); - - let mut bs = Vec::with_capacity(MAX_LENGTH); - let mut collects: VecDeque = VecDeque::with_capacity(2000); - - for i in 0..num_loops { - collects.clear(); - bs.clear(); - - let num_bits = rng.gen_range(0..=32); - let mask = 1u32.wrapping_shl(num_bits).wrapping_sub(1); - - let length = rng.gen_range(1..=MAX_LENGTH); - - unsafe { bs.set_len(length) }; - rng.fill(&mut bs[..]); - - let mut filled = 0; - while filled < bs.len() { - if rng.gen() { - let num_repeats = rng.gen_range(0..=(bs.len() - filled)); - let value = bs[filled] & mask; - for j in 0..num_repeats { - bs[filled + j] = value; - } - filled += num_repeats; - } else { - bs[filled] &= mask; - filled += 1; - } - } - - if rng.gen() { - let mut num_values = bs.len(); - while num_values > 0 { - let n = rng.gen_range(0..=num_values); - collects.push_back(n); - num_values -= n; - } - } else { - collects.resize(1, bs.len()); - } - - run_iteration( - &bs, - collects.iter().copied(), - &mut encoded, - &mut decoded, - num_bits, - )?; - - if decoded != bs { - minimize_failing_case(&mut bs, &mut collects, &mut encoded, &mut decoded, num_bits)?; - - eprintln!("Minimized case:"); - eprintln!("Expected: {bs:?}"); - eprintln!("Found: {decoded:?}"); - eprintln!("Collects: {collects:?}"); - eprintln!(); - - panic!("Found a failing case..."); - } - - if i % 512 == 0 { - eprintln!("{i} iterations done."); - } - } - - Ok(()) -} - -#[test] -fn small_fuzz() -> ParquetResult<()> { - fuzz_loops(2048) -} - -#[test] -#[ignore = "Large fuzz test. Too slow"] -fn large_fuzz() -> ParquetResult<()> { - fuzz_loops(1_000_000) -} - -#[test] -#[ignore = "Large fuzz test. Too slow"] -fn skip_fuzz() -> ParquetResult<()> { - let mut rng = rand::thread_rng(); - - const MAX_LENGTH: usize = 10_000; - - let mut encoded = Vec::with_capacity(10000); - - let mut bs: Vec = Vec::with_capacity(MAX_LENGTH); - let mut skips: VecDeque = VecDeque::with_capacity(2000); - - let num_loops = 100_000; - - for _ in 0..num_loops { - skips.clear(); - bs.clear(); - - let num_bits = rng.gen_range(0..=32); - let mask = 1u32.wrapping_shl(num_bits).wrapping_sub(1); - - let length = rng.gen_range(1..=MAX_LENGTH); - - unsafe { bs.set_len(length) }; - rng.fill(&mut bs[..]); - - let mut filled = 0; - while filled < bs.len() { - if rng.gen() { - let num_repeats = rng.gen_range(0..=(bs.len() - filled)); - let value = bs[filled] & mask; - for j in 0..num_repeats { - bs[filled + j] = value; - } - filled += num_repeats; - } else { - bs[filled] &= mask; - filled += 1; - } - } - - let mut num_done = 0; - while num_done < filled { - let num_skip = rng.gen_range(1..=filled - num_done); - num_done += num_skip; - skips.push_back(num_skip); - } - - encoder::encode(&mut encoded, bs.iter().copied(), num_bits).unwrap(); - let mut decoder = HybridRleDecoder::new(&encoded, num_bits, filled); - - for s in &skips { - decoder.skip_in_place(*s).unwrap(); - } - } - - Ok(()) -} diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/gatherer.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/gatherer.rs deleted file mode 100644 index 1548f6e50a02..000000000000 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/gatherer.rs +++ /dev/null @@ -1,545 +0,0 @@ -use crate::parquet::encoding::bitpacked::{Decoder, Unpackable, Unpacked}; -use crate::parquet::encoding::hybrid_rle::{BufferedBitpacked, HybridRleBuffered}; -use crate::parquet::error::{ParquetError, ParquetResult}; - -/// Trait that describes what to do with consumed Hybrid-RLE Encoded values. -/// -/// This is quite a general trait that provides a lot of open space as to how to handle the -/// Hybrid-RLE encoded values. There is also the [`Translator`] trait that is usually good enough -/// if you want just want to map values to another set of values and collect them into a vector. -/// -/// Although, this trait might seem quite over-engineered (it might be), it is very useful for -/// performance. This allows for usage of the properties that [`HybridRleDecoder`] provides and for -/// definition of efficient procedures for collecting slices, chunks, repeated elements and -/// bit-packed elements. -/// -/// The [`Translator`] doc-comment has a good description of why this trait is needed. -/// -/// [`HybridRleDecoder`]: super::HybridRleDecoder -pub trait HybridRleGatherer { - type Target; - - fn target_reserve(&self, target: &mut Self::Target, n: usize); - fn target_num_elements(&self, target: &Self::Target) -> usize; - - fn hybridrle_to_target(&self, value: u32) -> ParquetResult; - fn gather_one(&self, target: &mut Self::Target, value: O) -> ParquetResult<()>; - fn gather_repeated(&self, target: &mut Self::Target, value: O, n: usize) -> ParquetResult<()>; - fn gather_slice(&self, target: &mut Self::Target, source: &[u32]) -> ParquetResult<()> { - self.target_reserve(target, source.len()); - for v in source { - self.gather_one(target, self.hybridrle_to_target(*v)?)?; - } - Ok(()) - } - fn gather_chunk( - &self, - target: &mut Self::Target, - source: &::Unpacked, - ) -> ParquetResult<()> { - self.gather_slice(target, source) - } - fn gather_bitpacked_all( - &self, - target: &mut Self::Target, - mut decoder: Decoder, - ) -> ParquetResult<()> { - self.target_reserve(target, decoder.len()); - - let mut chunked = decoder.chunked(); - - for unpacked in &mut chunked { - self.gather_chunk(target, &unpacked)?; - } - - if let Some((last, last_length)) = chunked.remainder() { - self.gather_slice(target, &last[..last_length])?; - } - - Ok(()) - } - - fn gather_bitpacked_limited<'a>( - &self, - target: &mut Self::Target, - mut decoder: Decoder<'a, u32>, - limit: usize, - ) -> ParquetResult> { - assert!(limit < decoder.len()); - - const CHUNK_SIZE: usize = ::Unpacked::LENGTH; - - let mut chunked = decoder.chunked(); - - let num_full_chunks = limit / CHUNK_SIZE; - for unpacked in (&mut chunked).take(num_full_chunks) { - self.gather_chunk(target, &unpacked)?; - } - - let (unpacked, unpacked_length) = chunked.next_inexact().unwrap(); - let unpacked_offset = limit % CHUNK_SIZE; - debug_assert!(unpacked_offset < unpacked_length); - self.gather_slice(target, &unpacked[..unpacked_offset])?; - - Ok(BufferedBitpacked { - unpacked, - - unpacked_start: unpacked_offset, - unpacked_end: unpacked_length, - decoder, - }) - } - fn gather_bitpacked<'a>( - &self, - target: &mut Self::Target, - decoder: Decoder<'a, u32>, - limit: Option, - ) -> ParquetResult<(usize, Option>)> { - let length = decoder.len(); - - match limit { - None => self - .gather_bitpacked_all(target, decoder) - .map(|_| (length, None)), - Some(limit) if limit >= length => self - .gather_bitpacked_all(target, decoder) - .map(|_| (length, None)), - Some(limit) => self - .gather_bitpacked_limited(target, decoder, limit) - .map(|b| (limit, Some(HybridRleBuffered::Bitpacked(b)))), - } - } -} - -#[derive(Default, Clone, Copy)] -pub struct ZeroCount { - pub num_zero: usize, - pub num_nonzero: usize, -} -pub struct ZeroCountGatherer; - -impl HybridRleGatherer for ZeroCountGatherer { - type Target = ZeroCount; - - #[inline(always)] - fn target_reserve(&self, _target: &mut Self::Target, _n: usize) {} - - #[inline] - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.num_zero + target.num_nonzero - } - - #[inline] - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - Ok(ZeroCount { - num_zero: usize::from(value == 0), - num_nonzero: usize::from(value != 0), - }) - } - - #[inline] - fn gather_one(&self, target: &mut Self::Target, value: ZeroCount) -> ParquetResult<()> { - target.num_zero += value.num_zero; - target.num_nonzero += value.num_nonzero; - Ok(()) - } - - #[inline] - fn gather_repeated( - &self, - target: &mut Self::Target, - value: ZeroCount, - n: usize, - ) -> ParquetResult<()> { - target.num_zero += value.num_zero * n; - target.num_nonzero += value.num_nonzero * n; - Ok(()) - } - - #[inline] - fn gather_slice(&self, target: &mut Self::Target, source: &[u32]) -> ParquetResult<()> { - let mut num_zero = 0; - let mut num_nonzero = 0; - - for v in source { - num_zero += usize::from(*v == 0); - num_nonzero += usize::from(*v != 0); - } - - target.num_zero += num_zero; - target.num_nonzero += num_nonzero; - - Ok(()) - } -} - -/// A trait to describe a translation from a HybridRLE encoding to an another format. -/// -/// In essence, this is one method ([`Translator::translate`]) that maps an `u32` to the desired -/// output type `O`. There are several other methods that may provide optimized routines -/// for slices, chunks and decoders. -/// -/// # Motivation -/// -/// The [`HybridRleDecoder`] is used extensively during Parquet decoding because it is used for -/// Dremel decoding and dictionary decoding. We want to perform a transformation from this -/// space-efficient encoding to a buffer. Here, items might be skipped, might be mapped and only a -/// few items might be needed. There are 3 main ways to do this. -/// -/// 1. Element-by-element translation using iterator `map`, `filter`, `skip`, etc. This suffers -/// from the problem that is difficult to SIMD the translation and that a `collect` might need -/// to constantly poll the `next` function. Next to that monomorphization might need to generate -/// many, many variants. -/// 2. Buffer most everything, filter and translate later. This has high memory-consumption and -/// might suffer from cache-eviction problems. This is computationally the most efficient, but -/// probably still has a high runtime. Also, this fails to utilize run-length information and -/// needs to retranslate all repeated elements. -/// 3. Batched operations. Here, we try to utilize the run-length information and utilize SIMD to -/// process many bitpacked items. This can provide the best of both worlds. -/// -/// The [`HybridRleDecoder`][super::HybridRleDecoder] decoders utilizing both run-length encoding -/// and bitpacking. In both processes, this [`Translator`] trait allows for translation with (i) no -/// heap allocations and (ii) cheap buffering and can stop and start at any point. Consequently, -/// the memory consumption while doing these translations can be relatively low while still -/// processing items in batches. -/// -/// [`HybridRleDecoder`]: super::HybridRleDecoder -pub trait Translator { - /// Translate from a decoded value to the output format - fn translate(&self, value: u32) -> ParquetResult; - - /// Translate from a slice of decoded values to the output format and write them to a `target`. - /// - /// This can overwritten to be more optimized. - fn translate_slice(&self, target: &mut Vec, source: &[u32]) -> ParquetResult<()> { - target.reserve(source.len()); - for v in source { - target.push(self.translate(*v)?); - } - Ok(()) - } - - /// Translate from a chunk of unpacked items to the output format and write them to a `target`. - /// - /// This is the same as [`Translator::translate_slice`] but with a known slice size. This can - /// allow SIMD routines to better optimize the procedure. - /// - /// This can overwritten to be more optimized. - fn translate_chunk( - &self, - target: &mut Vec, - source: &::Unpacked, - ) -> ParquetResult<()> { - self.translate_slice(target, &source[..]) - } - - /// Translate and collect all the items in a [`Decoder`] to a `target`. - /// - /// This can overwritten to be more optimized. - fn translate_bitpacked_all( - &self, - target: &mut Vec, - mut decoder: Decoder, - ) -> ParquetResult<()> { - target.reserve(decoder.len()); - - let mut chunked = decoder.chunked(); - - for unpacked in &mut chunked { - self.translate_chunk(target, &unpacked)?; - } - - if let Some((last, last_length)) = chunked.remainder() { - self.translate_slice(target, &last[..last_length])?; - } - - Ok(()) - } - - /// Translate and collect a limited number of items in a [`Decoder`] to a `target`. - /// - /// This can overwritten to be more optimized. - /// - /// # Panics - /// - /// This method panics when `limit` is larger than the `decoder` length. - fn translate_bitpacked_limited<'a>( - &self, - target: &mut Vec, - mut decoder: Decoder<'a, u32>, - limit: usize, - ) -> ParquetResult> { - assert!(limit < decoder.len()); - - const CHUNK_SIZE: usize = ::Unpacked::LENGTH; - - let mut chunked = decoder.chunked(); - - let num_full_chunks = limit / CHUNK_SIZE; - for unpacked in (&mut chunked).take(num_full_chunks) { - self.translate_chunk(target, &unpacked)?; - } - - let (unpacked, unpacked_length) = chunked.next_inexact().unwrap(); - let unpacked_offset = limit % CHUNK_SIZE; - debug_assert!(unpacked_offset < unpacked_length); - self.translate_slice(target, &unpacked[..unpacked_offset])?; - - Ok(BufferedBitpacked { - unpacked, - - unpacked_start: unpacked_offset, - unpacked_end: unpacked_length, - decoder, - }) - } - - /// Translate and collect items in a [`Decoder`] to a `target`. - /// - /// This can overwritten to be more optimized. - fn translate_bitpacked<'a>( - &self, - target: &mut Vec, - decoder: Decoder<'a, u32>, - limit: Option, - ) -> ParquetResult<(usize, Option>)> { - let length = decoder.len(); - - match limit { - None => self - .translate_bitpacked_all(target, decoder) - .map(|_| (length, None)), - Some(limit) if limit >= length => self - .translate_bitpacked_all(target, decoder) - .map(|_| (length, None)), - Some(limit) => self - .translate_bitpacked_limited(target, decoder, limit) - .map(|b| (limit, Some(HybridRleBuffered::Bitpacked(b)))), - } - } -} - -impl> HybridRleGatherer for T { - type Target = Vec; - - #[inline(always)] - fn target_reserve(&self, target: &mut Self::Target, n: usize) { - target.reserve(n); - } - #[inline(always)] - fn target_num_elements(&self, target: &Self::Target) -> usize { - target.len() - } - - #[inline(always)] - fn hybridrle_to_target(&self, value: u32) -> ParquetResult { - self.translate(value) - } - - #[inline(always)] - fn gather_one(&self, target: &mut Self::Target, value: O) -> ParquetResult<()> { - target.push(value); - Ok(()) - } - - #[inline(always)] - fn gather_repeated(&self, target: &mut Self::Target, value: O, n: usize) -> ParquetResult<()> { - target.resize(target.len() + n, value); - Ok(()) - } - - #[inline(always)] - fn gather_slice(&self, target: &mut Self::Target, source: &[u32]) -> ParquetResult<()> { - self.translate_slice(target, source) - } - - #[inline(always)] - fn gather_chunk( - &self, - target: &mut Self::Target, - source: &::Unpacked, - ) -> ParquetResult<()> { - self.translate_chunk(target, source) - } - - #[inline(always)] - fn gather_bitpacked_all( - &self, - target: &mut Self::Target, - decoder: Decoder, - ) -> ParquetResult<()> { - self.translate_bitpacked_all(target, decoder) - } - - #[inline(always)] - fn gather_bitpacked_limited<'a>( - &self, - target: &mut Self::Target, - decoder: Decoder<'a, u32>, - limit: usize, - ) -> ParquetResult> { - self.translate_bitpacked_limited(target, decoder, limit) - } - - #[inline(always)] - fn gather_bitpacked<'a>( - &self, - target: &mut Self::Target, - decoder: Decoder<'a, u32>, - limit: Option, - ) -> ParquetResult<(usize, Option>)> { - self.translate_bitpacked(target, decoder, limit) - } -} - -/// This is a unit translation variant of [`Translator`]. This just maps all encoded values from a -/// [`HybridRleDecoder`] to themselves. -/// -/// [`HybridRleDecoder`]: super::HybridRleDecoder -pub struct UnitTranslator; - -impl Translator for UnitTranslator { - fn translate(&self, value: u32) -> ParquetResult { - Ok(value) - } - - fn translate_slice(&self, target: &mut Vec, source: &[u32]) -> ParquetResult<()> { - target.extend_from_slice(source); - Ok(()) - } - fn translate_chunk( - &self, - target: &mut Vec, - source: &::Unpacked, - ) -> ParquetResult<()> { - target.extend_from_slice(&source[..]); - Ok(()) - } - fn translate_bitpacked_all( - &self, - target: &mut Vec, - decoder: Decoder, - ) -> ParquetResult<()> { - decoder.collect_into(target); - Ok(()) - } -} - -/// This is a dictionary translation variant of [`Translator`]. -/// -/// All the [`HybridRleDecoder`] values are regarded as a offset into a dictionary. -/// -/// [`HybridRleDecoder`]: super::HybridRleDecoder -pub struct DictionaryTranslator<'a, T>(pub &'a [T]); - -impl<'a, T: Copy> Translator for DictionaryTranslator<'a, T> { - fn translate(&self, value: u32) -> ParquetResult { - self.0 - .get(value as usize) - .cloned() - .ok_or(ParquetError::oos("Dictionary index is out of range")) - } - - fn translate_slice(&self, target: &mut Vec, source: &[u32]) -> ParquetResult<()> { - let Some(source_max) = source.iter().copied().max() else { - return Ok(()); - }; - - if source_max as usize >= self.0.len() { - return Err(ParquetError::oos("Dictionary index is out of range")); - } - - // Safety: We have checked before that source only has indexes that are smaller than the - // dictionary length. - target.extend( - source - .iter() - .map(|&src_idx| unsafe { *self.0.get_unchecked(src_idx as usize) }), - ); - - Ok(()) - } - - fn translate_chunk( - &self, - target: &mut Vec, - source: &::Unpacked, - ) -> ParquetResult<()> { - let source_max: u32 = source.iter().copied().max().unwrap(); - - if source_max as usize >= self.0.len() { - return Err(ParquetError::oos("Dictionary index is out of range")); - } - - // Safety: We have checked before that source only has indexes that are smaller than the - // dictionary length. - target.extend( - source - .iter() - .map(|&src_idx| unsafe { *self.0.get_unchecked(src_idx as usize) }), - ); - - Ok(()) - } -} - -/// A closure-based translator -pub struct FnTranslator ParquetResult>(pub F); - -impl ParquetResult> Translator for FnTranslator { - fn translate(&self, value: u32) -> ParquetResult { - (self.0)(value) - } -} - -#[derive(Default)] -pub struct TryFromUsizeTranslator>(std::marker::PhantomData); - -impl> Translator for TryFromUsizeTranslator { - fn translate(&self, value: u32) -> ParquetResult { - O::try_from(value as usize).map_err(|_| ParquetError::oos("Invalid cast in translation")) - } -} - -pub struct SliceDictionaryTranslator<'a, T> { - dict: &'a [T], - size: usize, -} - -impl<'a, T> SliceDictionaryTranslator<'a, T> { - pub fn new(dict: &'a [T], size: usize) -> Self { - debug_assert_eq!(dict.len() % size, 0); - Self { dict, size } - } -} - -impl<'a, T> Translator<&'a [T]> for SliceDictionaryTranslator<'a, T> { - fn translate(&self, value: u32) -> ParquetResult<&'a [T]> { - let idx = value as usize; - - if idx >= self.dict.len() / self.size { - return Err(ParquetError::oos("Dictionary slice index is out of range")); - } - - Ok(&self.dict[idx * self.size..(idx + 1) * self.size]) - } - - fn translate_slice(&self, target: &mut Vec<&'a [T]>, source: &[u32]) -> ParquetResult<()> { - let Some(source_max) = source.iter().copied().max() else { - return Ok(()); - }; - - if source_max as usize >= self.dict.len() / self.size { - return Err(ParquetError::oos("Dictionary index is out of range")); - } - - // Safety: We have checked before that source only has indexes that are smaller than the - // dictionary length. - target.extend(source.iter().map(|&src_idx| unsafe { - self.dict - .get_unchecked((src_idx as usize) * self.size..(src_idx as usize + 1) * self.size) - })); - - Ok(()) - } -} diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs index 30fd4798ccb2..3bce3086bfa5 100644 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs @@ -1,37 +1,14 @@ // See https://github.com/apache/parquet-format/blob/master/Encodings.md#run-length-encoding--bit-packing-hybrid-rle--3 mod bitmap; -mod buffered; mod encoder; -pub mod gatherer; - -#[cfg(test)] -mod fuzz; pub use bitmap::{encode_bool as bitpacked_encode, BitmapIter}; -pub use buffered::BufferedBitpacked; pub use encoder::{encode, Encoder}; -pub use gatherer::{ - DictionaryTranslator, FnTranslator, Translator, TryFromUsizeTranslator, UnitTranslator, -}; use polars_utils::slice::GetSaferUnchecked; -use self::buffered::HybridRleBuffered; -use self::gatherer::HybridRleGatherer; -use super::{bitpacked, ceil8, uleb128}; -use crate::parquet::encoding::bitpacked::{Unpackable, Unpacked}; -use crate::parquet::encoding::hybrid_rle::buffered::BufferedRle; +use super::{bitpacked, uleb128}; use crate::parquet::error::{ParquetError, ParquetResult}; -/// The two possible states of an RLE-encoded run. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum HybridEncoded<'a> { - /// A bitpacked slice. The consumer must know its bit-width to unpack it. - Bitpacked(&'a [u8]), - /// A RLE-encoded slice. The first attribute corresponds to the slice (that can be interpreted) - /// the second attribute corresponds to the number of repetitions. - Rle(&'a [u8], usize), -} - /// A [`Iterator`] for Hybrid Run-Length Encoding /// /// The hybrid here means that each second is prepended by a bit that differentiates between two @@ -47,8 +24,6 @@ pub struct HybridRleDecoder<'a> { data: &'a [u8], num_bits: usize, num_values: usize, - - buffered: Option>, } pub struct HybridRleChunkIter<'a> { @@ -76,8 +51,6 @@ impl<'a> HybridRleDecoder<'a> { data, num_bits: num_bits as usize, num_values, - - buffered: None, } } @@ -94,7 +67,7 @@ impl<'a> HybridRleDecoder<'a> { } pub fn next_chunk(&mut self) -> ParquetResult>> { - if self.len() == 0 || self.data.is_empty() { + if self.len() == 0 { return Ok(None); } @@ -104,6 +77,10 @@ impl<'a> HybridRleDecoder<'a> { return Ok(Some(HybridRleChunk::Rle(0, num_values))); } + if self.data.is_empty() { + return Ok(None); + } + let (indicator, consumed) = uleb128::decode(self.data); self.data = unsafe { self.data.get_unchecked_release(consumed..) }; @@ -150,344 +127,21 @@ impl<'a> HybridRleDecoder<'a> { self.num_values = self.num_values.min(length); } - fn gather_limited_once>( - &mut self, - target: &mut G::Target, - limit: Option, - gatherer: &G, - ) -> ParquetResult { - if limit == Some(0) { - return Ok(0); - } - - let start_target_length = gatherer.target_num_elements(target); - let start_num_values = self.num_values; - - // @NOTE: - // This is basically a collapsed version of the `decoder::Decoder`. Any change here - // probably also applies there. In a microbenchmark this collapse did around 3x for this - // specific piece of code, but I think this actually also makes the code more readable. - - debug_assert!( - self.num_values > 0, - "{:?}", - gatherer.target_num_elements(target) - ); - debug_assert!(self.num_bits > 0); - - let (indicator, consumed) = uleb128::decode(self.data); - self.data = unsafe { self.data.get_unchecked_release(consumed..) }; - - if consumed == 0 { - let step_size = - limit.map_or(self.num_values, |limit| usize::min(self.num_values, limit)); - // In this case, we take the value encoded by `0`. For example, if the HybridRle - // encodes a dictionary. We should take the 0-th value. - let value = gatherer.hybridrle_to_target(0)?; - gatherer.gather_repeated(target, value, step_size)?; - self.num_values -= step_size; - - return Ok(step_size); - } - - let num_processed = if indicator & 1 == 1 { - // is bitpacking - let bytes = (indicator as usize >> 1) * self.num_bits; - let bytes = std::cmp::min(bytes, self.data.len()); - let (packed, remaining) = self.data.split_at(bytes); - self.data = remaining; - - let length = std::cmp::min(packed.len() * 8 / self.num_bits, self.num_values); - let decoder = bitpacked::Decoder::::try_new(packed, self.num_bits, length)?; - - let (num_processed, buffered) = gatherer.gather_bitpacked(target, decoder, limit)?; - debug_assert!(limit.map_or(true, |limit| limit >= num_processed)); - self.buffered = buffered; - - num_processed - } else { - // is rle - let run_length = indicator as usize >> 1; - // repeated-value := value that is repeated, using a fixed-width of round-up-to-next-byte(bit-width) - let rle_bytes = ceil8(self.num_bits); - let (pack, remaining) = self.data.split_at(rle_bytes); - self.data = remaining; - - if run_length == 0 { - 0 - } else { - let mut bytes = [0u8; size_of::()]; - pack.iter().zip(bytes.iter_mut()).for_each(|(src, dst)| { - *dst = *src; - }); - let value = u32::from_le_bytes(bytes); - - let num_elements = limit.map_or(run_length, |limit| usize::min(run_length, limit)); - - // Only translate once. Then, just do a memset. - let translated = gatherer.hybridrle_to_target(value)?; - gatherer.gather_repeated(target, translated, num_elements)?; - - if let Some(limit) = limit { - if run_length > limit { - self.buffered = (run_length != limit).then_some({ - HybridRleBuffered::Rle(BufferedRle { - value, - length: run_length - num_elements, - }) - }); - } - } - - num_elements + pub fn collect(self) -> ParquetResult> { + let mut target = Vec::with_capacity(self.len()); + + for chunk in self.into_chunk_iter() { + match chunk? { + HybridRleChunk::Rle(value, size) => { + target.resize(target.len() + size, value); + }, + HybridRleChunk::Bitpacked(decoder) => { + decoder.collect_into(&mut target); + }, } - }; - - self.num_values -= num_processed; - - debug_assert_eq!(num_processed, start_num_values - self.num_values); - debug_assert_eq!( - num_processed, - gatherer.target_num_elements(target) - start_target_length - ); - debug_assert!(limit.map_or(true, |limit| num_processed <= limit)); - - Ok(num_processed) - } - - #[inline(always)] - pub fn gather_into>( - mut self, - target: &mut G::Target, - gatherer: &G, - ) -> ParquetResult<()> { - if self.num_values == 0 { - return Ok(()); - } - - gatherer.target_reserve(target, self.num_values); - - if self.num_bits == 0 { - let value = gatherer.hybridrle_to_target(0)?; - gatherer.gather_repeated(target, value, self.num_values)?; - return Ok(()); } - if let Some(buffered) = self.buffered.take() { - let num_buffered = buffered.gather_into(target, gatherer)?; - self.num_values -= num_buffered; - } - - while self.num_values > 0 { - self.gather_limited_once(target, None, gatherer)?; - } - - Ok(()) - } - - pub fn gather_n_into>( - &mut self, - target: &mut G::Target, - n: usize, - gatherer: &G, - ) -> ParquetResult<()> { - if self.num_values == 0 || n == 0 { - return Ok(()); - } - - if self.num_bits == 0 { - let n = usize::min(n, self.num_values); - - let value = gatherer.hybridrle_to_target(0)?; - gatherer.gather_repeated(target, value, n)?; - self.num_values -= n; - return Ok(()); - } - - let target_length = gatherer.target_num_elements(target) + n; - gatherer.target_reserve(target, n); - - if let Some(buffered) = self.buffered.as_mut() { - let num_buffered = buffered.gather_limited_into(target, n, gatherer)?; - debug_assert!(num_buffered <= n); - self.num_values -= num_buffered; - - if num_buffered < n { - self.buffered = None; - } - } - - while gatherer.target_num_elements(target) < target_length && self.num_values > 0 { - self.gather_limited_once( - target, - Some(target_length - gatherer.target_num_elements(target)), - gatherer, - )?; - } - - Ok(()) - } - - #[inline(always)] - pub fn translate_and_collect_into>( - self, - target: &mut >::Target, - translator: &T, - ) -> ParquetResult<()> { - self.gather_into(target, translator) - } - - pub fn translate_and_collect_n_into>( - &mut self, - target: &mut >::Target, - n: usize, - translator: &T, - ) -> ParquetResult<()> { - self.gather_n_into(target, n, translator) - } - - #[inline(always)] - pub fn translate_and_collect>( - self, - translator: &T, - ) -> ParquetResult<>::Target> { - let mut vec = Vec::new(); - self.translate_and_collect_into(&mut vec, translator)?; - Ok(vec) - } - - #[inline(always)] - pub fn collect_into( - self, - target: &mut >::Target, - ) -> Result<(), ParquetError> { - self.translate_and_collect_into(target, &UnitTranslator) - } - - #[inline(always)] - pub fn collect_n_into( - &mut self, - target: &mut >::Target, - n: usize, - ) -> ParquetResult<()> { - self.translate_and_collect_n_into(target, n, &UnitTranslator) - } - - #[inline(always)] - pub fn collect(self) -> ParquetResult<>::Target> { - let mut vec = Vec::new(); - self.collect_into(&mut vec)?; - Ok(vec) - } - - pub fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if self.num_values == 0 || n == 0 { - return Ok(()); - } - - if n >= self.num_values { - self.data = &[]; - self.num_values = 0; - self.buffered = None; - return Ok(()); - } - - if self.num_bits == 0 { - self.num_values -= n; - return Ok(()); - } - - let mut n = n; - if let Some(buffered) = self.buffered.as_mut() { - let num_skipped = buffered.skip_in_place(n); - - if num_skipped < n { - self.buffered = None; - } - - self.num_values -= num_skipped; - n -= num_skipped; - } - - while n > 0 && self.num_values > 0 { - let start_num_values = self.num_values; - - let (indicator, consumed) = uleb128::decode(self.data); - self.data = unsafe { self.data.get_unchecked_release(consumed..) }; - - let num_skipped = if consumed == 0 { - n - } else if indicator & 1 == 1 { - // is bitpacking - let bytes = (indicator as usize >> 1) * self.num_bits; - let bytes = std::cmp::min(bytes, self.data.len()); - let (packed, remaining) = self.data.split_at(bytes); - self.data = remaining; - - let length = std::cmp::min(packed.len() * 8 / self.num_bits, self.num_values); - let mut decoder = - bitpacked::Decoder::::try_new(packed, self.num_bits, length)?; - - // Skip the whole decoder if it is possible - if decoder.len() <= n { - decoder.len() - } else { - const CHUNK_SIZE: usize = ::Unpacked::LENGTH; - - let num_full_chunks = n / CHUNK_SIZE; - decoder.skip_chunks(num_full_chunks); - - let (unpacked, unpacked_length) = decoder.chunked().next_inexact().unwrap(); - let unpacked_offset = n % CHUNK_SIZE; - debug_assert!(unpacked_offset < unpacked_length); - - self.buffered = Some(HybridRleBuffered::Bitpacked(BufferedBitpacked { - unpacked, - - unpacked_start: unpacked_offset, - unpacked_end: unpacked_length, - decoder, - })); - - n - } - } else { - // is rle - let run_length = indicator as usize >> 1; - // repeated-value := value that is repeated, using a fixed-width of round-up-to-next-byte(bit-width) - let rle_bytes = ceil8(self.num_bits); - let (pack, remaining) = self.data.split_at(rle_bytes); - self.data = remaining; - - // Skip the whole run-length encoded value if it is possible - if run_length <= n { - run_length - } else { - let mut bytes = [0u8; size_of::()]; - pack.iter().zip(bytes.iter_mut()).for_each(|(src, dst)| { - *dst = *src; - }); - let value = u32::from_le_bytes(bytes); - - self.buffered = Some(HybridRleBuffered::Rle(BufferedRle { - value, - length: run_length - n, - })); - - n - } - }; - - self.num_values -= num_skipped; - - debug_assert_eq!(num_skipped, start_num_values - self.num_values); - debug_assert!(num_skipped <= n, "{num_skipped} <= {n}"); - debug_assert!(indicator >> 1 == 0 || num_skipped > 0); - - n -= num_skipped; - } - - Ok(()) + Ok(target) } } diff --git a/crates/polars/tests/it/io/parquet/read/binary.rs b/crates/polars/tests/it/io/parquet/read/binary.rs index 4fcf27e173b7..086dd06bb9b9 100644 --- a/crates/polars/tests/it/io/parquet/read/binary.rs +++ b/crates/polars/tests/it/io/parquet/read/binary.rs @@ -1,11 +1,10 @@ -use polars_parquet::parquet::encoding::hybrid_rle::FnTranslator; use polars_parquet::parquet::error::ParquetResult; use polars_parquet::parquet::page::DataPage; use super::dictionary::BinaryPageDict; use super::utils::deserialize_optional; -use crate::io::parquet::read::hybrid_rle_iter; use crate::io::parquet::read::utils::FixedLenBinaryPageState; +use crate::io::parquet::read::{hybrid_rle_fn_collect, hybrid_rle_iter}; pub fn page_to_vec( page: &DataPage, @@ -25,9 +24,9 @@ pub fn page_to_vec( .map(|x| x.transpose()) .collect(), FixedLenBinaryPageState::RequiredDictionary(dict) => { - let dictionary = - FnTranslator(|v| dict.dict.value(v as usize).map(|v| Some(v.to_vec()))); - dict.indexes.translate_and_collect(&dictionary) + hybrid_rle_fn_collect(dict.indexes, |x| { + dict.dict.value(x as usize).map(<[u8]>::to_vec).map(Some) + }) }, FixedLenBinaryPageState::OptionalDictionary(validity, dict) => { let values = hybrid_rle_iter(dict.indexes)? diff --git a/crates/polars/tests/it/io/parquet/read/mod.rs b/crates/polars/tests/it/io/parquet/read/mod.rs index d671e085c86a..a9481ab0ab58 100644 --- a/crates/polars/tests/it/io/parquet/read/mod.rs +++ b/crates/polars/tests/it/io/parquet/read/mod.rs @@ -15,7 +15,7 @@ mod utils; use std::fs::File; use dictionary::DecodedDictPage; -use polars_parquet::parquet::encoding::hybrid_rle::HybridRleDecoder; +use polars_parquet::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use polars_parquet::parquet::error::{ParquetError, ParquetResult}; use polars_parquet::parquet::metadata::ColumnChunkMetadata; use polars_parquet::parquet::page::DataPage; @@ -32,6 +32,34 @@ pub fn hybrid_rle_iter(d: HybridRleDecoder) -> ParquetResult(d: HybridRleDecoder, mut f: impl FnMut(u32) -> ParquetResult) -> ParquetResult> { + let mut target = Vec::with_capacity(d.len()); + + for chunk in d.into_chunk_iter() { + match chunk? { + HybridRleChunk::Rle(value, size) => { + target.resize(target.len() + size, f(value)?); + }, + HybridRleChunk::Bitpacked(mut decoder) => { + let mut chunked = decoder.chunked(); + for dchunk in chunked.by_ref() { + for v in dchunk { + target.push(f(v)?); + } + } + + if let Some((dchunk, l)) = chunked.remainder() { + for &v in &dchunk[..l] { + target.push(f(v)?); + } + } + }, + } + } + + Ok(target) +} + pub fn get_path() -> PathBuf { let dir = env!("CARGO_MANIFEST_DIR"); PathBuf::from(dir).join("../../docs/assets/data") diff --git a/crates/polars/tests/it/io/parquet/read/primitive.rs b/crates/polars/tests/it/io/parquet/read/primitive.rs index 960c502fb82d..9e16b0240550 100644 --- a/crates/polars/tests/it/io/parquet/read/primitive.rs +++ b/crates/polars/tests/it/io/parquet/read/primitive.rs @@ -1,9 +1,10 @@ -use polars_parquet::parquet::encoding::hybrid_rle::FnTranslator; use polars_parquet::parquet::error::ParquetResult; use polars_parquet::parquet::page::DataPage; use polars_parquet::parquet::types::NativeType; use polars_parquet::read::ParquetError; +use crate::io::parquet::read::hybrid_rle_fn_collect; + use super::dictionary::PrimitivePageDict; use super::hybrid_rle_iter; use super::utils::{deserialize_optional, NativePageState}; @@ -44,8 +45,7 @@ pub fn page_to_vec( }, NativePageState::Required(values) => Ok(values.map(Some).collect()), NativePageState::RequiredDictionary(dict) => { - let dictionary = FnTranslator(|x| dict.dict.value(x as usize).copied().map(Some)); - dict.indexes.translate_and_collect(&dictionary) + hybrid_rle_fn_collect(dict.indexes, |x| dict.dict.value(x as usize).copied().map(Some)) }, NativePageState::OptionalDictionary(validity, dict) => { let values = diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py index f55dcb8f0ec6..709777c294bb 100644 --- a/py-polars/tests/unit/io/test_parquet.py +++ b/py-polars/tests/unit/io/test_parquet.py @@ -1098,6 +1098,7 @@ def test_hybrid_rle() -> None: @pytest.mark.slow def test_roundtrip_parametric(df: pl.DataFrame) -> None: f = io.BytesIO() + print(df) df.write_parquet(f) f.seek(0) result = pl.read_parquet(f) @@ -2195,3 +2196,23 @@ def test_filter_only_invalid(dtype: pl.DataType, filt: pl.Expr) -> None: out = pl.scan_parquet(f, parallel='prefiltered').filter(filt).collect() assert_frame_equal(df.filter(filt), out) + + +def test_nested_nulls() -> None: + df = pl.Series( + 'a', [ + [None, None], + None, + [None, 1], + [None, None], + [2, None], + ], + pl.Array(pl.Int32, 2), + ).to_frame() + + f = io.BytesIO() + df.write_parquet(f) + + f.seek(0) + out = pl.read_parquet(f) + assert_frame_equal(out, df) From aa35a3dd647c198559a88da8c9a7c1ba30f2fe5c Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 11:03:56 +0200 Subject: [PATCH 25/32] fix rebase problem --- .../arrow/read/deserialize/fixed_size_binary.rs | 14 +++++++------- .../arrow/read/deserialize/utils/array_chunks.rs | 4 ---- crates/polars-parquet/src/parquet/types.rs | 8 ++++++-- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs index 13a2cf390820..7ae39b3366ad 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary.rs @@ -90,13 +90,13 @@ impl FSBVec { pub fn into_bytes_buffer(self) -> Buffer { Buffer::from_storage(match self { - FSBVec::Size1(vec) => SharedStorage::bytes_from_aligned_bytes(vec), - FSBVec::Size2(vec) => SharedStorage::bytes_from_aligned_bytes(vec), - FSBVec::Size4(vec) => SharedStorage::bytes_from_aligned_bytes(vec), - FSBVec::Size8(vec) => SharedStorage::bytes_from_aligned_bytes(vec), - FSBVec::Size12(vec) => SharedStorage::bytes_from_aligned_bytes(vec), - FSBVec::Size16(vec) => SharedStorage::bytes_from_aligned_bytes(vec), - FSBVec::Size32(vec) => SharedStorage::bytes_from_aligned_bytes(vec), + FSBVec::Size1(vec) => SharedStorage::bytes_from_pod_vec(vec), + FSBVec::Size2(vec) => SharedStorage::bytes_from_pod_vec(vec), + FSBVec::Size4(vec) => SharedStorage::bytes_from_pod_vec(vec), + FSBVec::Size8(vec) => SharedStorage::bytes_from_pod_vec(vec), + FSBVec::Size12(vec) => SharedStorage::bytes_from_pod_vec(vec), + FSBVec::Size16(vec) => SharedStorage::bytes_from_pod_vec(vec), + FSBVec::Size32(vec) => SharedStorage::bytes_from_pod_vec(vec), FSBVec::Other(vec, _) => SharedStorage::from_vec(vec), }) } diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs index 2df4371460f2..534acde49369 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/array_chunks.rs @@ -16,11 +16,7 @@ impl<'a, B: AlignedBytes> ArrayChunks<'a, B> { /// /// This returns null if the `bytes` slice's length is not a multiple of the size of `P::Bytes`. pub(crate) fn new(bytes: &'a [u8]) -> Option { -<<<<<<< HEAD - if bytes.len() % size_of::() != 0 { -======= if bytes.len() % B::SIZE != 0 { ->>>>>>> 26e3481c1d (massively reduce monomorphisations for kernels) return None; } diff --git a/crates/polars-parquet/src/parquet/types.rs b/crates/polars-parquet/src/parquet/types.rs index 3d74ba7e0563..74c5da2adc2f 100644 --- a/crates/polars-parquet/src/parquet/types.rs +++ b/crates/polars-parquet/src/parquet/types.rs @@ -1,9 +1,13 @@ -use arrow::types::{AlignedBytes, Bytes12Alignment4, Bytes4Alignment4, Bytes8Alignment8, SameBytes}; +use arrow::types::{ + AlignedBytes, AlignedBytesCast, Bytes12Alignment4, Bytes4Alignment4, Bytes8Alignment8, +}; use crate::parquet::schema::types::PhysicalType; /// A physical native representation of a Parquet fixed-sized type. -pub trait NativeType: std::fmt::Debug + Send + Sync + 'static + Copy + Clone + SameBytes { +pub trait NativeType: + std::fmt::Debug + Send + Sync + 'static + Copy + Clone + AlignedBytesCast +{ type Bytes: AsRef<[u8]> + bytemuck::Pod + IntoIterator From 36b8d51fb709f195552256faf6ff3e867e6da7c8 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 11:07:13 +0200 Subject: [PATCH 26/32] format --- crates/polars-arrow/src/types/aligned_bytes.rs | 2 +- .../polars-parquet/src/arrow/read/deserialize/boolean.rs | 3 +-- .../src/arrow/read/deserialize/primitive/integer.rs | 8 +++++--- .../src/arrow/read/deserialize/utils/dict_encoded.rs | 6 +++++- crates/polars/tests/it/io/parquet/read/mod.rs | 5 ++++- crates/polars/tests/it/io/parquet/read/primitive.rs | 9 ++++----- 6 files changed, 20 insertions(+), 13 deletions(-) diff --git a/crates/polars-arrow/src/types/aligned_bytes.rs b/crates/polars-arrow/src/types/aligned_bytes.rs index 3e0351089788..95aa84cfd349 100644 --- a/crates/polars-arrow/src/types/aligned_bytes.rs +++ b/crates/polars-arrow/src/types/aligned_bytes.rs @@ -50,7 +50,7 @@ macro_rules! impl_aligned_bytes { ) => { $( /// Bytes with a size and alignment. - /// + /// /// This is used to reduce the monomorphizations for routines that solely rely on the size /// and alignment of types. #[derive(Debug, Copy, Clone, PartialEq, Eq, Default, Pod, Zeroable)] diff --git a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs index c567a2e46a01..11019b3ab614 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/boolean.rs @@ -7,8 +7,7 @@ use polars_compute::filter::filter_boolean_kernel; use super::utils::dict_encoded::{append_validity, constrain_page_validity}; use super::utils::{ - self, decode_hybrid_rle_into_bitmap, filter_from_range, freeze_validity, - Decoder, ExactSize, + self, decode_hybrid_rle_into_bitmap, filter_from_range, freeze_validity, Decoder, ExactSize, }; use super::Filter; use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; diff --git a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs index ba3c249ae772..087fc1c447d5 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/primitive/integer.rs @@ -5,14 +5,16 @@ use arrow::types::NativeType; use super::super::utils; use super::{ - AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, - IntoDecoderFunction, PrimitiveDecoder, UnitDecoderFunction, + AsDecoderFunction, ClosureDecoderFunction, DecoderFunction, IntoDecoderFunction, + PrimitiveDecoder, UnitDecoderFunction, }; use crate::parquet::encoding::{byte_stream_split, delta_bitpacked, hybrid_rle, Encoding}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage, DictPage}; use crate::parquet::types::{decode, NativeType as ParquetNativeType}; -use crate::read::deserialize::utils::{dict_indices_decoder, freeze_validity, unspecialized_decode}; +use crate::read::deserialize::utils::{ + dict_indices_decoder, freeze_validity, unspecialized_decode, +}; use crate::read::Filter; #[allow(clippy::large_enum_variant)] diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs index 36e65e82ede4..69c1cd6c549d 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/dict_encoded.rs @@ -48,7 +48,11 @@ pub(crate) fn append_validity( } } -pub(crate) fn constrain_page_validity(values_len: usize, page_validity: Option<&Bitmap>, filter: Option<&Filter>) -> Option { +pub(crate) fn constrain_page_validity( + values_len: usize, + page_validity: Option<&Bitmap>, + filter: Option<&Filter>, +) -> Option { let num_unfiltered_rows = match (filter.as_ref(), page_validity) { (None, None) => values_len, (None, Some(pv)) => { diff --git a/crates/polars/tests/it/io/parquet/read/mod.rs b/crates/polars/tests/it/io/parquet/read/mod.rs index a9481ab0ab58..448134a8d955 100644 --- a/crates/polars/tests/it/io/parquet/read/mod.rs +++ b/crates/polars/tests/it/io/parquet/read/mod.rs @@ -32,7 +32,10 @@ pub fn hybrid_rle_iter(d: HybridRleDecoder) -> ParquetResult(d: HybridRleDecoder, mut f: impl FnMut(u32) -> ParquetResult) -> ParquetResult> { +pub fn hybrid_rle_fn_collect( + d: HybridRleDecoder, + mut f: impl FnMut(u32) -> ParquetResult, +) -> ParquetResult> { let mut target = Vec::with_capacity(d.len()); for chunk in d.into_chunk_iter() { diff --git a/crates/polars/tests/it/io/parquet/read/primitive.rs b/crates/polars/tests/it/io/parquet/read/primitive.rs index 9e16b0240550..96ed3f4eac11 100644 --- a/crates/polars/tests/it/io/parquet/read/primitive.rs +++ b/crates/polars/tests/it/io/parquet/read/primitive.rs @@ -3,11 +3,10 @@ use polars_parquet::parquet::page::DataPage; use polars_parquet::parquet::types::NativeType; use polars_parquet::read::ParquetError; -use crate::io::parquet::read::hybrid_rle_fn_collect; - use super::dictionary::PrimitivePageDict; use super::hybrid_rle_iter; use super::utils::{deserialize_optional, NativePageState}; +use crate::io::parquet::read::hybrid_rle_fn_collect; /// The deserialization state of a `DataPage` of `Primitive` parquet primitive type #[derive(Debug)] @@ -44,9 +43,9 @@ pub fn page_to_vec( deserialize_optional(validity, values.by_ref().map(Ok)) }, NativePageState::Required(values) => Ok(values.map(Some).collect()), - NativePageState::RequiredDictionary(dict) => { - hybrid_rle_fn_collect(dict.indexes, |x| dict.dict.value(x as usize).copied().map(Some)) - }, + NativePageState::RequiredDictionary(dict) => hybrid_rle_fn_collect(dict.indexes, |x| { + dict.dict.value(x as usize).copied().map(Some) + }), NativePageState::OptionalDictionary(validity, dict) => { let values = hybrid_rle_iter(dict.indexes)?.map(|x| dict.dict.value(x as usize).copied()); From 474baf29826b2b1a7d1dc7c6a57cafb9b4818647 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 12:54:42 +0200 Subject: [PATCH 27/32] fix Storage try_transmute --- crates/polars-arrow/src/buffer/immutable.rs | 2 ++ crates/polars-arrow/src/storage.rs | 6 ++++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/polars-arrow/src/buffer/immutable.rs b/crates/polars-arrow/src/buffer/immutable.rs index eb27a62341b9..a3ad6721a6f8 100644 --- a/crates/polars-arrow/src/buffer/immutable.rs +++ b/crates/polars-arrow/src/buffer/immutable.rs @@ -164,6 +164,8 @@ impl Buffer { #[inline] #[must_use] pub unsafe fn sliced_unchecked(mut self, offset: usize, length: usize) -> Self { + debug_assert!(offset + length <= self.len()); + self.slice_unchecked(offset, length); self } diff --git a/crates/polars-arrow/src/storage.rs b/crates/polars-arrow/src/storage.rs index 76116f917316..1d4288eed2ee 100644 --- a/crates/polars-arrow/src/storage.rs +++ b/crates/polars-arrow/src/storage.rs @@ -283,10 +283,12 @@ impl SharedStorage { return Err(self); } - Ok(SharedStorage { + let storage = SharedStorage { inner: self.inner.cast(), phantom: PhantomData, - }) + }; + std::mem::forget(self); + Ok(storage) } } From cce38f5c395d6cc8f10c925230834c414874078d Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 13:22:46 +0200 Subject: [PATCH 28/32] add some docs --- .../src/arrow/read/deserialize/binview.rs | 43 +++++-- .../arrow/read/deserialize/nested_utils.rs | 99 ++++++++++++---- .../src/arrow/read/deserialize/utils/mod.rs | 112 ------------------ 3 files changed, 113 insertions(+), 141 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs index 6bc9831cd240..62a0e7f2df4e 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/binview.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/binview.rs @@ -307,6 +307,23 @@ pub fn decode_plain_generic( verify_utf8: bool, ) -> ParquetResult<()> { + // Since the offset in the buffer is decided by the interleaved lengths, every value has to be + // walked no matter what. This makes decoding rather inefficient in general. + // + // There are three cases: + // 1. All inlinable values + // - Most time is spend in decoding + // - No additional buffer has to be formed + // - Possible UTF-8 verification is fast because the len_below_128 trick + // 2. All non-inlinable values + // - Little time is spend in decoding + // - Most time is spend in buffer memcopying (we remove the interleaved lengths) + // - Possible UTF-8 verification is fast because the continuation byte trick + // 3. Mixed inlinable and non-inlinable values + // - Time shared between decoding and buffer forming + // - UTF-8 verification might still use len_below_128 trick, but might need to fall back to + // slow path. + target.finish_in_progress(); unsafe { target.views_mut() }.reserve(num_rows); @@ -375,16 +392,30 @@ pub fn decode_plain_generic( } if verify_utf8 { + // This is a trick that allows us to check the resulting buffer which allows to batch the + // UTF-8 verification. + // + // This is allowed if none of the strings start with a UTF-8 continuation byte, so we keep + // track of that during the decoding. if num_inlined == 0 { if !none_starting_with_continuation_byte || simdutf8::basic::from_utf8(&buffer).is_err() { return Err(invalid_utf8_err()); } + + // This is a small trick that allows us to check the Parquet buffer instead of the view + // buffer. Batching the UTF-8 verification is more performant. For this to be allowed, + // all the interleaved lengths need to be valid UTF-8. + // + // Every strings prepended by 4 bytes (L, 0, 0, 0), since we check here L < 128. L is + // only a valid first byte of a UTF-8 code-point and (L, 0, 0, 0) is valid UTF-8. + // Consequently, it is valid to just check the whole buffer. } else if all_len_below_128 { if simdutf8::basic::from_utf8(values).is_err() { return Err(invalid_utf8_err()); } } else { + // We check all the non-inlined values here. if !none_starting_with_continuation_byte || simdutf8::basic::from_utf8(&buffer).is_err() { return Err(invalid_utf8_err()); @@ -392,16 +423,14 @@ pub fn decode_plain_generic( let mut all_inlined_are_ascii = true; - const MASK: [u128; 2] = [ - 0x0000_0000_8080_8080_8080_8080_8080_8080, - 0x0000_0000_0000_0000_0000_0000_0000_0000, - ]; - + // @NOTE: This is only valid because we initialize our inline View's to be zeroes on + // non-included bytes. for view in &target.views()[target.len() - num_seen..] { - all_inlined_are_ascii &= - view.as_u128() & MASK[usize::from(view.length > View::MAX_INLINE_SIZE)] == 0; + all_inlined_are_ascii &= (view.length > View::MAX_INLINE_SIZE) + | (view.as_u128() & 0x0000_0000_8080_8080_8080_8080_8080_8080 == 0); } + // This is the very slow path. if !all_inlined_are_ascii { let mut is_valid = true; for view in &target.views()[target.len() - num_seen..] { diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index c1a6a5b75507..287f6ce88440 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -2,13 +2,12 @@ use arrow::bitmap::utils::BitmapIter; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; -use super::utils::{self, BatchableCollector}; +use super::utils; use super::{BasicDecompressor, Filter}; use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage}; use crate::parquet::read::levels::get_bit_width; -use crate::read::deserialize::utils::BatchedCollector; #[derive(Debug)] pub struct Nested { @@ -209,31 +208,79 @@ impl Nested { } } +/// Utility structure to create a `Filter` and `Validity` mask for the leaf values. +/// +/// This batches the extending. pub struct BatchedNestedDecoder<'a> { + pub(crate) num_waiting_valids: usize, + pub(crate) num_waiting_invalids: usize, + filter: &'a mut MutableBitmap, validity: &'a mut MutableBitmap, } -impl<'a> BatchableCollector<(), ()> for BatchedNestedDecoder<'a> { - fn push_n(&mut self, _target: &mut (), n: usize) -> ParquetResult<()> { - self.filter.extend_constant(n, true); - self.validity.extend_constant(n, true); - Ok(()) +impl<'a> BatchedNestedDecoder<'a> { + fn push_valid(&mut self) -> ParquetResult<()> { + self.push_n_valids(1) } - fn push_n_nulls(&mut self, _target: &mut (), n: usize) -> ParquetResult<()> { - self.filter.extend_constant(n, true); - self.validity.extend_constant(n, false); + fn push_invalid(&mut self) -> ParquetResult<()> { + self.push_n_invalids(1) + } + + fn push_n_valids(&mut self, n: usize) -> ParquetResult<()> { + if self.num_waiting_invalids == 0 { + self.num_waiting_valids += n; + return Ok(()); + } + + self.filter.extend_constant(self.num_waiting_valids, true); + self.validity.extend_constant(self.num_waiting_valids, true); + + self.filter.extend_constant(self.num_waiting_invalids, true); + self.validity + .extend_constant(self.num_waiting_invalids, false); + + self.num_waiting_valids = n; + self.num_waiting_invalids = 0; + + Ok(()) + } + fn push_n_invalids(&mut self, n: usize) -> ParquetResult<()> { + self.num_waiting_invalids += n; Ok(()) } fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { + if self.num_waiting_valids > 0 { + self.filter.extend_constant(self.num_waiting_valids, true); + self.validity.extend_constant(self.num_waiting_valids, true); + self.num_waiting_valids = 0; + } + if self.num_waiting_invalids > 0 { + self.filter.extend_constant(self.num_waiting_invalids, true); + self.validity + .extend_constant(self.num_waiting_invalids, false); + self.num_waiting_invalids = 0; + } + self.filter.extend_constant(n, false); self.validity.extend_constant(n, true); Ok(()) } + + fn finalize(self) -> ParquetResult<()> { + self.filter.extend_constant(self.num_waiting_valids, true); + self.validity.extend_constant(self.num_waiting_valids, true); + + self.filter.extend_constant(self.num_waiting_invalids, true); + self.validity + .extend_constant(self.num_waiting_invalids, false); + + Ok(()) + } } /// The initial info of nested data types. @@ -342,6 +389,16 @@ fn collect_level_values( Ok(()) } +/// State to keep track of how many top-level values (i.e. rows) still need to be skipped and +/// collected. +/// +/// This state should be kept between pages because a top-level value / row value may span several +/// pages. +/// +/// - `num_skips = Some(n)` means that it will skip till the `n + 1`-th occurrance of the repetition +/// level of `0` (i.e. the start of a top-level value / row value). +/// - `num_collects = Some(n)` means that it will collect values till the `n + 1`-th occurrance of +/// the repetition level of `0` (i.e. the start of a top-level value / row value). struct DecodingState { num_skips: Option, num_collects: Option, @@ -352,7 +409,7 @@ fn decode_nested( mut current_def_levels: &[u16], mut current_rep_levels: &[u16], - batched_collector: &mut BatchedCollector<'_, (), (), BatchedNestedDecoder<'_>>, + batched_collector: &mut BatchedNestedDecoder<'_>, nested: &mut [Nested], state: &mut DecodingState, @@ -455,7 +512,7 @@ fn decode_nested( } } - batched_collector.push_n_invalids(num_elements); + batched_collector.push_n_invalids(num_elements)?; break; } @@ -470,7 +527,7 @@ fn decode_nested( if is_valid { batched_collector.push_valid()?; } else { - batched_collector.push_invalid(); + batched_collector.push_invalid()?; } } } @@ -608,16 +665,14 @@ impl PageNestedDecoder { let mut leaf_filter = MutableBitmap::new(); let mut leaf_validity = MutableBitmap::new(); - let mut _tgt = (); - // @TODO: move this to outside the loop. - let mut batched_collector = BatchedCollector::new( - BatchedNestedDecoder { - filter: &mut leaf_filter, - validity: &mut leaf_validity, - }, - &mut _tgt, - ); + let mut batched_collector = BatchedNestedDecoder { + num_waiting_valids: 0, + num_waiting_invalids: 0, + + filter: &mut leaf_filter, + validity: &mut leaf_validity, + }; decode_nested( ¤t_def_levels, diff --git a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs index c067a5f257b1..057870db828d 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/utils/mod.rs @@ -108,96 +108,6 @@ pub fn not_implemented(page: &DataPage) -> ParquetError { )) } -pub trait BatchableCollector { - fn push_n(&mut self, target: &mut T, n: usize) -> ParquetResult<()>; - fn push_n_nulls(&mut self, target: &mut T, n: usize) -> ParquetResult<()>; - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()>; -} - -/// This batches sequential collect operations to try and prevent unnecessary buffering and -/// `Iterator::next` polling. -#[must_use] -pub struct BatchedCollector<'a, I, T, C: BatchableCollector> { - pub(crate) num_waiting_valids: usize, - pub(crate) num_waiting_invalids: usize, - - target: &'a mut T, - collector: C, - _pd: std::marker::PhantomData, -} - -impl<'a, I, T, C: BatchableCollector> BatchedCollector<'a, I, T, C> { - pub fn new(collector: C, target: &'a mut T) -> Self { - Self { - num_waiting_valids: 0, - num_waiting_invalids: 0, - target, - collector, - _pd: Default::default(), - } - } - - #[inline] - pub fn push_valid(&mut self) -> ParquetResult<()> { - self.push_n_valids(1) - } - - #[inline] - pub fn push_invalid(&mut self) { - self.push_n_invalids(1) - } - - #[inline] - pub fn push_n_valids(&mut self, n: usize) -> ParquetResult<()> { - if self.num_waiting_invalids == 0 { - self.num_waiting_valids += n; - return Ok(()); - } - - self.collector - .push_n(self.target, self.num_waiting_valids)?; - self.collector - .push_n_nulls(self.target, self.num_waiting_invalids)?; - - self.num_waiting_valids = n; - self.num_waiting_invalids = 0; - - Ok(()) - } - - #[inline] - pub fn push_n_invalids(&mut self, n: usize) { - self.num_waiting_invalids += n; - } - - #[inline] - pub fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if self.num_waiting_valids > 0 { - self.collector - .push_n(self.target, self.num_waiting_valids)?; - self.num_waiting_valids = 0; - } - if self.num_waiting_invalids > 0 { - self.collector - .push_n_nulls(self.target, self.num_waiting_invalids)?; - self.num_waiting_invalids = 0; - } - - self.collector.skip_in_place(n)?; - - Ok(()) - } - - #[inline] - pub fn finalize(mut self) -> ParquetResult<()> { - self.collector - .push_n(self.target, self.num_waiting_valids)?; - self.collector - .push_n_nulls(self.target, self.num_waiting_invalids)?; - Ok(()) - } -} - pub(crate) type PageValidity<'a> = HybridRleDecoder<'a>; pub(crate) fn page_validity_decoder(page: &DataPage) -> ParquetResult { let validity = split_buffer(page)?.def; @@ -347,28 +257,6 @@ pub(crate) fn unspecialized_decode( Ok(()) } -impl, I: Iterator> BatchableCollector for I { - #[inline] - fn push_n(&mut self, target: &mut P, n: usize) -> ParquetResult<()> { - target.extend_n(n, self); - Ok(()) - } - - #[inline] - fn push_n_nulls(&mut self, target: &mut P, n: usize) -> ParquetResult<()> { - target.extend_null_constant(n); - Ok(()) - } - - #[inline] - fn skip_in_place(&mut self, n: usize) -> ParquetResult<()> { - if n > 0 { - _ = self.nth(n - 1); - } - Ok(()) - } -} - /// An item with a known size pub(super) trait ExactSize { /// The number of items in the container From 0f88e782f9f16fe4ba630583192f8eee022023ce Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 13:41:16 +0200 Subject: [PATCH 29/32] pyfmt --- py-polars/tests/unit/io/test_parquet.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/py-polars/tests/unit/io/test_parquet.py b/py-polars/tests/unit/io/test_parquet.py index 709777c294bb..833a089db387 100644 --- a/py-polars/tests/unit/io/test_parquet.py +++ b/py-polars/tests/unit/io/test_parquet.py @@ -2055,7 +2055,7 @@ def test_conserve_sortedness( [list(range(i, i + size)) for i in range(13)], [list(range(i, i + size)) if i % 3 < 2 else None for i in range(13)], ] - ] + ], ) @pytest.mark.parametrize( "filt", @@ -2144,7 +2144,7 @@ def test_decode_f16() -> None: def test_invalid_utf8_binary() -> None: - a = pl.Series('a', [b"\x80"], pl.Binary).to_frame() + a = pl.Series("a", [b"\x80"], pl.Binary).to_frame() f = io.BytesIO() a.write_parquet(f) @@ -2164,11 +2164,11 @@ def test_invalid_utf8_binary() -> None: pl.String, pl.Binary, pl.Boolean, - pl.Struct({ 'x': pl.Int32 }), + pl.Struct({"x": pl.Int32}), pl.List(pl.Int32), pl.Array(pl.Int32, 0), pl.Array(pl.Int32, 2), - ] + ], ) @pytest.mark.parametrize( "filt", @@ -2181,26 +2181,29 @@ def test_invalid_utf8_binary() -> None: pl.col.f != 2, pl.col.f == 3, pl.col.f != 3, - ] + ], ) def test_filter_only_invalid(dtype: pl.DataType, filt: pl.Expr) -> None: - df = pl.DataFrame([ - pl.Series('a', [None, None, None], dtype), - pl.Series('f', range(3), pl.Int32), - ]) + df = pl.DataFrame( + [ + pl.Series("a", [None, None, None], dtype), + pl.Series("f", range(3), pl.Int32), + ] + ) f = io.BytesIO() df.write_parquet(f) f.seek(0) - out = pl.scan_parquet(f, parallel='prefiltered').filter(filt).collect() + out = pl.scan_parquet(f, parallel="prefiltered").filter(filt).collect() assert_frame_equal(df.filter(filt), out) def test_nested_nulls() -> None: df = pl.Series( - 'a', [ + "a", + [ [None, None], None, [None, 1], From 07ffdc887013beaa8c2d68d4011a5bd4d19112ae Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 13:52:13 +0200 Subject: [PATCH 30/32] fix doc and broken test --- .../polars-parquet/src/arrow/read/deserialize/nested_utils.rs | 4 ++-- crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index 287f6ce88440..f865966272af 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -396,9 +396,9 @@ fn collect_level_values( /// pages. /// /// - `num_skips = Some(n)` means that it will skip till the `n + 1`-th occurrance of the repetition -/// level of `0` (i.e. the start of a top-level value / row value). +/// level of `0` (i.e. the start of a top-level value / row value). /// - `num_collects = Some(n)` means that it will collect values till the `n + 1`-th occurrance of -/// the repetition level of `0` (i.e. the start of a top-level value / row value). +/// the repetition level of `0` (i.e. the start of a top-level value / row value). struct DecodingState { num_skips: Option, num_collects: Option, diff --git a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs index 3bce3086bfa5..b862cc56b36d 100644 --- a/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs +++ b/crates/polars-parquet/src/parquet/encoding/hybrid_rle/mod.rs @@ -281,7 +281,7 @@ mod tests { fn empty_values() -> ParquetResult<()> { let data = []; - let num_bits = 1; + let num_bits = 0; let decoder = HybridRleDecoder::new(&data, num_bits, 100); From 5040cafe29467da11d7c1dbfa4216d06362ddee0 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 13:53:40 +0200 Subject: [PATCH 31/32] typo --- .../polars-parquet/src/arrow/read/deserialize/nested_utils.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index f865966272af..a316317d1bae 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -395,9 +395,9 @@ fn collect_level_values( /// This state should be kept between pages because a top-level value / row value may span several /// pages. /// -/// - `num_skips = Some(n)` means that it will skip till the `n + 1`-th occurrance of the repetition +/// - `num_skips = Some(n)` means that it will skip till the `n + 1`-th occurrence of the repetition /// level of `0` (i.e. the start of a top-level value / row value). -/// - `num_collects = Some(n)` means that it will collect values till the `n + 1`-th occurrance of +/// - `num_collects = Some(n)` means that it will collect values till the `n + 1`-th occurrence of /// the repetition level of `0` (i.e. the start of a top-level value / row value). struct DecodingState { num_skips: Option, From e395c93132c6d9a881032d2ec81fe4e360207837 Mon Sep 17 00:00:00 2001 From: coastalwhite Date: Thu, 24 Oct 2024 13:54:55 +0200 Subject: [PATCH 32/32] rustfmt --- .../src/arrow/read/deserialize/nested_utils.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs index a316317d1bae..f4a53a123089 100644 --- a/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs +++ b/crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs @@ -2,8 +2,7 @@ use arrow::bitmap::utils::BitmapIter; use arrow::bitmap::{Bitmap, MutableBitmap}; use arrow::datatypes::ArrowDataType; -use super::utils; -use super::{BasicDecompressor, Filter}; +use super::{utils, BasicDecompressor, Filter}; use crate::parquet::encoding::hybrid_rle::{HybridRleChunk, HybridRleDecoder}; use crate::parquet::error::ParquetResult; use crate::parquet::page::{split_buffer, DataPage}; @@ -396,9 +395,9 @@ fn collect_level_values( /// pages. /// /// - `num_skips = Some(n)` means that it will skip till the `n + 1`-th occurrence of the repetition -/// level of `0` (i.e. the start of a top-level value / row value). +/// level of `0` (i.e. the start of a top-level value / row value). /// - `num_collects = Some(n)` means that it will collect values till the `n + 1`-th occurrence of -/// the repetition level of `0` (i.e. the start of a top-level value / row value). +/// the repetition level of `0` (i.e. the start of a top-level value / row value). struct DecodingState { num_skips: Option, num_collects: Option,