From 7598605a0546a9ca687a001d051f0f6134cb0fcc Mon Sep 17 00:00:00 2001 From: Pierre Lacave Date: Thu, 2 Jul 2026 09:53:18 +0200 Subject: [PATCH 1/3] perf: Optimize array_has() for array needle `array_has(array, element)` returns, for each row, whether the array contains the element. When the `element` (needle) is an array rather than a scalar -- the needle argument is a column with one value per row, e.g. `array_has(t1.tags, t2.key)` in a join filter -- execution goes through `array_has_dispatch_for_array` (the `ColumnarValue::Array` needle branch), which compared each row by invoking the Arrow `eq` kernel once per row. That kernel allocates a `BooleanArray` and pays downcast and dispatch overhead on every row. (The scalar-needle branch was optimized separately in #20374.) Add a fast path for primitive and string element types, preserving the Arrow `eq` kernel semantics (total-order float equality; null elements never match). With all-valid elements each row is a single branchless OR-reduction over the native values. When primitive elements contain nulls -- whose backing values are arbitrary -- the per-element equality bitmap is ANDed with the validity bitmap (one word-parallel op, no per-element branch) before the per-row reduction, so null slots never match regardless of their value. Past a moderate average list length (`NULL_FAST_PATH_MAX_LEN`) this bitmap's extra passes lose to the per-row kernel, so the element-null branch bails to it there; that length is measured over the visible (sliced) region, so a sliced array's hidden child elements don't route a small window to the slow path. The all-valid fold has no such crossover. For string elements each row is a single pass over the row's values; for `Utf8View` this is view-aware -- the byte length and 4-byte prefix packed in the 128-bit view reject non-matches before touching the data buffer (what the `eq` kernel does, but without its per-row allocation), and an inline needle (<= 12 bytes) is matched by full-view equality with no materialization at all. Nested and other element types keep using the per-row `eq` kernel. What this removes is the fixed per-row kernel overhead, not the element comparison itself, so the gain is largest for short lists and shrinks as lists grow. Criterion microbenchmark, array needle on a 100k-row batch vs the per-row `eq` kernel (NOT end-to-end query time): i64, all-valid elements: ~15x (64-elem lists), ~1.9x at 1024-elem i64, ~30% null elements: ~9x (8-elem) ... ~1x (512-elem); falls back beyond Utf8 / LargeUtf8: ~2.5x Utf8View (view-aware): ~5x on short/inline strings Every shape improves with no regression. End-to-end impact depends on how much of a query `array_has` accounts for. For a query dominated by an array-needle `array_has` join filter (a `NestedLoopJoinExec` with `filter=array_has(tags, key)` over 3000x3000 rows of 8-element lists) total time drops from 0.95s to 0.059s (~16x, identical results). For a workload where `array_has` is a smaller fraction -- e.g. the ~6% of profile that motivated this (see #18070 / #18161, which fixed the join's deep-copy but left the per-row `array_has` cost) -- the overall speedup is single-digit percent. The array-needle criterion benchmark used for these numbers is added separately (so it can be run against `main` to capture the before baseline). Part of #18727. Co-Authored-By: Claude Opus 4.8 --- datafusion/functions-nested/src/array_has.rs | 522 ++++++++++++++++++- 1 file changed, 516 insertions(+), 6 deletions(-) diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index 04818258f040b..a146dd04a4f21 100644 --- a/datafusion/functions-nested/src/array_has.rs +++ b/datafusion/functions-nested/src/array_has.rs @@ -18,11 +18,13 @@ //! [`ScalarUDFImpl`] definitions for array_has, array_has_all and array_has_any functions. use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, Datum, Scalar, - StringArrayType, + Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray, + BooleanBufferBuilder, Datum, PrimitiveArray, Scalar, StringArrayType, + StringViewArray, }; use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::DataType; +use arrow::downcast_primitive_array; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::cast::{as_fixed_size_list_array, as_generic_list_array}; use datafusion_common::utils::string_utils::string_array_to_vec; @@ -323,11 +325,26 @@ impl<'a> ArrayWrapper<'a> { } } +/// Evaluate `array_has` with an array (per-row) needle. +/// +/// The straightforward implementation compares each row with the Arrow `eq` +/// kernel, which allocates a `BooleanArray` and pays dispatch on every row -- +/// overhead that dominates for short lists. Primitive and string element types +/// therefore take [`array_has_array_fast_path`]; nested (and any other) types +/// fall back to the per-row kernel. fn array_has_dispatch_for_array<'a>( haystack: ArrayWrapper<'a>, needle: &ArrayRef, ) -> Result { let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls()); + + if let Some(values) = + array_has_array_fast_path(&haystack, needle, combined_nulls.as_ref()) + { + return Ok(Arc::new(BooleanArray::new(values, combined_nulls))); + } + + // Fallback: per-row `eq` kernel (nested element types, or a type mismatch). let mut result = BooleanBufferBuilder::new(haystack.len()); for (i, arr) in haystack.iter().enumerate() { if combined_nulls.as_ref().is_some_and(|n| n.is_null(i)) { @@ -344,6 +361,229 @@ fn array_has_dispatch_for_array<'a>( Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls))) } +/// Average list length past which the element-null fast path (case 2 of +/// [`array_has_array_primitive`]) stops beating the per-row `eq` kernel and is +/// bailed out of -- an empirically measured crossover. +const NULL_FAST_PATH_MAX_LEN: usize = 512; + +/// Per-element fast path for primitive and string element types. Returns `None` +/// for any other type (and on a needle/element type mismatch), so the caller +/// falls back to the per-row `eq` kernel. +fn array_has_array_fast_path( + haystack: &ArrayWrapper<'_>, + needle: &ArrayRef, + combined_nulls: Option<&NullBuffer>, +) -> Option { + let needle = needle.as_ref(); + + // Normalize for sliced arrays (like `array_has_dispatch_for_scalar`): slice + // to the visible region so `offsets[i] - first_offset` indexes the values. + let offsets: Vec = haystack.offsets().collect(); + let first_offset = offsets[0]; + let visible_values = haystack + .values() + .slice(first_offset, offsets[offsets.len() - 1] - first_offset); + let visible_values = visible_values.as_ref(); + + // The needle shares the haystack's element type after coercion; defer any + // mismatch to the generic path rather than panicking in the downcasts. + if visible_values.data_type() != needle.data_type() { + return None; + } + + downcast_primitive_array! { + visible_values => { + // The element-null branch of `array_has_array_primitive` makes + // several passes over the values; past a moderate average list + // length the per-row `eq` kernel wins, so bail to it there. Measured + // over the *visible* region -- the offset span and `visible_values` + // nulls, not the full backing child -- so a sliced array's hidden + // elements can't skew the decision. The average check short-circuits, + // so the all-valid win path never pays for `null_count`. Strings + // (single-pass) and nested types have no such crossover and never + // reach this arm. + let num_rows = offsets.len() - 1; + if num_rows > 0 + && (offsets[num_rows] - first_offset) / num_rows > NULL_FAST_PATH_MAX_LEN + && visible_values.null_count() > 0 + { + return None; + } + Some(array_has_array_primitive( + visible_values, needle, &offsets, first_offset, combined_nulls, + )) + }, + DataType::Utf8 => Some(array_has_array_string( + visible_values.as_string::(), + needle.as_string::(), + &offsets, + first_offset, + combined_nulls, + )), + DataType::LargeUtf8 => Some(array_has_array_string( + visible_values.as_string::(), + needle.as_string::(), + &offsets, + first_offset, + combined_nulls, + )), + DataType::Utf8View => Some(array_has_array_string_view( + visible_values.as_string_view(), + needle.as_string_view(), + &offsets, + first_offset, + combined_nulls, + )), + _ => None, + } +} + +/// Primitive fast path, with two branches on element validity: +/// +/// 1. No element nulls: a branchless OR-reduction over the raw value slice -- +/// auto-vectorizes, and is the common, fastest case. +/// 2. Element nulls: a null slot's backing value is arbitrary, so validity must +/// be consulted. Compare into an equality bitmap and AND it with the validity +/// bitmap -- one word-parallel op, no per-element branch -- then reduce each +/// row to "any bit set". Chunked so the expanded-needle scratch stays bounded +/// regardless of batch size. +fn array_has_array_primitive( + values: &PrimitiveArray, + needle: &dyn Array, + offsets: &[usize], + first_offset: usize, + combined_nulls: Option<&NullBuffer>, +) -> BooleanBuffer +where + T::Native: ArrowNativeTypeOp, +{ + let needle = needle.as_primitive::(); + let num_rows = offsets.len() - 1; + let value_slice = values.values(); + let needle_slice = needle.values(); + + let Some(element_nulls) = values.nulls() else { + return BooleanBuffer::collect_bool(num_rows, |i| { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + return false; + } + // `needle[i]` is non-null here: combined_nulls covers the needle nulls. + let needle_val = needle_slice[i]; + let start = offsets[i] - first_offset; + let end = offsets[i + 1] - first_offset; + value_slice[start..end] + .iter() + .fold(false, |acc, &v| acc | v.is_eq(needle_val)) + }); + }; + + // Case 2 (see fn doc), chunked like the all/any kernels. + let mut result = BooleanBufferBuilder::new(num_rows); + let mut needle_expanded: Vec = Vec::new(); + for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) { + let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows); + let elem_start = offsets[chunk_start] - first_offset; + let elem_end = offsets[chunk_end] - first_offset; + + // Expand the per-row needle across this chunk's elements (reused scratch), + // then compare in one vectorizable pass and mask out null elements. + needle_expanded.clear(); + for i in chunk_start..chunk_end { + needle_expanded + .resize(offsets[i + 1] - first_offset - elem_start, needle_slice[i]); + } + let chunk_values = &value_slice[elem_start..elem_end]; + let eq_bits = BooleanBuffer::collect_bool(chunk_values.len(), |k| { + chunk_values[k].is_eq(needle_expanded[k]) + }); + let matched = &eq_bits + & &element_nulls + .inner() + .slice(elem_start, elem_end - elem_start); + + for i in chunk_start..chunk_end { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + result.append(false); + continue; + } + let start = offsets[i] - first_offset - elem_start; + let end = offsets[i + 1] - first_offset - elem_start; + result.append(matched.slice(start, end - start).has_true()); + } + } + result.finish() +} + +/// String implementation of `array_has_array_fast_path`, generic over the +/// concrete string array type (`Utf8`, `LargeUtf8`, `Utf8View`). +fn array_has_array_string<'a, S: StringArrayType<'a> + Copy>( + values: S, + needle: S, + offsets: &[usize], + first_offset: usize, + combined_nulls: Option<&NullBuffer>, +) -> BooleanBuffer { + let num_rows = offsets.len() - 1; + BooleanBuffer::collect_bool(num_rows, |i| { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + return false; + } + // `needle[i]` is non-null here: combined_nulls covers the needle nulls. + let needle_val = needle.value(i); + let start = offsets[i] - first_offset; + let end = offsets[i + 1] - first_offset; + // Compare the value first and only consult validity on a match (see the + // primitive path for why this is correct and faster on no-match scans). + (start..end).any(|k| values.value(k) == needle_val && !values.is_null(k)) + }) +} + +/// View-aware `Utf8View` variant of [`array_has_array_string`]. A `StringView` +/// packs the byte length and a 4-byte prefix into its 128-bit view; Arrow's +/// per-row `eq` kernel compares those before ever touching the data buffer, +/// which the generic `value(k) == needle` path gives up by materializing every +/// element. Here we compare the raw views directly: an inline needle (<= 12 +/// bytes, whose view holds the whole value zero-padded) matches iff the full +/// views are equal -- no materialization at all -- and a longer needle is +/// rejected on length + prefix and only materialized to confirm a candidate. A +/// null slot's backing view is arbitrary, so (as in the primitive path) +/// validity is consulted only on a view match. +fn array_has_array_string_view( + values: &StringViewArray, + needle: &StringViewArray, + offsets: &[usize], + first_offset: usize, + combined_nulls: Option<&NullBuffer>, +) -> BooleanBuffer { + let num_rows = offsets.len() - 1; + let value_views = values.views(); + let needle_views = needle.views(); + BooleanBuffer::collect_bool(num_rows, |i| { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + return false; + } + // `needle[i]` is non-null here: combined_nulls covers the needle nulls. + let needle_view = needle_views[i]; + // Low 32 bits are the byte length; the next 32 are the inline prefix. + let needle_inline = (needle_view as u32) <= 12; + let needle_lo = needle_view as u64; + let needle_val = needle.value(i); + let start = offsets[i] - first_offset; + let end = offsets[i + 1] - first_offset; + (start..end).any(|k| { + let v = value_views[k]; + let matched = if needle_inline { + // Inline: the whole view is the canonical value (zero padded). + v == needle_view + } else { + // Longer: reject on length + prefix, then confirm the bytes. + (v as u64) == needle_lo && values.value(k) == needle_val + }; + matched && !values.is_null(k) + }) + }) +} + fn array_has_dispatch_for_scalar( haystack: ArrayWrapper<'_>, needle: &dyn Datum, @@ -972,13 +1212,14 @@ fn general_array_has_all_and_any_kernel( mod tests { use std::sync::Arc; - use arrow::datatypes::Int32Type; + use arrow::datatypes::{Float32Type, Float64Type, Int32Type}; use arrow::{ array::{ - Array, ArrayRef, AsArray, FixedSizeListArray, Int32Array, ListArray, - create_array, + Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, + FixedSizeListArray, Float32Array, Float64Array, Int32Array, LargeListArray, + ListArray, Scalar, StringArray, StringViewArray, create_array, }, - buffer::OffsetBuffer, + buffer::{NullBuffer, OffsetBuffer}, datatypes::{DataType, Field}, }; use datafusion_common::{ @@ -1311,4 +1552,273 @@ mod tests { &[Some(true), Some(true)], ); } + + /// Invoke `array_has` with the needle as an array (a column with one value + /// per row). This exercises `array_has_dispatch_for_array` and its fast path. + fn invoke_array_has_array(haystack: ArrayRef, needle: ArrayRef) -> ArrayRef { + let num_rows = haystack.len(); + let haystack_type = haystack.data_type().clone(); + let needle_type = needle.data_type().clone(); + ArrayHas::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(haystack), ColumnarValue::Array(needle)], + arg_fields: vec![ + Arc::new(Field::new("haystack", haystack_type, false)), + Arc::new(Field::new("needle", needle_type, false)), + ], + number_rows: num_rows, + return_field: Arc::new(Field::new("return", DataType::Boolean, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + .into_array(num_rows) + .unwrap() + } + + /// Reference implementation of the array-needle path using the original + /// per-row Arrow `eq` kernel + `has_true`, proving the fast path is + /// behavior-preserving. Handles `List`/`LargeList` (i32/i64 offsets) and + /// `FixedSizeList` haystacks. + fn reference_array_has_array(haystack: &ArrayRef, needle: &ArrayRef) -> BooleanArray { + let combined = NullBuffer::union(haystack.nulls(), needle.nulls()); + let mut builder = BooleanBufferBuilder::new(haystack.len()); + for i in 0..haystack.len() { + if combined.as_ref().is_some_and(|n| n.is_null(i)) { + builder.append(false); + continue; + } + let row = if let Some(l) = haystack.as_list_opt::() { + l.value(i) + } else if let Some(l) = haystack.as_list_opt::() { + l.value(i) + } else { + haystack.as_fixed_size_list().value(i) + }; + let needle_row = Scalar::new(needle.slice(i, 1)); + let eq = + super::compare_with_eq(&row, &needle_row, row.data_type().is_nested()) + .unwrap(); + builder.append(eq.has_true()); + } + BooleanArray::new(builder.finish(), combined) + } + + /// Build a `List` (i32 offsets) of string elements of `element_type` (`Utf8` + /// or `Utf8View`) from nested rows; a `None` row is a null list. + fn str_list( + rows: Vec>>>, + element_type: DataType, + ) -> ArrayRef { + let mut flat: Vec> = Vec::new(); + let mut offsets = vec![0i32]; + let mut nulls = Vec::new(); + for row in &rows { + match row { + Some(es) => { + flat.extend(es.iter().copied()); + nulls.push(true); + } + None => nulls.push(false), + } + offsets.push(flat.len() as i32); + } + let values: ArrayRef = match &element_type { + DataType::Utf8 => Arc::new(flat.into_iter().collect::()), + DataType::Utf8View => Arc::new(flat.into_iter().collect::()), + _ => unreachable!("str_list only supports Utf8/Utf8View"), + }; + Arc::new(ListArray::new( + Arc::new(Field::new_list_field(element_type, true)), + OffsetBuffer::new(offsets.into()), + values, + Some(NullBuffer::from(nulls)), + )) + } + + /// The fast path must produce exactly what the per-row `eq` kernel would, for + /// every element type and null shape. Each case is checked against the oracle, + /// which also pins total-order float equality (`NaN == NaN`, `+0.0 != -0.0`) + /// and the null-fill collision (a null slot's backing value must never match). + #[test] + fn test_array_has_array_matches_reference() { + let check = |haystack: ArrayRef, needle: ArrayRef| { + let got = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + got.as_boolean(), + &reference_array_has_array(&haystack, &needle) + ); + }; + + // Int32: element null, null needle, null row, empty row, and null-fill + // collision -- valid 0 matches (row 4), a null slot's backing 0 does not + // (row 5). + check( + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(7), None, Some(9)]), + Some(vec![Some(0), None]), + Some(vec![None, Some(5)]), + Some(vec![]), + ])), + Arc::new(Int32Array::from(vec![ + Some(2), + Some(5), + None, + Some(0), + Some(0), + Some(1), + ])), + ); + + // Float total order: NaN == NaN (same bits), +0.0 != -0.0 (f64 and f32). + check( + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(f64::NAN), Some(1.0)]), + Some(vec![Some(0.0)]), + Some(vec![Some(-0.0)]), + ])), + Arc::new(Float64Array::from(vec![f64::NAN, -0.0, -0.0])), + ); + check( + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(f32::NAN)]), + Some(vec![Some(0.0)]), + ])), + Arc::new(Float32Array::from(vec![f32::NAN, -0.0])), + ); + + // Utf8 and Utf8View: inline and long (> 12 byte) elements sharing a 4-byte + // prefix, element/needle nulls, and a "" needle against a null slot. + for element_type in [DataType::Utf8, DataType::Utf8View] { + let rows = vec![ + Some(vec![Some("a"), Some("bb")]), + Some(vec![Some("this_is_a_long_value_xyz")]), + Some(vec![Some("prefixAAAA_1111"), Some("prefixAAAA_2222")]), + Some(vec![Some("x"), None, Some("y")]), + Some(vec![None]), + None, + ]; + let needle_vals = vec![ + Some("bb"), + Some("this_is_a_long_value_xyz"), + Some("prefixAAAA_2222"), + Some("y"), + Some(""), + Some("q"), + ]; + let needle: ArrayRef = match &element_type { + DataType::Utf8 => { + Arc::new(needle_vals.into_iter().collect::()) + } + _ => Arc::new(needle_vals.into_iter().collect::()), + }; + check(str_list(rows, element_type), needle); + } + + // LargeList (i64 offsets): null branch + null-fill collision. + check( + Arc::new(LargeListArray::from_iter_primitive::( + vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![None, Some(3)]), + Some(vec![Some(4), None]), + ], + )), + Arc::new(Int32Array::from(vec![Some(2), Some(0), Some(4)])), + ); + } + + #[test] + fn test_array_has_array_sliced() { + // Offset normalization for sliced haystacks must keep element ranges and + // the needle column aligned, for both `List` (buffer offsets) and + // `FixedSizeList` (offsets computed as `i * value_length`), and must route + // by the *visible* average length, not the full backing child's. + let full = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(10), Some(20), Some(30)]), + Some(vec![Some(40)]), + Some(vec![Some(50), Some(60)]), + Some(vec![Some(70)]), + ]); + let hay: ArrayRef = Arc::new(full.slice(1, 3)); + let needle: ArrayRef = + Arc::new(Int32Array::from(vec![999, 20, 41, 60, 999]).slice(1, 3)); + let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); + assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); + + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let fsl_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32])); + let fsl: ArrayRef = + Arc::new(FixedSizeListArray::new(field, 2, fsl_values, None).slice(1, 2)); + let needle: ArrayRef = Arc::new(Int32Array::from(vec![11, 99])); + let got = invoke_array_has_array(Arc::clone(&fsl), Arc::clone(&needle)); + assert_eq!(got.as_boolean(), &reference_array_has_array(&fsl, &needle)); + + // Large null-bearing backing child sliced to a short 3-row window: the + // element-null bail heuristic must use the visible average (row_len = 20), + // not backing_rows * row_len / visible_rows. + let (backing_rows, row_len) = (200usize, 20usize); + let data: Vec>>> = (0..backing_rows) + .map(|i| { + Some( + (0..row_len) + .map(|j| (j % 4 != 0).then_some(((i + j) % 11) as i32)) + .collect(), + ) + }) + .collect(); + let needles: Vec> = + (0..backing_rows).map(|i| Some((i % 11) as i32)).collect(); + let hay: ArrayRef = Arc::new( + ListArray::from_iter_primitive::(data).slice(50, 3), + ); + let needle: ArrayRef = Arc::new(Int32Array::from(needles).slice(50, 3)); + let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); + assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); + } + + #[test] + fn test_array_has_array_chunk_boundary_and_bail() { + // More than 2 ROW_CONVERSION_CHUNK_SIZE (512) row chunks with element + // nulls, null rows and varying needles exercises chunk rebasing. + let n = 1500usize; + let rows: Vec>>> = (0..n) + .map(|i| { + (i % 13 != 0).then(|| { + let r = (i % 7) as i32; + vec![Some(r), None, Some((r + 1) % 5)] + }) + }) + .collect(); + let needles: Vec> = (0..n) + .map(|i| (i % 11 != 0).then_some((i % 5) as i32)) + .collect(); + let hay: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(rows)); + let needle: ArrayRef = Arc::new(Int32Array::from(needles)); + let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); + assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); + + // Average list length past NULL_FAST_PATH_MAX_LEN with element nulls + // routes to the per-row kernel; the result must still match the oracle. + let len = super::NULL_FAST_PATH_MAX_LEN + 64; + let rows = 20usize; + let data: Vec>>> = (0..rows) + .map(|i| { + Some( + (0..len) + .map(|j| (j % 5 != 0).then_some(((i + j) % 7) as i32)) + .collect(), + ) + }) + .collect(); + let needles: Vec> = (0..rows).map(|i| Some((i % 7) as i32)).collect(); + let hay: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(data)); + let needle: ArrayRef = Arc::new(Int32Array::from(needles)); + let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); + assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); + } } From d65997e2723648e3fd62ef95aaf4c24b6565dce1 Mon Sep 17 00:00:00 2001 From: Pierre Lacave Date: Mon, 6 Jul 2026 23:02:01 +0200 Subject: [PATCH 2/3] test: cover array_has array-needle path via sqllogictest Move the behavioral coverage for the array (column) needle path from Rust unit tests into datafusion/sqllogictest/test_files/array/array_has.slt (element and needle nulls, null/empty rows, the null-fill collision, found/not-found, over i64, Utf8, Utf8View, LargeList and LargeUtf8, plus a >512-row case that exercises the chunked element-null path), following the maintainer preference for covering UDFs with SLT. Keep a single unit test for sliced List / FixedSizeList offset handling, whose sliced-array shape SQL can't produce. Drops the reference oracle and the string-list test helper, no longer needed. Co-Authored-By: Claude Opus 4.8 --- datafusion/functions-nested/src/array_has.rs | 263 ++---------------- .../test_files/array/array_has.slt | 114 ++++++++ 2 files changed, 140 insertions(+), 237 deletions(-) diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index a146dd04a4f21..ebca49285a458 100644 --- a/datafusion/functions-nested/src/array_has.rs +++ b/datafusion/functions-nested/src/array_has.rs @@ -1212,14 +1212,13 @@ fn general_array_has_all_and_any_kernel( mod tests { use std::sync::Arc; - use arrow::datatypes::{Float32Type, Float64Type, Int32Type}; + use arrow::datatypes::Int32Type; use arrow::{ array::{ - Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, - FixedSizeListArray, Float32Array, Float64Array, Int32Array, LargeListArray, - ListArray, Scalar, StringArray, StringViewArray, create_array, + Array, ArrayRef, AsArray, FixedSizeListArray, Int32Array, ListArray, + create_array, }, - buffer::{NullBuffer, OffsetBuffer}, + buffer::OffsetBuffer, datatypes::{DataType, Field}, }; use datafusion_common::{ @@ -1575,250 +1574,40 @@ mod tests { .unwrap() } - /// Reference implementation of the array-needle path using the original - /// per-row Arrow `eq` kernel + `has_true`, proving the fast path is - /// behavior-preserving. Handles `List`/`LargeList` (i32/i64 offsets) and - /// `FixedSizeList` haystacks. - fn reference_array_has_array(haystack: &ArrayRef, needle: &ArrayRef) -> BooleanArray { - let combined = NullBuffer::union(haystack.nulls(), needle.nulls()); - let mut builder = BooleanBufferBuilder::new(haystack.len()); - for i in 0..haystack.len() { - if combined.as_ref().is_some_and(|n| n.is_null(i)) { - builder.append(false); - continue; - } - let row = if let Some(l) = haystack.as_list_opt::() { - l.value(i) - } else if let Some(l) = haystack.as_list_opt::() { - l.value(i) - } else { - haystack.as_fixed_size_list().value(i) - }; - let needle_row = Scalar::new(needle.slice(i, 1)); - let eq = - super::compare_with_eq(&row, &needle_row, row.data_type().is_nested()) - .unwrap(); - builder.append(eq.has_true()); - } - BooleanArray::new(builder.finish(), combined) - } - - /// Build a `List` (i32 offsets) of string elements of `element_type` (`Utf8` - /// or `Utf8View`) from nested rows; a `None` row is a null list. - fn str_list( - rows: Vec>>>, - element_type: DataType, - ) -> ArrayRef { - let mut flat: Vec> = Vec::new(); - let mut offsets = vec![0i32]; - let mut nulls = Vec::new(); - for row in &rows { - match row { - Some(es) => { - flat.extend(es.iter().copied()); - nulls.push(true); - } - None => nulls.push(false), - } - offsets.push(flat.len() as i32); - } - let values: ArrayRef = match &element_type { - DataType::Utf8 => Arc::new(flat.into_iter().collect::()), - DataType::Utf8View => Arc::new(flat.into_iter().collect::()), - _ => unreachable!("str_list only supports Utf8/Utf8View"), - }; - Arc::new(ListArray::new( - Arc::new(Field::new_list_field(element_type, true)), - OffsetBuffer::new(offsets.into()), - values, - Some(NullBuffer::from(nulls)), - )) - } - - /// The fast path must produce exactly what the per-row `eq` kernel would, for - /// every element type and null shape. Each case is checked against the oracle, - /// which also pins total-order float equality (`NaN == NaN`, `+0.0 != -0.0`) - /// and the null-fill collision (a null slot's backing value must never match). #[test] - fn test_array_has_array_matches_reference() { - let check = |haystack: ArrayRef, needle: ArrayRef| { - let got = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); - assert_eq!( - got.as_boolean(), - &reference_array_has_array(&haystack, &needle) - ); - }; - - // Int32: element null, null needle, null row, empty row, and null-fill - // collision -- valid 0 matches (row 4), a null slot's backing 0 does not - // (row 5). - check( - Arc::new(ListArray::from_iter_primitive::(vec![ - Some(vec![Some(1), Some(2), Some(3)]), - None, - Some(vec![Some(7), None, Some(9)]), - Some(vec![Some(0), None]), - Some(vec![None, Some(5)]), - Some(vec![]), - ])), - Arc::new(Int32Array::from(vec![ - Some(2), - Some(5), - None, - Some(0), - Some(0), - Some(1), - ])), - ); - - // Float total order: NaN == NaN (same bits), +0.0 != -0.0 (f64 and f32). - check( - Arc::new(ListArray::from_iter_primitive::(vec![ - Some(vec![Some(f64::NAN), Some(1.0)]), - Some(vec![Some(0.0)]), - Some(vec![Some(-0.0)]), - ])), - Arc::new(Float64Array::from(vec![f64::NAN, -0.0, -0.0])), - ); - check( - Arc::new(ListArray::from_iter_primitive::(vec![ - Some(vec![Some(f32::NAN)]), - Some(vec![Some(0.0)]), - ])), - Arc::new(Float32Array::from(vec![f32::NAN, -0.0])), - ); - - // Utf8 and Utf8View: inline and long (> 12 byte) elements sharing a 4-byte - // prefix, element/needle nulls, and a "" needle against a null slot. - for element_type in [DataType::Utf8, DataType::Utf8View] { - let rows = vec![ - Some(vec![Some("a"), Some("bb")]), - Some(vec![Some("this_is_a_long_value_xyz")]), - Some(vec![Some("prefixAAAA_1111"), Some("prefixAAAA_2222")]), - Some(vec![Some("x"), None, Some("y")]), - Some(vec![None]), - None, - ]; - let needle_vals = vec![ - Some("bb"), - Some("this_is_a_long_value_xyz"), - Some("prefixAAAA_2222"), - Some("y"), - Some(""), - Some("q"), - ]; - let needle: ArrayRef = match &element_type { - DataType::Utf8 => { - Arc::new(needle_vals.into_iter().collect::()) - } - _ => Arc::new(needle_vals.into_iter().collect::()), - }; - check(str_list(rows, element_type), needle); - } - - // LargeList (i64 offsets): null branch + null-fill collision. - check( - Arc::new(LargeListArray::from_iter_primitive::( - vec![ - Some(vec![Some(1), Some(2)]), - Some(vec![None, Some(3)]), - Some(vec![Some(4), None]), - ], - )), - Arc::new(Int32Array::from(vec![Some(2), Some(0), Some(4)])), - ); - } - - #[test] - fn test_array_has_array_sliced() { - // Offset normalization for sliced haystacks must keep element ranges and - // the needle column aligned, for both `List` (buffer offsets) and - // `FixedSizeList` (offsets computed as `i * value_length`), and must route - // by the *visible* average length, not the full backing child's. + fn test_array_has_array_needle_sliced() { + // Offset normalization for sliced haystacks must keep the element ranges + // and the needle column aligned, for both `List` (offsets from the + // buffer) and `FixedSizeList` (offsets computed as `i * value_length`). + // Slicing is an execution artifact SQL/SLT can't force, so this stays a + // unit test; value-level behavior is covered by `array/array_has.slt`. let full = ListArray::from_iter_primitive::(vec![ Some(vec![Some(1), Some(2)]), - Some(vec![Some(10), Some(20), Some(30)]), - Some(vec![Some(40)]), - Some(vec![Some(50), Some(60)]), + Some(vec![Some(10), Some(20), Some(30)]), // needle 20 -> true + Some(vec![Some(40)]), // needle 41 -> false + Some(vec![Some(50), Some(60)]), // needle 60 -> true Some(vec![Some(70)]), ]); - let hay: ArrayRef = Arc::new(full.slice(1, 3)); - let needle: ArrayRef = + let sliced_haystack: ArrayRef = Arc::new(full.slice(1, 3)); + let sliced_needle: ArrayRef = Arc::new(Int32Array::from(vec![999, 20, 41, 60, 999]).slice(1, 3)); - let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); - assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); + let result = invoke_array_has_array(sliced_haystack, sliced_needle); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true)] + ); + // Sliced FixedSizeList (width 2; rows 1..=2 of + // [[1,2],[11,12],[21,22],[31,32]] visible) with an aligned needle column. let field = Arc::new(Field::new("item", DataType::Int32, true)); let fsl_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32])); let fsl: ArrayRef = Arc::new(FixedSizeListArray::new(field, 2, fsl_values, None).slice(1, 2)); let needle: ArrayRef = Arc::new(Int32Array::from(vec![11, 99])); - let got = invoke_array_has_array(Arc::clone(&fsl), Arc::clone(&needle)); - assert_eq!(got.as_boolean(), &reference_array_has_array(&fsl, &needle)); - - // Large null-bearing backing child sliced to a short 3-row window: the - // element-null bail heuristic must use the visible average (row_len = 20), - // not backing_rows * row_len / visible_rows. - let (backing_rows, row_len) = (200usize, 20usize); - let data: Vec>>> = (0..backing_rows) - .map(|i| { - Some( - (0..row_len) - .map(|j| (j % 4 != 0).then_some(((i + j) % 11) as i32)) - .collect(), - ) - }) - .collect(); - let needles: Vec> = - (0..backing_rows).map(|i| Some((i % 11) as i32)).collect(); - let hay: ArrayRef = Arc::new( - ListArray::from_iter_primitive::(data).slice(50, 3), + let result = invoke_array_has_array(fsl, needle); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false)] ); - let needle: ArrayRef = Arc::new(Int32Array::from(needles).slice(50, 3)); - let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); - assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); - } - - #[test] - fn test_array_has_array_chunk_boundary_and_bail() { - // More than 2 ROW_CONVERSION_CHUNK_SIZE (512) row chunks with element - // nulls, null rows and varying needles exercises chunk rebasing. - let n = 1500usize; - let rows: Vec>>> = (0..n) - .map(|i| { - (i % 13 != 0).then(|| { - let r = (i % 7) as i32; - vec![Some(r), None, Some((r + 1) % 5)] - }) - }) - .collect(); - let needles: Vec> = (0..n) - .map(|i| (i % 11 != 0).then_some((i % 5) as i32)) - .collect(); - let hay: ArrayRef = - Arc::new(ListArray::from_iter_primitive::(rows)); - let needle: ArrayRef = Arc::new(Int32Array::from(needles)); - let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); - assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); - - // Average list length past NULL_FAST_PATH_MAX_LEN with element nulls - // routes to the per-row kernel; the result must still match the oracle. - let len = super::NULL_FAST_PATH_MAX_LEN + 64; - let rows = 20usize; - let data: Vec>>> = (0..rows) - .map(|i| { - Some( - (0..len) - .map(|j| (j % 5 != 0).then_some(((i + j) % 7) as i32)) - .collect(), - ) - }) - .collect(); - let needles: Vec> = (0..rows).map(|i| Some((i % 7) as i32)).collect(); - let hay: ArrayRef = - Arc::new(ListArray::from_iter_primitive::(data)); - let needle: ArrayRef = Arc::new(Int32Array::from(needles)); - let got = invoke_array_has_array(Arc::clone(&hay), Arc::clone(&needle)); - assert_eq!(got.as_boolean(), &reference_array_has_array(&hay, &needle)); } } diff --git a/datafusion/sqllogictest/test_files/array/array_has.slt b/datafusion/sqllogictest/test_files/array/array_has.slt index e343c1b1fae41..0fc246d24b928 100644 --- a/datafusion/sqllogictest/test_files/array/array_has.slt +++ b/datafusion/sqllogictest/test_files/array/array_has.slt @@ -894,4 +894,118 @@ statement ok DROP TABLE any_op_test; +# ------------------------------------------------------------------------- +# array_has with an array (column) needle -- one needle value per row, which +# goes through array_has_dispatch_for_array (the cases above use a scalar +# literal needle and take a different path). +# ------------------------------------------------------------------------- + +statement ok +create table array_has_int_needle (arr int[], needle int) as values + ([1, 2, 3], 2), -- found + ([4, 5, 6], 9), -- not found + (NULL, 5), -- null row + ([7, NULL, 9], NULL), -- null needle + ([7, NULL, 9], 7), -- element null skipped, found + ([0, NULL], 0), -- valid 0 matches + ([NULL, 5], 0), -- null-fill collision: a null slot must not match 0 + ([], 1), -- empty + ([NULL, NULL], 3); -- all null + +query B +select array_has(arr, needle) from array_has_int_needle; +---- +true +false +NULL +NULL +true +true +false +false +false + +# same over LargeList (i64) offsets +query B +select array_has(arrow_cast(arr, 'LargeList(Int32)'), needle) from array_has_int_needle; +---- +true +false +NULL +NULL +true +true +false +false +false + +statement ok +drop table array_has_int_needle; + +statement ok +create table array_has_str_needle (arr text[], needle text) as values + (['a', 'bb', 'ccc'], 'bb'), -- inline, found + (['short', 'tiny'], 'missing'), -- inline, not found + (['this_is_a_long_value_xyz'], 'this_is_a_long_value_xyz'), -- long, found + (['prefixAAAA_1111', 'prefixAAAA_2222'], 'prefixAAAA_2222'), -- long shared prefix + (['x', NULL, 'y'], 'y'), -- element null skipped + ([NULL], ''), -- null slot vs "" -> false + (NULL, 'q'); -- null row + +query B +select array_has(arr, needle) from array_has_str_needle; +---- +true +false +true +true +true +false +NULL + +# Utf8View exercises the view-aware fast path +query B +select array_has(arrow_cast(arr, 'List(Utf8View)'), arrow_cast(needle, 'Utf8View')) +from array_has_str_needle; +---- +true +false +true +true +true +false +NULL + +# LargeUtf8 elements +query B +select array_has(arrow_cast(arr, 'LargeList(LargeUtf8)'), arrow_cast(needle, 'LargeUtf8')) +from array_has_str_needle; +---- +true +false +true +true +true +false +NULL + +statement ok +drop table array_has_str_needle; + +# > ROW_CONVERSION_CHUNK_SIZE (512) rows with element nulls exercises the chunked +# element-null path. The needle equals an element that is always present, so all +# rows match; the second query shifts the needle out of range, so none do. +query I +select count(*) from generate_series(1, 2000) as t(v) +where array_has(make_array(v % 7, NULL, (v + 2) % 7), v % 7); +---- +2000 + +query I +select count(*) from generate_series(1, 2000) as t(v) +where array_has(make_array(v % 7, NULL, (v + 2) % 7), v % 7 + 100); +---- +0 + + include ./cleanup.slt.part From ba79a3816bd46520e430b903103d80d5c4eb7a2f Mon Sep 17 00:00:00 2001 From: Pierre Lacave Date: Fri, 10 Jul 2026 15:15:43 +0200 Subject: [PATCH 3/3] Address review feedback on the array-needle fast path. - Inline array_has_array_fast_path into array_has_dispatch_for_array so the dispatch reads in one pass. - Rebase offsets to 0 once and drop the first_offset parameter. - Use arrow's MAX_INLINE_VIEW_LEN instead of a hardcoded 12. - Build the expanded needle with std::iter::repeat_n instead of a cumulative Vec::resize. - Trim doc comments to the essentials. --- datafusion/functions-nested/src/array_has.rs | 214 ++++++++----------- 1 file changed, 93 insertions(+), 121 deletions(-) diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index ebca49285a458..3d51df36cd2ce 100644 --- a/datafusion/functions-nested/src/array_has.rs +++ b/datafusion/functions-nested/src/array_has.rs @@ -19,8 +19,8 @@ use arrow::array::{ Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray, - BooleanBufferBuilder, Datum, PrimitiveArray, Scalar, StringArrayType, - StringViewArray, + BooleanBufferBuilder, Datum, MAX_INLINE_VIEW_LEN, PrimitiveArray, Scalar, + StringArrayType, StringViewArray, }; use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::DataType; @@ -327,20 +327,75 @@ impl<'a> ArrayWrapper<'a> { /// Evaluate `array_has` with an array (per-row) needle. /// -/// The straightforward implementation compares each row with the Arrow `eq` -/// kernel, which allocates a `BooleanArray` and pays dispatch on every row -- -/// overhead that dominates for short lists. Primitive and string element types -/// therefore take [`array_has_array_fast_path`]; nested (and any other) types -/// fall back to the per-row kernel. +/// Primitive and string element types take a per-type fast path; nested (and any +/// other) element types fall back to the per-row `eq` kernel, which allocates a +/// `BooleanArray` per row. fn array_has_dispatch_for_array<'a>( haystack: ArrayWrapper<'a>, needle: &ArrayRef, ) -> Result { let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls()); + let needle = needle.as_ref(); + + // Slice to the visible region and rebase the offsets to 0 so `offsets[i]` + // indexes `visible_values` directly (the haystack may be a sliced list). + let mut offsets: Vec = haystack.offsets().collect(); + let first_offset = offsets[0]; + let visible_values = haystack + .values() + .slice(first_offset, offsets[offsets.len() - 1] - first_offset); + let visible_values = visible_values.as_ref(); + for offset in &mut offsets { + *offset -= first_offset; + } + + // Fast path for primitive/string elements whose (coerced) type matches the + // needle; a type mismatch or a nested type falls through to the per-row kernel. + let fast_path = if visible_values.data_type() != needle.data_type() { + None + } else { + downcast_primitive_array! { + visible_values => { + // The element-null path makes several passes over the values, so + // past a large average list length the per-row `eq` kernel is + // faster -- bail to it. The single-pass all-valid path has no such + // crossover, so only bail when elements are null. + let num_rows = offsets.len() - 1; + if num_rows > 0 + && offsets[num_rows] / num_rows > NULL_FAST_PATH_MAX_LEN + && visible_values.null_count() > 0 + { + None + } else { + Some(array_has_array_primitive( + visible_values, needle, &offsets, + combined_nulls.as_ref(), + )) + } + }, + DataType::Utf8 => Some(array_has_array_string( + visible_values.as_string::(), + needle.as_string::(), + &offsets, + combined_nulls.as_ref(), + )), + DataType::LargeUtf8 => Some(array_has_array_string( + visible_values.as_string::(), + needle.as_string::(), + &offsets, + combined_nulls.as_ref(), + )), + DataType::Utf8View => Some(array_has_array_string_view( + visible_values.as_string_view(), + needle.as_string_view(), + &offsets, + combined_nulls.as_ref(), + )), + _ => None, + } + }; - if let Some(values) = - array_has_array_fast_path(&haystack, needle, combined_nulls.as_ref()) - { + if let Some(values) = fast_path { return Ok(Arc::new(BooleanArray::new(values, combined_nulls))); } @@ -361,97 +416,20 @@ fn array_has_dispatch_for_array<'a>( Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls))) } -/// Average list length past which the element-null fast path (case 2 of -/// [`array_has_array_primitive`]) stops beating the per-row `eq` kernel and is -/// bailed out of -- an empirically measured crossover. +/// Average list length past which the element-null path loses to the per-row +/// `eq` kernel and bails to it (empirically measured). const NULL_FAST_PATH_MAX_LEN: usize = 512; -/// Per-element fast path for primitive and string element types. Returns `None` -/// for any other type (and on a needle/element type mismatch), so the caller -/// falls back to the per-row `eq` kernel. -fn array_has_array_fast_path( - haystack: &ArrayWrapper<'_>, - needle: &ArrayRef, - combined_nulls: Option<&NullBuffer>, -) -> Option { - let needle = needle.as_ref(); - - // Normalize for sliced arrays (like `array_has_dispatch_for_scalar`): slice - // to the visible region so `offsets[i] - first_offset` indexes the values. - let offsets: Vec = haystack.offsets().collect(); - let first_offset = offsets[0]; - let visible_values = haystack - .values() - .slice(first_offset, offsets[offsets.len() - 1] - first_offset); - let visible_values = visible_values.as_ref(); - - // The needle shares the haystack's element type after coercion; defer any - // mismatch to the generic path rather than panicking in the downcasts. - if visible_values.data_type() != needle.data_type() { - return None; - } - - downcast_primitive_array! { - visible_values => { - // The element-null branch of `array_has_array_primitive` makes - // several passes over the values; past a moderate average list - // length the per-row `eq` kernel wins, so bail to it there. Measured - // over the *visible* region -- the offset span and `visible_values` - // nulls, not the full backing child -- so a sliced array's hidden - // elements can't skew the decision. The average check short-circuits, - // so the all-valid win path never pays for `null_count`. Strings - // (single-pass) and nested types have no such crossover and never - // reach this arm. - let num_rows = offsets.len() - 1; - if num_rows > 0 - && (offsets[num_rows] - first_offset) / num_rows > NULL_FAST_PATH_MAX_LEN - && visible_values.null_count() > 0 - { - return None; - } - Some(array_has_array_primitive( - visible_values, needle, &offsets, first_offset, combined_nulls, - )) - }, - DataType::Utf8 => Some(array_has_array_string( - visible_values.as_string::(), - needle.as_string::(), - &offsets, - first_offset, - combined_nulls, - )), - DataType::LargeUtf8 => Some(array_has_array_string( - visible_values.as_string::(), - needle.as_string::(), - &offsets, - first_offset, - combined_nulls, - )), - DataType::Utf8View => Some(array_has_array_string_view( - visible_values.as_string_view(), - needle.as_string_view(), - &offsets, - first_offset, - combined_nulls, - )), - _ => None, - } -} - -/// Primitive fast path, with two branches on element validity: +/// Primitive fast path, two branches on element validity: /// -/// 1. No element nulls: a branchless OR-reduction over the raw value slice -- -/// auto-vectorizes, and is the common, fastest case. -/// 2. Element nulls: a null slot's backing value is arbitrary, so validity must -/// be consulted. Compare into an equality bitmap and AND it with the validity -/// bitmap -- one word-parallel op, no per-element branch -- then reduce each -/// row to "any bit set". Chunked so the expanded-needle scratch stays bounded -/// regardless of batch size. +/// 1. No nulls: branchless OR-reduction over the raw slice (auto-vectorizes). +/// 2. Nulls: AND the equality bitmap with validity (a null slot's value is +/// arbitrary), then reduce each row to "any bit set". Chunked to bound the +/// expanded needle. fn array_has_array_primitive( values: &PrimitiveArray, needle: &dyn Array, offsets: &[usize], - first_offset: usize, combined_nulls: Option<&NullBuffer>, ) -> BooleanBuffer where @@ -469,8 +447,8 @@ where } // `needle[i]` is non-null here: combined_nulls covers the needle nulls. let needle_val = needle_slice[i]; - let start = offsets[i] - first_offset; - let end = offsets[i + 1] - first_offset; + let start = offsets[i]; + let end = offsets[i + 1]; value_slice[start..end] .iter() .fold(false, |acc, &v| acc | v.is_eq(needle_val)) @@ -482,15 +460,17 @@ where let mut needle_expanded: Vec = Vec::new(); for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) { let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows); - let elem_start = offsets[chunk_start] - first_offset; - let elem_end = offsets[chunk_end] - first_offset; + let elem_start = offsets[chunk_start]; + let elem_end = offsets[chunk_end]; // Expand the per-row needle across this chunk's elements (reused scratch), // then compare in one vectorizable pass and mask out null elements. needle_expanded.clear(); for i in chunk_start..chunk_end { - needle_expanded - .resize(offsets[i + 1] - first_offset - elem_start, needle_slice[i]); + needle_expanded.extend(std::iter::repeat_n( + needle_slice[i], + offsets[i + 1] - offsets[i], + )); } let chunk_values = &value_slice[elem_start..elem_end]; let eq_bits = BooleanBuffer::collect_bool(chunk_values.len(), |k| { @@ -506,21 +486,19 @@ where result.append(false); continue; } - let start = offsets[i] - first_offset - elem_start; - let end = offsets[i + 1] - first_offset - elem_start; + let start = offsets[i] - elem_start; + let end = offsets[i + 1] - elem_start; result.append(matched.slice(start, end - start).has_true()); } } result.finish() } -/// String implementation of `array_has_array_fast_path`, generic over the -/// concrete string array type (`Utf8`, `LargeUtf8`, `Utf8View`). +/// String fast path, generic over the offset width (`Utf8` / `LargeUtf8`). fn array_has_array_string<'a, S: StringArrayType<'a> + Copy>( values: S, needle: S, offsets: &[usize], - first_offset: usize, combined_nulls: Option<&NullBuffer>, ) -> BooleanBuffer { let num_rows = offsets.len() - 1; @@ -530,29 +508,23 @@ fn array_has_array_string<'a, S: StringArrayType<'a> + Copy>( } // `needle[i]` is non-null here: combined_nulls covers the needle nulls. let needle_val = needle.value(i); - let start = offsets[i] - first_offset; - let end = offsets[i + 1] - first_offset; + let start = offsets[i]; + let end = offsets[i + 1]; // Compare the value first and only consult validity on a match (see the // primitive path for why this is correct and faster on no-match scans). (start..end).any(|k| values.value(k) == needle_val && !values.is_null(k)) }) } -/// View-aware `Utf8View` variant of [`array_has_array_string`]. A `StringView` -/// packs the byte length and a 4-byte prefix into its 128-bit view; Arrow's -/// per-row `eq` kernel compares those before ever touching the data buffer, -/// which the generic `value(k) == needle` path gives up by materializing every -/// element. Here we compare the raw views directly: an inline needle (<= 12 -/// bytes, whose view holds the whole value zero-padded) matches iff the full -/// views are equal -- no materialization at all -- and a longer needle is -/// rejected on length + prefix and only materialized to confirm a candidate. A -/// null slot's backing view is arbitrary, so (as in the primitive path) -/// validity is consulted only on a view match. +/// `Utf8View` variant of [`array_has_array_string`]: compare the packed 128-bit +/// views directly so the length + 4-byte prefix reject non-matches without +/// touching the data buffer, and an inline value matches on the view alone. A +/// longer view is only materialized to confirm a candidate; validity is +/// consulted only on a view match. fn array_has_array_string_view( values: &StringViewArray, needle: &StringViewArray, offsets: &[usize], - first_offset: usize, combined_nulls: Option<&NullBuffer>, ) -> BooleanBuffer { let num_rows = offsets.len() - 1; @@ -565,11 +537,11 @@ fn array_has_array_string_view( // `needle[i]` is non-null here: combined_nulls covers the needle nulls. let needle_view = needle_views[i]; // Low 32 bits are the byte length; the next 32 are the inline prefix. - let needle_inline = (needle_view as u32) <= 12; + let needle_inline = (needle_view as u32) <= MAX_INLINE_VIEW_LEN; let needle_lo = needle_view as u64; let needle_val = needle.value(i); - let start = offsets[i] - first_offset; - let end = offsets[i + 1] - first_offset; + let start = offsets[i]; + let end = offsets[i + 1]; (start..end).any(|k| { let v = value_views[k]; let matched = if needle_inline {