Add support for file row numbers in Parquet readers

[jkylling]

Which issue does this PR close?

Closes #7299.

What changes are included in this PR?

In this PR we:

• Add configuration to the ArrowReaderBuilder to set a row_number_column used to extend the read RecordBatches with an additional column with file row numbers.
• Keep track of the first row number in each row group in the file. This is computed from the file metadata.
• Add an ArrayReader to the vector of ArrayReaders reading columns from the Parquet file, if the row_number_column is set in the reader configuration. This is a RowNumberReader, which is a special ArrayReader. It reads no data from the Parquet pages, but uses the first row numbers in the RowGroupMetaData to keep track of progress.
• Add some basic tests and fuzz tests of the functionality.

The RowGroupMetaData::first_row_number is Option<i64>, since it is possible that the row number is unknown (I encountered an instance of this when trying to integrate this PR in delta-rs), and it's better if None is used instead of some special integer value.

The performance impact of this PR should be negligible when the row number column is not set. The only additional overhead would be the tracking of the first_row_number of each row group.

Are there any user-facing changes?

We add an additional public method:

• ArrowReaderBuilder::with_row_number_column

There are a few breaking changes as we touch a few public interfaces:

• RowGroupMetaData::from_thrift and RowGroupMetaData::from_thrift_encrypted takes an additional parameter first_row_number: Optional<i64>.
• The trait RowGroups has an additional method RowGroups::row_groups. Potentially this method could replace the RowGroups::num_rows method or provide a default implementation for it.
• An additional error variant ParquetError::RowGroupMetaDataMissingRowNumber.

I'm very open to suggestions on how to reduce the amount of breaking changes.

[etseidl]

Thanks for you submission @jkylling, I'll try to get a first pass review done this week. In the meantime please add the Apache license to row_number.rs and correct the other lint errors. 🙏

[jkylling]

Thanks for you submission @jkylling, I'll try to get a first pass review done this week. In the meantime please add the Apache license to row_number.rs and correct the other lint errors. 🙏

Updated. Looking forward to the first review!

