Wide-schema parquet: O(1) parquet_column + cacheable arrow-reader-metadata primitives

parquet/src/arrow/arrow_reader/mod.rs

@@ -29,7 +29,8 @@ use std::sync::Arc;
 pub use crate::arrow::array_reader::RowGroups;
 use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder};
 use crate::arrow::schema::{
-    ParquetField, parquet_to_arrow_schema_and_fields, virtual_type::is_virtual_column,
+    ParquetField, parquet_to_arrow_schema_and_field_levels_with_virtual,
+    virtual_type::is_virtual_column,
 };
 use crate::arrow::{FieldLevels, ProjectionMask, parquet_to_arrow_field_levels_with_virtual};
 use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash};
@@ -792,6 +793,24 @@ impl ArrowReaderOptions {
         })
     }
 
+    /// Returns the configured virtual columns, if any.
+    pub fn virtual_columns(&self) -> &[FieldRef] {
+        &self.virtual_columns
+    }
+
+    /// Returns the supplied arrow [`Schema`], if one was provided via
+    /// [`Self::with_schema`].
+    pub fn supplied_schema(&self) -> Option<&SchemaRef> {
+        self.supplied_schema.as_ref()
+    }
+
+    /// Returns whether the reader will skip the embedded arrow schema
+    /// metadata stored in the parquet file (set via
+    /// [`Self::with_skip_arrow_metadata`]).
+    pub fn skip_arrow_metadata(&self) -> bool {
+        self.skip_arrow_metadata
+    }
+
     #[deprecated(
         since = "57.2.0",
         note = "Use `column_index_policy` or `offset_index_policy` instead"
@@ -886,6 +905,43 @@ impl ArrowReaderMetadata {
         Self::try_new(Arc::new(metadata), options)
     }
 
+    /// Construct an [`ArrowReaderMetadata`] from precomputed
+    /// [`FieldLevels`].
+    ///
+    /// Use this when the caller has already computed the arrow [`Schema`]
+    /// and [`FieldLevels`] for `metadata` (e.g. by calling
+    /// [`parquet_to_arrow_field_levels`]) and wants to avoid recomputing
+    /// them. For wide schemas the field-levels conversion walks every leaf
+    /// in the parquet schema, so reusing it across reader builds for the
+    /// same metadata is significantly cheaper than calling [`Self::try_new`]
+    /// repeatedly.
+    ///
+    /// `arrow_schema` and `field_levels` must describe the same columns:
+    /// they are normally produced together by
+    /// [`parquet_to_arrow_schema_and_field_levels`]. `arrow_schema` is taken
+    /// separately only because it additionally carries the file's key/value
+    /// metadata, which [`FieldLevels`] does not.
+    ///
+    /// [`parquet_to_arrow_field_levels`]: crate::arrow::schema::parquet_to_arrow_field_levels
+    /// [`parquet_to_arrow_schema_and_field_levels`]: crate::arrow::schema::parquet_to_arrow_schema_and_field_levels
+    pub fn from_field_levels(
+        metadata: Arc<ParquetMetaData>,
+        arrow_schema: SchemaRef,
+        field_levels: FieldLevels,
+    ) -> Self {
+        let FieldLevels { fields, levels } = field_levels;
+        debug_assert_eq!(
+            arrow_schema.fields(),
+            &fields,
+            "arrow_schema fields must match field_levels fields"
+        );
+        Self {
+            metadata,
+            schema: arrow_schema,
+            fields: levels.map(Arc::new),
+        }
+    }
+
     /// Create a new [`ArrowReaderMetadata`] from a pre-existing
     /// [`ParquetMetaData`] and [`ArrowReaderOptions`].
     ///
@@ -906,18 +962,18 @@ impl ArrowReaderMetadata {
                     false => metadata.file_metadata().key_value_metadata(),
                 };
 
