diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index 04818258f040b..ebca49285a458 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, @@ -1311,4 +1551,63 @@ 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() + } + + #[test] + 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)]), // needle 20 -> true + Some(vec![Some(40)]), // needle 41 -> false + Some(vec![Some(50), Some(60)]), // needle 60 -> true + Some(vec![Some(70)]), + ]); + 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 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 result = invoke_array_has_array(fsl, needle); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false)] + ); + } } 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