I was very confused as to why cargo format did not work properly, but looks like you are already aware of this (#6179) :)

[etseidl]

Partial review, just a few nits for now.

[etseidl]

Thanks again @jkylling for taking this on. I've finished my first pass and have only one reservation. Otherwise it looks good and meets the criteria set forth in #7299 (comment).

[etseidl]

I'm finding myself a bit uneasy with adding the first row number to the RowGroupMetaData. Rather than that, could this bit here instead be changed to keep track of the first row number while populating the deque? Is there some wrinkle I'm missing? Might the row groups be filtered before instantiating the RowNumberReader?

[etseidl]

Answered my own question...it seems there's some complexity here at least when using the async reader.

[jkylling]

Yes, I believe we don't have access to all row groups when creating the array readers.

I took a quick look at the corresponding Parquet reader implementations for Trino and parquet-java.

Trino:

• Has a boolean to include a row number column, https://github.com/trinodb/trino/blob/a54d38a30e486a94a365c7f12a94e47beb30b0fa/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java#L112
• Includes this column when the boolean is set: https://github.com/trinodb/trino/blob/a54d38a30e486a94a365c7f12a94e47beb30b0fa/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java#L337
• Has a special block reader for reading row indexes https://github.com/trinodb/trino/blob/a54d38a30e486a94a365c7f12a94e47beb30b0fa/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java#L385-L393 I believe the positions play a similar role to our RowSelectors.
• Gets row indexes from RowGroupInfo, a pruned version of https://github.com/trinodb/trino/blob/a54d38a30e486a94a365c7f12a94e47beb30b0fa/lib/trino-parquet/src/main/java/io/trino/parquet/reader/ParquetReader.java#L456
• Populates the fileRowOffset by iterating through the row groups: https://github.com/trinodb/trino/blob/master/lib/trino-parquet/src/main/java/io/trino/parquet/metadata/ParquetMetadata.java#L107-L111

parquet-java:

• Has a method for tracking the current row index: https://github.com/apache/parquet-java/blob/7d1fe32c8c972710a9d780ec5e7d1f95d871374d/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetReader.java#L150-L155
• This row index is based on an iterator which starts form a row group row index, https://github.com/apache/parquet-java/blob/7d1fe32c8c972710a9d780ec5e7d1f95d871374d/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/InternalParquetRecordReader.java#L311-L339
• This row group row index is initialized by iterating through the row groups: https://github.com/apache/parquet-java/blob/7d1fe32c8c972710a9d780ec5e7d1f95d871374d/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java#L1654-L1656 (mapping obtained here: https://github.com/apache/parquet-java/blob/7d1fe32c8c972710a9d780ec5e7d1f95d871374d/parquet-hadoop/src/main/java/org/apache/parquet/format/converter/ParquetMetadataConverter.java#L1496-L1506)

Their approaches are rather similar to ours.

One take away is that the above implementations do not be keep the full RowGroupMetaDatas around as we do by requiring an iterator over RowGroupMetadata in the RowGroups trait. This is likely a good idea as this struct can be quite large. What do you think about changing the RowGroups trait to something like below?

/// A collection of row groups
pub trait RowGroups {
    /// Get the number of rows in this collection
    fn num_rows(&self) -> usize {
        self.row_group_infos.iter().map(|info| info.num_rows).sum()
    }

    /// Returns a [`PageIterator`] for the column chunks with the given leaf column index
    fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>>;

    /// Returns an iterator over the row groups in this collection
    fn row_group_infos(&self) -> Box<dyn Iterator<Item = &RowGroupInfo> + '_>;
}

struct RowGroupInfo {
    num_rows: usize,
    row_index: i64,
}

[etseidl]

I don't think this is necessary...the full ParquetMetaData is already available everywhere this trait is implemented, so I don't see a need to worry about adding another metadata structure here.

[etseidl]

Suggested change

	/// Returns the first row number in this row group.
	/// Returns the global index number for the first row in this row group.

And perhaps use first_row_index instead? That may be clearer.

[jkylling]

Agree. Updated.

[alamb]

Thanks @jkylling and @etseidl -- I think we need to be very careful to balance adding new features in the parquet reader with keeping it fast and maintainable. I haven't had a chance to look at this PR yet, but I do worry about performance and complexity

[jkylling]

@etseidl @alamb what would it take to move this PR forward? We could add benchmarks to settle concerns about performance, or refactor to try to reduce complexity?

[etseidl]

@etseidl @alamb what would it take to move this PR forward? We could add benchmarks to settle concerns about performance, or refactor to try to reduce complexity?

Sorry @jkylling, things have been rather hectic lately. I'll try to give it another look this week, along with some benchmarking (but I don't expect any perf hit). I'll just note that since this is a breaking change, it won't be able to be merged until the next major release (July-ish IIRC), so there's plenty of time to get this right. Also, I'll be deferring to those with more project history (e.g. @alamb @tustvold) as to whether the approach here is the best way to achieve the goal. Thank you for your contribution and your patience! 😄

[alamb]

Yeah, sorry I also have been slammed with many other projects. I'll try and find time to look but I suspect it may be a while

[jkylling]

Thank you for the update, and totally understand other responsibilities are taking up your time. I'll keep on being patient, and maybe do some minor improvements to this PR (use a smaller struct than the full RowGroupMetadata, and add some benchmarks for the RowNumberReader). Just want to make sure we have this PR ready before the next major release approaches.

[etseidl]

