perf: Add ArrayChunks to optimize codegen of BatchDecoder (#17632)

pola-rs · Jul 15, 2024 · 8bad2fd · 8bad2fd
1 parent 7a4a9d0
commit 8bad2fd
Show file tree

Hide file tree

Showing 14 changed files with 467 additions and 258 deletions.
diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binary/basic.rs b/crates/polars-parquet/src/arrow/read/deserialize/binary/basic.rs
@@ -75,7 +75,7 @@ impl<'a, O: Offset> StateTranslation<'a, BinaryDecoder<O>> for BinaryStateTransl
                 values,
                 page_values,
             )?,
-            (T::Dictionary(page), None) => {
+            (T::Dictionary(page, _), None) => {
                 // Already done on the dict.
                 validate_utf8 = false;
                 let page_dict = &page.dict;
@@ -90,7 +90,7 @@ impl<'a, O: Offset> StateTranslation<'a, BinaryDecoder<O>> for BinaryStateTransl
                 }
                 page.values.get_result()?;
             },
-            (T::Dictionary(page), Some(page_validity)) => {
+            (T::Dictionary(page, _), Some(page_validity)) => {
                 // Already done on the dict.
                 validate_utf8 = false;
                 let page_dict = &page.dict;

diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binary/decoders.rs b/crates/polars-parquet/src/arrow/read/deserialize/binary/decoders.rs
@@ -1,5 +1,5 @@
 use arrow::array::specification::try_check_utf8;
-use arrow::array::{BinaryArray, MutableBinaryValuesArray};
+use arrow::array::{BinaryArray, MutableBinaryValuesArray, View};
 use polars_error::PolarsResult;
 
 use super::super::utils;
@@ -140,7 +140,7 @@ impl<'a> ValuesDictionary<'a> {
 #[derive(Debug)]
 pub(crate) enum BinaryStateTranslation<'a> {
     Unit(BinaryIter<'a>),
-    Dictionary(ValuesDictionary<'a>),
+    Dictionary(ValuesDictionary<'a>, Option<Vec<View>>),
     Delta(Delta<'a>),
     DeltaBytes(DeltaBytes<'a>),
 }
@@ -160,6 +160,7 @@ impl<'a> BinaryStateTranslation<'a> {
                 }
                 Ok(BinaryStateTranslation::Dictionary(
                     ValuesDictionary::try_new(page, dict)?,
+                    None,
                 ))
             },
             (Encoding::Plain, _) => {
@@ -180,7 +181,7 @@ impl<'a> BinaryStateTranslation<'a> {
     pub(crate) fn len_when_not_nullable(&self) -> usize {
         match self {
             Self::Unit(v) => v.len_when_not_nullable(),
-            Self::Dictionary(v) => v.len(),
+            Self::Dictionary(v, _) => v.len(),
             Self::Delta(v) => v.len(),
             Self::DeltaBytes(v) => v.size_hint().0,
         }
@@ -193,7 +194,7 @@ impl<'a> BinaryStateTranslation<'a> {
 
         match self {
             Self::Unit(t) => _ = t.by_ref().nth(n - 1),
-            Self::Dictionary(t) => t.values.skip_in_place(n)?,
+            Self::Dictionary(t, _) => t.values.skip_in_place(n)?,
             Self::Delta(t) => _ = t.by_ref().nth(n - 1),
             Self::DeltaBytes(t) => _ = t.by_ref().nth(n - 1),
         }

diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binview/basic.rs b/crates/polars-parquet/src/arrow/read/deserialize/binview/basic.rs
@@ -69,13 +69,13 @@ impl<'a> StateTranslation<'a, BinViewDecoder> for BinaryStateTranslation<'a> {
                 values,
                 page_values,
             )?,
-            (Self::Dictionary(page), None) => {
+            (Self::Dictionary(page, views_dict), None) => {
                 // Already done on the dict.
                 validate_utf8 = false;
 
-                let page_dict = &page.dict;
-                let views_dict = binary_views_dict(values, page_dict);
-                let translator = DictionaryTranslator(&views_dict);
+                let views_dict =
+                    views_dict.get_or_insert_with(|| binary_views_dict(values, page.dict));
+                let translator = DictionaryTranslator(views_dict);
 
                 page.values.translate_and_collect_n_into(
                     values.views_mut(),
@@ -86,13 +86,13 @@ impl<'a> StateTranslation<'a, BinViewDecoder> for BinaryStateTranslation<'a> {
                     validity.extend_constant(additional, true);
                 }
             },
-            (Self::Dictionary(page), Some(page_validity)) => {
+            (Self::Dictionary(page, views_dict), Some(page_validity)) => {
                 // Already done on the dict.
                 validate_utf8 = false;
 
-                let page_dict = &page.dict;
-                let views_dict = binary_views_dict(values, page_dict);
-                let translator = DictionaryTranslator(&views_dict);
+                let views_dict =
+                    views_dict.get_or_insert_with(|| binary_views_dict(values, page.dict));
+                let translator = DictionaryTranslator(views_dict);
                 let collector = TranslatedHybridRle::new(&mut page.values, &translator);
 
                 extend_from_decoder(validity, page_validity, Some(additional), values, collector)?;

diff --git a/crates/polars-parquet/src/arrow/read/deserialize/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/mod.rs
@@ -1,4 +1,20 @@
 //! APIs to read from Parquet format.
+
+macro_rules! decoder_fn {
+    (($x:ident $(, $field:ident:$ty:ty)* $(,)?) => <$p:ty, $t:ty> => $expr:expr) => {{
+        #[derive(Clone, Copy)]
+        struct DecoderFn($($ty),*);
+        impl crate::arrow::read::deserialize::primitive::DecoderFunction<$p, $t> for DecoderFn {
+            #[inline(always)]
+            fn decode(self, $x: $p) -> $t {
+                let Self($($field),*) = self;
+                $expr
+            }
+        }
+        DecoderFn($($field),*)
+    }};
+}
+
 mod binary;
 mod binview;
 mod boolean;

diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested.rs
@@ -3,6 +3,7 @@ use arrow::match_integer_type;
 use ethnum::I256;
 use polars_error::polars_bail;
 
+use self::primitive::{AsDecoderFunction, IntoDecoderFunction, UnitDecoderFunction};
 use super::*;
 
 /// Converts an iterator of arrays to a trait object returning trait objects
@@ -80,7 +81,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i32| x as i8,
+                AsDecoderFunction::<i32, i8>::default(),
             ))
         },
         Primitive(Int16) => {
@@ -92,7 +93,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i32| x as i16,
+                AsDecoderFunction::<i32, i16>::default(),
             ))
         },
         Primitive(Int32) => {
@@ -104,7 +105,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i32| x,
+                UnitDecoderFunction::<i32>::default(),
             ))
         },
         Primitive(Int64) => {
@@ -116,7 +117,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i64| x,
+                UnitDecoderFunction::<i64>::default(),
             ))
         },
         Primitive(UInt8) => {
@@ -128,7 +129,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i32| x as u8,
+                AsDecoderFunction::<i32, u8>::default(),
             ))
         },
         Primitive(UInt16) => {
@@ -140,7 +141,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i32| x as u16,
+                AsDecoderFunction::<i32, u16>::default(),
             ))
         },
         Primitive(UInt32) => {
@@ -153,7 +154,7 @@ where
                     field.data_type().clone(),
                     num_rows,
                     chunk_size,
-                    |x: i32| x as u32,
+                    AsDecoderFunction::<i32, u32>::default(),
                 )),
                 // some implementations of parquet write arrow's u32 into i64.
                 PhysicalType::Int64 => primitive(primitive::NestedIter::new(
@@ -162,7 +163,7 @@ where
                     field.data_type().clone(),
                     num_rows,
                     chunk_size,
-                    |x: i64| x as u32,
+                    AsDecoderFunction::<i64, u32>::default(),
                 )),
                 other => {
                     polars_bail!(ComputeError:
@@ -180,7 +181,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: i64| x as u64,
+                AsDecoderFunction::<i64, u64>::default(),
             ))
         },
         Primitive(Float32) => {
@@ -192,7 +193,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: f32| x,
+                UnitDecoderFunction::<f32>::default(),
             ))
         },
         Primitive(Float64) => {
@@ -204,7 +205,7 @@ where
                 field.data_type().clone(),
                 num_rows,
                 chunk_size,
-                |x: f64| x,
+                UnitDecoderFunction::<f64>::default(),
             ))
         },
         BinaryView | Utf8View => {
@@ -283,15 +284,15 @@ where
                         field.data_type.clone(),
                         num_rows,
                         chunk_size,
-                        |x: i32| x as i128,
+                        IntoDecoderFunction::<i32, i128>::default(),
                     )),
                     PhysicalType::Int64 => primitive(primitive::NestedIter::new(
                         columns.pop().unwrap(),
                         init,
                         field.data_type.clone(),
                         num_rows,
                         chunk_size,
-                        |x: i64| x as i128,
+                        IntoDecoderFunction::<i64, i128>::default(),
                     )),
                     PhysicalType::FixedLenByteArray(n) if n > 16 => {
                         polars_bail!(
@@ -346,15 +347,15 @@ where
                         field.data_type.clone(),
                         num_rows,
                         chunk_size,
-                        |x: i32| i256(I256::new(x as i128)),
+                        decoder_fn!((x) => <i32, i256> => i256(I256::new(x as i128))),
                     )),
                     PhysicalType::Int64 => primitive(primitive::NestedIter::new(
                         columns.pop().unwrap(),
                         init,
                         field.data_type.clone(),
                         num_rows,
                         chunk_size,
-                        |x: i64| i256(I256::new(x as i128)),
+                        decoder_fn!((x) => <i64, i256> => i256(I256::new(x as i128))),
                     )),
                     PhysicalType::FixedLenByteArray(n) if n <= 16 => {
                         let iter = fixed_size_binary::NestedIter::new(
@@ -501,39 +502,39 @@ fn dict_read<'a, K: DictionaryKey, I: 'a + PagesIter>(
             data_type,
             num_rows,
             chunk_size,
-            |x: i32| x as u8,
+            AsDecoderFunction::<i32, u8>::default(),
         )),
         UInt16 => primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
             iter,
             init,
             data_type,
             num_rows,
             chunk_size,
-            |x: i32| x as u16,
+            AsDecoderFunction::<i32, u16>::default(),
         )),
         UInt32 => primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
             iter,
             init,
             data_type,
             num_rows,
             chunk_size,
