[parquet] Add row group index virtual column

[friendlymatthew]

Which issue does this PR close?

• Closes [Parquet] Support returning RowGroupIndex as a column #8641
• Related [Parquet] Virtual Column Support #8799

Rationale for this change

This PR adds support for storing row group indices as a virtual column, allowing users to determine which row group each row originated from

The usage pattern is quite simple, something like:

use parquet::arrow::RowGroupIndex;

let row_group_index_field = Arc::new(
    Field::new("row_group_index", DataType::Int64, false)
        .with_extension_type(RowGroupIndex)
);

let options = ArrowReaderOptions::new()
    .with_virtual_columns(vec![row_group_index_field])?;

let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)?
    .build()?;

[friendlymatthew]

@Jefffrey if you have some bandwidth, I'd be curious to get your thoughts

[friendlymatthew]

cc @alamb

[Jefffrey]

@Jefffrey if you have some bandwidth, I'd be curious to get your thoughts

I'll see if I can take a look at this PR soon, though I will say it's been a while since I looked at the parquet codebase 😅

[adriangb]

This would be sweet! I like the idea of using an extension type as a marker, and letting the caller customize the column name.

[alamb]

I think the extension type marker was pioneered by @jkylling and @vustef

I will also try and review this sooner rather than later

[scovich]

Nice! (a few nits)

[scovich]

nit: I know this is how the row index reader did it, but since that code merged I learned that Iterator::by_ref is a thing.

Suggested change

	.extend((&mut self.remaining_indices).take(batch_size));
	.extend(self.remaining_indices.by_ref().take(batch_size));

It's not shorter, but does seem more readable?

(more below)

[friendlymatthew]

TIL. Thank you @scovich

[scovich]

nit: is a match simpler?

Suggested change

	if metadata.is_some_and(str::is_empty) {
	Ok("")
	} else {
	Err(ArrowError::InvalidArgumentError(
	"Virtual column extension type expects an empty string as metadata".to_owned(),
	))
	}
	match metadata {
	Some(&"") => Ok(""),
	_ => Err(ArrowError::InvalidArgumentError(
	"Virtual column extension type expects an empty string as metadata".to_owned(),
	)),
	}

or even

Suggested change

	if metadata.is_some_and(str::is_empty) {
	Ok("")
	} else {
	Err(ArrowError::InvalidArgumentError(
	"Virtual column extension type expects an empty string as metadata".to_owned(),
	))
	}
	if let Some(&"") = metadata {
	return Ok("");
	};
	Err(ArrowError::InvalidArgumentError(
	"Virtual column extension type expects an empty string as metadata".to_owned(),
	))

[Jefffrey]

Looks reasonable to me 👍

[alamb]

Thanks @friendlymatthew and @scovich and @Jefffrey -- this is really cool. I can't wait to get this integrated into downstream projects

[alamb]

Most of the time each ArrayReader will be instantiated for pages from exactly one row group we can probably make this significantly faster by optimizing the case when reading from only a single row group.

The other thing would be to pre-calculate the arrays (where the row group values don't change) and return the same ArrayRef until the RowGroup actually changes

However, there is no reason we need to do that now.

[friendlymatthew]

Filed:

• perf: optimize RowGroupIndexReader for single row group reads #9180
• perf: cache ArrayRef in RowGroupIndexReader when row group doesn't change #9181

I wonder if we should have the same optimization for RowNumberReader 🤔

[alamb]

I don't see any test coverage for this method.

I ran code coverage

cargo llvm-cov --html --all-features -p parquet

And skip doesn't seem to be covered:

[alamb]

Thank you @friendlymatthew @Jefffrey and @scovich