Yes, a benchmark that shows minimal impact with no row numbers would be nice (and hopefully adding row numbers won't be bad either 😄).

[etseidl]

Are these sorts necessary? Would the ordinal ever be out of order? They shouldn't be if I understand the encryption spec correctly.

[scovich]

I think we need to be very careful to balance adding new features in the parquet reader with keeping it fast and maintainable. I haven't had a chance to look at this PR yet, but I do worry about performance and complexity

100% agreed that simplicity and maintainability are paramount... but row numbers are a pretty fundamental feature that's very hard to emulate in higher layers if the parquet reader doesn't support them. Back when https://github.com/delta-io/delta first took a dependency on row numbers, spark's parquet reader did not yet support them; we had to disable row group pruning and other optimizations in order to make it (mostly) safe to manually compute row numbers in the query engine. It was really painful.

AFAIK, most parquet readers now support row numbers. We can add DuckDB and Iceberg to the ones already mentioned above. I was actually surprised to trip over this PR and learn that arrow-parquet does not yet support row numbers.

[scovich]

Maybe a crazy idea, but wouldn't the implementation be simpler (and more flexible) with a RowNumber extension type? Then users could do e.g.

Field::new("row_index", DataType::Int64, false).with_extension_type(RowNumber))

and build_primitive_reader could just check for it, no matter where in the schema it hides, instead of implicitly adding an extra column to the schema?

[scovich]

Update: I don't think raw parquet types support metadata, so this may not be an option.

[jkylling]

This would simplify usage of the feature. Having to keep track of the additional row number column is quite cumbersome in clients of this API. One option could be to extend ParquetFieldType with an additional row number type and add it based on the extension type in ArrowReaderMetadata::with_supplied_metadata? @etseidl @alamb what do you think about this approach?

[scovich]

It seems like this whole RowGroupSizeIterator thing is a complicated and error-prone way of chaining several Range<i64>? Can we use standard iterator machinery instead?

pub(crate) struct RowNumberReader {
    buffered_row_numbers: Vec<i64>,
    remaining_row_numbers: std::iter::Flatten<std::vec::IntoIter<std::ops::Range<i64>>>,
}

impl RowNumberReader {
    pub(crate) fn try_new<'a>(
        row_groups: impl Iterator<Item = &'a RowGroupMetaData>,
    ) -> Result<Self> {
        let ranges = row_groups
            .map(|rg| {
                let first_row_number = rg.first_row_index().ok_or(ParquetError::General(
                    "Row group missing row number".to_string(),
                ))?;
                Ok(first_row_number..first_row_number + rg.num_rows())
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Self {
            buffered_row_numbers: Vec::new(),
            remaining_row_numbers: ranges.into_iter().flatten(),
        })
    }
    
    // Use `take` on a `&mut Iterator` to consume a number of elements without consuming the iterator.
    fn take(&mut self, batch_size: usize) -> impl Iterator<Item = i64> {
        (&mut self.remaining_row_numbers).take(batch_size)
    }
}

impl ArrayReader for RowNumberReader {
    fn read_records(&mut self, batch_size: usize) -> Result<usize> {
        let starting_len = self.buffered_row_numbers.len();
        self.buffered_row_numbers.extend(self.take(batch_size));
        Ok(self.buffered_row_numbers.len() - starting_len)
    }

    fn skip_records(&mut self, num_records: usize) -> Result<usize> {
        Ok(self.take(num_records).count())
    }

[jkylling]

This is much simpler. Thank you! I suspect we are missing out on some performance in skip_records with this, but the bulk of the data pruning will likely have happened by pruning Parquet row groups already.

[jkylling]

@scovich I see you are involved in the maintenance of delta-kernel-rs. If you are interested, I've started on an implementation of deletion vector read support in delta-rs in this branch, based on a back port of an early version of this PR to arrow-54.2.1. The PR is still very rough, but the read path has got okay test coverage and it's able to read tables with deletion vectors produced by Spark correctly. The write support for deletion vectors is rudimentary (deletion vectors are only used for deletes when configured, and deleting from the same file twice is unsupported), and is mostly there to be able to unit test the read support. Unfortunately, I've not had time to wokr on this lately.

[alamb]

FWIW @zhuqi-lucas and I are working on improvements to the filter application here, which may result in some additional API churn:

• Poc for adaptive parquet predicate pushdown(bitmap/range) with page cache(3 data pages) #7454