-            |x: i32| x as u32,
+            AsDecoderFunction::<i32, u32>::default(),
         )),
         Int8 => primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
             iter,
             init,
             data_type,
             num_rows,
             chunk_size,
-            |x: i32| x as i8,
+            AsDecoderFunction::<i32, i8>::default(),
         )),
         Int16 => primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
             iter,
             init,
             data_type,
             num_rows,
             chunk_size,
-            |x: i32| x as i16,
+            AsDecoderFunction::<i32, i16>::default(),
         )),
         Int32 | Date32 | Time32(_) | Interval(IntervalUnit::YearMonth) => {
             primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
@@ -542,7 +543,7 @@ fn dict_read<'a, K: DictionaryKey, I: 'a + PagesIter>(
                 data_type,
                 num_rows,
                 chunk_size,
-                |x: i32| x,
+                UnitDecoderFunction::<i32>::default(),
             ))
         },
         Int64 | Date64 | Time64(_) | Duration(_) => {
@@ -552,7 +553,7 @@ fn dict_read<'a, K: DictionaryKey, I: 'a + PagesIter>(
                 data_type,
                 num_rows,
                 chunk_size,
-                |x: i64| x as i32,
+                AsDecoderFunction::<i64, i32>::default(),
             ))
         },
         Float32 => primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
@@ -561,15 +562,15 @@ fn dict_read<'a, K: DictionaryKey, I: 'a + PagesIter>(
             data_type,
             num_rows,
             chunk_size,
-            |x: f32| x,
+            UnitDecoderFunction::<f32>::default(),
         )),
         Float64 => primitive(primitive::NestedDictIter::<K, _, _, _, _>::new(
             iter,
             init,
             data_type,
             num_rows,
             chunk_size,
-            |x: f64| x,
+            UnitDecoderFunction::<f64>::default(),
         )),
         LargeUtf8 | LargeBinary => primitive(binary::NestedDictIter::<K, i64, _>::new(
             iter, init, data_type, num_rows, chunk_size,