fix: work around DataFusion 54.1.0 Parquet page-index regression

native/core/Cargo.toml

@@ -36,6 +36,7 @@ publish = false
 
 [dependencies]
 arrow = { workspace = true }
+bytes = { workspace = true }
 parquet = { workspace = true, default-features = false, features = ["experimental", "arrow", "snap", "lz4", "zstd", "flate2-zlib-rs"] }
 futures = { workspace = true }
 mimalloc = { version = "*", default-features = false, optional = true }

native/core/src/parquet/eager_page_index_reader_factory.rs

@@ -0,0 +1,193 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A `ParquetFileReaderFactory` that always loads the Parquet page index into the shared
+//! `FileMetadataCache` on the first metadata fetch for a file, instead of deferring to
+//! DataFusion's opener.
+//!
+//! DataFusion's opener requests `PageIndexPolicy::Skip` on the initial metadata load and defers
+//! loading the page index until row-group pruning shows it is still needed
+//! (apache/datafusion#22857). That deferred load reads the page index directly off the
+//! `AsyncFileReader` (`load_page_index` in `datafusion-datasource-parquet`'s opener), bypassing
+//! `FileMetadataCache` entirely. For a predicate that never resolves to fully-matched by
+//! row-group statistics alone, such as `IS NOT NULL` on a column whose row groups don't carry
+//! `null_count` statistics, the skip never fires, and the page index is re-fetched, uncached, on
+//! every open. At the scale of a wide fact table scanned across many partitions, that is
+//! repeated, unbounded I/O for the same bytes.
+//!
+//! This factory forces `PageIndexPolicy::Optional` on every metadata fetch for files with no
+//! decryption properties, ignoring whatever policy the caller requests, so the page index is
+//! always present in the cached `ParquetMetaData` after the first load. DataFusion's opener
+//! checks whether the metadata it already has includes the page index before issuing its own
+//! fetch, so with this factory that check is always true and the uncached fetch never happens.
+//! The tradeoff: files where the opener's skip heuristic would have avoided the page index load
+//! entirely now load it anyway.
+//!
+//! Encrypted files are exempt from the override: `DFParquetMetadata::fetch_metadata` disables
+//! `FileMetadataCache` entirely whenever decryption properties are set, so nothing gets cached
+//! for them either way, and forcing eager loading would only add an unconditional page-index
+//! fetch to encrypted scans that have no pruning predicate at all. Encrypted opens get exactly
+//! the caller's requested policy, unchanged from stock behavior.
+//!
+//! Filed upstream as apache/datafusion#23978. Revert this once the opener merges its deferred
+//! page-index load back into `FileMetadataCache` instead of bypassing it.
+
+use bytes::Bytes;
+use datafusion::common::Result as DFResult;
+use datafusion::datasource::physical_plan::parquet::metadata::DFParquetMetadata;
+use datafusion::datasource::physical_plan::parquet::{
+    ParquetFileMetrics, ParquetFileReaderFactory,
+};
+use datafusion::execution::cache::cache_manager::FileMetadataCache;
+use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;
+use datafusion_datasource::PartitionedFile;
+use futures::future::BoxFuture;
+use futures::FutureExt;
+use object_store::ObjectStore;
+use parquet::arrow::arrow_reader::ArrowReaderOptions;
+use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader};
+use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
+use std::fmt::Debug;
+use std::ops::Range;
+use std::sync::Arc;
+
+#[derive(Debug)]
+pub struct EagerPageIndexReaderFactory {
+    store: Arc<dyn ObjectStore>,
+    metadata_cache: Arc<dyn FileMetadataCache>,
+}
+
+impl EagerPageIndexReaderFactory {
+    pub fn new(store: Arc<dyn ObjectStore>, metadata_cache: Arc<dyn FileMetadataCache>) -> Self {
+        Self {
+            store,
+            metadata_cache,
+        }
+    }
+}
+
+impl ParquetFileReaderFactory for EagerPageIndexReaderFactory {
+    fn create_reader(
+        &self,
+        partition_index: usize,
+        partitioned_file: PartitionedFile,
+        metadata_size_hint: Option<usize>,
+        metrics: &ExecutionPlanMetricsSet,
+    ) -> DFResult<Box<dyn AsyncFileReader + Send>> {
+        let file_metrics = ParquetFileMetrics::new(
+            partition_index,
+            partitioned_file.object_meta.location.as_ref(),
+            metrics,
+        );
+        let mut inner = ParquetObjectReader::new(
+            Arc::clone(&self.store),
+            partitioned_file.object_meta.location.clone(),
+        )
+        .with_file_size(partitioned_file.object_meta.size);
+        if let Some(hint) = metadata_size_hint {
+            inner = inner.with_footer_size_hint(hint);
+        }
+
+        Ok(Box::new(EagerPageIndexReader {
+            file_metrics,
+            store: Arc::clone(&self.store),
+            inner,
+            partitioned_file,
+            metadata_cache: Arc::clone(&self.metadata_cache),
+            metadata_size_hint,
+        }))
+    }
+}
+
+struct EagerPageIndexReader {
+    file_metrics: ParquetFileMetrics,
+    store: Arc<dyn ObjectStore>,
+    inner: ParquetObjectReader,
+    partitioned_file: PartitionedFile,
+    metadata_cache: Arc<dyn FileMetadataCache>,
+    metadata_size_hint: Option<usize>,
+}
+
+impl AsyncFileReader for EagerPageIndexReader {
+    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, parquet::errors::Result<Bytes>> {
+        let bytes_scanned = range.end - range.start;
+        self.file_metrics.bytes_scanned.add(bytes_scanned as usize);
+        self.inner.get_bytes(range)
+    }
+
+    fn get_byte_ranges(
+        &mut self,
+        ranges: Vec<Range<u64>>,
+    ) -> BoxFuture<'_, parquet::errors::Result<Vec<Bytes>>>
+    where
+        Self: Send,
+    {
+        let total: u64 = ranges.iter().map(|r| r.end - r.start).sum();
+        self.file_metrics.bytes_scanned.add(total as usize);
+        self.inner.get_byte_ranges(ranges)
+    }
+
+    fn get_metadata<'a>(
+        &'a mut self,
+        options: Option<&'a ArrowReaderOptions>,
+    ) -> BoxFuture<'a, parquet::errors::Result<Arc<ParquetMetaData>>> {
+        // Forward decryption properties like `CachedParquetFileReader` does. Only override the
+        // policy for non-encrypted opens; see module docs for why.
+        let object_meta = self.partitioned_file.object_meta.clone();
+        let metadata_cache = Arc::clone(&self.metadata_cache);
+        let store = Arc::clone(&self.store);
+        let metadata_size_hint = self.metadata_size_hint;
+        async move {
+            let file_decryption_properties = options
+                .and_then(|o| o.file_decryption_properties())
+                .map(Arc::clone);
+            let page_index_policy = if file_decryption_properties.is_none() {
+                Some(PageIndexPolicy::Optional)
+            } else {
+                options.map(|o| o.column_index_policy())
+            };
+
+            DFParquetMetadata::new(store.as_ref(), &object_meta)
+                .with_decryption_properties(file_decryption_properties)
+                .with_file_metadata_cache(Some(metadata_cache))
+                .with_metadata_size_hint(metadata_size_hint)
+                .with_page_index_policy(page_index_policy)
+                .fetch_metadata()
+                .await
+                .map_err(|e| {
+                    parquet::errors::ParquetError::General(format!(
+                        "Failed to fetch metadata for file {}: {e}",
+                        object_meta.location,
+                    ))
+                })
+        }
+        .boxed()
+    }
+}
+
+impl Drop for EagerPageIndexReader {
+    fn drop(&mut self) {
+        self.file_metrics
+            .scan_efficiency_ratio
+            .add_part(self.file_metrics.bytes_scanned.value());
+        // Multiple readers may run against the same file, so we set_total on every drop rather
+        // than accumulating it, to avoid adding the file's total size multiple times.
+        self.file_metrics
+            .scan_efficiency_ratio
+            .set_total(self.partitioned_file.object_meta.size as usize);
+    }
+}