-                let (schema, fields) = parquet_to_arrow_schema_and_fields(
+                let (schema, field_levels) = parquet_to_arrow_schema_and_field_levels_with_virtual(
                     metadata.file_metadata().schema_descr(),
                     ProjectionMask::all(),
                     kv_metadata,
                     &options.virtual_columns,
                 )?;
 
-                Ok(Self {
+                Ok(Self::from_field_levels(
                     metadata,
-                    schema: Arc::new(schema),
-                    fields: fields.map(Arc::new),
-                })
+                    Arc::new(schema),
+                    field_levels,
+                ))
             }
         }
     }

parquet/src/arrow/arrow_reader/statistics.rs

@@ -1487,12 +1487,36 @@ impl<'a> StatisticsConverter<'a> {
             None => None,
         };
 
-        Ok(Self {
-            parquet_column_index: parquet_index,
+        Ok(Self::from_arrow_field(
+            arrow_field,
+            parquet_schema,
+            parquet_index,
+        ))
+    }
+
+    /// Construct a [`StatisticsConverter`] from a precomputed
+    /// `(arrow_field, parquet_leaf_index)` pair.
+    ///
+    /// This is a low-overhead alternative to [`Self::try_new`] for callers
+    /// that need to build many converters against the same schemas — for
+    /// example when extracting per-column statistics for a wide schema.
+    /// [`Self::try_new`] performs O(N) name lookups against both the arrow
+    /// and parquet schemas; this constructor avoids those lookups by taking
+    /// the resolved field and column index directly.
+    ///
+    /// `parquet_leaf_index` should be `None` if the column does not exist in
+    /// the parquet file (schema evolution case).
+    pub fn from_arrow_field(
+        arrow_field: &'a Field,
+        parquet_schema: &'a SchemaDescriptor,
+        parquet_leaf_index: Option<usize>,
+    ) -> Self {
+        Self {
+            parquet_column_index: parquet_leaf_index,
             arrow_field,
             missing_null_counts_as_zero: true,
-            physical_type: parquet_index.map(|idx| parquet_schema.column(idx).physical_type()),
-        })
+            physical_type: parquet_leaf_index.map(|idx| parquet_schema.column(idx).physical_type()),
+        }
     }
 
     /// Extract the minimum values from row group statistics in [`RowGroupMetaData`]

parquet/src/arrow/async_reader/mod.rs

@@ -110,6 +110,31 @@ pub trait AsyncFileReader: Send {
         &'a mut self,
         options: Option<&'a ArrowReaderOptions>,
     ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>>;
+
+    /// Return a future resolving to a fully-built [`ArrowReaderMetadata`].
+    ///
+    /// The default implementation calls [`Self::get_metadata`] and then
+    /// constructs the [`ArrowReaderMetadata`] via
+    /// [`ArrowReaderMetadata::try_new`], which walks every leaf in the
+    /// parquet schema to derive the matching arrow [`Schema`] and field
+    /// levels.
+    ///
+    /// Implementations that already cache the per-file metadata (e.g.
+    /// because the same file is opened many times across queries) can
+    /// override this method to also cache the derived arrow schema view.
+    /// For wide schemas this can dominate per-file CPU overhead, since the
+    /// per-leaf walk is `O(N_columns)` and is otherwise repeated on every
+    /// reader build for the same file.
+    fn get_arrow_reader_metadata<'a>(
+        &'a mut self,
+        options: ArrowReaderOptions,
+    ) -> BoxFuture<'a, Result<ArrowReaderMetadata>> {
+        async move {
+            let metadata = self.get_metadata(Some(&options)).await?;
+            ArrowReaderMetadata::try_new(metadata, options)
+        }
+        .boxed()
+    }
 }
 
 /// This allows Box<dyn AsyncFileReader + '_> to be used as an AsyncFileReader,
@@ -128,6 +153,13 @@ impl AsyncFileReader for Box<dyn AsyncFileReader + '_> {
     ) -> BoxFuture<'a, Result<Arc<ParquetMetaData>>> {
         self.as_mut().get_metadata(options)
     }
