Test out pushdecoder

Author: Dandandan

datafusion/core/tests/datasource/object_store_access.rs

@@ -502,7 +502,7 @@ async fn query_single_parquet_file_with_single_predicate() {
     RequestCountingObjectStore()
     Total Requests: 2
     - GET  (opts) path=parquet_table.parquet head=true
-    - GET  (ranges) path=parquet_table.parquet ranges=1064-1594,1594-2124
+    - GET  (ranges) path=parquet_table.parquet ranges=1064-1481,1481-1594,1594-2011,2011-2124
     "
     );
 }
@@ -526,8 +526,8 @@ async fn query_single_parquet_file_multi_row_groups_multiple_predicates() {
     RequestCountingObjectStore()
     Total Requests: 3
     - GET  (opts) path=parquet_table.parquet head=true
-    - GET  (ranges) path=parquet_table.parquet ranges=4-534,534-1064
-    - GET  (ranges) path=parquet_table.parquet ranges=1064-1594,1594-2124
+    - GET  (ranges) path=parquet_table.parquet ranges=4-421,421-534,534-951,951-1064
+    - GET  (ranges) path=parquet_table.parquet ranges=1064-1481,1481-1594,1594-2011,2011-2124
     "
     );
 }

datafusion/datasource-parquet/src/opener.rs

