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
128 changes: 89 additions & 39 deletions datafusion/physical-expr/benches/in_list_strategy.rs
Original file line number Diff line number Diff line change
Expand Up @@ -35,16 +35,16 @@
//! |------|-------|-----------------|-------------------|
//! | Narrow integer cases | UInt8 | small value domain | 4, 16 |
//! | Narrow integer cases | Int16, Float16 | larger value domain | 4, 64, 256 |
//! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 64, 256 |
//! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 32, 128 |
//! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 33, 64, 256, 1024, 10000 |
//! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 17, 32, 128, 1024, 10000 |
//! | 128-bit interval cases | IntervalMonthDayNano | small lists | 4 |
//! | Utf8 short-string cases | Utf8 | 8-byte strings | 4, 64, 256 |
//! | Utf8 long-string cases | Utf8 | 24-byte strings | 4, 64, 256 |
//! | Utf8View short-string cases | Utf8View | 8-byte strings | 4, 16, 64, 256 |
//! | Utf8View length-12 cases | Utf8View | 12-byte strings | 16, 64 |
//! | Utf8View long-string cases | Utf8View | 24-byte strings | 4, 16, 64, 256 |
//! | Shared-prefix string cases | Utf8, Utf8View | same prefix, different suffix | 16, 32, 64 |
//! | Fixed-size binary cases | FixedSizeBinary(16) | fixed-width binary values | 4, 64, 256, 10000 |
//! | Fixed-size binary cases | FixedSizeBinary(1/2/4/8/16) | fixed-width binary values | representative thresholds and large lists |

