fix(rust, python): fix bug in batched parquet reader that dropped dfs… (

#5506)
pola-rs · Nov 14, 2022 · 319aaaa · 319aaaa
1 parent a7a7e0f
commit 319aaaa
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/polars/polars-io/src/parquet/read_impl.rs b/polars/polars-io/src/parquet/read_impl.rs
@@ -356,7 +356,7 @@ impl BatchedParquetReader {
 
     pub fn next_batches(&mut self, n: usize) -> PolarsResult<Option<Vec<DataFrame>>> {
         // fill up fifo stack
-        if self.row_group_offset < self.n_row_groups && self.chunks_fifo.len() < n {
+        if self.row_group_offset <= self.n_row_groups && self.chunks_fifo.len() < n {
             let dfs = match self.parallel {
                 ParallelStrategy::Columns => {
                     let dfs = rg_to_dfs(
@@ -416,12 +416,11 @@ impl BatchedParquetReader {
             let mut chunks = Vec::with_capacity(n);
             let mut i = 0;
             while let Some(df) = self.chunks_fifo.pop_front() {
+                chunks.push(df);
+                i += 1;
                 if i == n {
                     break;
                 }
-
-                chunks.push(df);
-                i += 1;
             }
 
             Ok(Some(chunks))

diff --git a/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs b/polars/polars-lazy/polars-pipe/src/executors/sinks/groupby/generic.rs
@@ -371,7 +371,7 @@ impl Sink for GenericGroupbySink {
                                     )
                                 };
 
-                                // extend the keys buffer with the new keys from othger
+                                // extend the keys buffer with the new keys from other
                                 keys_buffer_self.extend_from_slice(keys_other);
 
                                 // insert the keys and values_offset