IN LIST: reuse primitive filters for FixedSizeBinary

datafusion/physical-expr/benches/in_list_strategy.rs

@@ -44,14 +44,14 @@
 //! | 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, 16) | fixed-width binary values | 16 (1 byte), 64 (2 bytes), 4/64/256/10000 (16 bytes) |
 
 use arrow::array::types::IntervalMonthDayNano;
 use arrow::array::*;
 use arrow::datatypes::{Field, Int32Type, IntervalMonthDayNanoType, Schema};
 use arrow::record_batch::RecordBatch;
 use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
-use datafusion_common::ScalarValue;
+use datafusion_common::{HashSet, ScalarValue};
 use datafusion_physical_expr::expressions::{col, in_list, lit};
 use half::f16;
 use rand::distr::Alphanumeric;
@@ -996,40 +996,51 @@ 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];
+fn random_fixed_binary(rng: &mut StdRng, width: i32) -> Vec<u8> {
+    let mut buf = vec![0u8; width as usize];
     rng.fill(&mut buf[..]);
     buf
 }
 
-/// 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))
-        .collect();
+    // Keep the haystack unique so each configured list size reaches the
+    // intended filter strategy.
+    let mut haystack_set = HashSet::with_capacity(list_size);
+    let mut haystack = Vec::with_capacity(list_size);
+    while haystack.len() < list_size {
+        let value = random_fixed_binary(&mut rng, width);
+        if haystack_set.insert(value.clone()) {
+            haystack.push(value);
+        }
+    }
 
     // Generate array with controlled match rate
     let values: Vec<Vec<u8>> = (0..ARRAY_SIZE)
         .map(|_| {
             if !haystack.is_empty() && rng.random_bool(match_rate) {
                 haystack.choose(&mut rng).unwrap().clone()
             } else {
-                random_fixed_binary_16(&mut rng)
+                loop {
+                    let value = random_fixed_binary(&mut rng, width);
+                    if !haystack_set.contains(&value) {
+                        break value;
+                    }
+                }
             }
         })
         .collect();
@@ -1040,28 +1051,28 @@ 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),
+        ),
         &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), (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);
         }
     }
 }

datafusion/physical-expr/src/expressions/in_list.rs

@@ -38,6 +38,7 @@ use datafusion_expr::{ColumnarValue, expr_vec_fmt};
 
 mod array_static_filter;
 mod branchless_filter;
+mod fixed_size_binary_filter;
 mod primitive_filter;
 mod result;
 mod static_filter;
@@ -3548,6 +3549,39 @@ mod tests {
             );
         }
 
+        // 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(Arc::clone(&fsb_needle)),
+                Arc::clone(&fsb_in),
+            )?
+        );
+        assert_eq!(
+            expected,
+            eval_in_list_from_array(wrap_in_dict(fsb_needle), wrap_in_dict(fsb_in))?
+        );
+
         // Utf8 (falls through to ArrayStaticFilter)
         let utf8_in = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
         let utf8_needle = Arc::new(StringArray::from(vec!["a", "d", "b"])) as ArrayRef;