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

Commit

Permalink
Do not initiallize the iterator on creation, instead handle None on r…
Browse files Browse the repository at this point in the history
…ow groups
  • Loading branch information
Aron Hansen Berggren authored and Aron Hansen Berggren committed Jun 4, 2022
1 parent 6398333 commit 753a4ce
Showing 1 changed file with 11 additions and 9 deletions.
20 changes: 11 additions & 9 deletions src/io/parquet/read/file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,21 +80,20 @@ impl<R: Read + Seek> FileReader<R> {
metadata: schema_metadata,
};

let mut row_groups = RowGroupReader::new(
let 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,
current_row_group: None,
})
}

Expand All @@ -116,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)
}
}
Expand Down

0 comments on commit 753a4ce

Please sign in to comment.