@@ -445,7 +445,7 @@ impl FileOpener for ParquetOpener {
             // Determine which row groups to actually read. The idea is to skip
             // as many row groups as possible based on the metadata and query
             let file_metadata = Arc::clone(reader_metadata.metadata());
-            let predicate_ref = pruning_predicate.as_ref().map(|p| p.as_ref());
+            let pruning_pred = pruning_predicate.as_ref().map(|p| p.as_ref());
             let rg_metadata = file_metadata.row_groups();
             // track which row groups to actually read
             let access_plan =
@@ -457,13 +457,13 @@ impl FileOpener for ParquetOpener {
             }
 
             // If there is a predicate that can be evaluated against the metadata
-            if let Some(predicate) = predicate_ref.as_ref() {
+            if let Some(pruning_pred) = pruning_pred.as_ref() {
                 if enable_row_group_stats_pruning {
                     row_groups.prune_by_statistics(
                         &physical_file_schema,
                         reader_metadata.parquet_schema(),
                         rg_metadata,
-                        predicate,
+                        pruning_pred,
                         &file_metrics,
                     );
                 } else {
@@ -474,16 +474,18 @@ impl FileOpener for ParquetOpener {
                         .add_matched(row_groups.remaining_row_group_count());
                 }
 
-                // Bloom filter pruning: create a separate reader for bloom
-                // filter I/O since the push decoder doesn't own the reader.
                 if enable_bloom_filter && !row_groups.is_empty() {
-                    let bf_reader: Box<dyn AsyncFileReader> = parquet_file_reader_factory
-                        .create_reader(
+                    // Use the existing reader for bloom filter I/O;
+                    // replace with a fresh reader for decoding below.
+                    let bf_reader = std::mem::replace(
+                        &mut async_file_reader,
+                        parquet_file_reader_factory.create_reader(
                             partition_index,
                             partitioned_file.clone(),
                             metadata_size_hint,
                             &metrics,
-                        )?;
+                        )?,
+                    );
                     let mut bf_builder =
                         ParquetRecordBatchStreamBuilder::new_with_metadata(
                             bf_reader,
@@ -493,7 +495,7 @@ impl FileOpener for ParquetOpener {
                         .prune_by_bloom_filters(
                             &physical_file_schema,
                             &mut bf_builder,
-                            predicate,
+                            pruning_pred,
                             &file_metrics,
                         )
                         .await;
@@ -558,7 +560,7 @@ impl FileOpener for ParquetOpener {
             }
 
             // ---------------------------------------------------------
-            // Step: construct the push decoder builder
+            // Step: construct builder for the final RecordBatch stream
             // ---------------------------------------------------------
 
             let mut builder =
@@ -619,7 +621,10 @@ impl FileOpener for ParquetOpener {
             let arrow_reader_metrics = ArrowReaderMetrics::enabled();
 
             let indices = projection.column_indices();
-            let mask = ProjectionMask::roots(reader_metadata.parquet_schema(), indices);
+            let mask = ProjectionMask::roots(
+                reader_metadata.parquet_schema(),
+                indices.clone(),
+            );
 
             let decoder = builder
                 .with_projection(mask)
@@ -632,13 +637,22 @@ impl FileOpener for ParquetOpener {
                 file_metrics.predicate_cache_inner_records.clone();
             let predicate_cache_records = file_metrics.predicate_cache_records.clone();
 
-            // Create a stream that drives the decoder by fetching data as needed
+            // Rebase column indices to match the narrowed stream schema.
+            // The projection expressions have indices based on physical_file_schema,
+            // but the stream only contains the columns selected by the ProjectionMask.
+            let stream_schema =
+                Arc::new(physical_file_schema.project(&indices)?);
+            let replace_schema = stream_schema != output_schema;
+            let projection = projection
+                .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?;
+            let projector = projection.make_projector(&stream_schema)?;
             let stream = futures::stream::unfold(
                 PushDecoderStreamState {
                     decoder,
                     reader: async_file_reader,
-                    projection,
+                    projector,
                     output_schema,
+                    replace_schema,
                     arrow_reader_metrics,
                     predicate_cache_inner_records,
                     predicate_cache_records,
@@ -672,11 +686,7 @@ impl FileOpener for ParquetOpener {
                                     &state.predicate_cache_inner_records,
                                     &state.predicate_cache_records,
                                 );
-                                let result = apply_projection(
-                                    batch,
-                                    &state.projection,
-                                    &state.output_schema,
-                                );
+                                let result = state.project_batch(batch);
                                 return Some((result, state));
                             }
                             Ok(DecodeResult::Finished) => {
@@ -694,7 +704,6 @@ impl FileOpener for ParquetOpener {
             // Step: wrap the stream so a dynamic filter can stop the file scan early
             // ----------------------------------------------------------------------
             if let Some(file_pruner) = file_pruner {
-                // Box the stream first to satisfy Unpin bound on EarlyStoppingStream
                 let boxed_stream = stream.boxed();
                 Ok(EarlyStoppingStream::new(
                     boxed_stream,
@@ -709,41 +718,31 @@ impl FileOpener for ParquetOpener {
     }
 }
 
-/// State carried through the unfold stream
 struct PushDecoderStreamState {
     decoder: parquet::arrow::push_decoder::ParquetPushDecoder,
     reader: Box<dyn AsyncFileReader>,
-    projection: ProjectionExprs,
+    projector: datafusion_physical_expr::projection::Projector,
     output_schema: Arc<arrow::datatypes::Schema>,
+    replace_schema: bool,
     arrow_reader_metrics: ArrowReaderMetrics,
     predicate_cache_inner_records: Gauge,
     predicate_cache_records: Gauge,
 }
 
-/// Apply projection expressions to the decoded batch and ensure correct schema
-fn apply_projection(
-    batch: RecordBatch,
-    projection: &ProjectionExprs,
-    output_schema: &Arc<arrow::datatypes::Schema>,
-) -> Result<RecordBatch> {
-    let stream_schema = batch.schema();
-    let projection = projection
-        .clone()
-        .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?;
-    let projector = projection.make_projector(&stream_schema)?;
-    let mut batch = projector.project_batch(&batch)?;
-
-    if batch.schema() != *output_schema {
-        let (_schema, arrays, num_rows) = batch.into_parts();
-        let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
-        batch = RecordBatch::try_new_with_options(
-            Arc::clone(output_schema),
-            arrays,
-            &options,
-        )?;
+impl PushDecoderStreamState {
+    fn project_batch(&self, batch: RecordBatch) -> Result<RecordBatch> {
+        let mut batch = self.projector.project_batch(&batch)?;
+        if self.replace_schema {
+            let (_schema, arrays, num_rows) = batch.into_parts();
+            let options = RecordBatchOptions::new().with_row_count(Some(num_rows));
+            batch = RecordBatch::try_new_with_options(
+                Arc::clone(&self.output_schema),
+                arrays,
+                &options,
+            )?;
+        }
+        Ok(batch)
     }
-
-    Ok(batch)
 }
 
 /// Copies metrics from ArrowReaderMetrics (the metrics collected by the