Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion parquet/src/arrow/arrow_reader/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1318,7 +1318,12 @@ impl<T: ChunkReader + 'static> ReaderPageIterator<T> {
// To avoid `i[rg_idx][self.column_idx`] panic, we need to filter out empty `i[rg_idx]`.
let page_locations = offset_index
.filter(|i| !i[rg_idx].is_empty())
.map(|i| i[rg_idx][self.column_idx].page_locations.clone());
.map(|i| {
i[rg_idx][self.column_idx]
.as_ref()
.map(|o| o.page_locations.clone())
})
.unwrap_or(None);
let total_rows = rg.num_rows() as usize;
let reader = self.reader.clone();

Expand Down
167 changes: 87 additions & 80 deletions parquet/src/arrow/arrow_reader/statistics.rs

Large diffs are not rendered by default.

20 changes: 11 additions & 9 deletions parquet/src/arrow/arrow_writer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2964,7 +2964,7 @@ mod tests {
assert!(reader.metadata().offset_index().is_some());
let offset_indexes = &reader.metadata().offset_index().unwrap()[0];

let page_locations = offset_indexes[0].page_locations.clone();
let page_locations = offset_indexes[0].as_ref().unwrap().page_locations.clone();

// We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes
// so we expect one dictionary encoded page and then a page per row thereafter.
Expand Down Expand Up @@ -3372,6 +3372,8 @@ mod tests {
if let Some(col_indexes) = file_meta_data.column_index() {
for rg_idx in col_indexes {
for idx in rg_idx {
assert!(idx.is_some());
let idx = idx.as_ref().unwrap();
assert!(idx.nan_counts().is_some());
let float_idx = match idx {
ColumnIndexMetaData::DOUBLE(idx) => idx,
Expand Down Expand Up @@ -3449,11 +3451,11 @@ mod tests {

assert!(file_meta_data.column_index().is_some());
let col_idx = &file_meta_data.column_index().as_ref().unwrap()[0][0];
assert_eq!(col_idx.num_pages(), 4);
assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);

// test each page
let float_idx = match col_idx {
ColumnIndexMetaData::DOUBLE(idx) => idx,
Some(ColumnIndexMetaData::DOUBLE(idx)) => idx,
_ => panic!("expected double statistics"),
};

Expand Down Expand Up @@ -4947,8 +4949,8 @@ mod tests {

assert_eq!(index.len(), 1);
assert_eq!(index[0].len(), 2); // 2 columns
assert_eq!(index[0][0].page_locations().len(), 1); // 1 page
assert_eq!(index[0][1].page_locations().len(), 1); // 1 page
assert_eq!(index[0][0].as_ref().unwrap().page_locations().len(), 1); // 1 page
assert_eq!(index[0][1].as_ref().unwrap().page_locations().len(), 1); // 1 page
}

#[test]
Expand Down Expand Up @@ -5019,11 +5021,11 @@ mod tests {

let a_idx = &column_index[0][0];
assert!(
matches!(a_idx, ColumnIndexMetaData::BYTE_ARRAY(_)),
matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
"{a_idx:?}"
);
let b_idx = &column_index[0][1];
assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
assert!(b_idx.is_none(), "{b_idx:?}");
}

#[test]
Expand Down Expand Up @@ -5089,9 +5091,9 @@ mod tests {
assert_eq!(column_index[0].len(), 2); // 2 columns

let a_idx = &column_index[0][0];
assert!(matches!(a_idx, ColumnIndexMetaData::NONE), "{a_idx:?}");
assert!(a_idx.is_none(), "{a_idx:?}");
let b_idx = &column_index[0][1];
assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
assert!(b_idx.is_none(), "{b_idx:?}");
}

#[test]
Expand Down
18 changes: 11 additions & 7 deletions parquet/src/arrow/in_memory_row_group.rs
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ use std::sync::Arc;
/// An in-memory collection of column chunks
#[derive(Debug)]
pub(crate) struct InMemoryRowGroup<'a> {
pub(crate) offset_index: Option<&'a [OffsetIndexMetaData]>,
pub(crate) offset_index: Option<&'a [Option<OffsetIndexMetaData>]>,
/// Column chunks for this row group
pub(crate) column_chunks: Vec<Option<Arc<ColumnChunkData>>>,
pub(crate) row_count: usize,
Expand Down Expand Up @@ -85,7 +85,13 @@ impl InMemoryRowGroup<'_> {
// then we need to also fetch a dictionary page.
let mut ranges: Vec<Range<u64>> = vec![];
let (start, _len) = chunk_meta.byte_range();
match offset_index[idx].page_locations.first() {
let Some(offset_idx) = offset_index[idx].as_ref() else {
// No offset index for this column, fetch the entire column
ranges.push(start..start + _len);
return ranges;
};

match offset_idx.page_locations.first() {
Some(first) if first.offset as u64 != start => {
ranges.push(start..first.offset as u64);
}
Expand All @@ -96,11 +102,9 @@ impl InMemoryRowGroup<'_> {
// (see doc comment for this function for details on `cache_mask`)
let use_expanded = cache_mask.map(|m| m.leaf_included(idx)).unwrap_or(false);
if use_expanded {
ranges.extend(
expanded_selection.scan_ranges(&offset_index[idx].page_locations),
);
ranges.extend(expanded_selection.scan_ranges(&offset_idx.page_locations));
} else {
ranges.extend(selection.scan_ranges(&offset_index[idx].page_locations));
ranges.extend(selection.scan_ranges(&offset_idx.page_locations));
}
page_start_offsets.push(ranges.iter().map(|range| range.start).collect());

Expand Down Expand Up @@ -203,7 +207,7 @@ impl RowGroups for InMemoryRowGroup<'_> {
.offset_index
// filter out empty offset indexes (old versions specified Some(vec![]) when no present)
.filter(|index| !index.is_empty())
.map(|index| index[i].page_locations.clone());
.and_then(|index| index[i].as_ref().map(|idx| idx.page_locations.clone()));
let column_chunk_metadata = self.metadata.row_group(self.row_group_idx).column(i);
let page_reader = SerializedPageReader::new(
data.clone(),
Expand Down
2 changes: 1 addition & 1 deletion parquet/src/arrow/push_decoder/reader_builder/data.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,7 @@ impl<'a> DataRequestBuilder<'a> {
fn get_offset_index(
parquet_metadata: &ParquetMetaData,
row_group_idx: usize,
) -> Option<&[OffsetIndexMetaData]> {
) -> Option<&[Option<OffsetIndexMetaData>]> {
parquet_metadata
.offset_index()
// filter out empty offset indexes (old versions specified Some(vec![]) when no present)
Expand Down
40 changes: 23 additions & 17 deletions parquet/src/arrow/push_decoder/reader_builder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -849,7 +849,10 @@ impl RowGroupReaderBuilder {
}

/// Get the offset index for the specified row group, if any
fn row_group_offset_index(&self, row_group_idx: usize) -> Option<&[OffsetIndexMetaData]> {
fn row_group_offset_index(
&self,
row_group_idx: usize,
) -> Option<&[Option<OffsetIndexMetaData>]> {
self.metadata
.offset_index()
.filter(|index| !index.is_empty())
Expand Down Expand Up @@ -881,7 +884,7 @@ impl RowGroupReaderBuilder {
fn prepare_selection_for_page_skipping(
plan_builder: ReadPlanBuilder,
projection_mask: &ProjectionMask,
offset_index: Option<&[OffsetIndexMetaData]>,
offset_index: Option<&[Option<OffsetIndexMetaData>]>,
total_rows: usize,
) -> ReadPlanBuilder {
match plan_builder.resolve_selection_strategy() {
Expand All @@ -906,7 +909,7 @@ fn prepare_selection_for_page_skipping(
fn loaded_row_ranges_for_projection(
selection: Option<&RowSelection>,
projection_mask: &ProjectionMask,
offset_index: Option<&[OffsetIndexMetaData]>,
offset_index: Option<&[Option<OffsetIndexMetaData>]>,
total_rows: usize,
) -> Option<LoadedRowRanges> {
let selection = selection?;
Expand All @@ -916,7 +919,8 @@ fn loaded_row_ranges_for_projection(
.iter()
.enumerate()
.filter_map(|(leaf_idx, column)| {
let pages = column.page_locations();
let column_metadata = column.as_ref()?;
let pages = column_metadata.page_locations();
(projection_mask.leaf_included(leaf_idx) && !pages.is_empty()).then(|| {
RowSelection::from_consecutive_ranges(
selection
Expand Down Expand Up @@ -945,17 +949,19 @@ mod tests {

#[test]
fn test_loaded_row_ranges_intersect_column_page_boundaries() {
let column = |first_rows: &[i64]| OffsetIndexMetaData {
page_locations: first_rows
.iter()
.enumerate()
.map(|(idx, first_row_index)| PageLocation {
offset: (idx * 10) as i64,
compressed_page_size: 10,
first_row_index: *first_row_index,
})
.collect(),
unencoded_byte_array_data_bytes: None,
let column = |first_rows: &[i64]| {
Some(OffsetIndexMetaData {
page_locations: first_rows
.iter()
.enumerate()
.map(|(idx, first_row_index)| PageLocation {
offset: (idx * 10) as i64,
compressed_page_size: 10,
first_row_index: *first_row_index,
})
.collect(),
unencoded_byte_array_data_bytes: None,
})
};
let columns = vec![column(&[0, 4, 8]), column(&[0, 6, 10])];
let selection = RowSelection::from(vec![
Expand All @@ -978,7 +984,7 @@ mod tests {

#[test]
fn test_auto_keeps_mask_when_page_pruning_skips_pages() {
let columns = vec![OffsetIndexMetaData {
let columns = vec![Some(OffsetIndexMetaData {
page_locations: [0, 2, 4, 6, 8, 10]
.into_iter()
.enumerate()
Expand All @@ -989,7 +995,7 @@ mod tests {
})
.collect(),
unencoded_byte_array_data_bytes: None,
}];
})];
let selection = RowSelection::from(vec![
RowSelector::select(1),
RowSelector::skip(10),
Expand Down
8 changes: 6 additions & 2 deletions parquet/src/bin/parquet-concat.rs
Original file line number Diff line number Diff line change
Expand Up @@ -109,9 +109,13 @@ impl Args {
let mut rg_out = writer.next_row_group()?;
for (col_idx, column) in rg.columns().iter().enumerate() {
let bloom_filter = read_bloom_filter(column, &input);
let column_index = rg_column_indexes.and_then(|row| row.get(col_idx)).cloned();
let column_index = rg_column_indexes
.and_then(|row| row.get(col_idx))
.and_then(|opt| opt.clone());

let offset_index = rg_offset_indexes.and_then(|row| row.get(col_idx)).cloned();
let offset_index = rg_offset_indexes
.and_then(|row| row.get(col_idx))
.and_then(|opt| opt.clone());

let result = ColumnCloseResult {
bytes_written: column.compressed_size() as _,
Expand Down
34 changes: 19 additions & 15 deletions parquet/src/bin/parquet-index.rs
Original file line number Diff line number Diff line change
Expand Up @@ -99,24 +99,28 @@ impl Args {
ParquetError::General(format!(
"No offset index for row group {row_group_idx} column chunk {column_idx}"
))
})?.as_ref().ok_or_else(|| {
ParquetError::General(format!(
"Offset index is None for row group {row_group_idx} column chunk {column_idx}"
))
})?;

let row_counts =
compute_row_counts(offset_index.page_locations.as_slice(), row_group.num_rows());
match &column_indices[column_idx] {
ColumnIndexMetaData::NONE => println!("NO INDEX"),
ColumnIndexMetaData::BOOLEAN(v) => {
compute_row_counts(offset_index.page_locations(), row_group.num_rows());
match column_indices[column_idx].as_ref() {
None => println!("NO INDEX"),
Some(ColumnIndexMetaData::BOOLEAN(v)) => {
print_index::<bool>(v, offset_index, &row_counts)?
}
ColumnIndexMetaData::INT32(v) => print_index(v, offset_index, &row_counts)?,
ColumnIndexMetaData::INT64(v) => print_index(v, offset_index, &row_counts)?,
ColumnIndexMetaData::INT96(v) => print_index(v, offset_index, &row_counts)?,
ColumnIndexMetaData::FLOAT(v) => print_index(v, offset_index, &row_counts)?,
ColumnIndexMetaData::DOUBLE(v) => print_index(v, offset_index, &row_counts)?,
ColumnIndexMetaData::BYTE_ARRAY(v) => {
Some(ColumnIndexMetaData::INT32(v)) => print_index(v, offset_index, &row_counts)?,
Some(ColumnIndexMetaData::INT64(v)) => print_index(v, offset_index, &row_counts)?,
Some(ColumnIndexMetaData::INT96(v)) => print_index(v, offset_index, &row_counts)?,
Some(ColumnIndexMetaData::FLOAT(v)) => print_index(v, offset_index, &row_counts)?,
Some(ColumnIndexMetaData::DOUBLE(v)) => print_index(v, offset_index, &row_counts)?,
Some(ColumnIndexMetaData::BYTE_ARRAY(v)) => {
print_bytes_index(v, offset_index, &row_counts)?
}
ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(v) => {
Some(ColumnIndexMetaData::FIXED_LEN_BYTE_ARRAY(v)) => {
print_bytes_index(v, offset_index, &row_counts)?
}
}
Expand Down Expand Up @@ -147,11 +151,11 @@ fn print_index<T: std::fmt::Display>(
offset_index: &OffsetIndexMetaData,
row_counts: &[i64],
) -> Result<()> {
if column_index.num_pages() as usize != offset_index.page_locations.len() {
if column_index.num_pages() as usize != offset_index.page_locations().len() {
return Err(ParquetError::General(format!(
"Index length mismatch, got {} and {}",
column_index.num_pages(),
offset_index.page_locations.len()
offset_index.page_locations().len()
)));
}

Expand Down Expand Up @@ -186,11 +190,11 @@ fn print_bytes_index(
offset_index: &OffsetIndexMetaData,
row_counts: &[i64],
) -> Result<()> {
if column_index.num_pages() as usize != offset_index.page_locations.len() {
if column_index.num_pages() as usize != offset_index.page_locations().len() {
return Err(ParquetError::General(format!(
"Index length mismatch, got {} and {}",
column_index.num_pages(),
offset_index.page_locations.len()
offset_index.page_locations().len()
)));
}

Expand Down
1 change: 0 additions & 1 deletion parquet/src/file/metadata/memory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -242,7 +242,6 @@ impl HeapSize for OffsetIndexMetaData {
impl HeapSize for ColumnIndexMetaData {
fn heap_size(&self) -> usize {
match self {
Self::NONE => 0,
Self::BOOLEAN(native_index) => native_index.heap_size(),
Self::INT32(native_index) => native_index.heap_size(),
Self::INT64(native_index) => native_index.heap_size(),
Expand Down
22 changes: 13 additions & 9 deletions parquet/src/file/metadata/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,29 +141,31 @@ pub(crate) use writer::ThriftMetadataWriter;
/// documentation]. Each [`ColumnIndex`] holds statistics about all the pages in a
/// particular column chunk.
///
/// `column_index[row_group_number][column_number]` holds the
/// `column_index[row_group_number][column_number]` holds the optional
/// [`ColumnIndex`] corresponding to column `column_number` of row group
/// `row_group_number`.
/// `row_group_number`. This will be `None` if no index is present for the given
/// column chunk.
///
/// For example `column_index[2][3]` holds the [`ColumnIndex`] for the fourth
/// column in the third row group of the parquet file.
///
/// [PageIndex documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
/// [`ColumnIndex`]: crate::file::page_index::column_index::ColumnIndexMetaData
pub type ParquetColumnIndex = Vec<Vec<ColumnIndexMetaData>>;
pub type ParquetColumnIndex = Vec<Vec<Option<ColumnIndexMetaData>>>;

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changes here are the big change...the rest is dealing with the consequences

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have always found this structure to be very confusing (as it is a Vec of Vecs). Adding Option makes it even more confusing in my mind.

What would you think about at least encapsulating the PageIndex into a structure of its own (rather than two parallel structure)?

struct PageIndex {
  column_indexes: Vec<Vec<Option<ColumnIndexMetaData>>>,
  offset_indexes:  Vec<Vec<Option<OffsetIndexMetaData>>>,
}

🤔

Then we could add accessors like

if let Some(offset_index) = page_index.offset_index(rg_idx) { 
  // use offset index for rg_idx
}

That might also allow us to tweak the internal representation of these indexes to support options, etc without breaking the structure again


/// [`OffsetIndexMetaData`] for each data page of each row group of each column
/// [`OffsetIndexMetaData`] for each column chunk of each row group
///
/// This structure is the parsed representation of the [`OffsetIndex`] from the
/// Parquet file footer, as described in the Parquet [PageIndex documentation].
///
/// `offset_index[row_group_number][column_number]` holds
/// the [`OffsetIndexMetaData`] corresponding to column
/// `column_number`of row group `row_group_number`.
/// the optional [`OffsetIndexMetaData`] corresponding to column
/// `column_number`of row group `row_group_number`. This will be `None` if no index
/// is present for the given column chunk.
///
/// [PageIndex documentation]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
/// [`OffsetIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
pub type ParquetOffsetIndex = Vec<Vec<OffsetIndexMetaData>>;
pub type ParquetOffsetIndex = Vec<Vec<Option<OffsetIndexMetaData>>>;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we are going to change the page index representation (and force a breaking change downstream on the users) I wonder if we can think bigger than just adding an Option here and making it align with the parquet-format names

For example what do you think about making it a struct so that we have a better chance of evolving it over time (and make it easier to document)?

For example:

https://github.com/apache/parquet-format/blob/2076361bb64e2de9ca6a8d06eda025a6fa4e9df6/src/main/thrift/parquet.thrift#L1261

struct ParquetOffsetIndex {
  page_locations: Vec<ParquetPageLocation>,
  unencoded_byte_array_data_bytes: Option<Vec<i64>,
}

And https://github.com/apache/parquet-format/blob/2076361bb64e2de9ca6a8d06eda025a6fa4e9df6/src/main/thrift/parquet.thrift#L1236

struct ParquetPageLocation {
  offset: i64,
  compressed_page_size: i32,
  first_row_index: i64,
 }

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could also just do something slightly more encapsulated rather than a typedef

struct ParquetOffsetIndex {
  inner: Vec<Vec<Option<OffsetIndexMetaData>>>;
}

🤔


/// Parsed metadata for a single Parquet file
///
Expand Down Expand Up @@ -2117,11 +2119,13 @@ mod tests {
offset_index.append_row_count(1);
offset_index.append_offset_and_size(2, 3);
offset_index.append_unencoded_byte_array_data_bytes(Some(10));
let offset_index = offset_index.build();
let offset_index = Some(offset_index.build());

let parquet_meta = ParquetMetaDataBuilder::new(file_metadata)
.set_row_groups(row_group_meta)
.set_column_index(Some(vec![vec![ColumnIndexMetaData::BOOLEAN(native_index)]]))
.set_column_index(Some(vec![vec![Some(ColumnIndexMetaData::BOOLEAN(
native_index,
))]]))
.set_offset_index(Some(vec![vec![offset_index]]))
.build();

Expand Down
Loading
Loading