use arrow::array::types::IntervalMonthDayNano;
use arrow::array::*;
Expand Down Expand Up @@ -166,11 +166,11 @@ fn random_string(rng: &mut StdRng, len: usize) -> String {
fn strings_with_shared_prefix(
rng: &mut StdRng,
count: usize,
prefix_len: usize,
prefix: &str,
discriminator: char,
) -> Vec<String> {
let prefix = random_string(rng, prefix_len);
(0..count)
.map(|_| format!("{}{}", prefix, random_string(rng, 8))) // prefix + random 8-char suffix
.map(|_| format!("{prefix}{}{discriminator}", random_string(rng, 7)))
.collect()
}

Expand Down Expand Up @@ -275,12 +275,14 @@ fn bench_string_shared_prefix<A>(
.wrapping_add(prefix_len as u64 * 0x4444);
let mut rng = StdRng::seed_from_u64(seed);

// Generate IN list with a shared prefix.
let haystack = strings_with_shared_prefix(&mut rng, list_size, prefix_len);
// Use the same prefix and equal-length, disjoint suffixes for both pools so
// misses exercise length/prefix collisions rather than immediate rejection.
let prefix = random_string(&mut rng, prefix_len);
let haystack = strings_with_shared_prefix(&mut rng, list_size, &prefix, 'h');

// Generate non-matching strings with the same prefix to keep misses close
// to the matching set.
let non_match_pool = strings_with_shared_prefix(&mut rng, 100, prefix_len);
let non_match_pool = strings_with_shared_prefix(&mut rng, 100, &prefix, 'm');

// Generate array with controlled match rate
let values: A = (0..ARRAY_SIZE)
Expand Down Expand Up @@ -313,19 +315,29 @@ fn bench_string_mixed_lengths<A>(
name: &str,
list_size: usize,
match_rate: f64,
inline_rate: f64,
to_scalar: fn(String) -> ScalarValue,
) where
A: Array + FromIterator<Option<String>> + 'static,
{
let seed = 0xABCD_EF01_u64.wrapping_add(list_size as u64 * 0x5555);
let mut rng = StdRng::seed_from_u64(seed);

// Mixed lengths: some short (<= 12), some long (> 12)
let lengths = [4, 8, 12, 16, 20, 24];
let inline_lengths = [4, 8, 12];
let long_lengths = [16, 20, 24];
let inline_count = ((list_size as f64 * inline_rate).round() as usize)
.max(1)
.min(list_size - 1);

// Generate IN list with mixed lengths
let haystack: Vec<String> = (0..list_size)
.map(|_| {
.map(|idx| {
let inline = idx < inline_count;
let lengths = if inline {
&inline_lengths
} else {
&long_lengths
};
let len = *lengths.choose(&mut rng).unwrap();
random_string(&mut rng, len)
})
Expand All @@ -337,6 +349,11 @@ fn bench_string_mixed_lengths<A>(
Some(if !haystack.is_empty() && rng.random_bool(match_rate) {
haystack.choose(&mut rng).unwrap().clone()
} else {
let lengths = if rng.random_bool(inline_rate) {
&inline_lengths
} else {
&long_lengths
};
let len = *lengths.choose(&mut rng).unwrap();
random_string(&mut rng, len)
})
Expand Down Expand Up @@ -420,7 +437,7 @@ fn bench_narrow_integer(c: &mut Criterion) {

fn bench_primitive(c: &mut Criterion) {
// Int32: small and larger list sizes
for list_size in [4, 32, 64, 256] {
for list_size in [4, 32, 33, 64, 256, 1024, 10_000] {
let list_case = if list_size <= 32 {
"small_list"
} else {
Expand All @@ -442,7 +459,7 @@ fn bench_primitive(c: &mut Criterion) {
}

// Int64: small and larger list sizes
for list_size in [4, 16, 32, 128] {
for list_size in [4, 16, 17, 32, 128, 1024, 10_000] {
let list_case = if list_size <= 16 {
"small_list"
} else {
Expand Down Expand Up @@ -602,6 +619,7 @@ fn bench_utf8(c: &mut Criterion) {
&format!("mixed_len/list={list_size}/match={match_pct}%"),
list_size,
match_pct as f64 / 100.0,
0.5,
to_scalar,
);
}
Expand Down Expand Up @@ -694,6 +712,23 @@ fn bench_utf8view(c: &mut Criterion) {
&format!("mixed_len/list={list_size}/match={match_pct}%"),
list_size,
match_pct as f64 / 100.0,
0.5,
to_scalar,
);
}
}

// Strongly skewed mixed lists exercise routing near the all-inline and
// all-long boundaries while retaining both representations.
for inline_pct in [2, 98] {
for match_pct in MATCH_RATES {
bench_string_mixed_lengths::<StringViewArray>(
c,
"utf8view",
&format!("mixed_len/inline={inline_pct}%/list=64/match={match_pct}%"),
64,
match_pct as f64 / 100.0,
inline_pct as f64 / 100.0,
to_scalar,
);
}
Expand Down Expand Up @@ -996,40 +1031,43 @@ fn bench_nulls(c: &mut Criterion) {
}

// =============================================================================
// FIXED SIZE BINARY BENCHMARKS (FixedSizeBinary<16>, e.g. UUIDs)
// FIXED SIZE BINARY BENCHMARKS
// =============================================================================

/// Generates a random 16-byte value (UUID-sized).
fn random_fixed_binary_16(rng: &mut StdRng) -> Vec<u8> {
let mut buf = vec![0u8; 16];
rng.fill(&mut buf[..]);
buf
/// Generates deterministic, disjoint hit and miss values. The high bit of the
/// final byte distinguishes misses, including for the one-byte domain.
fn fixed_binary_value(width: i32, index: usize, miss: bool) -> Vec<u8> {
let mut value = (index as u128).to_le_bytes()[..width as usize].to_vec();
let last = value.last_mut().unwrap();
if miss {
*last |= 0x80;
} else {
*last &= 0x7f;
}
value
}

/// Benchmarks FixedSizeBinary(16) IN list evaluation.
/// FixedSizeBinary doesn't use the generic numeric helpers since its array
/// construction differs from primitive types.
fn bench_fixed_size_binary_inner(
c: &mut Criterion,
name: &str,
width: i32,
list_size: usize,
match_rate: f64,
) {
let seed = 0xF1ED_B1A7_u64.wrapping_add(list_size as u64 * 0x6666);
let seed = 0xF1ED_B1A7_u64
.wrapping_add(list_size as u64 * 0x6666)
.wrapping_add(width as u64 * 0x7777);
let mut rng = StdRng::seed_from_u64(seed);

// Generate IN list values (16-byte each)
let haystack: Vec<Vec<u8>> = (0..list_size)
.map(|_| random_fixed_binary_16(&mut rng))
.map(|index| fixed_binary_value(width, index, false))
.collect();

// Generate array with controlled match rate
let values: Vec<Vec<u8>> = (0..ARRAY_SIZE)
.map(|_| {
.map(|index| {
if !haystack.is_empty() && rng.random_bool(match_rate) {
haystack.choose(&mut rng).unwrap().clone()
} else {
random_fixed_binary_16(&mut rng)
fixed_binary_value(width, index, true)
}
})
.collect();
Expand All @@ -1040,28 +1078,40 @@ fn bench_fixed_size_binary_inner(
let schema = Schema::new(vec![Field::new("a", array.data_type().clone(), true)]);
let exprs: Vec<_> = haystack
.iter()
.map(|v| lit(ScalarValue::FixedSizeBinary(16, Some(v.clone()))))
.map(|v| lit(ScalarValue::FixedSizeBinary(width, Some(v.clone()))))
.collect();
let expr = in_list(col("a", &schema).unwrap(), exprs, &false, &schema).unwrap();
let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array) as ArrayRef])
.unwrap();

c.bench_with_input(
BenchmarkId::new("fixed_size_binary", name),
BenchmarkId::new(
"fixed_size_binary",
format!(
"fsb{width}/list={list_size}/match={}%",
(match_rate * 100.0) as u32
),
),
&batch,
|b, batch| b.iter(|| expr.evaluate(batch).unwrap()),
);
}

fn bench_fixed_size_binary(c: &mut Criterion) {
for list_size in [4, 64, 256, 10000] {
for (width, list_size) in [
(1, 16),
(2, 64),
(4, 4),
(4, 64),
(8, 4),
(8, 64),
(16, 4),
(16, 64),
(16, 256),
(16, 10000),
] {
for match_pct in MATCH_RATES {
bench_fixed_size_binary_inner(
c,
&format!("fsb16/list={list_size}/match={match_pct}%"),
list_size,
match_pct as f64 / 100.0,
);
bench_fixed_size_binary_inner(c, width, list_size, match_pct as f64 / 100.0);
}
}
}
Expand Down
56 changes: 51 additions & 5 deletions datafusion/physical-expr/src/expressions/in_list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt};

mod array_static_filter;
mod branchless_filter;
mod byte_view_filter;
mod fixed_size_binary_filter;
mod frozen_set;
mod primitive_filter;
mod result;
mod static_filter;
Expand Down Expand Up @@ -215,15 +218,15 @@ impl InListExpr {
expr,
list,
negated,
Some(instantiate_static_filter(array)?),
Some(instantiate_static_filter(array, &expr_data_type)?),
))
}

/// Create a new InList expression, using a static filter when possible.
///
/// This validates data types and attempts to create a static filter for constant
/// list expressions. Uses specialized StaticFilter implementations for better
/// performance (e.g., Int32StaticFilter for Int32).
/// list expressions. Uses specialized `StaticFilter` implementations for better
/// performance (for example, primitive filters for fixed-width values).
///
/// Returns an error if data types don't match. If the list contains non-constant
/// expressions, falls back to dynamic evaluation at runtime.
Expand All @@ -242,7 +245,7 @@ impl InListExpr {

// Try to create a static filter if all list expressions are constants
let static_filter = match try_evaluate_constant_list(&list, schema)? {
Some(in_array) => Some(instantiate_static_filter(in_array)?),
Some(in_array) => Some(instantiate_static_filter(in_array, &expr_data_type)?),
None => None, // Non-constant expressions, fall back to dynamic evaluation
};

Expand Down Expand Up @@ -2592,7 +2595,7 @@ mod tests {
// Create IN list with Int32 literals: (100, 200, 300)
let list = vec![lit(100i32), lit(200i32), lit(300i32)];

// Create InListExpr via in_list() - this uses Int32StaticFilter for Int32 lists
// Create InListExpr via in_list() - this uses a specialized primitive filter
let expr = in_list(col_a, list, &false, &schema)?;

// Create dictionary-encoded batch with values [100, 200, 500]
Expand Down Expand Up @@ -3576,6 +3579,49 @@ mod tests {
)?
);

// Utf8View in_array, Utf8View and Dict(Utf8View) needles
let utf8view_in =
Arc::new(StringViewArray::from(vec!["a", "b", "c"])) as ArrayRef;
let utf8view_needle =
Arc::new(StringViewArray::from(vec!["a", "d", "b"])) as ArrayRef;
assert_eq!(
expected,
eval_in_list_from_array(
Arc::clone(&utf8view_needle),
Arc::clone(&utf8view_in),
)?
);
assert_eq!(
expected,
eval_in_list_from_array(wrap_in_dict(utf8view_needle), utf8view_in)?
);

// FixedSizeBinary in_array, FixedSizeBinary and Dictionary needles
let fsb_in = Arc::new(FixedSizeBinaryArray::try_from_iter(
[
[1, 2, 3, 4].as_slice(),
[5, 6, 7, 8].as_slice(),
[9, 10, 11, 12].as_slice(),
]
.into_iter(),
)?) as ArrayRef;
let fsb_needle = Arc::new(FixedSizeBinaryArray::try_from_iter(
[
[1, 2, 3, 4].as_slice(),
[13, 14, 15, 16].as_slice(),
[5, 6, 7, 8].as_slice(),
]
.into_iter(),
)?) as ArrayRef;
assert_eq!(
expected,
eval_in_list_from_array(Arc::clone(&fsb_needle), Arc::clone(&fsb_in))?
);
assert_eq!(
expected,
eval_in_list_from_array(wrap_in_dict(fsb_needle), fsb_in)?
);

// Struct in_array, Struct needle: multi-column join
let struct_fields = Fields::from(vec![
Field::new("c0", DataType::Utf8, true),
Expand Down
Loading
Loading