perf(rust): faster decode on Parquet HybridRLE (#17208)

Co-authored-by: Ritchie Vink <[email protected]>

crates/polars-parquet/src/arrow/read/deserialize/binary/decoders.rs

@@ -179,13 +179,13 @@ impl<'a> FilteredDelta<'a> {
 
 #[derive(Debug)]
 pub(crate) struct RequiredDictionary<'a> {
-    pub values: hybrid_rle::HybridRleDecoder<'a>,
+    pub values: hybrid_rle::BufferedHybridRleDecoderIter<'a>,
     pub dict: &'a BinaryDict,
 }
 
 impl<'a> RequiredDictionary<'a> {
     pub fn try_new(page: &'a DataPage, dict: &'a BinaryDict) -> PolarsResult<Self> {
-        let values = utils::dict_indices_decoder(page)?;
+        let values = utils::dict_indices_decoder(page)?.into_iter();
 
         Ok(Self { dict, values })
     }
@@ -198,13 +198,13 @@ impl<'a> RequiredDictionary<'a> {
 
 #[derive(Debug)]
 pub(crate) struct FilteredRequiredDictionary<'a> {
-    pub values: SliceFilteredIter<hybrid_rle::HybridRleDecoder<'a>>,
+    pub values: SliceFilteredIter<hybrid_rle::BufferedHybridRleDecoderIter<'a>>,
     pub dict: &'a BinaryDict,
 }
 
 impl<'a> FilteredRequiredDictionary<'a> {
     pub fn try_new(page: &'a DataPage, dict: &'a BinaryDict) -> PolarsResult<Self> {
-        let values = utils::dict_indices_decoder(page)?;
+        let values = utils::dict_indices_decoder(page)?.into_iter();
 
         let rows = get_selected_rows(page);
         let values = SliceFilteredIter::new(values, rows);
@@ -220,20 +220,20 @@ impl<'a> FilteredRequiredDictionary<'a> {
 
 #[derive(Debug)]
 pub(crate) struct ValuesDictionary<'a> {
-    pub values: hybrid_rle::HybridRleDecoder<'a>,
+    pub values: hybrid_rle::BufferedHybridRleDecoderIter<'a>,
     pub dict: &'a BinaryDict,
 }
 
 impl<'a> ValuesDictionary<'a> {
     pub fn try_new(page: &'a DataPage, dict: &'a BinaryDict) -> PolarsResult<Self> {
-        let values = utils::dict_indices_decoder(page)?;
+        let values = utils::dict_indices_decoder(page)?.into_iter();
 
         Ok(Self { dict, values })
     }
 
     #[inline]
     pub fn len(&self) -> usize {
-        self.values.size_hint().0
+        self.values.len()
     }
 }
 

crates/polars-parquet/src/arrow/read/deserialize/dictionary/mod.rs

@@ -5,14 +5,17 @@ use std::collections::VecDeque;
 use arrow::array::{Array, DictionaryArray, DictionaryKey, PrimitiveArray};
 use arrow::bitmap::MutableBitmap;
 use arrow::datatypes::ArrowDataType;
+pub use nested::next_dict as nested_next_dict;
+use polars_error::{polars_err, PolarsResult};
+use polars_utils::iter::FallibleIterator;
 
 use super::utils::{
     self, dict_indices_decoder, extend_from_decoder, get_selected_rows, DecodedState, Decoder,
     FilteredOptionalPageValidity, MaybeNext, OptionalPageValidity,
 };
 use super::PagesIter;
 use crate::parquet::deserialize::SliceFilteredIter;
-use crate::parquet::encoding::hybrid_rle::HybridRleDecoder;
+use crate::parquet::encoding::hybrid_rle::BufferedHybridRleDecoderIter;
 use crate::parquet::encoding::Encoding;
 use crate::parquet::page::{DataPage, DictPage, Page};
 use crate::parquet::schema::Repetition;
@@ -23,29 +26,32 @@ pub enum State<'a> {
     Optional(Optional<'a>),
     Required(Required<'a>),
     FilteredRequired(FilteredRequired<'a>),
-    FilteredOptional(FilteredOptionalPageValidity<'a>, HybridRleDecoder<'a>),
+    FilteredOptional(
+        FilteredOptionalPageValidity<'a>,
+        BufferedHybridRleDecoderIter<'a>,
+    ),
 }
 
 #[derive(Debug)]
 pub struct Required<'a> {
-    values: HybridRleDecoder<'a>,
+    values: BufferedHybridRleDecoderIter<'a>,
 }
 
 impl<'a> Required<'a> {
     fn try_new(page: &'a DataPage) -> PolarsResult<Self> {
-        let values = dict_indices_decoder(page)?;
+        let values = dict_indices_decoder(page)?.into_iter();
         Ok(Self { values })
     }
 }
 
 #[derive(Debug)]
 pub struct FilteredRequired<'a> {
-    values: SliceFilteredIter<HybridRleDecoder<'a>>,
+    values: SliceFilteredIter<BufferedHybridRleDecoderIter<'a>>,
 }
 
 impl<'a> FilteredRequired<'a> {
     fn try_new(page: &'a DataPage) -> PolarsResult<Self> {
-        let values = dict_indices_decoder(page)?;
+        let values = dict_indices_decoder(page)?.into_iter();
 
         let rows = get_selected_rows(page);
         let values = SliceFilteredIter::new(values, rows);
@@ -56,13 +62,13 @@ impl<'a> FilteredRequired<'a> {
 
 #[derive(Debug)]
 pub struct Optional<'a> {
-    values: HybridRleDecoder<'a>,
+    values: BufferedHybridRleDecoderIter<'a>,
     validity: OptionalPageValidity<'a>,
 }
 
 impl<'a> Optional<'a> {
     fn try_new(page: &'a DataPage) -> PolarsResult<Self> {
-        let values = dict_indices_decoder(page)?;
+        let values = dict_indices_decoder(page)?.into_iter();
 
         Ok(Self {
             values,
@@ -132,7 +138,7 @@ where
             (Encoding::PlainDictionary | Encoding::RleDictionary, true, true) => {
                 Ok(State::FilteredOptional(
                     FilteredOptionalPageValidity::try_new(page)?,
-                    dict_indices_decoder(page)?,
+                    dict_indices_decoder(page)?.into_iter(),
                 ))
             },
             _ => Err(utils::not_implemented(page)),
@@ -316,7 +322,3 @@ pub(super) fn next_dict<K: DictionaryKey, I: PagesIter, F: Fn(&DictPage) -> Box<
         },
     }
 }
-
-pub use nested::next_dict as nested_next_dict;
-use polars_error::{polars_err, PolarsResult};
-use polars_utils::iter::FallibleIterator;

crates/polars-parquet/src/arrow/read/deserialize/dictionary/nested.rs

@@ -10,21 +10,21 @@ use super::super::super::PagesIter;
 use super::super::nested_utils::*;
 use super::super::utils::{dict_indices_decoder, not_implemented, MaybeNext, PageState};
 use super::finish_key;
-use crate::parquet::encoding::hybrid_rle::HybridRleDecoder;
+use crate::parquet::encoding::hybrid_rle::BufferedHybridRleDecoderIter;
 use crate::parquet::encoding::Encoding;
 use crate::parquet::page::{DataPage, DictPage, Page};
 use crate::parquet::schema::Repetition;
 
 // The state of a required DataPage with a boolean physical type
 #[derive(Debug)]
 pub struct Required<'a> {
-    values: HybridRleDecoder<'a>,
+    values: BufferedHybridRleDecoderIter<'a>,
     length: usize,
 }
 
 impl<'a> Required<'a> {
     fn try_new(page: &'a DataPage) -> PolarsResult<Self> {
-        let values = dict_indices_decoder(page)?;
+        let values = dict_indices_decoder(page)?.into_iter();
         let length = page.num_values();
         Ok(Self { values, length })
     }
@@ -34,7 +34,7 @@ impl<'a> Required<'a> {
 #[allow(clippy::large_enum_variant)]
 #[derive(Debug)]
 pub enum State<'a> {
-    Optional(HybridRleDecoder<'a>),
+    Optional(BufferedHybridRleDecoderIter<'a>),
     Required(Required<'a>),
 }
 
@@ -89,7 +89,9 @@ impl<'a, K: DictionaryKey> NestedDecoder<'a> for DictionaryDecoder<K> {
 
         match (page.encoding(), is_optional, is_filtered) {
             (Encoding::RleDictionary | Encoding::PlainDictionary, true, false) => {
-                dict_indices_decoder(page).map(State::Optional)
+                dict_indices_decoder(page)
+                    .map(|v| v.into_iter())
+                    .map(State::Optional)
             },
             (Encoding::RleDictionary | Encoding::PlainDictionary, false, false) => {
                 Required::try_new(page).map(State::Required)
@@ -161,7 +163,7 @@ pub fn next_dict<K: DictionaryKey, I: PagesIter, F: Fn(&DictPage) -> Box<dyn Arr
     match iter.next() {
         Err(e) => MaybeNext::Some(Err(e.into())),
         Ok(Some(page)) => {
-            let (page, dict) = match (&dict, page) {
+            let (page, dict) = match (&dict, &page) {
                 (None, Page::Data(_)) => {
                     return MaybeNext::Some(Err(polars_err!(ComputeError:
                         "not implemented: dictionary arrays from non-dict-encoded pages",

crates/polars-parquet/src/arrow/read/deserialize/fixed_size_binary/basic.rs

@@ -84,13 +84,13 @@ impl<'a> FilteredRequired<'a> {
 
 #[derive(Debug)]
 pub(super) struct RequiredDictionary<'a> {
-    pub values: hybrid_rle::HybridRleDecoder<'a>,
+    pub values: hybrid_rle::BufferedHybridRleDecoderIter<'a>,
     pub dict: &'a Dict,
 }
 
 impl<'a> RequiredDictionary<'a> {
     pub(super) fn try_new(page: &'a DataPage, dict: &'a Dict) -> PolarsResult<Self> {
-        let values = dict_indices_decoder(page)?;
+        let values = dict_indices_decoder(page)?.into_iter();
 
         Ok(Self { dict, values })
     }
@@ -103,14 +103,14 @@ impl<'a> RequiredDictionary<'a> {
 
 #[derive(Debug)]
 pub(super) struct OptionalDictionary<'a> {
-    pub(super) values: hybrid_rle::HybridRleDecoder<'a>,
+    pub(super) values: hybrid_rle::BufferedHybridRleDecoderIter<'a>,
     pub(super) validity: OptionalPageValidity<'a>,
     pub(super) dict: &'a Dict,
 }
 
 impl<'a> OptionalDictionary<'a> {
     pub(super) fn try_new(page: &'a DataPage, dict: &'a Dict) -> PolarsResult<Self> {
-        let values = dict_indices_decoder(page)?;
+        let values = dict_indices_decoder(page)?.into_iter();
 
         Ok(Self {
             values,

crates/polars-parquet/src/arrow/read/deserialize/nested_utils.rs

@@ -1,4 +1,5 @@
 use std::collections::VecDeque;
+use std::iter::{Peekable, Zip};
 
 use arrow::array::Array;
 use arrow::bitmap::MutableBitmap;
@@ -7,7 +8,7 @@ use polars_utils::slice::GetSaferUnchecked;
 
 use super::super::PagesIter;
 use super::utils::{DecodedState, MaybeNext, PageState};
-use crate::parquet::encoding::hybrid_rle::HybridRleDecoder;
+use crate::parquet::encoding::hybrid_rle::{BufferedHybridRleDecoderIter, HybridRleDecoder};
 use crate::parquet::page::{split_buffer, DataPage, DictPage, Page};
 use crate::parquet::read::levels::get_bit_width;
 
@@ -238,7 +239,7 @@ pub fn init_nested(init: &[InitNested], capacity: usize) -> NestedState {
 }
 
 pub struct NestedPage<'a> {
-    iter: std::iter::Peekable<std::iter::Zip<HybridRleDecoder<'a>, HybridRleDecoder<'a>>>,
+    iter: Peekable<Zip<BufferedHybridRleDecoderIter<'a>, BufferedHybridRleDecoderIter<'a>>>,
 }
 
 impl<'a> NestedPage<'a> {
@@ -251,9 +252,12 @@ impl<'a> NestedPage<'a> {
         let max_def_level = page.descriptor.max_def_level;
 
         let reps =
-            HybridRleDecoder::try_new(rep_levels, get_bit_width(max_rep_level), page.num_values())?;
+            HybridRleDecoder::new(rep_levels, get_bit_width(max_rep_level), page.num_values());
         let defs =
-            HybridRleDecoder::try_new(def_levels, get_bit_width(max_def_level), page.num_values())?;
+            HybridRleDecoder::new(def_levels, get_bit_width(max_def_level), page.num_values());
+
+        let reps = reps.into_iter();
+        let defs = defs.into_iter();
 
         let iter = reps.zip(defs).peekable();
 
@@ -523,7 +527,7 @@ where
             }
         },
         Ok(Some(page)) => {
-            let page = match page {
+            let page = match &page {
                 Page::Data(page) => page,
                 Page::Dict(dict_page) => {
                     *dict = Some(decoder.deserialize_dict(dict_page));

crates/polars-parquet/src/arrow/read/deserialize/primitive/basic.rs

@@ -65,7 +65,7 @@ pub(super) struct ValuesDictionary<'a, T>
 where
     T: NativeType,
 {
-    pub values: hybrid_rle::HybridRleDecoder<'a>,
+    pub values: hybrid_rle::BufferedHybridRleDecoderIter<'a>,
     pub dict: &'a Vec<T>,
 }
 
@@ -74,7 +74,7 @@ where
     T: NativeType,
 {
     pub fn try_new(page: &'a DataPage, dict: &'a Vec<T>) -> PolarsResult<Self> {
-        let values = utils::dict_indices_decoder(page)?;
+        let values = utils::dict_indices_decoder(page)?.into_iter();
 
         Ok(Self { dict, values })
     }

crates/polars-parquet/src/arrow/read/deserialize/utils.rs

@@ -3,7 +3,7 @@ use std::collections::VecDeque;
 use arrow::bitmap::utils::BitmapIter;
 use arrow::bitmap::MutableBitmap;
 use arrow::pushable::Pushable;
-use polars_error::{polars_err, to_compute_err, PolarsError, PolarsResult};
+use polars_error::{polars_err, PolarsError, PolarsResult};
 
 use super::super::PagesIter;
 use crate::parquet::deserialize::{
@@ -424,8 +424,8 @@ pub(super) fn next<'a, I: PagesIter, D: Decoder<'a>>(
         Err(e) => MaybeNext::Some(Err(e.into())),
         Ok(Some(page)) => {
             let page = match page {
-                Page::Data(page) => page,
-                Page::Dict(dict_page) => {
+                Page::Data(ref page) => page,
+                Page::Dict(ref dict_page) => {
                     *dict = Some(decoder.deserialize_dict(dict_page));
                     return MaybeNext::More;
                 },
@@ -472,8 +472,11 @@ pub(super) fn dict_indices_decoder(page: &DataPage) -> PolarsResult<hybrid_rle::
     let bit_width = indices_buffer[0];
     let indices_buffer = &indices_buffer[1..];
 
-    hybrid_rle::HybridRleDecoder::try_new(indices_buffer, bit_width as u32, page.num_values())
-        .map_err(to_compute_err)
+    Ok(hybrid_rle::HybridRleDecoder::new(
+        indices_buffer,
+        bit_width as u32,
+        page.num_values(),
+    ))
 }
 
 pub(super) fn page_is_optional(page: &DataPage) -> bool {