native/core/src/parquet/mod.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+pub mod eager_page_index_reader_factory;
 pub mod encryption_support;
 
 pub mod parquet_exec;

native/core/src/parquet/parquet_exec.rs

@@ -16,13 +16,13 @@
 // under the License.
 
 use crate::execution::operators::ExecutionError;
+use crate::parquet::eager_page_index_reader_factory::EagerPageIndexReaderFactory;
 use crate::parquet::encryption_support::{CometEncryptionConfig, ENCRYPTION_FACTORY_ID};
 use crate::parquet::parquet_support::SparkParquetOptions;
 use crate::parquet::schema_adapter::SparkPhysicalExprAdapterFactory;
 use arrow::datatypes::{Field, SchemaRef};
 use datafusion::config::TableParquetOptions;
 use datafusion::datasource::listing::PartitionedFile;
-use datafusion::datasource::physical_plan::parquet::CachedParquetFileReaderFactory;
 use datafusion::datasource::physical_plan::{
     FileGroup, FileScanConfigBuilder, FileSource, ParquetSource,
 };
@@ -142,11 +142,15 @@ pub(crate) fn init_datasource_exec(
         );
     }
 
-    // DataFusion's metadata-caching reader factory: loads each file's footer metadata once into the
-    // per-task RuntimeEnv cache (bounded LRU, `metadata_cache_limit`) and reuses it across that
-    // file's row-group splits. The page index is loaded lazily by the opener, only when row-group
-    // statistics cannot already prove no further pruning is possible (apache/datafusion#22857), and
-    // that load bypasses this factory, so only footer metadata is cached.
+    // DataFusion's opener requests `PageIndexPolicy::Skip` on the initial metadata load and
+    // defers loading the page index until row-group pruning shows it is still needed
+    // (apache/datafusion#22857). That deferred load bypasses `FileMetadataCache` entirely, so on
+    // a predicate that never resolves to fully-matched by row-group statistics alone (e.g. `IS
+    // NOT NULL` on a column whose row groups don't carry `null_count` stats, as with TPC-DS
+    // `store_sales`), the page index is re-fetched, uncached, on every open (comet#3978).
+    // `EagerPageIndexReaderFactory` forces the page index to load on the first fetch and be
+    // cached with the footer, at the cost of losing the skip's benefit when it would have
+    // applied. Filed upstream as apache/datafusion#23978; revert this once that's fixed.
     //
     // TODO: metadata I/O is invisible in metrics. `fetch_metadata` reads via `ObjectStore::get_ranges`,
     // bypassing the `get_bytes` path where `bytes_scanned` is counted. A byte-counting ObjectStore
