extract more

Author: alamb

datafusion/datasource-parquet/src/opener.rs

@@ -211,6 +211,19 @@ struct MetadataLoadedParquetOpen {
     options: ArrowReaderOptions,
 }
 
+/// Result of CPU-only preparation after metadata has been loaded.
+///
+/// This captures the file schema coercions and file-specific pruning predicates
+/// so the next async step only has to fetch any missing page index data.
+struct FiltersPreparedParquetOpen {
+    loaded: MetadataLoadedParquetOpen,
+    physical_file_schema: SchemaRef,
+    projection: ProjectionExprs,
+    predicate: Option<Arc<dyn PhysicalExpr>>,
+    pruning_predicate: Option<Arc<PruningPredicate>>,
+    page_pruning_predicate: Option<Arc<PagePruningAccessPlanFilter>>,
+}
+
 impl ParquetMorselizerState {
     /// Perform the CPU-only setup for opening a parquet file.
     fn prepare_open_file(
@@ -399,51 +412,20 @@ impl ParquetMorselizerState {
         })
     }
 
-    /// Continue opening once parquet metadata has been loaded.
-    async fn execute_metadata_loaded_open(
+    /// Prepare file-specific filters and schema coercions after metadata is loaded.
+    fn prepare_loaded_metadata(
         loaded: MetadataLoadedParquetOpen,
-    ) -> Result<BoxStream<'static, Result<RecordBatch>>> {
+    ) -> Result<FiltersPreparedParquetOpen> {
         let MetadataLoadedParquetOpen {
-            prepared,
+            mut prepared,
             mut reader_metadata,
             mut options,
         } = loaded;
-        let PreparedParquetOpen {
-            state,
-            partitioned_file,
-            file_range,
-            extensions,
-            file_metrics,
-            file_pruner,
-            metadata_size_hint,
-            mut async_file_reader,
-            logical_file_schema,
-            output_schema,
-            mut projection,
-            mut predicate,
-            #[cfg(feature = "parquet_encryption")]
-                file_decryption_properties: _,
-        } = prepared;
-
-        let file_name = partitioned_file.object_meta.location.to_string();
-        let batch_size = state.batch_size;
-        let reorder_predicates = state.reorder_filters;
-        let pushdown_filters = state.pushdown_filters;
-        let force_filter_selections = state.force_filter_selections;
+        let state = Arc::clone(&prepared.state);
         let coerce_int96 = state.coerce_int96;
-        let enable_bloom_filter = state.enable_bloom_filter;
-        let enable_row_group_stats_pruning = state.enable_row_group_stats_pruning;
-        let limit = state.limit;
-        let parquet_file_reader_factory = Arc::clone(&state.parquet_file_reader_factory);
-        let partition_index = state.partition_index;
-        let metrics = state.metrics.clone();
         let predicate_creation_errors = MetricBuilder::new(&state.metrics)
             .global_counter("num_predicate_creation_errors");
         let expr_adapter_factory = Arc::clone(&state.expr_adapter_factory);
-        let enable_page_index = state.enable_page_index;
-        let max_predicate_cache_size = state.max_predicate_cache_size;
-        let reverse_row_groups = state.reverse_row_groups;
-        let preserve_order = state.preserve_order;
 
         // Note about schemas: we are actually dealing with **3 different schemas** here:
         // - The table schema as defined by the TableProvider.
@@ -452,6 +434,7 @@ impl ParquetMorselizerState {
         //   This is what the physical file schema is coerced to.
         // - The physical file schema: this is the schema that the arrow-rs
         //   parquet reader will actually produce.
+        let logical_file_schema = Arc::clone(&prepared.logical_file_schema);
         let mut physical_file_schema = Arc::clone(reader_metadata.schema());
 
         // The schema loaded from the file may not be the same as the
@@ -495,32 +478,109 @@ impl ParquetMorselizerState {
             Arc::clone(&physical_file_schema),
         )?;
         let simplifier = PhysicalExprSimplifier::new(&physical_file_schema);
-        predicate = predicate
+        prepared.predicate = prepared
+            .predicate
             .map(|p| simplifier.simplify(rewriter.rewrite(p)?))
             .transpose()?;
         // Adapt projections to the physical file schema as well
-        projection =
-            projection.try_map_exprs(|p| simplifier.simplify(rewriter.rewrite(p)?))?;
+        prepared.projection = prepared
+            .projection
+            .try_map_exprs(|p| simplifier.simplify(rewriter.rewrite(p)?))?;
 
         // Build predicates for this specific file
         let (pruning_predicate, page_pruning_predicate) = build_pruning_predicates(
-            predicate.as_ref(),
+            prepared.predicate.as_ref(),
             &physical_file_schema,
             &predicate_creation_errors,
         );
 
-        // The page index is not stored inline in the parquet footer so the
-        // code above may not have read the page index structures yet. If we
-        // need them for reading and they aren't yet loaded, we need to load them now.
-        if should_enable_page_index(enable_page_index, &page_pruning_predicate) {
-            reader_metadata = load_page_index(
+        let projection = prepared.projection.clone();
+        let predicate = prepared.predicate.clone();
+
+        Ok(FiltersPreparedParquetOpen {
+            loaded: MetadataLoadedParquetOpen {
+                prepared,
                 reader_metadata,
-                &mut async_file_reader,
+                options,
+            },
+            physical_file_schema,
+            projection,
+            predicate,
+            pruning_predicate,
+            page_pruning_predicate,
+        })
+    }
+
+    /// Fetch the page index if it is needed and missing from the loaded metadata.
+    async fn load_prepared_page_index(
+        mut prepared: FiltersPreparedParquetOpen,
+    ) -> Result<FiltersPreparedParquetOpen> {
+        let enable_page_index = prepared.loaded.prepared.state.enable_page_index;
+        if should_enable_page_index(enable_page_index, &prepared.page_pruning_predicate)
+        {
+            prepared.loaded.reader_metadata = load_page_index(
+                prepared.loaded.reader_metadata,
+                &mut prepared.loaded.prepared.async_file_reader,
                 // Since we're manually loading the page index the option here should not matter but we pass it in for consistency
-                options.with_page_index_policy(PageIndexPolicy::Optional),
+                prepared
+                    .loaded
+                    .options
+                    .clone()
+                    .with_page_index_policy(PageIndexPolicy::Optional),
             )
             .await?;
         }
+        Ok(prepared)
+    }
+
+    /// Continue opening once metadata-derived filters and page indexes are ready.
+    async fn execute_metadata_loaded_open(
+        prepared: FiltersPreparedParquetOpen,
+    ) -> Result<BoxStream<'static, Result<RecordBatch>>> {
+        let FiltersPreparedParquetOpen {
+            loaded,
+            physical_file_schema,
+            projection,
+            predicate,
+            pruning_predicate,
+            page_pruning_predicate,
+        } = prepared;
+        let MetadataLoadedParquetOpen {
+            prepared,
+            reader_metadata,
+            options: _,
+        } = loaded;
+        let PreparedParquetOpen {
+            state,
+            partitioned_file,
+            file_range,
+            extensions,
+            file_metrics,
+            file_pruner,
+            metadata_size_hint,
+            mut async_file_reader,
+            logical_file_schema: _,
+            output_schema,
+            projection: _,
+            predicate: _,
+            #[cfg(feature = "parquet_encryption")]
+            file_decryption_properties: _,
+        } = prepared;
+
+        let file_name = partitioned_file.object_meta.location.to_string();
+        let batch_size = state.batch_size;
+        let reorder_predicates = state.reorder_filters;
+        let pushdown_filters = state.pushdown_filters;
+        let force_filter_selections = state.force_filter_selections;
+        let enable_bloom_filter = state.enable_bloom_filter;
+        let enable_row_group_stats_pruning = state.enable_row_group_stats_pruning;
+        let limit = state.limit;
+        let parquet_file_reader_factory = Arc::clone(&state.parquet_file_reader_factory);
+        let partition_index = state.partition_index;
+        let metrics = state.metrics.clone();
+        let max_predicate_cache_size = state.max_predicate_cache_size;
+        let reverse_row_groups = state.reverse_row_groups;
+        let preserve_order = state.preserve_order;
 
         // ------------------------------------------------------------
         // Step: prune row groups by range, predicate and bloom filter
@@ -612,8 +672,7 @@ impl ParquetMorselizerState {
         // be ruled using page metadata, rows from other columns
         // with that range can be skipped as well
         // --------------------------------------------------------
-        if enable_page_index
-            && !access_plan.is_empty()
+        if !access_plan.is_empty()
             && let Some(p) = page_pruning_predicate
         {
             access_plan = p.prune_plan_with_page_index(
@@ -775,8 +834,10 @@ enum ReadyState {
     PrepareEncryption(Box<PreparedParquetOpen>),
     /// Planner can do file-level pruning before requesting parquet metadata.
     PruneFiles(Box<PreparedParquetOpen>),
-    /// Planner has loaded parquet metadata and is ready for further planning.
-    Prepared(Box<MetadataLoadedParquetOpen>),
+    /// Planner has loaded parquet metadata and can do CPU-only filter preparation.
+    PrepareFilters(Box<MetadataLoadedParquetOpen>),
+    /// Planner has prepared filters and can request any missing page index data.
+    Prepared(Box<FiltersPreparedParquetOpen>),
     /// Planner has a fully prepared stream ready to wrap as a morsel.
     EmitMorsel(BoxStream<'static, Result<RecordBatch>>),
 }
@@ -796,6 +857,10 @@ impl ReadyState {
     }
 
     fn prepared(prepared: MetadataLoadedParquetOpen) -> Self {
+        Self::PrepareFilters(prepared.into())
+    }
+
+    fn filters_prepared(prepared: FiltersPreparedParquetOpen) -> Self {
         Self::Prepared(prepared.into())
     }
 
@@ -845,6 +910,7 @@ impl ParquetMorselPlannerState {
                 #[cfg(feature = "parquet_encryption")]
                 ReadyState::PrepareEncryption(_) => "Ready(PrepareEncryption)",
                 ReadyState::PruneFiles(_) => "Ready(PruneFiles)",
+                ReadyState::PrepareFilters(_) => "Ready(PrepareFilters)",
                 ReadyState::Prepared(_) => "Ready(Prepared)",
                 ReadyState::EmitMorsel(_) => "Ready(EmitMorsel)",
             },
@@ -961,9 +1027,20 @@ impl MorselPlanner for ParquetMorselPlanner {
                         Ok(ReadyState::prepared(loaded))
                     }))
                 }
+                ReadyState::PrepareFilters(prepared) => {
+                    let prepared =
+                        ParquetMorselizerState::prepare_loaded_metadata(*prepared)?;
+                    self.state = ParquetMorselPlannerState::ready(
+                        ReadyState::filters_prepared(prepared),
+                    );
+                    Ok(Some(MorselPlan::new()))
+                }
                 ReadyState::Prepared(prepared) => Ok(self.schedule_io(async move {
+                    let prepared =
+                        ParquetMorselizerState::load_prepared_page_index(*prepared)
+                            .await?;
                     let stream =
-                        ParquetMorselizerState::execute_metadata_loaded_open(*prepared)
+                        ParquetMorselizerState::execute_metadata_loaded_open(prepared)
                             .await?;
                     Ok(ReadyState::emit_morsel(stream))
                 })),