From 3a025da35ae97ea5f88f8c1dc9b3a1352b102102 Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Mon, 27 Jun 2022 08:12:06 +0000
Subject: [PATCH] Fixed error in reading chunked parquet

---
 src/io/parquet/read/deserialize/utils.rs | 12 ++++++------
 src/io/parquet/read/row_group.rs         |  2 +-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/io/parquet/read/deserialize/utils.rs b/src/io/parquet/read/deserialize/utils.rs
index 9e560e43014..dd2174c4180 100644
--- a/src/io/parquet/read/deserialize/utils.rs
+++ b/src/io/parquet/read/deserialize/utils.rs
@@ -230,6 +230,7 @@ impl<'a> OptionalPageValidity<'a> {
         }
     }
 
+    /// Number of items remaining
     pub fn len(&self) -> usize {
         self.iter.len()
             + self
@@ -300,10 +301,9 @@ pub(super) fn extend_from_decoder<'a, T: Default, P: Pushable<T>, I: Iterator<It
 ) {
     let limit = limit.unwrap_or(usize::MAX);
 
-    // todo: remove `consumed_here` and compute next limit from `consumed`
-    let mut consumed_here = 0;
-    while consumed_here < limit {
-        let run = page_validity.next_limited(limit);
+    let mut remaining = limit;
+    while remaining > 0 {
+        let run = page_validity.next_limited(remaining);
         let run = if let Some(run) = run { run } else { break };
 
         match run {
@@ -325,7 +325,7 @@ pub(super) fn extend_from_decoder<'a, T: Default, P: Pushable<T>, I: Iterator<It
                 }
                 validity.extend_from_slice(values, offset, length);
 
-                consumed_here += length;
+                remaining -= length;
             }
             FilteredHybridEncoded::Repeated { is_set, length } => {
                 validity.extend_constant(length, is_set);
@@ -335,7 +335,7 @@ pub(super) fn extend_from_decoder<'a, T: Default, P: Pushable<T>, I: Iterator<It
                     pushable.extend_constant(length, T::default());
                 }
 
-                consumed_here += length;
+                remaining -= length;
             }
             FilteredHybridEncoded::Skipped(valids) => for _ in values_iter.by_ref().take(valids) {},
         };
diff --git a/src/io/parquet/read/row_group.rs b/src/io/parquet/read/row_group.rs
index 25c21658190..0c7222343c8 100644
--- a/src/io/parquet/read/row_group.rs
+++ b/src/io/parquet/read/row_group.rs
@@ -74,7 +74,7 @@ impl Iterator for RowGroupDeserializer {
                 })
             })
             .collect::<Result<Vec<_>>>()
-            .map(Chunk::new);
+            .and_then(Chunk::try_new);
         self.remaining_rows = self.remaining_rows.saturating_sub(
             chunk
                 .as_ref()