Skip to content

Commit bc6e058

Browse files
perf: Optimize array_has() for array needle (#23337)
## Which issue does this PR close? - Part of #23334. > The numbers below come from the committed criterion benchmark added in #23335 (`cargo bench --bench array_has`) — **origin** = the per-row `eq` kernel (unoptimized `main` / #23335), **now** = with this optimization applied. Run the bench on `main` and on this branch to reproduce. Full disclosure - this was heavily assisted by AI, and I did my best to understand and justify every change here before submitting. ## Rationale for this change `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.) 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. All numbers below are from the committed criterion benchmark (`cargo bench --bench array_has`, groups `array_has_array_null_patterns` / `array_has_array_by_size` / `array_has_array_by_rows`): the `array_has` UDF evaluated in isolation with an array needle, **origin** (the per-row `eq` kernel) vs **now**. "list length" is the number of elements in each row's array (not the row count). Not end-to-end query time. ### By data type and null pattern (list length 64, 10K rows) | element | element len | null pattern | origin | now | speedup | |-----------|----------------|----------------------|---------|---------|---------| | i64 | - | no nulls, found | 1.10 ms | 73 µs | 15.1x | | i64 | - | no nulls, not found | 1.07 ms | 72 µs. | 14.9x | | i64 | - | 30% nulls, found | 1.17 ms | 315 µs | 3.7x | | i64 | - | 30% nulls, not found | 1.10 ms | 274 µs | 4.0x | | i64 | - | all null | 1.10 ms | 272 µs | 4.0x | | i64 | - | collision | 1.10 ms | 270 µs | 4.1x | | Utf8 | short (inline) | no nulls | 2.57 ms | 1.01 ms | 2.5x | | Utf8 | short (inline) | 30% nulls | 3.37 ms | 1.52 ms | 2.2x | | Utf8 | long (>12B) | no nulls | 2.61 ms | 1.04 ms | 2.5x | | Utf8 | long (>12B) | 30% nulls | 3.31 ms | 1.52 ms | 2.2x | | Utf8 | - | all null | 1.26 ms | 256 µs | 4.9x | | LargeUtf8 | short (inline) | no nulls | 2.56 ms | 1.02 ms | 2.5x | | LargeUtf8 | short (inline) | 30% nulls | 3.20 ms | 1.54 ms | 2.1x | | LargeUtf8 | long (>12B) | no nulls | 2.67 ms | 1.05 ms | 2.6x | | LargeUtf8 | long (>12B) | 30% nulls | 3.42 ms | 1.59 ms | 2.2x | | LargeUtf8 | - | all null | 1.31 ms | 263 µs | 5.0x | | Utf8View | short (inline) | no nulls | 1.18 ms | 239 µs | 4.9x | | Utf8View | short (inline) | 30% nulls | 1.26 ms | 246 µs | 5.1x | | Utf8View | long (>12B) | no nulls | 2.86 ms | 1.17 ms | 2.4x | | Utf8View | long (>12B) | 30% nulls | 3.51 ms | 1.66 ms | 2.1x | | Utf8View | - | all null | 1.20 ms | 267 µs | 4.5x | The i64 null cases are uniform (~4x) whether the match is present, absent, the whole list is null, or the needle collides with a null slot's backing fill value — validity is folded in with one word-parallel op, so there is no per-row rescan and no null slot can match. Strings win ~2.1–2.5x mainly by dropping the per-row `BooleanArray` allocation. `Utf8View` additionally uses a view-aware compare: the byte length and 4-byte prefix packed into the 128-bit view reject non-matches before touching the data buffer, and an inline value (≤ 12 bytes) is matched by whole-view equality with no materialization at all — hence ~5x on short/inline strings. When long strings share a prefix (e.g. ARNs) the prefix can't reject, so `Utf8View` falls in line with the other string types (~2.1–2.4x). No string case regresses. ### By list length (i64, 30% element nulls, not found, 10K rows) | elems/row | origin | now | speedup | |-----------|---------|---------|--------------------------------------| | 8 | 1.03 ms | 111 µs | 9.3x | | 32 | 1.07 ms | 197 µs | 5.5x | | 128 | 1.18 ms | 446 µs | 2.6x | | 256 | 1.28 ms | 780 µs | 1.6x | | 512 | 1.54 ms | 1.44 ms | 1.1x | | 1024 | 2.17 ms | 2.15 ms | 1.0x (falls back to per-row kernel) | The element-null branch makes a few passes over the values; past a moderate average list length (`NULL_FAST_PATH_MAX_LEN`) the per-row kernel wins, so it bails to it there — no meaningful regression. That average is measured over the visible (sliced) region, so a sliced array's hidden child elements can't route a small window to the slow path. The all-valid fold has no such crossover. ### By row count (i64, 8 elems/row, 30% nulls, not found) | rows | origin | now | speedup | |------|-----------|----------|---------| | 10K | 1.04 ms | 111 µs | 9.4x | | 100K | 10.42 ms | 1.09 ms | 9.6x | | 1M | 102.68 ms | 10.91 ms | 9.4x | Invariant to the number of rows — the per-row overhead removed is a fixed cost, so absolute savings scale linearly with the column height. The remaining benchmarks in the suite (scalar `array_has`, `array_has_all`, `array_has_any` — paths this PR does not touch) are unchanged (median 0.99x, within measurement noise), confirming no regression outside the array-needle path. ### End-to-end (context) 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. ## What changes are included in this PR? A fast path for primitive and string element types in `array_has_dispatch_for_array`, preserving the Arrow `eq` kernel semantics (total-order float equality; null elements never match): - **All-valid elements:** each row is a single branchless OR-reduction over the raw native value slice (auto-vectorizes; the common case). - **Element nulls:** a null slot's backing value is arbitrary, so the per-element equality bitmap is ANDed with the validity bitmap (one word-parallel op, no per-element branch) before reducing each row to "any bit set", a null slot can never match regardless of its value. This branch is processed in row chunks so the scratch buffer stays bounded, and past `NULL_FAST_PATH_MAX_LEN` average elements/row a length check over the visible (sliced) region bails to the per-row kernel (see the list-length table). - **String elements:** each row is a single pass over the row's values (compare, then consult validity only on a match). `Utf8View` compares the packed 128-bit views directly — length + 4-byte prefix reject non-matches before any data-buffer access, and an inline value (≤ 12 bytes) matches by whole-view equality with no materialization. - **Nested (and any other) element types** keep using the per-row `eq` kernel. The array-needle benchmarks used for the numbers above are added in #3 (null patterns, list length, and row count). ## Are these changes tested? Yes: - New unit tests for the array-needle path covering element nulls, the null-fill collision (needle equal to a null slot's backing value), total-order float equality (`NaN` / `-0.0`), sliced arrays (including a small visible window over a large backing child), `LargeList` offsets, empty rows, a multi-chunk input, and a long-list input that exercises the per-row fallback, each cross-checked against the original per-row `eq` kernel as an oracle. - Existing `array_has` / `array_contains` / `join_lists` sqllogictest suites pass. ## Are there any user-facing changes? No. --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 7f6cc60 commit bc6e058

2 files changed

Lines changed: 392 additions & 3 deletions

File tree

datafusion/functions-nested/src/array_has.rs

Lines changed: 278 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,11 +18,13 @@
1818
//! [`ScalarUDFImpl`] definitions for array_has, array_has_all and array_has_any functions.
1919
2020
use arrow::array::{
21-
Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, Datum, Scalar,
22-
StringArrayType,
21+
Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray,
22+
BooleanBufferBuilder, Datum, MAX_INLINE_VIEW_LEN, PrimitiveArray, Scalar,
23+
StringArrayType, StringViewArray,
2324
};
24-
use arrow::buffer::{BooleanBuffer, NullBuffer};
25+
use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer};
2526
use arrow::datatypes::DataType;
27+
use arrow::downcast_primitive_array;
2628
use arrow::row::{RowConverter, Rows, SortField};
2729
use datafusion_common::cast::{as_fixed_size_list_array, as_generic_list_array};
2830
use datafusion_common::utils::string_utils::string_array_to_vec;
@@ -323,11 +325,85 @@ impl<'a> ArrayWrapper<'a> {
323325
}
324326
}
325327

