Fixed using lower limit than size of first parquet row group (#1046)

jorgecarleitao · Jun 5, 2022 · 745c199 · 745c199
1 parent 287dae3
commit 745c199
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 10 deletions.
diff --git a/src/io/parquet/read/file.rs b/src/io/parquet/read/file.rs
@@ -115,12 +115,15 @@ impl<R: Read + Seek> FileReader<R> {
     fn next_row_group(&mut self) -> Result<Option<RowGroupDeserializer>> {
         let result = self.row_groups.next().transpose()?;
 
-        self.remaining_rows = self.remaining_rows.saturating_sub(
-            result
-                .as_ref()
-                .map(|x| x.num_rows())
-                .unwrap_or(self.remaining_rows),
-        );
+        // If current_row_group is None, then there will be no elements to remove.
+        if self.current_row_group.is_some() {
+            self.remaining_rows = self.remaining_rows.saturating_sub(
+                result
+                    .as_ref()
+                    .map(|x| x.num_rows())
+                    .unwrap_or(self.remaining_rows),
+            );
+        }
         Ok(result)
     }
 }

diff --git a/src/io/parquet/read/row_group.rs b/src/io/parquet/read/row_group.rs
@@ -78,10 +78,12 @@ impl Iterator for RowGroupDeserializer {
             })
             .collect::<Result<Vec<_>>>()
             .map(Chunk::new);
-        self.remaining_rows -= chunk
-            .as_ref()
-            .map(|x| x.len())
-            .unwrap_or(self.remaining_rows);
+        self.remaining_rows = self.remaining_rows.saturating_sub(
+            chunk
+                .as_ref()
+                .map(|x| x.len())
+                .unwrap_or(self.remaining_rows),
+        );
 
         Some(chunk)
     }