+
+    fn get_arrow_reader_metadata<'a>(
+        &'a mut self,
+        options: ArrowReaderOptions,
+    ) -> BoxFuture<'a, Result<ArrowReaderMetadata>> {
+        self.as_mut().get_arrow_reader_metadata(options)
+    }
 }
 
 impl<T: AsyncFileReader + MetadataFetch + AsyncRead + AsyncSeek + Unpin> MetadataSuffixFetch for T {
@@ -194,8 +226,11 @@ impl ArrowReaderMetadata {
         input: &mut T,
         options: ArrowReaderOptions,
     ) -> Result<Self> {
-        let metadata = input.get_metadata(Some(&options)).await?;
-        Self::try_new(metadata, options)
+        // Delegate to AsyncFileReader::get_arrow_reader_metadata so
+        // implementations that cache derived state (e.g. an arrow Schema
+        // and field-levels view) can short-circuit the per-leaf walk
+        // performed by Self::try_new.
+        input.get_arrow_reader_metadata(options).await
     }
 }
 

parquet/src/arrow/mod.rs

@@ -210,7 +210,8 @@ use arrow_schema::{FieldRef, Schema};
 pub use self::schema::{
     ArrowSchemaConverter, FieldLevels, add_encoded_arrow_schema_to_metadata, encode_arrow_schema,
     parquet_to_arrow_field_levels, parquet_to_arrow_field_levels_with_virtual,
-    parquet_to_arrow_schema, parquet_to_arrow_schema_by_columns, virtual_type::*,
+    parquet_to_arrow_schema, parquet_to_arrow_schema_and_field_levels,
+    parquet_to_arrow_schema_by_columns, virtual_type::*,
 };
 
 /// Schema metadata key used to store serialized Arrow schema
@@ -475,9 +476,11 @@ pub fn parquet_column<'a>(
         return None;
     }
 
-    // This could be made more efficient (#TBD)
-    let parquet_idx = (0..parquet_schema.columns().len())
-        .find(|x| parquet_schema.get_column_root_idx(*x) == root_idx)?;
+    // For non-nested fields the parquet column ordering matches the arrow
+    // root field ordering, so the first leaf descended from `root_idx` is
+    // exactly the parquet column we want. SchemaDescriptor caches this
+    // mapping for O(1) lookup.
+    let parquet_idx = parquet_schema.root_first_leaf_index(root_idx)?;
     Some((parquet_idx, field))
 }
 

parquet/src/arrow/schema/mod.rs

@@ -67,17 +67,40 @@ pub fn parquet_to_arrow_schema_by_columns(
     Ok(parquet_to_arrow_schema_and_fields(parquet_schema, mask, key_value_metadata, &[])?.0)
 }
 
-/// Determines the Arrow Schema from a Parquet schema
+/// Convert a parquet [`SchemaDescriptor`] into both the arrow [`Schema`]
+/// and matching [`FieldLevels`] in a single walk.
 ///
