Skip to content

Commit bebaf2e

Browse files
committed
perf: prune ListingTable files before grouping
1 parent 5f21582 commit bebaf2e

18 files changed

Lines changed: 795 additions & 136 deletions

‎Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎datafusion/catalog-listing/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ datafusion-physical-expr = { workspace = true }
4242
datafusion-physical-expr-adapter = { workspace = true }
4343
datafusion-physical-expr-common = { workspace = true }
4444
datafusion-physical-plan = { workspace = true }
45+
datafusion-pruning = { workspace = true }
4546
futures = { workspace = true }
4647
itertools = { workspace = true }
4748
log = { workspace = true }

‎datafusion/catalog-listing/src/table.rs‎

Lines changed: 221 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ use crate::{ListingOptions, ListingTableConfig};
2323
use arrow::datatypes::{Field, Metadata, Schema, SchemaBuilder, SchemaRef};
2424
use async_trait::async_trait;
2525
use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider};
26+
use datafusion_common::pruning::PrunableStatistics;
2627
use datafusion_common::stats::{Precision, is_known_empty};
2728
use datafusion_common::{
2829
Constraints, DFSchema, SchemaExt, Statistics, internal_datafusion_err, plan_err,
@@ -54,6 +55,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
5455
use datafusion_physical_plan::ExecutionPlan;
5556
use datafusion_physical_plan::empty::EmptyExec;
5657
use datafusion_physical_plan::limit::GlobalLimitExec;
58+
use datafusion_pruning::PruningPredicateBuilder;
5759
use futures::future::BoxFuture;
5860
use futures::{Stream, StreamExt, TryStreamExt, future, stream};
5961
use object_store::ObjectStore;
@@ -669,7 +671,12 @@ impl ListingTable {
669671
statistics,
670672
grouped_by_partition: partitioned_by_file_group,
671673
} = self
672-
.list_files_for_scan(state, &partition_filters, statistic_file_limit)
674+
.list_files_for_scan_inner(
675+
state,
676+
&partition_filters,
677+
&filters,
678+
statistic_file_limit,
679+
)
673680
.await?;
674681

675682
// if no files need to be read, return an `EmptyExec`
@@ -921,15 +928,28 @@ impl ListingTable {
921928
filters: &'a [Expr],
922929
limit: Option<usize>,
923930
) -> datafusion_common::Result<ListFilesResult> {
931+
self.list_files_for_scan_inner(ctx, filters, &[], limit)
932+
.await
933+
}
934+
935+
async fn list_files_for_scan_inner<'a>(
936+
&'a self,
937+
ctx: &'a dyn Session,
938+
partition_filters: &'a [Expr],
939+
data_filters: &'a [Expr],
940+
limit: Option<usize>,
941+
) -> datafusion_common::Result<ListFilesResult> {
942+
// Declared layouts must retain their partition identities, including empty groups.
924943
if let Some(output_partitioning) = self.options.output_partitioning.as_ref() {
925944
self.list_files_for_declared_output_partitioning(
926945
ctx,
927946
output_partitioning,
928-
filters,
947+
partition_filters,
929948
)
930949
.await
931950
} else {
932-
self.list_files_for_regular_scan(ctx, filters, limit).await
951+
self.list_files_for_regular_scan(ctx, partition_filters, data_filters, limit)
952+
.await
933953
}
934954
}
935955

@@ -985,10 +1005,47 @@ impl ListingTable {
9851005
get_files_with_limit(files, file_limit, ctx.config().collect_statistics()).await
9861006
}
9871007

