jayzhan211 commented on code in PR #26441:
URL: https://github.com/apache/datafusion/pull/26441#discussion_r4504870290
##########
datafusion/catalog-listing/src/table.rs:
##########
@@ -1431,4 +1534,137 @@ mod tests {
let result = derive_common_ordering_from_files(&file_groups);
assert_eq!(result, Some(ordering));
}
+
+ fn prune_test_files(expr: Expr, statistics: Vec<Option<Statistics>>) ->
Vec<String> {
+ use arrow::datatypes::DataType;
+ let schema = Arc::new(Schema::new(vec![Field::new("d",
DataType::Int64, true)]));
+ let physical = datafusion_physical_expr::create_physical_expr(
+ &expr,
+ &DFSchema::try_from(schema.as_ref().clone()).unwrap(),
+ &ExecutionProps::new(),
+ &PhysicalPlanningContext::default(),
+ )
+ .unwrap();
+ let predicate = PruningPredicateBuilder::new()
+ .with_file_schema(Arc::clone(&schema))
+ .try_build(physical)
+ .unwrap();
+ let files = FileGroup::new(
+ statistics
+ .into_iter()
+ .enumerate()
+ .map(|(i, stats)| {
+ let mut file = PartitionedFile::new(format!("f{i}"), 10);
+ file.statistics = stats.map(Arc::new);
+ file
+ })
+ .collect(),
+ );
+ prune_collected_files(&predicate, &schema, files)
+ .iter()
+ .map(|f| f.object_meta.location.to_string())
+ .collect()
+ }
+
+ fn bounds(
+ min: Precision<datafusion_common::ScalarValue>,
+ max: Precision<datafusion_common::ScalarValue>,
+ ) -> Statistics {
+ Statistics {
+ num_rows: Precision::Exact(10),
+ column_statistics: vec![datafusion_common::ColumnStatistics {
+ min_value: min,
+ max_value: max,
+ ..Default::default()
+ }],
+ ..Default::default()
+ }
+ }
+
+ #[test]
+ fn statistics_pruning_keeps_missing_and_inexact_bounds() {
+ use datafusion_expr::{col, lit};
+ let exact = |v: i64| Precision::Exact(v.into());
+ let inexact = |v: i64| Precision::Inexact(v.into());
+ assert_eq!(
+ prune_test_files(
+ col("d").gt_eq(lit(20_i64)),
+ vec![
+ None,
+ Some(bounds(Precision::Absent, Precision::Absent)),
+ Some(bounds(inexact(0), inexact(9))),
+ Some(bounds(exact(0), exact(9))),
+ Some(bounds(exact(20), exact(29)))
+ ]
+ ),
+ ["f0", "f1", "f2", "f4"]
+ );
+ }
+
+ #[test]
+ fn statistics_pruning_uses_independent_bounds() {
+ use datafusion_expr::{col, lit};
+ assert_eq!(
+ prune_test_files(
+ col("d").lt(lit(20_i64)),
+ vec![
+ Some(bounds(Precision::Exact(30_i64.into()),
Precision::Absent)),
+ Some(bounds(Precision::Absent,
Precision::Exact(10_i64.into())))
+ ]
+ ),
+ ["f1"]
+ );
+ }
+
+ #[test]
+ fn statistics_pruning_preserves_null_semantics() {
+ use datafusion_expr::{col, lit};
+ let mut all_null = bounds(Precision::Absent, Precision::Absent);
+ all_null.column_statistics[0].null_count = Precision::Exact(10);
+ let mut no_null = bounds(
+ Precision::Exact(0_i64.into()),
+ Precision::Exact(9_i64.into()),
+ );
+ no_null.column_statistics[0].null_count = Precision::Exact(0);
+ assert_eq!(
+ prune_test_files(
+ col("d").is_null(),
+ vec![Some(all_null.clone()), Some(no_null.clone())]
+ ),
+ ["f0"]
+ );
+ assert_eq!(
+ prune_test_files(
+ col("d").is_not_null(),
+ vec![Some(all_null.clone()), Some(no_null.clone())]
+ ),
+ ["f1"]
+ );
+ assert!(
+ prune_test_files(col("d").eq(lit(5_i64)),
vec![Some(all_null)]).is_empty()
+ );
+ assert_eq!(
+ prune_test_files(
+ col("d").eq(lit(5_i64)).or(col("d").is_null()),
+ vec![Some(no_null)]
+ ),
+ ["f0"]
+ );
+ }
+
+ #[test]
+ fn statistics_pruning_keeps_file_on_evaluation_error() {
+ use datafusion_expr::{col, lit};
+ // A malformed statistics type cannot establish that the file is empty.
+ assert_eq!(
+ prune_test_files(
+ col("d").gt(lit(10_i64)),
+ vec![Some(bounds(
+ Precision::Exact("invalid".into()),
+ Precision::Exact("invalid".into())
+ ))]
+ ),
Review Comment:
Optional cleanup: this test doesn't reach the error branch yet. A Utf8 bound
doesn't make `prune` fail, so the test still passes if `.unwrap_or(true)`
becomes `.unwrap_or(false)`. Binary bounds can't be cast to Int64, so they do
error, and the test then fails with `unwrap_or(false)`:
```suggestion
use datafusion_common::ScalarValue;
use datafusion_expr::{col, lit};
// Binary bounds cannot be cast to Int64, so `prune` returns an
error.
// A Utf8 bound does not error, so it would not reach that branch.
let binary = || Precision::Exact(ScalarValue::Binary(Some(vec![1])));
assert_eq!(
prune_test_files(
col("d").gt(lit(10_i64)),
vec![Some(bounds(binary(), binary()))]
),
```
##########
datafusion/catalog-listing/src/table.rs:
##########
@@ -1028,17 +1090,30 @@ impl ListingTable {
let (grouped, split_points) = file_group
.group_by_partition_values_with_split_points(file_group_count);
if grouped.len() >= threshold {
+ // Preserve the groups and split points derived from all
files so
+ // pruning cannot change the declared range partitioning.
+ let grouped = grouped
+ .into_iter()
+ .map(|group| {
+ self.prune_files_by_statistics(
+ pruning_predicate.as_deref(),
+ group,
+ )
+ })
+ .collect();
(grouped, true, split_points)
} else {
let all_files: Vec<_> =
grouped.into_iter().flat_map(|g|
g.into_inner()).collect();
- (
-
FileGroup::new(all_files).split_files(file_group_count),
- false,
- vec![],
- )
+ let file_group = self.prune_files_by_statistics(
+ pruning_predicate.as_deref(),
+ FileGroup::new(all_files),
+ );
Review Comment:
Optional: no test covers this branch, where `preserve_file_partitions` is
above the number of distinct partition values. If the prune here is removed,
every test still passes. With two `p` values and a threshold of 3, the test
below expects `[[f2], [f3]]`. It passes on this PR and fails when the prune is
removed.
<details><summary>Test for <code>file_statistics_pruning.rs</code></summary>
```rs
#[tokio::test]
async fn preserve_partitions_below_threshold_still_prunes() -> Result<()> {
let dir = fixture(
(0..8)
.map(|i| (i * 10..i * 10 + 10).map(Some).collect())
.collect(),
true,
true,
);
let mut config = SessionConfig::new()
.with_target_partitions(4)
.with_collect_statistics(true);
// Two distinct `p` values are below the threshold, so files are split
normally.
config.options_mut().optimizer.preserve_file_partitions = 3;
let ctx = SessionContext::new_with_config(config);
let table = table(&dir, None, true, false);
let plan = table.scan(&ctx.state(), None, &[predicate()], None).await?;
assert_eq!(
names(scan_config(&plan)),
vec![vec!["f2.parquet"], vec!["f3.parquet"]]
);
Ok(())
}
```
</details>
##########
datafusion/catalog-listing/src/table.rs:
##########
@@ -1270,6 +1345,34 @@ async fn get_files_with_limit(
Ok((file_group, inexact_stats))
}
+/// Keep a file unless the pruning predicate proves it cannot contain a match.
+fn prune_collected_files(
+ predicate: &datafusion_pruning::PruningPredicate,
+ file_schema: &SchemaRef,
+ files: FileGroup,
+) -> FileGroup {
+ FileGroup::new(
+ files
+ .into_inner()
+ .into_iter()
+ .filter(|file| {
+ let Some(statistics) = &file.statistics else {
+ return true;
+ };
Review Comment:
Optional: files whose statistics are all `Absent` can't be pruned, but each
one still gets a `PrunableStatistics` and a `prune` call. That covers every
file of a CSV or JSON table, whose `infer_stats` returns unknown statistics.
`FilePruner::try_new` skips these files with the same check:
```suggestion
// Same check as `FilePruner::try_new`: all-`Absent`
statistics
// cannot prune, so skip building `PrunableStatistics` for
them.
if !file.has_statistics() {
return true;
}
let Some(statistics) = &file.statistics else {
return true;
};
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
