Skip to content

Commit

Permalink
parquet writer: Raise an error when the row_group_index overflows i16 (
Browse files Browse the repository at this point in the history
…#6378)

This caused confusing panics down the line because 'ordinal' is
negative.
  • Loading branch information
progval committed Sep 13, 2024
1 parent 60ec869 commit f80bc5f
Showing 1 changed file with 14 additions and 3 deletions.
17 changes: 14 additions & 3 deletions parquet/src/file/writer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -190,14 +190,25 @@ impl<W: Write + Send> SerializedFileWriter<W> {
/// Creates new row group from this file writer.
/// In case of IO error or Thrift error, returns `Err`.
///
/// There is no limit on a number of row groups in a file; however, row groups have
/// There can be at most 2^15 row groups in a file; and row groups have
/// to be written sequentially. Every time the next row group is requested, the
/// previous row group must be finalised and closed using `RowGroupWriter::close` method.
pub fn next_row_group(&mut self) -> Result<SerializedRowGroupWriter<'_, W>> {
self.assert_previous_writer_closed()?;
let ordinal = self.row_group_index;

self.row_group_index += 1;
let ordinal: i16 = ordinal.try_into().map_err(|_| {
ParquetError::General(format!(
"Parquet does not support more than {} row groups per file (currently: {})",
i16::MAX,
ordinal
))
})?;

self.row_group_index = self
.row_group_index
.checked_add(1)
.expect("SerializedFileWriter::row_group_index overflowed");

let bloom_filter_position = self.properties().bloom_filter_position();
let row_groups = &mut self.row_groups;
Expand Down Expand Up @@ -227,7 +238,7 @@ impl<W: Write + Send> SerializedFileWriter<W> {
self.descr.clone(),
self.props.clone(),
&mut self.buf,
ordinal as i16,
ordinal,
Some(Box::new(on_close)),
);
Ok(row_group_writer)
Expand Down

0 comments on commit f80bc5f

Please sign in to comment.