Jefffrey commented on code in PR #23337:
URL: https://github.com/apache/datafusion/pull/23337#discussion_r3555887074
##########
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))
Review Comment:
is it better to use `any()` here or does LLVM generate equivalent code anyway
##########
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();
Review Comment:
might consider using the new `OffsetBuffer::subtract` to normalize the
offsets instead of needing to handle the offset from the first offset
everywhere downstream, see
- https://github.com/apache/datafusion/pull/23424
##########
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:
we can omit detail about null_count and string/nested types here. null_count
check is cheap anyway (its always precomputed on array creation), and
string/other types are apparent in the below arms
##########
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]);
Review Comment:
```suggestion
needle_expanded.extend(std::iter::repeat_n(
needle_slice[i],
offsets[i + 1] - offsets[i],
));
```
thoughts on if this is easier to read?
--
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]
