perf: Batch nested embed parquet decoding (#17549)

pola-rs · Jul 10, 2024 · 8fd0c5b · 8fd0c5b
1 parent a55d9bc
commit 8fd0c5b
Show file tree

Hide file tree

Showing 5 changed files with 137 additions and 130 deletions.
diff --git a/crates/polars-parquet/src/arrow/read/deserialize/binary/utils.rs b/crates/polars-parquet/src/arrow/read/deserialize/binary/utils.rs
@@ -108,7 +108,8 @@ impl<'a> Iterator for BinaryIter<'a> {
 
     #[inline]
     fn next(&mut self) -> Option<Self::Item> {
-        if self.values.is_empty() || self.num_values == 0 {
+        if self.num_values == 0 {
+            assert!(self.values.is_empty());
             return None;
         }
 

diff --git a/crates/polars-parquet/src/arrow/read/deserialize/mod.rs b/crates/polars-parquet/src/arrow/read/deserialize/mod.rs
@@ -46,7 +46,7 @@ pub fn create_list(
     nested: &mut NestedState,
     values: Box<dyn Array>,
 ) -> Box<dyn Array> {
-    let (mut offsets, validity) = nested.nested.pop().unwrap().take();
+    let (mut offsets, validity) = nested.pop().unwrap();
     match data_type.to_logical_type() {
         ArrowDataType::List(_) => {
             offsets.push(values.len() as i64);
@@ -89,7 +89,7 @@ pub fn create_map(
     nested: &mut NestedState,
     values: Box<dyn Array>,
 ) -> Box<dyn Array> {
-    let (mut offsets, validity) = nested.nested.pop().unwrap().take();
+    let (mut offsets, validity) = nested.pop().unwrap();
     match data_type.to_logical_type() {
         ArrowDataType::Map(_, _) => {
             offsets.push(values.len() as i64);
@@ -147,7 +147,7 @@ where
                 chunk_size,
                 num_rows,
             )?
-            .map(|x| Ok((NestedState::new(vec![]), x?))),
+            .map(|x| Ok((NestedState::default(), x?))),
         ));
     }
 

diff --git a/crates/polars-parquet/src/arrow/read/deserialize/nested.rs b/crates/polars-parquet/src/arrow/read/deserialize/nested.rs
@@ -13,7 +13,7 @@ where
 {
     Box::new(iter.map(|x| {
         x.map(|(mut nested, array)| {
-            let _ = nested.nested.pop().unwrap(); // the primitive
+            let _ = nested.pop().unwrap(); // the primitive
             (nested, array)
         })
     }))
@@ -28,7 +28,7 @@ where
 {
     Box::new(iter.map(|x| {
         x.map(|(mut nested, array)| {
-            let _ = nested.nested.pop().unwrap(); // the primitive
+            let _ = nested.pop().unwrap(); // the primitive
             (nested, Box::new(array) as _)
         })
     }))
@@ -322,7 +322,7 @@ where
                                 validity,
                             )?);
 
-                            let _ = nested.nested.pop().unwrap(); // the primitive
+                            let _ = nested.pop().unwrap(); // the primitive
 
                             Ok((nested, array))
                         });
@@ -380,7 +380,7 @@ where
                                 validity,
                             )?);
 
-                            let _ = nested.nested.pop().unwrap(); // the primitive
+                            let _ = nested.pop().unwrap(); // the primitive
 
                             Ok((nested, array))
                         });
@@ -411,7 +411,7 @@ where
                                 validity,
                             )?);
 
-                            let _ = nested.nested.pop().unwrap(); // the primitive
+                            let _ = nested.pop().unwrap(); // the primitive
 
                             Ok((nested, array))
                         });