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
21 changes: 12 additions & 9 deletions datafusion/pruning/benches/string_in_list_pruning.rs
Original file line number Diff line number Diff line change
Expand Up @@ -47,12 +47,10 @@ use datafusion_common::{Column, ScalarValue};
use datafusion_expr_common::operator::Operator;
use datafusion_physical_expr::PhysicalExprRef;
use datafusion_physical_expr::expressions::{BinaryExpr, col, in_list, lit};
use datafusion_pruning::{
MAX_IN_LIST_SIZE, PruningPredicate, PruningPredicateBuilder, PruningStatistics,
};
use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder, PruningStatistics};

const DOMAIN_SIZES: [usize; 9] = [1, 2, 4, 8, 16, 20, 21, 256, 1024];
const CONTAINER_COUNTS: [usize; 3] = [16, 256, 4096];
const CONTAINER_COUNTS: [usize; 4] = [1, 16, 256, 4096];
const BASELINE_CONTAINER_COUNT: usize = 4096;

fn value(index: usize) -> String {
Expand Down Expand Up @@ -332,7 +330,12 @@ fn assert_equivalent_results(case: &BenchmarkCase, statistics: &IntervalStatisti
.not_in_list_with_null_predicate
.prune(statistics)
.unwrap();
if case.size + 1 > MAX_IN_LIST_SIZE {
if case
.in_list_with_null_predicate
.predicate_expr()
.to_string()
.contains("IN_SET_INTERSECTS")
{
// Compact pruning applies filter semantics: NULL does not change which
// rows can satisfy IN, while NOT IN (..., NULL) can never be true.
assert_eq!(in_with_null, expected);
Expand Down Expand Up @@ -435,11 +438,11 @@ fn criterion_benchmark(criterion: &mut Criterion) {
criterion.benchmark_group("string_in_list_pruning/not_in_distributions");
distributions.throughput(Throughput::Elements(BASELINE_CONTAINER_COUNT as u64));
let singleton = IntervalStatistics::uniform_singleton();
for case in cases.iter().filter(|case| case.size > 20) {
let compact = case.not_in_list_predicate.prune(&singleton).unwrap();
assert!(compact.iter().all(|keep| !keep));
for case in &cases {
let actual = case.not_in_list_predicate.prune(&singleton).unwrap();
assert!(actual.iter().all(|keep| !keep));
assert_eq!(
compact,
actual,
case.expanded_and_predicate.prune(&singleton).unwrap()
);
for (name, predicate) in [
Expand Down
52 changes: 44 additions & 8 deletions datafusion/pruning/src/pruning_predicate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1636,14 +1636,18 @@ fn build_compact_in_list_expr(
}

/// Default maximum number of entries in an `IN (...)` list eligible for
/// statistics pruning. Eligible literal lists of supported ordered types above
/// this threshold use a compact sorted domain instead of per-value min/max
/// checks, for both `IN` and `NOT IN`.
/// statistics pruning. See [`PruningPredicateBuilder::with_max_in_list_size`]
/// for the representation used within this cap.
/// Callers can raise the cap via [`PredicateRewriter::with_max_in_list_size`], and
/// query engines can wire it from the
/// `datafusion.execution.parquet.max_in_list_size` config option.
pub const MAX_IN_LIST_SIZE: usize = 20;

// Keep the representation threshold independent of the configurable pruning cap.
// Small lists can be faster with vectorized per-value comparisons than with
// compact per-container searches, particularly for large statistics batches.
const MIN_COMPACT_IN_LIST_SIZE: usize = 21;

/// Rewrite a predicate expression in terms of statistics (min/max/null_counts)
/// for use as a [`PruningPredicate`].
pub struct PredicateRewriter {
Expand Down Expand Up @@ -1770,11 +1774,8 @@ fn build_predicate_expression(
}
}
if let Some(in_list) = expr.downcast_ref::<phys_expr::InListExpr>() {
// Keep the existing expression shape for lists of at most 20 values.
// This lower bound is a scope/compatibility choice, not a measured
// performance threshold; compact pruning is opt-in via a raised cap.
// The compact form covers both `IN` and `NOT IN`.
if in_list.list().len() > MAX_IN_LIST_SIZE
if in_list.list().len() >= MIN_COMPACT_IN_LIST_SIZE
&& in_list.list().len() <= max_in_list_size
&& let Some(pruning_expr) =
build_compact_in_list_expr(in_list, schema, required_columns, properties)
Expand Down Expand Up @@ -4247,7 +4248,7 @@ mod tests {
predicate.required_columns().single_column().unwrap().name(),
"c1"
);
if count > MAX_IN_LIST_SIZE {
if count >= MIN_COMPACT_IN_LIST_SIZE {
// The expression and statistics schema do not grow with the domain.
assert_eq!(predicate.required_columns.columns.len(), 4);
let mut nodes = 0;
Expand Down Expand Up @@ -4385,6 +4386,41 @@ mod tests {
Ok(())
}

#[test]
fn string_in_list_representation_respects_cap_and_threshold() -> Result<()> {
let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)]));
for count in [1, 2, 4, 8, 20, 21] {
let values = (0..count)
.map(|i| lit(format!("a{i:03}")))
.collect::<Vec<_>>();
for (negated, marker) in
[(false, "IN_SET_INTERSECTS"), (true, "NOT_IN_SET_MAY_MATCH")]
{
let expr = logical2physical(
&col("c1").in_list(values.clone(), negated),
&schema,
);
for limit in [0, count - 1, count, 32] {
let predicate = PruningPredicateBuilder::new()
.with_file_schema(Arc::clone(&schema))
.with_max_in_list_size(limit)
.try_build(Arc::clone(&expr))?;
assert_eq!(
predicate.predicate_expr().to_string().contains(marker),
count >= 21 && count <= limit,
"count={count}, limit={limit}, negated={negated}"
);
assert_eq!(
is_always_true(predicate.predicate_expr()),
count > limit,
"count={count}, limit={limit}, negated={negated}"
);
}
}
}
Ok(())
}

#[test]
fn large_string_in_list_respects_configured_limit() -> Result<()> {
let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Utf8, true)]));
Expand Down