1008+
/// Remove only files whose collected statistics prove the filter cannot match.
1009+
/// This performs no I/O and leaves the normal row filters in place.
1010+
fn prune_files_by_statistics(
1011+
&self,
1012+
ctx: &dyn Session,
1013+
filters: &[Expr],
1014+
files: FileGroup,
1015+
) -> FileGroup {
1016+
if !ctx.config().collect_statistics() || self.expr_adapter_factory.is_some() {
1017+
return files;
1018+
}
1019+
// Custom adapters can change values relative to file statistics, and volatile
1020+
// expressions must not be evaluated while planning.
1021+
let Some(filter) = filters
1022+
.iter()
1023+
.filter(|f| !f.is_volatile())
1024+
.cloned()
1025+
.reduce(Expr::and)
1026+
else {
1027+
return files;
1028+
};
1029+
let Ok(schema) = DFSchema::try_from(self.table_schema.as_ref().clone()) else {
1030+
return files;
1031+
};
1032+
let Ok(predicate) = ctx.create_physical_expr(filter, &schema) else {
1033+
return files;
1034+
};
1035+
let Some(predicate) = PruningPredicateBuilder::new()
1036+
.with_file_schema(Arc::clone(&self.table_schema))
1037+
.build(predicate)
1038+
else {
1039+
return files;
1040+
};
1041+
prune_collected_files(&predicate, &self.file_schema, files)
1042+
}
1043+
9881044
async fn list_files_for_regular_scan<'a>(
9891045
&'a self,
9901046
ctx: &'a dyn Session,
9911047
filters: &'a [Expr],
1048+
data_filters: &'a [Expr],
9921049
limit: Option<usize>,
9931050
) -> datafusion_common::Result<ListFilesResult> {
9941051
let file_group_count = ctx.config().target_partitions();
@@ -1010,6 +1067,7 @@ impl ListingTable {
10101067
let (file_group, inexact_stats) = self
10111068
.collect_files_for_scan(ctx, &store, filters, limit)
10121069
.await?;
1070+
let file_group = self.prune_files_by_statistics(ctx, data_filters, file_group);
10131071

10141072
// Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N
10151073
//
@@ -1259,6 +1317,34 @@ async fn get_files_with_limit(
12591317
Ok((file_group, inexact_stats))
12601318
}
12611319

1320+
/// Keep a file unless the pruning predicate proves it cannot contain a match.
1321+
fn prune_collected_files(
1322+
predicate: &datafusion_pruning::PruningPredicate,
1323+
file_schema: &SchemaRef,
1324+
files: FileGroup,
1325+
) -> FileGroup {
1326+
FileGroup::new(
1327+
files
1328+
.into_inner()
1329+
.into_iter()
1330+
.filter(|file| {
1331+
let Some(statistics) = &file.statistics else {
1332+
return true;
1333+
};
1334+
// One unknown file must not hide the exact bounds of another file.
1335+
let statistics = PrunableStatistics::new(
1336+
vec![Arc::clone(statistics)],
1337+
Arc::clone(file_schema),
1338+
);
1339+
predicate
1340+
.prune(&statistics)
1341+
.map(|keep| keep != [false])
1342+
.unwrap_or(true)
1343+
})
1344+
.collect(),
1345+
)
1346+
}
1347+
12621348
#[cfg(test)]
12631349
mod tests {
12641350
use super::*;
@@ -1420,4 +1506,136 @@ mod tests {
14201506
let result = derive_common_ordering_from_files(&file_groups);
14211507
assert_eq!(result, Some(ordering));
14221508
}
1509+
fn prune_test_files(expr: Expr, statistics: Vec<Option<Statistics>>) -> Vec<String> {
1510+
use arrow::datatypes::DataType;
1511+
let schema = Arc::new(Schema::new(vec![Field::new("d", DataType::Int64, true)]));
1512+
let physical = datafusion_physical_expr::create_physical_expr(
1513+
&expr,
1514+
&DFSchema::try_from(schema.as_ref().clone()).unwrap(),
1515+
&ExecutionProps::new(),
1516+
&PhysicalPlanningContext::default(),
1517+
)
1518+
.unwrap();
1519+
let predicate = PruningPredicateBuilder::new()
1520+
.with_file_schema(Arc::clone(&schema))
1521+
.try_build(physical)
1522+
.unwrap();
1523+
let files = FileGroup::new(
1524+
statistics
1525+
.into_iter()
1526+
.enumerate()
1527+
.map(|(i, stats)| {
1528+
let mut file = PartitionedFile::new(format!("f{i}"), 10);
1529+
file.statistics = stats.map(Arc::new);
1530+
file
1531+
})
1532+
.collect(),
1533+
);
1534+
prune_collected_files(&predicate, &schema, files)
1535+
.iter()
1536+
.map(|f| f.object_meta.location.to_string())
1537+
.collect()
1538+
}
1539+
1540+
fn bounds(
1541+
min: Precision<datafusion_common::ScalarValue>,
1542+
max: Precision<datafusion_common::ScalarValue>,
1543+
) -> Statistics {
1544+
Statistics {
1545+
num_rows: Precision::Exact(10),
1546+
column_statistics: vec![datafusion_common::ColumnStatistics {
1547+
min_value: min,
1548+
max_value: max,
1549+
..Default::default()
1550+
}],
1551+
..Default::default()
1552+
}
1553+
}
1554+
1555+
#[test]
1556+
fn statistics_pruning_keeps_missing_and_inexact_bounds() {
1557+
use datafusion_expr::{col, lit};
1558+
let exact = |v: i64| Precision::Exact(v.into());
1559+
let inexact = |v: i64| Precision::Inexact(v.into());
1560+
assert_eq!(
1561+
prune_test_files(
1562+
col("d").gt_eq(lit(20_i64)),
1563+
vec![
1564+
None,
1565+
Some(bounds(Precision::Absent, Precision::Absent)),
1566+
Some(bounds(inexact(0), inexact(9))),
1567+
Some(bounds(exact(0), exact(9))),
1568+
Some(bounds(exact(20), exact(29)))
1569+
]
1570+
),
1571+
["f0", "f1", "f2", "f4"]
1572+
);
1573+
}
1574+
1575+
#[test]
1576+
fn statistics_pruning_uses_independent_bounds() {
1577+
use datafusion_expr::{col, lit};
1578+
assert_eq!(
1579+
prune_test_files(
1580+
col("d").lt(lit(20_i64)),
1581+
vec![
1582+
Some(bounds(Precision::Exact(30_i64.into()), Precision::Absent)),
1583+
Some(bounds(Precision::Absent, Precision::Exact(10_i64.into())))
1584+
]
1585+
),
1586+
["f1"]
1587+
);
1588+
}
1589+
1590+
#[test]
1591+
fn statistics_pruning_preserves_null_semantics() {
1592+
use datafusion_expr::{col, lit};
1593+
let mut all_null = bounds(Precision::Absent, Precision::Absent);
1594+
all_null.column_statistics[0].null_count = Precision::Exact(10);
1595+
let mut no_null = bounds(
1596+
Precision::Exact(0_i64.into()),
1597+
Precision::Exact(9_i64.into()),
1598+
);
1599+
no_null.column_statistics[0].null_count = Precision::Exact(0);
1600+
assert_eq!(
1601+
prune_test_files(
1602+
col("d").is_null(),
1603+
vec![Some(all_null.clone()), Some(no_null.clone())]
1604+
),
1605+
["f0"]
1606+
);
1607+
assert_eq!(
1608+
prune_test_files(
1609+
col("d").is_not_null(),
1610+
vec![Some(all_null.clone()), Some(no_null.clone())]
1611+
),
1612+
["f1"]
1613+
);
1614+
assert!(
1615+
prune_test_files(col("d").eq(lit(5_i64)), vec![Some(all_null)]).is_empty()
1616+
);
1617+
assert_eq!(
1618+
prune_test_files(
1619+
col("d").eq(lit(5_i64)).or(col("d").is_null()),
1620+
vec![Some(no_null)]
1621+
),
1622+
["f0"]
1623+
);
1624+
}
1625+
1626+
#[test]
1627+
fn statistics_pruning_keeps_file_on_evaluation_error() {
1628+
use datafusion_expr::{col, lit};
1629+
// A malformed statistics type cannot establish that the file is empty.
1630+
assert_eq!(
1631+
prune_test_files(
1632+
col("d").gt(lit(10_i64)),
1633+
vec![Some(bounds(
1634+
Precision::Exact("invalid".into()),
1635+
Precision::Exact("invalid".into())
1636+
))]
1637+
),
1638+
["f0"]
1639+
);
1640+
}
14231641
}

0 commit comments

Comments
 (0)