328+
/// Evaluate `array_has` with an array (per-row) needle.
329+
///
330+
/// Primitive and string element types take a per-type fast path; nested (and any
331+
/// other) element types fall back to the per-row `eq` kernel, which allocates a
332+
/// `BooleanArray` per row.
326333
fn array_has_dispatch_for_array<'a>(
327334
haystack: ArrayWrapper<'a>,
328335
needle: &ArrayRef,
329336
) -> Result<ArrayRef> {
330337
let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls());
338+
let needle = needle.as_ref();
339+
340+
// Rebase offsets to 0 with `OffsetBuffer::subtract` so `offsets[i]` indexes
341+
// `visible_values` directly (the haystack may be a sliced list).
342+
let raw = OffsetBuffer::new(
343+
haystack
344+
.offsets()
345+
.map(|o| o as i64)
346+
.collect::<Vec<_>>()
347+
.into(),
348+
);
349+
let first = raw[0];
350+
let visible_values = haystack
351+
.values()
352+
.slice(first as usize, (raw[raw.len() - 1] - first) as usize);
353+
let visible_values = visible_values.as_ref();
354+
let offsets: Vec<usize> = raw.subtract(first).iter().map(|&o| o as usize).collect();
355+
356+
// Fast path for primitive/string elements whose (coerced) type matches the
357+
// needle; a type mismatch or a nested type falls through to the per-row kernel.
358+
let fast_path = if visible_values.data_type() != needle.data_type() {
359+
None
360+
} else {
361+
downcast_primitive_array! {
362+
visible_values => {
363+
// The element-null path makes several passes over the values, so
364+
// past a large average list length the per-row `eq` kernel is
365+
// faster -- bail to it. The single-pass all-valid path has no such
366+
// crossover, so only bail when elements are null.
367+
let num_rows = offsets.len() - 1;
368+
if num_rows > 0
369+
&& offsets[num_rows] / num_rows > NULL_FAST_PATH_MAX_LEN
370+
&& visible_values.null_count() > 0
371+
{
372+
None
373+
} else {
374+
Some(array_has_array_primitive(
375+
visible_values, needle, &offsets,
376+
combined_nulls.as_ref(),
377+
))
378+
}
379+
},
380+
DataType::Utf8 => Some(array_has_array_string(
381+
visible_values.as_string::<i32>(),
382+
needle.as_string::<i32>(),
383+
&offsets,
384+
combined_nulls.as_ref(),
385+
)),
386+
DataType::LargeUtf8 => Some(array_has_array_string(
387+
visible_values.as_string::<i64>(),
388+
needle.as_string::<i64>(),
389+
&offsets,
390+
combined_nulls.as_ref(),
391+
)),
392+
DataType::Utf8View => Some(array_has_array_string_view(
393+
visible_values.as_string_view(),
394+
needle.as_string_view(),
395+
&offsets,
396+
combined_nulls.as_ref(),
397+
)),
398+
_ => None,
399+
}
400+
};
401+
402+
if let Some(values) = fast_path {
403+
return Ok(Arc::new(BooleanArray::new(values, combined_nulls)));
404+
}
405+
406+
// Fallback: per-row `eq` kernel (nested element types, or a type mismatch).
331407
let mut result = BooleanBufferBuilder::new(haystack.len());
332408
for (i, arr) in haystack.iter().enumerate() {
333409
if combined_nulls.as_ref().is_some_and(|n| n.is_null(i)) {
@@ -344,6 +420,146 @@ fn array_has_dispatch_for_array<'a>(
344420
Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls)))
345421
}
346422

