Skip to content
Merged
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
10 changes: 6 additions & 4 deletions datafusion/common/src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1373,10 +1373,12 @@ config_namespace! {
/// rewrite; other predicates and Bloom-filter pruning remain available.
///
/// Within the cap, nonempty lists of at most 20 values use the existing
/// per-value rewrite. Larger non-null literal string lists on a string
/// column use a compact sorted domain, for both `IN` and `NOT IN`.
/// Other lists retain the existing per-value rewrite, so raising the cap
/// can make those predicates expensive to build and evaluate.
/// per-value rewrite. Larger literal string lists on a string column use
/// a compact representation, for both `IN` and `NOT IN`, including lists
/// with NULL members. `NOT IN` with NULL and all-NULL `IN` lists cannot
/// match any rows. Other lists retain the existing per-value rewrite, so
/// raising the cap can make those predicates expensive to build and
/// evaluate.
///
/// Defaults to 20.
pub max_in_list_size: usize, default = 20
Expand Down
132 changes: 81 additions & 51 deletions datafusion/core/tests/parquet/string_in_list_pruning.rs
Original file line number Diff line number Diff line change
Expand Up @@ -440,23 +440,24 @@ async fn string_not_in_list_with_truncated_bounds() {
}

#[tokio::test]
async fn string_not_in_list_with_null_does_not_bypass_row_filter() {
async fn string_in_list_with_null_preserves_filter_semantics() {
let mut file = tempfile::Builder::new()
.prefix("string_not_in_list_pruning")
.prefix("string_in_list_null_pruning")
.suffix(".parquet")
.tempfile()
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)]));
// The first row group has a known zero null count, and every value lies
// in a gap in the IN list. Dropping the NULL list member while inverting
// NOT IN would incorrectly prove that this entire row group matches.
// in a gap in the IN list. The matching value in the second row group is
// deliberately not first, so incorrectly bypassing the row filter changes
// the result when the scan has a limit.
let values = vec![
Some("v000001"),
Some("v000001"),
Some("v000001"),
Some("v000001"),
Some("v000000"),
Some("v000001"),
Some("v000000"),
None,
Some("v999999"),
];
Expand All @@ -475,13 +476,11 @@ async fn string_not_in_list_with_null_does_not_bypass_row_filter() {
assert_eq!(writer.close().unwrap().num_row_groups(), 2);

// Build the physical source directly so a logical optimizer cannot fold
// the SQL NOT IN (..., NULL) filter to an empty relation before the scan.
// NOT IN (..., NULL) to an empty relation before the scan.
let mut list = (0..21)
.map(|index| lit(format!("v{:06}", index * 10)))
.collect::<Vec<_>>();
list.push(lit(ScalarValue::Utf8(None)));
let predicate =
in_list(col("value", &schema).unwrap(), list, &true, &schema).unwrap();
let location = Path::from_filesystem_path(file.path()).unwrap();
let partitioned_file = PartitionedFile::new(
location.to_string(),
Expand All @@ -490,48 +489,79 @@ async fn string_not_in_list_with_null_does_not_bypass_row_filter() {
let ctx =
SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1));

for max_in_list_size in [0, 32] {
let mut options = TableParquetOptions::default();
options.global.max_in_list_size = max_in_list_size;
let source = Arc::new(
ParquetSource::new(Arc::clone(&schema))
.with_table_parquet_options(options)
.with_predicate(Arc::clone(&predicate))
.with_pushdown_filters(true)
.with_enable_page_index(false)
.with_bloom_filter_on_read(false),
);
let config =
FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
.with_file(partitioned_file.clone())
.with_limit(Some(1))
.build();
let plan: Arc<dyn ExecutionPlan> =
Arc::new(DataSourceExec::new(Arc::new(config)));
let plan_text = displayable(plan.as_ref()).indent(true).to_string();
assert!(plan_text.contains("NOT IN"), "{plan_text}");
let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
let output = ScanOutput {
batches,
plan: plan_text,
metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(),
};

assert_eq!(
output
.batches
.iter()
.map(RecordBatch::num_rows)
.sum::<usize>(),
0,
"cap={max_in_list_size}, plan={}, metrics={}",
output.plan,
output.metrics
);
assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0);
assert_eq!(output.pruned("row_groups_pruned_statistics"), 0);
assert_eq!(output.pruned("limit_pruned_row_groups"), 0);
assert_eq!(output.counter("pushdown_rows_pruned"), 8);
assert_eq!(output.counter("predicate_evaluation_errors"), 0);
for negated in [false, true] {
let predicate = in_list(
col("value", &schema).unwrap(),
list.clone(),
&negated,
&schema,
)
.unwrap();
for max_in_list_size in [0, 32] {
let mut options = TableParquetOptions::default();
options.global.max_in_list_size = max_in_list_size;
let source = Arc::new(
ParquetSource::new(Arc::clone(&schema))
.with_table_parquet_options(options)
.with_predicate(Arc::clone(&predicate))
.with_pushdown_filters(true)
.with_enable_page_index(false)
.with_bloom_filter_on_read(false),
);
let config =
FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
.with_file(partitioned_file.clone())
.with_limit(Some(1))
.build();
let plan: Arc<dyn ExecutionPlan> =
Arc::new(DataSourceExec::new(Arc::new(config)));
let plan_text = displayable(plan.as_ref()).indent(true).to_string();
assert!(plan_text.contains("IN"), "{plan_text}");
let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await.unwrap();
let output = ScanOutput {
batches,
plan: plan_text,
metrics: MetricsFinder::find_metrics(plan.as_ref()).unwrap(),
};

if negated {
assert!(output.batches.iter().all(|batch| batch.num_rows() == 0));
} else {
assert_batches_eq!(
[
"+---------+",
"| value |",
"+---------+",
"| v000000 |",
"+---------+",
],
&output.batches
);
}
assert_eq!(
output.counter("pushdown_rows_pruned"),
match (negated, max_in_list_size) {
(false, 0) => 7,
(false, _) => 3,
(true, 0) => 8,
(true, _) => 0,
}
);
assert_eq!(output.fully_matched("row_groups_pruned_statistics"), 0);
assert_eq!(
output.pruned("row_groups_pruned_statistics"),
if max_in_list_size == 0 {
0
} else if negated {
2
} else {
1
},
"negated={negated}, cap={max_in_list_size}, metrics={}",
output.metrics
);
assert_eq!(output.pruned("limit_pruned_row_groups"), 0);
assert_eq!(output.counter("predicate_evaluation_errors"), 0);
}
}
}
6 changes: 6 additions & 0 deletions datafusion/datasource-parquet/src/row_group_filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -395,6 +395,11 @@ impl RowGroupAccessPlanFilter {
if candidate_row_group_indices.is_empty() {
return;
}
// Some pruning rewrites preserve only whether rows can be TRUE, while
// full-match inference must distinguish FALSE from UNKNOWN.
if !predicate.can_be_inverted_for_full_match() {
return;
}
let arrow_schema = pruning_stats.arrow_schema;

let mut inverted_expr: Arc<dyn PhysicalExpr> =
Expand Down Expand Up @@ -432,6 +437,7 @@ impl RowGroupAccessPlanFilter {

let Ok(inverted_predicate) = PruningPredicateBuilder::new()
.with_file_schema(Arc::clone(predicate.schema()))
.with_max_in_list_size(predicate.max_in_list_size())
.try_build(inverted_expr)
else {
return;
Expand Down
85 changes: 68 additions & 17 deletions datafusion/pruning/benches/string_in_list_pruning.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,10 @@
//! without making the baseline depend on a deeply nested expression.
//!
//! The main cases alternate intervals pinned to a domain member with intervals
//! that span a sparse gap. Supplemental `NOT IN` cases measure uniform singleton
//! containers, long bounds, and long literals that share bounds' prefixes. Bloom
//! filters are not involved.
//! that span a sparse gap. NULL-containing variants cover the optimized
//! filter-semantics paths, including the constant-false `NOT IN` rewrite.
//! Supplemental `NOT IN` cases measure uniform singleton containers, long bounds,
//! and long literals that share bounds' prefixes. Bloom filters are not involved.
//!
//! Run with `cargo bench -p datafusion-pruning --bench string_in_list_pruning`.
//! The construction benchmarks reuse their input physical expressions; the
Expand Down Expand Up @@ -188,12 +189,16 @@ struct BenchmarkCase {
size: usize,
schema: SchemaRef,
in_list: PhysicalExprRef,
in_list_with_null: PhysicalExprRef,
expanded_or: PhysicalExprRef,
not_in_list: PhysicalExprRef,
not_in_list_with_null: PhysicalExprRef,
expanded_and: PhysicalExprRef,
in_list_predicate: PruningPredicate,
in_list_with_null_predicate: PruningPredicate,
expanded_or_predicate: PruningPredicate,
not_in_list_predicate: PruningPredicate,
not_in_list_with_null_predicate: PruningPredicate,
expanded_and_predicate: PruningPredicate,
statistics: IntervalStatistics,
}
Expand All @@ -209,15 +214,30 @@ impl BenchmarkCase {
let values = (0..size)
.map(|index| lit(ScalarValue::new_utf8view(value(index * 10))))
.collect::<Vec<_>>();
let mut values_with_null = values.clone();
values_with_null.push(lit(ScalarValue::Utf8View(None)));
let not_in_list =
in_list(Arc::clone(&column), values.clone(), &true, &schema).unwrap();
let not_in_list_with_null = in_list(
Arc::clone(&column),
values_with_null.clone(),
&true,
&schema,
)
.unwrap();
let in_list_with_null =
in_list(Arc::clone(&column), values_with_null, &false, &schema).unwrap();
let in_list =
in_list(Arc::clone(&column), values.clone(), &false, &schema).unwrap();
let expanded_or = expanded(&column, &values, Operator::Eq, Operator::Or);
let expanded_and = expanded(&column, &values, Operator::NotEq, Operator::And);
let in_list_predicate = build_predicate(&in_list, &schema, size);
let in_list_with_null_predicate =
build_predicate(&in_list_with_null, &schema, size + 1);
let expanded_or_predicate = build_predicate(&expanded_or, &schema, size);
let not_in_list_predicate = build_predicate(&not_in_list, &schema, size);
let not_in_list_with_null_predicate =
build_predicate(&not_in_list_with_null, &schema, size + 1);
let expanded_and_predicate = build_predicate(&expanded_and, &schema, size);
eprintln!(
"string_in_list_pruning: {size} values, compact in={}, compact not in={}",
Expand All @@ -239,6 +259,10 @@ impl BenchmarkCase {
.collect::<Vec<_>>();
let negated_kept = kept.iter().map(|keep| !keep).collect::<Vec<_>>();
assert_eq!(in_list_predicate.prune(&statistics).unwrap(), kept);
assert_eq!(
in_list_with_null_predicate.prune(&statistics).unwrap(),
kept
);
assert_eq!(expanded_or_predicate.prune(&statistics).unwrap(), kept);
assert_eq!(
not_in_list_predicate.prune(&statistics).unwrap(),
Expand All @@ -248,17 +272,28 @@ impl BenchmarkCase {
expanded_and_predicate.prune(&statistics).unwrap(),
negated_kept
);
assert!(
not_in_list_with_null_predicate
.prune(&statistics)
.unwrap()
.iter()
.all(|keep| !keep)
);

Self {
size,
schema,
in_list,
in_list_with_null,
expanded_or,
not_in_list,
not_in_list_with_null,
expanded_and,
in_list_predicate,
in_list_with_null_predicate,
expanded_or_predicate,
not_in_list_predicate,
not_in_list_with_null_predicate,
expanded_and_predicate,
statistics,
}
Expand All @@ -269,22 +304,28 @@ fn criterion_benchmark(criterion: &mut Criterion) {
let cases = DOMAIN_SIZES.map(BenchmarkCase::new);
let mut construction = criterion.benchmark_group("string_in_list_pruning/construct");
for case in &cases {
construction.throughput(Throughput::Elements(case.size as u64));
for (name, expression) in [
("in_list", &case.in_list),
("expanded_or", &case.expanded_or),
("not_in_list", &case.not_in_list),
("expanded_and", &case.expanded_and),
for (name, expression, max_in_list_size) in [
("in_list", &case.in_list, case.size),
("in_list_with_null", &case.in_list_with_null, case.size + 1),
("expanded_or", &case.expanded_or, case.size),
("not_in_list", &case.not_in_list, case.size),
(
"not_in_list_with_null_constant_false",
&case.not_in_list_with_null,
case.size + 1,
),
("expanded_and", &case.expanded_and, case.size),
] {
construction.throughput(Throughput::Elements(max_in_list_size as u64));
construction.bench_with_input(
BenchmarkId::new(name, case.size),
BenchmarkId::new(name, max_in_list_size),
expression,
|bencher, expression| {
bencher.iter(|| {
black_box(build_predicate(
black_box(expression),
&case.schema,
case.size,
max_in_list_size,
))
});
},
Expand All @@ -296,14 +337,24 @@ fn criterion_benchmark(criterion: &mut Criterion) {
let mut evaluation = criterion.benchmark_group("string_in_list_pruning/evaluate");
evaluation.throughput(Throughput::Elements(CONTAINERS as u64));
for case in &cases {
for (name, predicate) in [
("in_list", &case.in_list_predicate),
("expanded_or", &case.expanded_or_predicate),
("not_in_list", &case.not_in_list_predicate),
("expanded_and", &case.expanded_and_predicate),
for (name, predicate, list_size) in [
("in_list", &case.in_list_predicate, case.size),
(
"in_list_with_null",
&case.in_list_with_null_predicate,
case.size + 1,
),
("expanded_or", &case.expanded_or_predicate, case.size),
("not_in_list", &case.not_in_list_predicate, case.size),
(
"not_in_list_with_null_constant_false",
&case.not_in_list_with_null_predicate,
case.size + 1,
),
("expanded_and", &case.expanded_and_predicate, case.size),
] {
evaluation.bench_with_input(
BenchmarkId::new(name, case.size),
BenchmarkId::new(name, list_size),
predicate,
|bencher, predicate| {
bencher.iter(|| {
Expand Down
Loading