@@ -155,7 +159,7 @@ pub(crate) fn init_datasource_exec(
     let store = runtime_env.object_store(&object_store_url)?;
     let metadata_cache = runtime_env.cache_manager.get_file_metadata_cache();
     parquet_source = parquet_source.with_parquet_file_reader_factory(Arc::new(
-        CachedParquetFileReaderFactory::new(store, metadata_cache),
+        EagerPageIndexReaderFactory::new(store, metadata_cache),
     ));
 
     // Route data filters through `try_pushdown_filters` rather than calling
@@ -259,16 +263,18 @@ mod tests {
     use parquet::file::properties::{EnabledStatistics, WriterProperties};
     use std::fs::File;
 
-    // The scan's reader factory caches each file's footer metadata in the per-task RuntimeEnv
-    // metadata cache, reused across the file's row-group splits (#3978). The page index is loaded
-    // lazily and only when row-group statistics cannot already prune the surviving row groups
-    // (apache/datafusion#22857); that load bypasses the caching factory, so the cache holds footer
-    // metadata only. Without a pruning predicate here the page index is never loaded.
+    // Regression test for #3978: DataFusion's opener requests `PageIndexPolicy::Skip` on the
+    // initial metadata load and only loads the page index later, on demand, when row-group
+    // pruning shows it is still needed (apache/datafusion#22857). That on-demand load bypasses
+    // `FileMetadataCache` entirely, so on a predicate that never resolves to fully-matched (e.g.
+    // non-null `IS NOT NULL` join keys without row-group `null_count` stats), the page index is
+    // re-fetched, uncached, on every open. `EagerPageIndexReaderFactory` ignores the requested
+    // policy and always loads the page index into the cache on the first fetch, so it asserts
+    // present here even though this scan has no pruning predicate to request it.
     #[tokio::test]
-    async fn caches_footer_metadata() {
+    async fn caches_full_metadata_with_page_index() {
         // A file with a page index (page-level statistics and a small data page row limit) so we can
-        // assert the cached metadata deliberately excludes the page index rather than lacking it
-        // because none was written.
+        // assert the cached metadata includes it rather than lacking it because none was written.
         let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
         let batch = RecordBatch::try_new(
             Arc::clone(&schema),
@@ -322,8 +328,8 @@ mod tests {
             batch.unwrap();
         }
 
-        // The per-task RuntimeEnv metadata cache holds this file's footer metadata. The page index
-        // is not loaded (no pruning predicate), so the cached metadata excludes it.
+        // The per-task RuntimeEnv metadata cache holds this file's metadata with the page index
+        // (column + offset index) loaded, even though this scan pushed no pruning predicate.
         let cache = session_ctx
             .runtime_env()
             .cache_manager
@@ -338,8 +344,8 @@ mod tests {
             .expect("cached entry should hold Parquet metadata")
             .parquet_metadata();
         assert!(
-            parquet_meta.column_index().is_none() && parquet_meta.offset_index().is_none(),
-            "footer-only cache must not include the page index"
+            parquet_meta.column_index().is_some() && parquet_meta.offset_index().is_some(),
+            "cached metadata must include the page index"
         );
     }
 }