423+
/// Average list length past which the element-null path loses to the per-row
424+
/// `eq` kernel and bails to it (empirically measured).
425+
const NULL_FAST_PATH_MAX_LEN: usize = 512;
426+
427+
/// Primitive fast path, two branches on element validity:
428+
///
429+
/// 1. No nulls: branchless OR-reduction over the raw slice (auto-vectorizes).
430+
/// 2. Nulls: AND the equality bitmap with validity (a null slot's value is
431+
/// arbitrary), then reduce each row to "any bit set". Chunked to bound the
432+
/// expanded needle.
433+
fn array_has_array_primitive<T: ArrowPrimitiveType>(
434+
values: &PrimitiveArray<T>,
435+
needle: &dyn Array,
436+
offsets: &[usize],
437+
combined_nulls: Option<&NullBuffer>,
438+
) -> BooleanBuffer
439+
where
440+
T::Native: ArrowNativeTypeOp,
441+
{
442+
let needle = needle.as_primitive::<T>();
443+
let num_rows = offsets.len() - 1;
444+
let value_slice = values.values();
445+
let needle_slice = needle.values();
446+
447+
let Some(element_nulls) = values.nulls() else {
448+
return BooleanBuffer::collect_bool(num_rows, |i| {
449+
if combined_nulls.is_some_and(|n| n.is_null(i)) {
450+
return false;
451+
}
452+
// `needle[i]` is non-null here: combined_nulls covers the needle nulls.
453+
let needle_val = needle_slice[i];
454+
let start = offsets[i];
455+
let end = offsets[i + 1];
456+
value_slice[start..end]
457+
.iter()
458+
.fold(false, |acc, &v| acc | v.is_eq(needle_val))
459+
});
460+
};
461+
462+
// Case 2 (see fn doc), chunked like the all/any kernels.
463+
let mut result = BooleanBufferBuilder::new(num_rows);
464+
let mut needle_expanded: Vec<T::Native> = Vec::new();
465+
for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) {
466+
let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows);
467+
let elem_start = offsets[chunk_start];
468+
let elem_end = offsets[chunk_end];
469+
470+
// Expand the per-row needle across this chunk's elements (reused scratch),
471+
// then compare in one vectorizable pass and mask out null elements.
472+
needle_expanded.clear();
473+
for i in chunk_start..chunk_end {
474+
needle_expanded.extend(std::iter::repeat_n(
475+
needle_slice[i],
476+
offsets[i + 1] - offsets[i],
477+
));
478+
}
479+
let chunk_values = &value_slice[elem_start..elem_end];
480+
let eq_bits = BooleanBuffer::collect_bool(chunk_values.len(), |k| {
481+
chunk_values[k].is_eq(needle_expanded[k])
482+
});
483+
let matched = &eq_bits
484+
& &element_nulls
485+
.inner()
486+
.slice(elem_start, elem_end - elem_start);
487+
488+
for i in chunk_start..chunk_end {
489+
if combined_nulls.is_some_and(|n| n.is_null(i)) {
490+
result.append(false);
491+
continue;
492+
}
493+
let start = offsets[i] - elem_start;
494+
let end = offsets[i + 1] - elem_start;
495+
result.append(matched.slice(start, end - start).has_true());
496+
}
497+
}
498+
result.finish()
499+
}
500+
501+
/// String fast path, generic over the offset width (`Utf8` / `LargeUtf8`).
502+
fn array_has_array_string<'a, S: StringArrayType<'a> + Copy>(
503+
values: S,
504+
needle: S,
505+
offsets: &[usize],
506+
combined_nulls: Option<&NullBuffer>,
507+
) -> BooleanBuffer {
508+
let num_rows = offsets.len() - 1;
509+
BooleanBuffer::collect_bool(num_rows, |i| {
510+
if combined_nulls.is_some_and(|n| n.is_null(i)) {
511+
return false;
512+
}
513+
// `needle[i]` is non-null here: combined_nulls covers the needle nulls.
514+
let needle_val = needle.value(i);
515+
let start = offsets[i];
516+
let end = offsets[i + 1];
517+
// Compare the value first and only consult validity on a match (see the
518+
// primitive path for why this is correct and faster on no-match scans).
519+
(start..end).any(|k| values.value(k) == needle_val && !values.is_null(k))
520+
})
521+
}
522+
523+
/// `Utf8View` variant of [`array_has_array_string`]: compare the packed 128-bit
524+
/// views directly so the length + 4-byte prefix reject non-matches without
525+
/// touching the data buffer, and an inline value matches on the view alone. A
526+
/// longer view is only materialized to confirm a candidate; validity is
527+
/// consulted only on a view match.
528+
fn array_has_array_string_view(
529+
values: &StringViewArray,
530+
needle: &StringViewArray,
531+
offsets: &[usize],
532+
combined_nulls: Option<&NullBuffer>,
533+
) -> BooleanBuffer {
534+
let num_rows = offsets.len() - 1;
535+
let value_views = values.views();
536+
let needle_views = needle.views();
537+
BooleanBuffer::collect_bool(num_rows, |i| {
538+
if combined_nulls.is_some_and(|n| n.is_null(i)) {
539+
return false;
540+
}
541+
// `needle[i]` is non-null here: combined_nulls covers the needle nulls.
542+
let needle_view = needle_views[i];
543+
// Low 32 bits are the byte length; the next 32 are the inline prefix.
544+
let needle_inline = (needle_view as u32) <= MAX_INLINE_VIEW_LEN;
545+
let needle_lo = needle_view as u64;
546+
let needle_val = needle.value(i);
547+
let start = offsets[i];
548+
let end = offsets[i + 1];
549+
(start..end).any(|k| {
550+
let v = value_views[k];
551+
let matched = if needle_inline {
552+
// Inline: the whole view is the canonical value (zero padded).
553+
v == needle_view
554+
} else {
555+
// Longer: reject on length + prefix, then confirm the bytes.
556+
(v as u64) == needle_lo && values.value(k) == needle_val
557+
};
558+
matched && !values.is_null(k)
559+
})
560+
})
561+
}
562+
347563
fn array_has_dispatch_for_scalar(
348564
haystack: ArrayWrapper<'_>,
349565
needle: &dyn Datum,
@@ -1311,4 +1527,63 @@ mod tests {
13111527
&[Some(true), Some(true)],
13121528
);
13131529
}
1530+
1531+
/// Invoke `array_has` with the needle as an array (a column with one value
1532+
/// per row). This exercises `array_has_dispatch_for_array` and its fast path.
1533+
fn invoke_array_has_array(haystack: ArrayRef, needle: ArrayRef) -> ArrayRef {
1534+
let num_rows = haystack.len();
1535+
let haystack_type = haystack.data_type().clone();
1536+
let needle_type = needle.data_type().clone();
1537+
ArrayHas::new()
1538+
.invoke_with_args(ScalarFunctionArgs {
1539+
args: vec![ColumnarValue::Array(haystack), ColumnarValue::Array(needle)],
1540+
arg_fields: vec![
1541+
Arc::new(Field::new("haystack", haystack_type, false)),
1542+
Arc::new(Field::new("needle", needle_type, false)),
1543+
],
1544+
number_rows: num_rows,
1545+
return_field: Arc::new(Field::new("return", DataType::Boolean, true)),
1546+
config_options: Arc::new(ConfigOptions::default()),
1547+
})
1548+
.unwrap()
1549+
.into_array(num_rows)
1550+
.unwrap()
1551+
}
1552+
1553+
#[test]
1554+
fn test_array_has_array_needle_sliced() {
1555+
// Offset normalization for sliced haystacks must keep the element ranges
1556+
// and the needle column aligned, for both `List` (offsets from the
1557+
// buffer) and `FixedSizeList` (offsets computed as `i * value_length`).
1558+
// Slicing is an execution artifact SQL/SLT can't force, so this stays a
1559+
// unit test; value-level behavior is covered by `array/array_has.slt`.
1560+
let full = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
1561+
Some(vec![Some(1), Some(2)]),
1562+
Some(vec![Some(10), Some(20), Some(30)]), // needle 20 -> true
1563+
Some(vec![Some(40)]), // needle 41 -> false
1564+
Some(vec![Some(50), Some(60)]), // needle 60 -> true
1565+
Some(vec![Some(70)]),
1566+
]);
1567+
let sliced_haystack: ArrayRef = Arc::new(full.slice(1, 3));
1568+
let sliced_needle: ArrayRef =
1569+
Arc::new(Int32Array::from(vec![999, 20, 41, 60, 999]).slice(1, 3));
1570+
let result = invoke_array_has_array(sliced_haystack, sliced_needle);
1571+
assert_eq!(
1572+
result.as_boolean().iter().collect::<Vec<_>>(),
1573+
vec![Some(true), Some(false), Some(true)]
1574+
);
1575+
1576+
// Sliced FixedSizeList (width 2; rows 1..=2 of
1577+
// [[1,2],[11,12],[21,22],[31,32]] visible) with an aligned needle column.
1578+
let field = Arc::new(Field::new("item", DataType::Int32, true));
1579+
let fsl_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32]));
1580+
let fsl: ArrayRef =
1581+
Arc::new(FixedSizeListArray::new(field, 2, fsl_values, None).slice(1, 2));
1582+
let needle: ArrayRef = Arc::new(Int32Array::from(vec![11, 99]));
1583+
let result = invoke_array_has_array(fsl, needle);
1584+
assert_eq!(
1585+
result.as_boolean().iter().collect::<Vec<_>>(),
1586+
vec![Some(true), Some(false)]
1587+
);
1588+
}
13141589
}

0 commit comments

Comments
 (0)