gabotechs commented on code in PR #23337:
URL: https://github.com/apache/datafusion/pull/23337#discussion_r3554443554
##########
datafusion/functions-nested/src/array_has.rs:
##########
@@ -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<BooleanBuffer> {
+ 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<usize> = 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::<i32>(),
+ needle.as_string::<i32>(),
+ &offsets,
+ first_offset,
+ combined_nulls,
+ )),
+ DataType::LargeUtf8 => Some(array_has_array_string(
+ visible_values.as_string::<i64>(),
+ needle.as_string::<i64>(),
+ &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<T: ArrowPrimitiveType>(
+ values: &PrimitiveArray<T>,
+ needle: &dyn Array,
+ offsets: &[usize],
+ first_offset: usize,
+ combined_nulls: Option<&NullBuffer>,
+) -> BooleanBuffer
+where
+ T::Native: ArrowNativeTypeOp,
+{
+ let needle = needle.as_primitive::<T>();
+ 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<T::Native> = 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;
Review Comment:
I spent a fair amount of time looking at this not knowing what's happening
here.
It might be better to not hardcode the 12 and just use the exported constant
from arrow:
```rust
let needle_inline = (needle_view as u32) <=
arrow::array::MAX_INLINE_VIEW_LEN;
```
##########
datafusion/functions-nested/src/array_has.rs:
##########
@@ -323,11 +325,26 @@ impl<'a> ArrayWrapper<'a> {
}
}
+/// Evaluate `array_has` with an array (per-row) needle.
Review Comment:
After reading this, my first instinct is that this should not be contributed
to this repo.
I see that this PR is pretty much implement an `array_has` arrow kernel from
scratch, and the code contributed here is not even related to DataFusion, it's
just plain Arrow.
However, looking at the rest of this file... I think that ship sailed some
time ago. This file was already filled with code that looks like `arrow-rs`
could have been a better fit.
So nothing to do here I guess.
##########
datafusion/functions-nested/src/array_has.rs:
##########
@@ -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<BooleanBuffer> {
+ 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<usize> = 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;
Review Comment:
I've been starting for longer that I'd like to admit at this comment and I
have to say I don't understand it. For example, I struggle to understand what
"backing child" means, what's a "hidden element", or what is "The average
check", or what's the "all-valid win path".
Any chance of maybe rewording this comment? I think it's trying to make the
justification of what the `(offsets[num_rows] - first_offset) / num_rows >
NULL_FAST_PATH_MAX_LEN`, which is actually a question that I had that I was not
able to understand with this comment.
##########
datafusion/functions-nested/src/array_has.rs:
##########
@@ -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<ArrayRef> {
let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls());
+
+ if let Some(values) =
+ array_has_array_fast_path(&haystack, needle, combined_nulls.as_ref())
Review Comment:
`_fast_path` is probably not the best name for this.
`array_has_array_fast_path` describes the current motivation/performance
role, but not what the function actually does.
I'd even go beyond this and avoid some callstack nesting by restructuring
this function to something that requires less jumps from readers to understand.
##########
datafusion/functions-nested/src/array_has.rs:
##########
@@ -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<ArrayRef> {
let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls());
+
+ if let Some(values) =
+ array_has_array_fast_path(&haystack, needle, combined_nulls.as_ref())
Review Comment:
I'd probably even remove the `array_has_array_fast_path` extra function jump
at all, and integrate its body inside this `array_has_dispatch_for_array`
function.
In the end, it's not possible to understand what
`array_has_dispatch_for_array` without navigating to the body of
`array_has_array_fast_path`, as the `array_has_array_fast_path` name and
signature are not informative enough to know what they do without reading its
contents.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
