Commit beced87

use crate::opener::build_pruning_predicates;

use crate::row_filter;

use crate::row_group_filter::RowGroupAccessPlanFilter;

use crate::{

ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory,

apply_file_schema_type_coercions,

use arrow::record_batch::RecordBatch;

use datafusion_common::{DataFusionError, Result};

use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};

use datafusion_datasource::{PartitionedFile, TableSchema};

use datafusion_physical_expr::projection::ProjectionExprs;

use datafusion_physical_expr::utils::reassign_expr_columns;

use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;

use datafusion_physical_expr_common::physical_expr::PhysicalExpr;

use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder};

use futures::StreamExt;

use parquet::DecodeResult;

use parquet::arrow::ProjectionMask;

use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions};

use parquet::arrow::async_reader::AsyncFileReader;

use parquet::arrow::push_decoder::ParquetPushDecoderBuilder;

use std::sync::Arc;

pub(super) struct PushParquetOpener {

pub partition_index: usize,

pub projection: ProjectionExprs,

pub batch_size: usize,

pub limit: Option<usize>,

pub predicate: Option<Arc<dyn PhysicalExpr>>,

pub table_schema: TableSchema,

pub metadata_size_hint: Option<usize>,

pub metrics: ExecutionPlanMetricsSet,

pub parquet_file_reader_factory: Arc<dyn ParquetFileReaderFactory>,

pub pushdown_filters: bool,

pub reorder_filters: bool,

pub enable_row_group_stats_pruning: bool,

pub expr_adapter_factory: Arc<dyn PhysicalExprAdapterFactory>,

impl FileOpener for PushParquetOpener {

fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {

let file_name = partitioned_file.object_meta.location.to_string();

let file_metrics =

ParquetFileMetrics::new(self.partition_index, &file_name, &self.metrics);

let metadata_size_hint = partitioned_file

.metadata_size_hint

.or(self.metadata_size_hint);

let mut async_file_reader: Box<dyn AsyncFileReader> =

self.parquet_file_reader_factory.create_reader(

self.partition_index,

partitioned_file.clone(),

metadata_size_hint,

&self.metrics,

)?;

let batch_size = self.batch_size;

let limit = self.limit;

let projection = self.projection.clone();

let logical_file_schema = Arc::clone(self.table_schema.file_schema());

let predicate = self.predicate.clone();

let pushdown_filters = self.pushdown_filters;

let reorder_filters = self.reorder_filters;

let enable_row_group_stats_pruning = self.enable_row_group_stats_pruning;

let expr_adapter_factory = Arc::clone(&self.expr_adapter_factory);

let extensions = partitioned_file.extensions.clone();

let predicate_creation_errors = MetricBuilder::new(&self.metrics)

.global_counter("num_predicate_creation_errors");

let output_schema = Arc::new(

self.projection

.project_schema(self.table_schema.table_schema())?,

);

Ok(Box::pin(async move {

// Load metadata

let options = ArrowReaderOptions::new();

let reader_metadata =

ArrowReaderMetadata::load_async(&mut async_file_reader, options.clone())

.await?;

let mut physical_file_schema = Arc::clone(reader_metadata.schema());

// Apply schema coercions (e.g. Utf8 -> Utf8View)

if let Some(merged) = apply_file_schema_type_coercions(

&logical_file_schema,

&physical_file_schema,

) {

physical_file_schema = Arc::new(merged);

}

let options = options.with_schema(Arc::clone(&physical_file_schema));

let reader_metadata = ArrowReaderMetadata::try_new(

Arc::clone(reader_metadata.metadata()),

options,

)?;

// Adapt predicate and projection to the physical file schema

let rewriter = expr_adapter_factory.create(

Arc::clone(&logical_file_schema),

Arc::clone(&physical_file_schema),

)?;

let predicate = predicate.map(|p| rewriter.rewrite(p)).transpose()?;

let projection = projection.try_map_exprs(|p| rewriter.rewrite(p))?;

// Build pruning predicates

let (pruning_predicate, _page_pruning_predicate) = build_pruning_predicates(

predicate.as_ref(),

&physical_file_schema,

&predicate_creation_errors,

);

// Determine which row groups to read

let file_metadata = Arc::clone(reader_metadata.metadata());

let rg_metadata = file_metadata.row_groups();

let num_row_groups = rg_metadata.len();

let access_plan = match extensions {

Some(ext) => ext

.downcast_ref::<ParquetAccessPlan>()

.cloned()

.unwrap_or_else(|| ParquetAccessPlan::new_all(num_row_groups)),

None => ParquetAccessPlan::new_all(num_row_groups),

};

let mut row_groups_filter = RowGroupAccessPlanFilter::new(access_plan);

// Prune row groups by statistics

if let Some(ref pred) = pruning_predicate {

if enable_row_group_stats_pruning {

row_groups_filter.prune_by_statistics(

&physical_file_schema,

reader_metadata.parquet_schema(),

rg_metadata,

pred,

&file_metrics,

);

}

}

let access_plan = row_groups_filter.build();

let row_group_indexes = access_plan.row_group_indexes();

// Build the push decoder

let mut builder =

ParquetPushDecoderBuilder::new_with_metadata(reader_metadata.clone())

.with_batch_size(batch_size);

// Set row filter for predicate pushdown

if let Some(ref predicate) = predicate {

if pushdown_filters {

let row_filter = row_filter::build_row_filter(

predicate,

&physical_file_schema,

file_metadata.as_ref(),

reorder_filters,

&file_metrics,

);

match row_filter {

Ok(Some(filter)) => {

builder = builder.with_row_filter(filter);

}

Ok(None) => {}

Err(e) => {

log::debug!(

"Ignoring error building row filter for '{predicate:?}': {e}"

);

}

}

}

}

if !row_group_indexes.is_empty() {

builder = builder.with_row_groups(row_group_indexes);

} else {

// All row groups pruned - return empty stream

return Ok(futures::stream::empty().boxed());

}

// Set row selection from the access plan

let row_selection = access_plan.into_overall_row_selection(rg_metadata)?;

if let Some(selection) = row_selection {

builder = builder.with_row_selection(selection);

}

if let Some(limit) = limit {

builder = builder.with_limit(limit);

}

// Apply projection

let projection = projection.clone();

let indices = projection.column_indices();

let mask = ProjectionMask::roots(reader_metadata.parquet_schema(), indices);

builder = builder.with_projection(mask);

let decoder = builder.build()?;

// Create a stream that drives the decoder by fetching data as needed

let stream = futures::stream::unfold(

(decoder, async_file_reader, projection, output_schema),

|(mut decoder, mut reader, projection, output_schema)| async move {

loop {

match decoder.try_decode() {

Ok(DecodeResult::NeedsData(ranges)) => {

match reader.get_byte_ranges(ranges.clone()).await {

Ok(data) => {

if let Err(e) = decoder.push_ranges(ranges, data)

{

return Some((

Err(DataFusionError::from(e)),

(

decoder,

reader,

projection,

output_schema,

),

));

}

}

Err(e) => {

return Some((

Err(DataFusionError::from(e)),

(decoder, reader, projection, output_schema),

));

}

}

}

Ok(DecodeResult::Data(batch)) => {

// Apply projection expressions and schema

let result =

apply_projection(batch, &projection, &output_schema);

return Some((

result,

(decoder, reader, projection, output_schema),

));

}

Ok(DecodeResult::Finished) => {

return None;

}

Err(e) => {

return Some((

Err(DataFusionError::from(e)),

(decoder, reader, projection, output_schema),

));

}

}

}

},

);

Ok(stream.boxed())

}))

}

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 =

arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(num_rows));

batch = RecordBatch::try_new_with_options(

Arc::clone(output_schema),

arrays,

&options,

)?;

}

Ok(batch)