-/// Looks for an Arrow schema metadata "hint" (see
-/// [`parquet_to_arrow_field_levels`]), and uses it if present to ensure
-/// lossless round trips.
-pub(crate) fn parquet_to_arrow_schema_and_fields(
+/// This is a convenience around [`parquet_to_arrow_field_levels`] that
+/// also produces a [`Schema`] (with the parquet file's key/value metadata
+/// attached). Useful for callers that want to pre-compute and reuse both
+/// for [`ArrowReaderMetadata::from_field_levels`].
+///
+/// [`ArrowReaderMetadata::from_field_levels`]: crate::arrow::arrow_reader::ArrowReaderMetadata::from_field_levels
+pub fn parquet_to_arrow_schema_and_field_levels(
+    parquet_schema: &SchemaDescriptor,
+    mask: ProjectionMask,
+    key_value_metadata: Option<&Vec<KeyValue>>,
+) -> Result<(Schema, FieldLevels)> {
+    parquet_to_arrow_schema_and_field_levels_with_virtual(
+        parquet_schema,
+        mask,
+        key_value_metadata,
+        &[],
+    )
+}
+
+/// As [`parquet_to_arrow_schema_and_field_levels`] but also appends the
+/// given `virtual_columns` to the resulting schema and field levels.
+///
+/// This is the single schema-walk used by [`parquet_to_arrow_schema_and_fields`]
+/// (which discards the [`FieldLevels`] wrapper) and by
+/// [`parquet_to_arrow_schema_and_field_levels`] (which keeps it).
+pub(crate) fn parquet_to_arrow_schema_and_field_levels_with_virtual(
     parquet_schema: &SchemaDescriptor,
     mask: ProjectionMask,
     key_value_metadata: Option<&Vec<KeyValue>>,
     virtual_columns: &[FieldRef],
-) -> Result<(Schema, Option<ParquetField>)> {
+) -> Result<(Schema, FieldLevels)> {
     let mut metadata = parse_key_value_metadata(key_value_metadata).unwrap_or_default();
     let maybe_schema = metadata
         .remove(super::ARROW_SCHEMA_META_KEY)
@@ -94,7 +117,27 @@ pub(crate) fn parquet_to_arrow_schema_and_fields(
     let hint = maybe_schema.as_ref().map(|s| s.fields());
     let field_levels =
         parquet_to_arrow_field_levels_with_virtual(parquet_schema, mask, hint, virtual_columns)?;
-    let schema = Schema::new_with_metadata(field_levels.fields, metadata);
+    let schema = Schema::new_with_metadata(field_levels.fields.clone(), metadata);
+    Ok((schema, field_levels))
+}
+
+/// Determines the Arrow Schema from a Parquet schema
+///
+/// Looks for an Arrow schema metadata "hint" (see
+/// [`parquet_to_arrow_field_levels`]), and uses it if present to ensure
+/// lossless round trips.
+pub(crate) fn parquet_to_arrow_schema_and_fields(
+    parquet_schema: &SchemaDescriptor,
+    mask: ProjectionMask,
+    key_value_metadata: Option<&Vec<KeyValue>>,
+    virtual_columns: &[FieldRef],
+) -> Result<(Schema, Option<ParquetField>)> {
+    let (schema, field_levels) = parquet_to_arrow_schema_and_field_levels_with_virtual(
+        parquet_schema,
+        mask,
+        key_value_metadata,
+        virtual_columns,
+    )?;
     Ok((schema, field_levels.levels))
 }
 

parquet/src/file/metadata/mod.rs

@@ -2046,10 +2046,12 @@ mod tests {
             .set_row_groups(row_group_meta_with_stats)
             .build();
 
+        // Sizes include the `root_to_first_leaf: Vec<usize>` cache that
+        // SchemaDescriptor adds to make `parquet_column` lookups O(1).
         #[cfg(not(feature = "encryption"))]
-        let base_expected_size = 2734;
+        let base_expected_size = 2774;
         #[cfg(feature = "encryption")]
-        let base_expected_size = 2902;
+        let base_expected_size = 2942;
 
         assert_eq!(parquet_meta.memory_size(), base_expected_size);
 
@@ -2078,9 +2080,9 @@ mod tests {
             .build();
 
         #[cfg(not(feature = "encryption"))]
-        let bigger_expected_size = 3160;
+        let bigger_expected_size = 3200;
         #[cfg(feature = "encryption")]
-        let bigger_expected_size = 3328;
+        let bigger_expected_size = 3368;
 
         // more set fields means more memory usage
         assert!(bigger_expected_size > base_expected_size);
@@ -2127,7 +2129,7 @@ mod tests {
             .set_row_groups(row_group_meta.clone())
             .build();
 
-        let base_expected_size = 2042;
+        let base_expected_size = 2082;
         assert_eq!(parquet_meta_data.memory_size(), base_expected_size);
 
         let footer_key = "0123456789012345".as_bytes();
@@ -2153,7 +2155,7 @@ mod tests {
             .set_file_decryptor(Some(decryptor))
             .build();
 
-        let expected_size_with_decryptor = 3056;
+        let expected_size_with_decryptor = 3096;
         assert!(expected_size_with_decryptor > base_expected_size);
 
         assert_eq!(