Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 74 additions & 0 deletions datafusion/datasource-parquet/src/opener/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -812,6 +812,7 @@ impl ParquetMorselizer {
let file_pruner = predicate.as_ref().and_then(|p| {
FilePruner::try_new(
Arc::clone(p),
self.table_schema.table_schema(),
&logical_file_schema,
&partitioned_file,
predicate_creation_errors.clone(),
Expand Down Expand Up @@ -3761,4 +3762,77 @@ mod test {
assert_eq!(rows, 5);
}
}

#[tokio::test]
async fn test_prune_on_constant_columns() {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the original test from #23272

let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;

let batch = record_batch!((
"a",
Int32,
vec![Some(5), Some(5), Some(5), Some(5), Some(5)]
))
.unwrap();
let data_size =
write_parquet(Arc::clone(&store), "constant.parquet", batch.clone()).await;
let schema = batch.schema();

let file = PartitionedFile::new(
"constant.parquet".to_string(),
u64::try_from(data_size).unwrap(),
)
.with_statistics(Arc::new(Statistics {
num_rows: Precision::Exact(5),
total_byte_size: Precision::Absent,
column_statistics: vec![
ColumnStatistics::new_unknown()
.with_min_value(Precision::Exact(ScalarValue::Int32(Some(5))))
.with_max_value(Precision::Exact(ScalarValue::Int32(Some(5))))
.with_null_count(Precision::Exact(0)),
],
}));

let make_opener = |predicate| {
let metrics = ExecutionPlanMetricsSet::new();
let morselizer = ParquetMorselizerBuilder::new()
.with_store(Arc::clone(&store))
.with_schema(Arc::clone(&schema))
.with_projection_indices(&[0])
.with_predicate(predicate)
.with_metrics(metrics.clone())
.build();
(morselizer, metrics)
};

// constant expression `a > 100`
let expr = col("a").gt(lit(100));
let predicate = logical2physical(&expr, &schema);
let (opener, metrics) = make_opener(predicate);
let stream = open_file(&opener, file).await.unwrap();
let (_num_batches, _num_rows) = count_batches_and_rows(stream).await;

// The bug: `files_ranges_pruned_statistics` pruned count is 0 because
// `constant_columns_from_stats` folds the column reference into a literal,
// and `FilePruner` cannot prune a constant expression.
let pruned = {
use datafusion_physical_plan::metrics::MetricValue;
metrics
.clone_inner()
.iter()
.filter_map(|m| match m.value() {
MetricValue::PruningMetrics {
name,
pruning_metrics,
} if name.as_ref() == "files_ranges_pruned_statistics" => {
Some(pruning_metrics.pruned())
}
_ => None,
})
.sum::<usize>()
};
assert!(
pruned > 0,
"file with constant column should be pruned at file level, got pruned={pruned}"
);
}
}
81 changes: 79 additions & 2 deletions datafusion/pruning/src/file_pruner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ use std::sync::Arc;
use arrow::datatypes::{FieldRef, SchemaRef};
use datafusion_common::{Result, internal_datafusion_err, pruning::PrunableStatistics};
use datafusion_datasource::PartitionedFile;
use datafusion_physical_expr::DynamicFilterTracking;
use datafusion_physical_expr::{DynamicFilterTracking, PhysicalExprSimplifier};
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
use datafusion_physical_plan::metrics::Count;
use log::debug;
Expand All @@ -45,6 +45,7 @@ pub struct FilePruner {
checked_once: bool,
/// Schema used for pruning (the logical file schema).
file_schema: SchemaRef,
table_schema: SchemaRef,
file_stats_pruning: PrunableStatistics,
predicate_creation_errors: Count,
}
Expand All @@ -57,13 +58,15 @@ impl FilePruner {
#[expect(clippy::needless_pass_by_value)]
pub fn new(
predicate: Arc<dyn PhysicalExpr>,
table_schema: &SchemaRef,
logical_file_schema: &SchemaRef,
_partition_fields: Vec<FieldRef>,
partitioned_file: PartitionedFile,
predicate_creation_errors: Count,
) -> Result<Self> {
Self::try_new(
predicate,
table_schema,
logical_file_schema,
&partitioned_file,
predicate_creation_errors,
Expand All @@ -87,6 +90,7 @@ impl FilePruner {
/// partition-value folding even without column statistics.
pub fn try_new(
predicate: Arc<dyn PhysicalExpr>,
table_schema: &SchemaRef,
file_schema: &SchemaRef,
partitioned_file: &PartitionedFile,
predicate_creation_errors: Count,
Expand All @@ -103,13 +107,15 @@ impl FilePruner {
if !partitioned_file.has_statistics() && !tracking.contains_dynamic_filter() {
return None;
}

let file_stats_pruning =
PrunableStatistics::new(vec![file_stats.clone()], Arc::clone(file_schema));
Some(Self {
predicate,
tracking,
checked_once: false,
file_schema: Arc::clone(file_schema),
table_schema: Arc::clone(table_schema),
file_stats_pruning,
predicate_creation_errors,
})
Expand Down Expand Up @@ -143,8 +149,17 @@ impl FilePruner {
if !should_build {
return Ok(false);
}
// If there is no dynamic-filter-expression involved, convert constant expression
// such as `a > 100` to `true`/`false` so `PruningPredicate` can use to prune the file.
let predicate = if !self.tracking.contains_dynamic_filter() {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This works, but i am not sure if this the best place to do it. The expression will eventually end up in
https://github.com/apache/datafusion/blob/main/datafusion/pruning/src/pruning_predicate.rs#L476
but the PhysicalExprSimplifier won't be applied because it's not a dynamic filter in this case.

let simplifier = PhysicalExprSimplifier::new(&self.table_schema);
simplifier.simplify(Arc::clone(&self.predicate))?
} else {
Arc::clone(&self.predicate)
};

let pruning_predicate = build_pruning_predicate(
Arc::clone(&self.predicate),
predicate,
&self.file_schema,
&self.predicate_creation_errors,
);
Expand All @@ -169,3 +184,65 @@ impl FilePruner {
Ok(false)
}
}

#[cfg(test)]
mod tests {
use super::*;
use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::{
ColumnStatistics, ScalarValue, Statistics, stats::Precision,
};
use datafusion_expr::{col, lit};
use datafusion_physical_expr::planner::logical2physical;

#[test]
fn test_pruning_on_various_predicates() {

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

More additional tests

let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));

let file = PartitionedFile::new("test.parquet".to_string(), 100).with_statistics(
Arc::new(Statistics {
num_rows: Precision::Exact(5),
total_byte_size: Precision::Absent,
column_statistics: vec![
ColumnStatistics::new_unknown()
.with_min_value(Precision::Exact(ScalarValue::Int32(Some(5))))
.with_max_value(Precision::Exact(ScalarValue::Int32(Some(5))))
.with_null_count(Precision::Exact(0)),
],
}),
);

let predicate = logical2physical(&lit(5i32).gt(lit(100i32)), &schema);

let mut pruner =
FilePruner::try_new(predicate, &schema, &schema, &file, Count::new())
.unwrap();

assert!(
pruner.should_prune().unwrap(),
"constant predicate `5 > 100` should simplify to `false` and prune the file"
);

let predicate = logical2physical(&lit(5i32).gt(lit(1i32)), &schema);

let mut pruner =
FilePruner::try_new(predicate, &schema, &schema, &file, Count::new())
.unwrap();

assert!(
!pruner.should_prune().unwrap(),
"constant predicate `5 > 1` should simplify to `true` and not prune the file"
);

let predicate = logical2physical(&col("a").gt(lit(100i32)), &schema);

let mut pruner =
FilePruner::try_new(predicate, &schema, &schema, &file, Count::new())
.unwrap();

assert!(
pruner.should_prune().unwrap(),
"`a > 100` with max(a) = 5 should prune via column statistics"
);
}
}
Loading