diff --git a/datafusion/functions-nested/benches/array_has.rs b/datafusion/functions-nested/benches/array_has.rs index f5e66d56c0efe..52cc05bf93f34 100644 --- a/datafusion/functions-nested/benches/array_has.rs +++ b/datafusion/functions-nested/benches/array_has.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, Int64Array, ListArray, StringArray}; +use arrow::array::{ + ArrayRef, Int64Array, LargeStringArray, ListArray, StringArray, StringViewArray, +}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field}; use criterion::{ @@ -43,17 +45,224 @@ fn criterion_benchmark(c: &mut Criterion) { for &size in &array_sizes { bench_array_has(c, size); + bench_array_has_array(c, size); bench_array_has_all(c, size); bench_array_has_any(c, size); } // Specific benchmarks for string arrays (common use case) bench_array_has_strings(c); + bench_array_has_array_strings(c); bench_array_has_all_strings(c); bench_array_has_any_strings(c); // Benchmark for array_has_any with one scalar arg bench_array_has_any_scalar(c); + + // Array-needle fast-path profile: null patterns, list length, row height. + bench_array_has_array_null_patterns(c); + bench_array_has_array_by_size(c); + bench_array_has_array_by_rows(c); +} + +/// Invoke `array_has` once with an array (column) needle -- exercises the +/// `array_has_dispatch_for_array` fast path. +fn run_array_needle_case( + group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>, + id: String, + haystack: ArrayRef, + needle: ArrayRef, + rows: usize, +) { + let config_options = Arc::new(ConfigOptions::default()); + let return_field: Arc = Field::new("result", DataType::Boolean, true).into(); + let arg_fields: Vec> = vec![ + Field::new("arr", haystack.data_type().clone(), false).into(), + Field::new("el", needle.data_type().clone(), false).into(), + ]; + let args = vec![ColumnarValue::Array(haystack), ColumnarValue::Array(needle)]; + group.bench_function(id, |b| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: rows, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }); + }); +} + +/// Build a `List` of `array_size` string elements per row (`{prefix}{i}`) with +/// the given element type (`Utf8` / `LargeUtf8` / `Utf8View`) and null density. +/// The prefix controls element length: a short one stays inline in a `Utf8View` +/// (<= 12 bytes), a long one spills to the data buffer. +fn string_list_array( + num_rows: usize, + array_size: usize, + null_density: f64, + prefix: &str, + element_type: &DataType, +) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED); + let data = (0..num_rows * array_size).map(|_| { + if rng.random::() < null_density { + None + } else { + Some(format!("{prefix}{}", rng.random_range(0..array_size))) + } + }); + let values: ArrayRef = match element_type { + DataType::Utf8 => Arc::new(data.collect::()), + DataType::LargeUtf8 => Arc::new(data.collect::()), + DataType::Utf8View => Arc::new(data.collect::()), + other => panic!("unsupported string element type: {other}"), + }; + let offsets = (0..=num_rows) + .map(|i| (i * array_size) as i32) + .collect::>(); + Arc::new( + ListArray::try_new( + Arc::new(Field::new("item", element_type.clone(), true)), + OffsetBuffer::new(offsets.into()), + values, + None, + ) + .unwrap(), + ) +} + +/// Build a string needle column (one value per row) of the given element type. +fn string_value_array( + num_rows: usize, + range: usize, + prefix: &str, + element_type: &DataType, +) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED + 2); + let data = + (0..num_rows).map(|_| Some(format!("{prefix}{}", rng.random_range(0..range)))); + match element_type { + DataType::Utf8 => Arc::new(data.collect::()), + DataType::LargeUtf8 => Arc::new(data.collect::()), + DataType::Utf8View => Arc::new(data.collect::()), + other => panic!("unsupported string element type: {other}"), + } +} + +/// Array needle, fixed list length (64), across null patterns. i64 covers no +/// nulls / found / not-found / all-null / null-fill collision. Each string +/// element type (`Utf8`, `LargeUtf8`, `Utf8View`) covers no nulls / 30% nulls at +/// both short (inline, <= 12 byte) and long (> 12 byte, shared-prefix) element +/// lengths, plus all-null. `not_found` shifts the needle out of the value range. +fn bench_array_has_array_null_patterns(c: &mut Criterion) { + let (rows, size) = (10_000usize, 64usize); + let s = size as i64; + let mut group = c.benchmark_group("array_has_array_null_patterns"); + run_array_needle_case( + &mut group, + "i64/no_nulls".to_string(), + create_int64_list_array(rows, size, 0.0), + create_int64_value_array(rows, s, 0), + rows, + ); + run_array_needle_case( + &mut group, + "i64/nulls30_found".to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, 0), + rows, + ); + run_array_needle_case( + &mut group, + "i64/nulls30_not_found".to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, s), + rows, + ); + run_array_needle_case( + &mut group, + "i64/all_null".to_string(), + create_int64_list_array(rows, size, 1.0), + create_int64_value_array(rows, s, 0), + rows, + ); + run_array_needle_case( + &mut group, + "i64/collision".to_string(), + create_int64_list_array(rows, size, 1.0), + create_int64_value_array(rows, 1, 0), + rows, + ); + // Short elements stay inline in a `Utf8View` (<= 12 bytes); long elements + // share a 4-byte prefix (the realistic case where the view prefix can't + // reject, forcing a buffer compare). `_short` / `_long` labels distinguish + // them; all-null has no content so it is length-independent. + let short = "value_"; // "value_0".."value_63": <= 8 bytes, inline + let long = "long_element_string_value_"; // ~28 bytes, spills to the buffer + for (type_label, element_type) in [ + ("utf8", DataType::Utf8), + ("largeutf8", DataType::LargeUtf8), + ("utf8view", DataType::Utf8View), + ] { + for (len_label, prefix) in [("short", short), ("long", long)] { + for (pat_label, density) in [("no_nulls", 0.0), ("nulls30", 0.3)] { + run_array_needle_case( + &mut group, + format!("{type_label}_{len_label}/{pat_label}"), + string_list_array(rows, size, density, prefix, &element_type), + string_value_array(rows, size, prefix, &element_type), + rows, + ); + } + } + run_array_needle_case( + &mut group, + format!("{type_label}/all_null"), + string_list_array(rows, size, 1.0, short, &element_type), + string_value_array(rows, size, short, &element_type), + rows, + ); + } + group.finish(); +} + +/// Array needle, i64, 30% element nulls, not found, across list lengths. +fn bench_array_has_array_by_size(c: &mut Criterion) { + let rows = 10_000usize; + let mut group = c.benchmark_group("array_has_array_by_size"); + for size in [8usize, 32, 128, 256, 512, 1024] { + let s = size as i64; + run_array_needle_case( + &mut group, + size.to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, s), + rows, + ); + } + group.finish(); +} + +/// Array needle, i64, 8 elems/row, 30% nulls, not found, across row counts. +fn bench_array_has_array_by_rows(c: &mut Criterion) { + let (size, s) = (8usize, 8i64); + let mut group = c.benchmark_group("array_has_array_by_rows"); + for rows in [10_000usize, 100_000, 1_000_000] { + run_array_needle_case( + &mut group, + rows.to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, s), + rows, + ); + } + group.finish(); } fn bench_array_has(c: &mut Criterion, array_size: usize) { @@ -119,6 +328,136 @@ fn bench_array_has(c: &mut Criterion, array_size: usize) { group.finish(); } +/// Benchmarks array_has where the needle is an array (a column with one value +/// per row) rather than a scalar. +fn bench_array_has_array(c: &mut Criterion, array_size: usize) { + let mut group = c.benchmark_group("array_has_array_i64"); + let haystack = create_int64_list_array(NUM_ROWS, array_size, NULL_DENSITY); + let config_options = Arc::new(ConfigOptions::default()); + let return_field: Arc = Field::new("result", DataType::Boolean, true).into(); + let arg_fields: Vec> = vec![ + Field::new("arr", haystack.data_type().clone(), false).into(), + Field::new("el", DataType::Int64, false).into(), + ]; + + // Needle values drawn from the same range as the haystack values, so many + // rows find a match (and the inner loop can short-circuit). + let needle_found = create_int64_value_array(NUM_ROWS, array_size as i64, 0); + let args_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_found), + ]; + group.bench_with_input( + BenchmarkId::new("found", array_size), + &array_size, + |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }, + ); + + // Needle values outside the haystack range: never matches, so every row + // scans its full element list (worst case for the inner loop). + let needle_not_found = + create_int64_value_array(NUM_ROWS, array_size as i64, array_size as i64); + let args_not_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_not_found), + ]; + group.bench_with_input( + BenchmarkId::new("not_found", array_size), + &array_size, + |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_not_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }, + ); + + group.finish(); +} + +fn bench_array_has_array_strings(c: &mut Criterion) { + let mut group = c.benchmark_group("array_has_array_strings"); + let config_options = Arc::new(ConfigOptions::default()); + let return_field: Arc = Field::new("result", DataType::Boolean, true).into(); + + let sizes = vec![10, 100, 500]; + + for &size in &sizes { + let haystack = create_string_list_array(NUM_ROWS, size, NULL_DENSITY); + let arg_fields: Vec> = vec![ + Field::new("arr", haystack.data_type().clone(), false).into(), + Field::new("el", DataType::Utf8, false).into(), + ]; + + let needle_found = create_string_value_array(NUM_ROWS, size, "value_"); + let args_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_found), + ]; + group.bench_with_input(BenchmarkId::new("found", size), &size, |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }); + + let needle_not_found = create_string_value_array(NUM_ROWS, size, "missing_"); + let args_not_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_not_found), + ]; + group.bench_with_input(BenchmarkId::new("not_found", size), &size, |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_not_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }); + } + + group.finish(); +} + fn bench_array_has_all(c: &mut Criterion, array_size: usize) { let mut group = c.benchmark_group("array_has_all"); let haystack = create_int64_list_array(NUM_ROWS, array_size, NULL_DENSITY); @@ -659,6 +998,30 @@ fn create_int64_list_array( ) } +/// Create an `Int64Array` of `num_rows` non-null values in `[offset, offset + +/// range)`, used as an array needle for `array_has`. +fn create_int64_value_array(num_rows: usize, range: i64, offset: i64) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED + 2); + let values = (0..num_rows) + .map(|_| Some(rng.random_range(0..range) + offset)) + .collect::(); + Arc::new(values) +} + +/// Create a `StringArray` of `num_rows` non-null values like "{prefix}{idx}" +/// where `idx` is drawn from `[0, range)`, used as an array needle for +/// `array_has`. +fn create_string_value_array(num_rows: usize, range: usize, prefix: &str) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED + 2); + let values = (0..num_rows) + .map(|_| { + let idx = rng.random_range(0..range); + Some(format!("{prefix}{idx}")) + }) + .collect::(); + Arc::new(values) +} + /// Like `create_int64_list_array` but values are offset so they won't /// appear in a standard list array (useful for "not found" benchmarks). fn create_int64_list_array_with_offset( diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index 04818258f040b..1b449cce6e556 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,230 @@ 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; + // `BooleanBuffer::has_true()` is arrow >= 59 only; use count_set_bits. + result.append(matched.slice(start, end - start).count_set_bits() > 0); + } + } + 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 +1213,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 +1553,448 @@ 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`. Used to prove the fast path is + /// behavior-preserving. Handles `List` haystacks (Int32 offsets). + fn reference_array_has_array( + haystack: &ListArray, + 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 = haystack.value(i); + let is_nested = row.data_type().is_nested(); + let needle_row = Scalar::new(needle.slice(i, 1)); + let eq = super::compare_with_eq(&row, &needle_row, is_nested).unwrap(); + builder.append(eq.has_true()); + } + BooleanArray::new(builder.finish(), combined) + } + + #[test] + fn test_array_has_array_needle_element_nulls() { + // Element nulls, needle nulls, null rows and empty rows, for a primitive + // (Int32) and a string (Utf8) element type. Verified against the + // original per-row `eq` + `has_true` reference. + // + // [1, 2, 3] needle 2 -> true + // [4, 5, 6] needle 9 -> false + // null row needle 5 -> null + // [7, null, 9] needle null -> null (null needle) + // [7, null, 9] needle 7 -> true + // [null] needle 1 -> false (null element never matches) + // [] needle 1 -> false (empty) + let haystack: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(4), Some(5), Some(6)]), + None, + Some(vec![Some(7), None, Some(9)]), + Some(vec![Some(7), None, Some(9)]), + Some(vec![None]), + Some(vec![]), + ])); + let needle: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(2), + Some(9), + Some(5), + None, + Some(7), + Some(1), + Some(1), + ])); + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![ + Some(true), + Some(false), + None, + None, + Some(true), + Some(false), + Some(false), + ] + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + + // ["a", "b"] needle "b" -> true + // [null] needle "z" -> false (null element never matches) + // ["c", "d"] needle "x" -> false + let values = + StringArray::from(vec![Some("a"), Some("b"), None, Some("c"), Some("d")]); + let offsets = OffsetBuffer::new(vec![0, 2, 3, 5].into()); + let haystack: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + offsets, + Arc::new(values), + None, + )); + let needle: ArrayRef = + Arc::new(StringArray::from(vec![Some("b"), Some("z"), Some("x")])); + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(false)] + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + } + + #[test] + fn test_array_has_array_needle_null_fill_collision() { + // A null primitive slot keeps the type default (0) in the value buffer, + // so a needle of 0 collides with it. The null element must still not + // match; only a valid 0 does. + // [0, null] needle 0 -> true (valid 0 matches) + // [null, 5] needle 0 -> false (null slot's backing 0 must not match) + // [null] needle 0 -> false + let haystack: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(0), None]), + Some(vec![None, Some(5)]), + Some(vec![None]), + ])); + let needle: ArrayRef = + Arc::new(Int32Array::from(vec![Some(0), Some(0), Some(0)])); + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(false)] + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + } + + #[test] + fn test_array_has_array_needle_chunk_boundary() { + // The element-null path is processed in ROW_CONVERSION_CHUNK_SIZE (512) + // row chunks; use > 2 chunks with element nulls, null rows and varying + // needles to exercise the chunk rebasing, verified against the oracle. + let n = 1500usize; + let mut rows: Vec>>> = Vec::with_capacity(n); + let mut needles: Vec> = Vec::with_capacity(n); + for i in 0..n { + if i % 13 == 0 { + rows.push(None); // null row + } else { + let r = (i % 7) as i32; + rows.push(Some(vec![Some(r), None, Some((r + 1) % 5)])); + } + needles.push(if i % 11 == 0 { + None + } else { + Some((i % 5) as i32) + }); + } + let haystack: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(rows)); + let needle: ArrayRef = Arc::new(Int32Array::from(needles)); + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + } + + #[test] + fn test_array_has_array_needle_long_lists_with_nulls_bail() { + // Average list length past NULL_FAST_PATH_MAX_LEN with element nulls + // routes to the per-row `eq` kernel; the result must still match the + // oracle. + let rows = 20usize; + let len = super::NULL_FAST_PATH_MAX_LEN + 64; // > threshold + let mut data: Vec>>> = Vec::with_capacity(rows); + let mut needles: Vec> = Vec::with_capacity(rows); + for i in 0..rows { + let row = (0..len) + .map(|j| { + if j % 5 == 0 { + None + } else { + Some(((i + j) % 7) as i32) + } + }) + .collect::>(); + data.push(Some(row)); + needles.push(Some((i % 7) as i32)); + } + let haystack: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(data)); + let needle: ArrayRef = Arc::new(Int32Array::from(needles)); + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + } + + #[test] + fn test_array_has_array_needle_large_list() { + // The fast path also covers LargeList (i64 offsets); exercise it through + // the null branch, including a null-fill collision (needle 0). + // [1, 2] needle 2 -> true + // [null, 3] needle 0 -> false (null slot's backing 0 must not match) + // [4, null] needle 4 -> true + let haystack: ArrayRef = + Arc::new(LargeListArray::from_iter_primitive::( + vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![None, Some(3)]), + Some(vec![Some(4), None]), + ], + )); + let needle: ArrayRef = + Arc::new(Int32Array::from(vec![Some(2), Some(0), Some(4)])); + let result = invoke_array_has_array(haystack, needle); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true)] + ); + } + + #[test] + fn test_array_has_array_needle_string_view() { + // Utf8View has its own view-aware fast path. Exercise inline (<= 12 byte) + // and long (> 12 byte) elements, two long strings sharing a 4-byte prefix + // (the view prefix matches but the full bytes must be confirmed), element + // nulls, an empty row, a null row, and a null-fill collision (needle "" + // against a null slot). Cross-checked against the per-row `eq` oracle. + let rows: Vec>>> = vec![ + Some(vec![Some("a"), Some("bb"), Some("ccc")]), // inline, match "bb" + Some(vec![Some("short"), Some("tiny")]), // inline, no match + Some(vec![Some("this_is_a_long_value_xyz")]), // long, match + Some(vec![Some("prefixAAAA_1111"), Some("prefixAAAA_2222")]), // long, shared prefix + Some(vec![Some("prefixAAAA_1111")]), // long, shared prefix but no match + Some(vec![Some("x"), None, Some("y")]), // element null skipped + Some(vec![None]), // null slot, needle "" must not match + Some(vec![]), // empty row + None, // null row + ]; + let needle_vals = vec![ + Some("bb"), + Some("missing"), + Some("this_is_a_long_value_xyz"), + Some("prefixAAAA_2222"), + Some("prefixAAAA_9999"), + Some("y"), + Some(""), + Some("z"), + Some("q"), + ]; + + let mut flat: Vec> = Vec::new(); + let mut offsets: Vec = vec![0]; + let mut list_nulls: Vec = Vec::new(); + for row in &rows { + match row { + Some(elems) => { + flat.extend(elems.iter().copied()); + list_nulls.push(true); + } + None => list_nulls.push(false), + } + offsets.push(flat.len() as i32); + } + let values = flat.into_iter().collect::(); + let haystack: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new("item", DataType::Utf8View, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(values), + Some(NullBuffer::from(list_nulls)), + )); + let needle: ArrayRef = + Arc::new(needle_vals.into_iter().collect::()); + + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![ + Some(true), // "bb" inline match + Some(false), // no match + Some(true), // long match + Some(true), // shared-prefix long match (second element) + Some(false), // shared prefix, different bytes -> no match + Some(true), // element null skipped, "y" matches + Some(false), // null slot vs "" -> not a match + Some(false), // empty row + None, // null row + ] + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + } + + #[test] + fn test_array_has_array_needle_float_total_order() { + // Arrow `eq` uses total-order float equality: NaN == NaN (same bits) is + // true, and +0.0 == -0.0 is false. The fast path must match this. + let haystack: ArrayRef = + 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)]), + Some(vec![Some(2.0)]), + ])); + let needle: ArrayRef = Arc::new(Float64Array::from(vec![ + f64::NAN, // matches NaN element + -0.0, // does NOT match +0.0 + -0.0, // matches -0.0 + 2.0, // matches + ])); + + let result = invoke_array_has_array(Arc::clone(&haystack), Arc::clone(&needle)); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true), Some(true)] + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(haystack.as_list::(), &needle) + ); + + // Float32 too, to cover the f32 total-order path. + let haystack32: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(f32::NAN), Some(1.0)]), + Some(vec![Some(0.0)]), + Some(vec![Some(-0.0)]), + Some(vec![Some(2.0)]), + ])); + let needle32: ArrayRef = + Arc::new(Float32Array::from(vec![f32::NAN, -0.0, -0.0, 2.0])); + let result32 = invoke_array_has_array(haystack32, needle32); + assert_eq!( + result32.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true), Some(true)] + ); + } + + #[test] + fn test_array_has_array_needle_sliced() { + // Offset normalization for sliced haystacks must keep the element ranges + // and the needle column aligned. Covers both `ArrayWrapper` offset + // branches: `List` (offsets from the buffer) and `FixedSizeList` + // (offsets computed as `i * value_length`). + + // Sliced List haystack AND a sliced needle column (rows 1..=3 visible). + // Verified against the original per-row reference. + 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( + Arc::clone(&sliced_haystack), + Arc::clone(&sliced_needle), + ); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true)] + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(sliced_haystack.as_list::(), &sliced_needle) + ); + + // Sliced FixedSizeList haystack (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)] + ); + } + + #[test] + fn test_array_has_array_needle_sliced_big_backing() { + // A sliced haystack must be routed by its *visible* average length, not + // the full backing child's. Here the backing list has a large child but + // only a few short rows are visible after slicing: the element-null bail + // heuristic must measure the visible region (short rows -> fast path), + // not the full child (`backing_rows * row_len / visible_rows` would + // exceed NULL_FAST_PATH_MAX_LEN and wrongly bail). Result verified + // against the per-row `eq` oracle regardless of which path is taken. + let backing_rows = 200usize; + let row_len = 20usize; // visible rows stay well under the threshold + let mut data: Vec>>> = Vec::with_capacity(backing_rows); + let mut needles: Vec> = Vec::with_capacity(backing_rows); + for i in 0..backing_rows { + let row = (0..row_len) + .map(|j| { + if j % 4 == 0 { + None // element nulls, so the null branch is engaged + } else { + Some(((i + j) % 11) as i32) + } + }) + .collect::>(); + data.push(Some(row)); + needles.push(Some((i % 11) as i32)); + } + let full = ListArray::from_iter_primitive::(data); + let needle_full = Int32Array::from(needles); + + // Visible window: 3 short rows from the middle (first_offset != 0). A + // full-child heuristic sees 200*20 / 3 > NULL_FAST_PATH_MAX_LEN and bails; + // the visible average is row_len (20), which keeps the fast path. + let (start, window) = (50, 3); + let sliced_haystack: ArrayRef = Arc::new(full.slice(start, window)); + let sliced_needle: ArrayRef = Arc::new(needle_full.slice(start, window)); + let result = invoke_array_has_array( + Arc::clone(&sliced_haystack), + Arc::clone(&sliced_needle), + ); + assert_eq!( + result.as_boolean(), + &reference_array_has_array(sliced_haystack.as_list::(), &sliced_needle) + ); + } }