General virtual columns support + row numbers as a first use-case (#8715)

Based on #7307.

# Which issue does this PR close?

- Closes #7299

# Rationale for this change

We need row numbers for many of the downstream features, e.g. computing
unique row identifier in iceberg.

# What changes are included in this PR?

New API to get row numbers as a virtual column:
```
let file = File::open(path).unwrap();
let row_number_field = Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber);
let options = ArrowReaderOptions::new().with_virtual_columns(vec![row_number_field]);
let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)
    .unwrap()
    .build()
    .expect("Could not create reader");
reader
    .collect::<Result<Vec<_>, _>>()
    .expect("Could not read")
    ```
```
This column is defined as an extension type.
Parquet metadata is propagated to the array builder to compute first row
indexes.
New Virtual column is included in addition to Primitive and Group.

# Are these changes tested?

Yes

# Are there any user-facing changes?

This is user facing feature, and has added docstrings.
No breaking changes, at least I tried not to, by creating a duplicate of
public method to add more parameters.

---------

Co-authored-by: Jonas Irgens Kylling <[email protected]>
Co-authored-by: scovich <[email protected]>
Co-authored-by: Andrew Lamb <[email protected]>

Author: 4 people

parquet/examples/read_with_rowgroup.rs

@@ -22,7 +22,7 @@ use parquet::arrow::async_reader::AsyncFileReader;
 use parquet::arrow::{ProjectionMask, parquet_to_arrow_field_levels};
 use parquet::column::page::{PageIterator, PageReader};
 use parquet::errors::{ParquetError, Result};
-use parquet::file::metadata::RowGroupMetaData;
+use parquet::file::metadata::{ParquetMetaData, RowGroupMetaData};
 use parquet::file::reader::{ChunkReader, Length};
 use parquet::file::serialized_reader::SerializedPageReader;
 use std::sync::Arc;
@@ -35,10 +35,11 @@ async fn main() -> Result<()> {
     let mut file = File::open(&path).await.unwrap();
 
     // The metadata could be cached in other places, this example only shows how to read
-    let metadata = file.get_metadata(None).await?;
+    let metadata = Arc::try_unwrap(file.get_metadata(None).await?).unwrap();
 
-    for rg in metadata.row_groups() {
-        let mut rowgroup = InMemoryRowGroup::create(rg.clone(), ProjectionMask::all());
+    for row_group_idx in 0..metadata.row_groups().len() {
+        let mut rowgroup =
+            InMemoryRowGroup::create(metadata.clone(), row_group_idx, ProjectionMask::all());
         rowgroup.async_fetch_data(&mut file, None).await?;
         let reader = rowgroup.build_reader(1024, None)?;
 
@@ -100,14 +101,15 @@ impl ChunkReader for ColumnChunkData {
 
 #[derive(Clone)]
 pub struct InMemoryRowGroup {
-    pub metadata: RowGroupMetaData,
+    metadata: ParquetMetaData,
+    row_group_idx: usize,
     mask: ProjectionMask,
     column_chunks: Vec<Option<Arc<ColumnChunkData>>>,
 }
 
 impl RowGroups for InMemoryRowGroup {
     fn num_rows(&self) -> usize {
-        self.metadata.num_rows() as usize
+        self.row_group_metadata().num_rows() as usize
     }
 
     fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>> {
@@ -118,7 +120,7 @@ impl RowGroups for InMemoryRowGroup {
             Some(data) => {
                 let page_reader: Box<dyn PageReader> = Box::new(SerializedPageReader::new(
                     data.clone(),
-                    self.metadata.column(i),
+                    self.row_group_metadata().column(i),
                     self.num_rows(),
                     None,
                 )?);
@@ -129,26 +131,44 @@ impl RowGroups for InMemoryRowGroup {
             }
         }
     }
+
+    fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_> {
+        Box::new(std::iter::once(self.row_group_metadata()))
+    }
+
+    fn metadata(&self) -> &ParquetMetaData {
+        &self.metadata
+    }
 }
 
 impl InMemoryRowGroup {
-    pub fn create(metadata: RowGroupMetaData, mask: ProjectionMask) -> Self {
-        let column_chunks = metadata.columns().iter().map(|_| None).collect::<Vec<_>>();
+    pub fn create(metadata: ParquetMetaData, row_group_idx: usize, mask: ProjectionMask) -> Self {
+        let column_chunks = metadata
+            .row_group(row_group_idx)
+            .columns()
+            .iter()
+            .map(|_| None)
+            .collect::<Vec<_>>();
 
         Self {
             metadata,
+            row_group_idx,
             mask,
             column_chunks,
         }
     }
 
+    pub fn row_group_metadata(&self) -> &RowGroupMetaData {
+        self.metadata.row_group(self.row_group_idx)
+    }
+
     pub fn build_reader(
         &self,
         batch_size: usize,
         selection: Option<RowSelection>,
     ) -> Result<ParquetRecordBatchReader> {
         let levels = parquet_to_arrow_field_levels(
-            &self.metadata.schema_descr_ptr(),
+            &self.row_group_metadata().schema_descr_ptr(),
             self.mask.clone(),
             None,
         )?;
@@ -163,7 +183,7 @@ impl InMemoryRowGroup {
         _selection: Option<&RowSelection>,
     ) -> Result<()> {
         let mut vs = std::mem::take(&mut self.column_chunks);
-        for (leaf_idx, meta) in self.metadata.columns().iter().enumerate() {
+        for (leaf_idx, meta) in self.row_group_metadata().columns().iter().enumerate() {
             if self.mask.leaf_included(leaf_idx) {
                 let (start, len) = meta.byte_range();
                 let data = reader.get_bytes(start..(start + len)).await?;

parquet/src/arrow/array_reader/builder.rs

@@ -26,16 +26,18 @@ use crate::arrow::array_reader::cached_array_reader::CachedArrayReader;
 use crate::arrow::array_reader::empty_array::make_empty_array_reader;
 use crate::arrow::array_reader::fixed_len_byte_array::make_fixed_len_byte_array_reader;
 use crate::arrow::array_reader::row_group_cache::RowGroupCache;
+use crate::arrow::array_reader::row_number::RowNumberReader;
 use crate::arrow::array_reader::{
     ArrayReader, FixedSizeListArrayReader, ListArrayReader, MapArrayReader, NullArrayReader,
     PrimitiveArrayReader, RowGroups, StructArrayReader, make_byte_array_dictionary_reader,
     make_byte_array_reader,
 };
 use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics;
-use crate::arrow::schema::{ParquetField, ParquetFieldType};
+use crate::arrow::schema::{ParquetField, ParquetFieldType, VirtualColumnType};
 use crate::basic::Type as PhysicalType;
 use crate::data_type::{BoolType, DoubleType, FloatType, Int32Type, Int64Type, Int96Type};
 use crate::errors::{ParquetError, Result};
+use crate::file::metadata::ParquetMetaData;
 use crate::schema::types::{ColumnDescriptor, ColumnPath, Type};
 
 /// Builder for [`CacheOptions`]
@@ -89,6 +91,8 @@ pub struct ArrayReaderBuilder<'a> {
     row_groups: &'a dyn RowGroups,
     /// Optional cache options for the array reader
     cache_options: Option<&'a CacheOptions<'a>>,
+    /// Parquet metadata for computing virtual column values
+    parquet_metadata: Option<&'a ParquetMetaData>,
     /// metrics
     metrics: &'a ArrowReaderMetrics,
 }
@@ -98,6 +102,7 @@ impl<'a> ArrayReaderBuilder<'a> {
         Self {
             row_groups,
             cache_options: None,
+            parquet_metadata: None,
             metrics,
         }
     }
@@ -108,6 +113,12 @@ impl<'a> ArrayReaderBuilder<'a> {
         self
     }
 
+    /// Add parquet metadata to the builder for computing virtual column values
+    pub fn with_parquet_metadata(mut self, parquet_metadata: &'a ParquetMetaData) -> Self {
+        self.parquet_metadata = Some(parquet_metadata);
+        self
+    }
+
     /// Create [`ArrayReader`] from parquet schema, projection mask, and parquet file reader.
     pub fn build_array_reader(
         &self,
@@ -153,6 +164,13 @@ impl<'a> ArrayReaderBuilder<'a> {
                     Ok(Some(reader))
                 }
             }
+            ParquetFieldType::Virtual(virtual_type) => {
+                // Virtual columns don't have data in the parquet file
+                // They need to be built by specialized readers
+                match virtual_type {
+                    VirtualColumnType::RowNumber => Ok(Some(self.build_row_number_reader()?)),
+                }
+            }
             ParquetFieldType::Group { .. } => match &field.arrow_type {
                 DataType::Map(_, _) => self.build_map_reader(field, mask),
                 DataType::Struct(_) => self.build_struct_reader(field, mask),
@@ -164,6 +182,18 @@ impl<'a> ArrayReaderBuilder<'a> {
         }
     }
 
+    fn build_row_number_reader(&self) -> Result<Box<dyn ArrayReader>> {
+        let parquet_metadata = self.parquet_metadata.ok_or_else(|| {
+            ParquetError::General(
+                "ParquetMetaData is required to read virtual row number columns.".to_string(),
+            )
+        })?;
+        Ok(Box::new(RowNumberReader::try_new(
+            parquet_metadata,
+            self.row_groups.row_groups(),
+        )?))
+    }
+
     /// Build array reader for map type.
     fn build_map_reader(
         &self,
@@ -439,6 +469,7 @@ impl<'a> ArrayReaderBuilder<'a> {
 mod tests {
     use super::*;
     use crate::arrow::schema::parquet_to_arrow_schema_and_fields;
+    use crate::arrow::schema::virtual_type::RowNumber;
     use crate::file::reader::{FileReader, SerializedFileReader};
     use crate::util::test_common::file_util::get_test_file;
     use arrow::datatypes::Field;
@@ -455,6 +486,7 @@ mod tests {
             file_metadata.schema_descr(),
             ProjectionMask::all(),
             file_metadata.key_value_metadata(),
+            &[],
         )
         .unwrap();
 
@@ -472,4 +504,41 @@ mod tests {
 
         assert_eq!(array_reader.get_data_type(), &arrow_type);
     }
+
+    #[test]
+    fn test_create_array_reader_with_row_numbers() {
+        let file = get_test_file("nulls.snappy.parquet");
+        let file_reader: Arc<dyn FileReader> = Arc::new(SerializedFileReader::new(file).unwrap());
+
+        let file_metadata = file_reader.metadata().file_metadata();
+        let mask = ProjectionMask::leaves(file_metadata.schema_descr(), [0]);
+        let row_number_field = Arc::new(
+            Field::new("row_number", DataType::Int64, false).with_extension_type(RowNumber),
+        );
+        let (_, fields) = parquet_to_arrow_schema_and_fields(
+            file_metadata.schema_descr(),
+            ProjectionMask::all(),
+            file_metadata.key_value_metadata(),
+            std::slice::from_ref(&row_number_field),
+        )
+        .unwrap();
+
+        let metrics = ArrowReaderMetrics::disabled();
+        let array_reader = ArrayReaderBuilder::new(&file_reader, &metrics)
+            .with_parquet_metadata(file_reader.metadata())
+            .build_array_reader(fields.as_ref(), &mask)
+            .unwrap();
+
+        // Create arrow types
+        let arrow_type = DataType::Struct(Fields::from(vec![
+            Field::new(
+                "b_struct",
+                DataType::Struct(vec![Field::new("b_c_int", DataType::Int32, true)].into()),
+                true,
+            ),
+            (*row_number_field).clone(),
+        ]));
+
+        assert_eq!(array_reader.get_data_type(), &arrow_type);
+    }
 }

parquet/src/arrow/array_reader/list_array.rs

@@ -561,6 +561,7 @@ mod tests {
             schema,
             ProjectionMask::all(),
             file_metadata.key_value_metadata(),
+            &[],
         )
         .unwrap();
 

parquet/src/arrow/array_reader/mod.rs

@@ -27,6 +27,7 @@ use crate::arrow::record_reader::GenericRecordReader;
 use crate::arrow::record_reader::buffer::ValuesBuffer;
 use crate::column::page::PageIterator;
 use crate::column::reader::decoder::ColumnValueDecoder;
+use crate::file::metadata::ParquetMetaData;
 use crate::file::reader::{FilePageIterator, FileReader};
 
 mod builder;
@@ -42,12 +43,14 @@ mod map_array;
 mod null_array;
 mod primitive_array;
 mod row_group_cache;
+mod row_number;
 mod struct_array;
 
 #[cfg(test)]
 mod test_util;
 
 // Note that this crate is public under the `experimental` feature flag.
+use crate::file::metadata::RowGroupMetaData;
 pub use builder::{ArrayReaderBuilder, CacheOptions, CacheOptionsBuilder};
 pub use byte_array::make_byte_array_reader;
 pub use byte_array_dictionary::make_byte_array_dictionary_reader;
@@ -139,17 +142,35 @@ pub trait RowGroups {
     /// Returns a [`PageIterator`] for all pages in the specified column chunk
     /// across all row groups in this collection.
     fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>>;
+
+    /// Returns an iterator over the row groups in this collection
+    ///
+    /// Note this may not include all row groups in [`Self::metadata`].
+    fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_>;
+
+    /// Returns the parquet metadata
+    fn metadata(&self) -> &ParquetMetaData;
 }
 
 impl RowGroups for Arc<dyn FileReader> {
     fn num_rows(&self) -> usize {
-        self.metadata().file_metadata().num_rows() as usize
+        FileReader::metadata(self.as_ref())
+            .file_metadata()
+            .num_rows() as usize
     }
 
     fn column_chunks(&self, column_index: usize) -> Result<Box<dyn PageIterator>> {
         let iterator = FilePageIterator::new(column_index, Arc::clone(self))?;
         Ok(Box::new(iterator))
     }
+
+    fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_> {
+        Box::new(FileReader::metadata(self.as_ref()).row_groups().iter())
+    }
+
+    fn metadata(&self) -> &ParquetMetaData {
+        FileReader::metadata(self.as_ref())
+    }
 }
 
 /// Uses `record_reader` to read up to `batch_size` records from `pages`