Optimize IN LIST for byte view arrays

Author: geoffreyclaude

datafusion/physical-expr/benches/in_list_strategy.rs

@@ -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()
 }
 
@@ -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)
@@ -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)
         })
@@ -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)
             })
@@ -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,
             );
         }
@@ -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,
             );
         }

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

@@ -37,6 +37,7 @@ use datafusion_common::{
 use datafusion_expr::{ColumnarValue, expr_vec_fmt};
 
 mod array_static_filter;
+mod byte_view_filter;
 mod frozen_set;
 mod primitive_filter;
 mod result;
@@ -215,7 +216,7 @@ impl InListExpr {
             expr,
             list,
             negated,
-            Some(instantiate_static_filter(array)?),
+            Some(instantiate_static_filter(array, &expr_data_type)?),
         ))
     }
 
@@ -242,7 +243,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
         };
 
@@ -3576,6 +3577,23 @@ 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)?
+        );
+
         // Struct in_array, Struct needle: multi-column join
         let struct_fields = Fields::from(vec![
             Field::new("c0", DataType::Utf8, true),