Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Fixed using lower limit than size of first parquet row group
Browse files Browse the repository at this point in the history
  • Loading branch information
Aron Hansen Berggren authored and Aron Hansen Berggren committed Jun 3, 2022
1 parent 06f8f36 commit 6398333
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 4 deletions.
5 changes: 3 additions & 2 deletions src/io/parquet/read/file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,20 +80,21 @@ impl<R: Read + Seek> FileReader<R> {
metadata: schema_metadata,
};

let row_groups = RowGroupReader::new(
let mut row_groups = RowGroupReader::new(
reader,
schema,
groups_filter,
metadata.row_groups.clone(),
chunk_size,
limit,
);
let current_row_group = row_groups.next().transpose()?;

Ok(Self {
row_groups,
metadata,
remaining_rows: limit.unwrap_or(usize::MAX),
current_row_group: None,
current_row_group,
})
}

Expand Down
5 changes: 3 additions & 2 deletions src/io/parquet/read/row_group.rs
Original file line number Diff line number Diff line change
Expand Up @@ -78,10 +78,11 @@ impl Iterator for RowGroupDeserializer {
})
.collect::<Result<Vec<_>>>()
.map(Chunk::new);
self.remaining_rows -= chunk
self.remaining_rows = self.remaining_rows.saturating_sub(chunk
.as_ref()
.map(|x| x.len())
.unwrap_or(self.remaining_rows);
.unwrap_or(self.remaining_rows)
);

Some(chunk)
}
Expand Down

0 comments on commit 6398333

Please sign in to comment.