You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/05/08 11:57:57 UTC
(datafusion) branch gh-readonly-queue/main/pr-22070-bdb1c83a74c89d53fc6531967f6618f87fa11dd2 created (now 53517af674)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22070-bdb1c83a74c89d53fc6531967f6618f87fa11dd2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 53517af674 perf: Use `NullBuffer::union_many` (#22070)
This branch includes the following new commits:
new 53517af674 perf: Use `NullBuffer::union_many` (#22070)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: perf: Use `NullBuffer::union_many` (#22070)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22070-bdb1c83a74c89d53fc6531967f6618f87fa11dd2
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 53517af674e2540805dab0f688f324537ada3220
Author: Neil Conway <ne...@gmail.com>
AuthorDate: Fri May 8 07:57:34 2026 -0400
perf: Use `NullBuffer::union_many` (#22070)
## Which issue does this PR close?
- Closes #21880.
## Rationale for this change
Now that `NullBuffer::union_many` has landed in arrow-rs 58.2, we can
use it to simplify and speedup various places in DataFusion that were
union'ing together > 2 null buffers. In such cases, `union_many` avoids
the intermediate allocations that are done by repeated `union` calls.
## What changes are included in this PR?
* Adopt `NullBuffer::union_many` throughout the codebase
* Along the way, clean up `compute_null_mask` in the Spark crate to be
more readable and avoid needlessly padding scalars out to arrays just to
compute the null mask
## Are these changes tested?
Yes, covered by existing tests.
## Are there any user-facing changes?
No.
---
datafusion/functions-nested/src/extract.rs | 6 +-
datafusion/functions/src/datetime/make_time.rs | 7 ++-
datafusion/functions/src/string/replace.rs | 14 +++--
datafusion/functions/src/unicode/substr.rs | 21 ++++---
datafusion/functions/src/unicode/substrindex.rs | 14 +++--
datafusion/spark/src/function/array/repeat.rs | 2 +-
datafusion/spark/src/function/map/str_to_map.rs | 15 +++--
datafusion/spark/src/function/null_utils.rs | 80 ++++++++++---------------
datafusion/spark/src/function/string/concat.rs | 2 +-
9 files changed, 76 insertions(+), 85 deletions(-)
diff --git a/datafusion/functions-nested/src/extract.rs b/datafusion/functions-nested/src/extract.rs
index 12cca0e52a..fe1e31e8d5 100644
--- a/datafusion/functions-nested/src/extract.rs
+++ b/datafusion/functions-nested/src/extract.rs
@@ -602,14 +602,12 @@ fn combine_input_nulls(
to_array: &Int64Array,
stride: Option<&Int64Array>,
) -> Option<NullBuffer> {
- [
+ NullBuffer::union_many([
array.nulls(),
from_array.nulls(),
to_array.nulls(),
stride.and_then(|s| s.nulls()),
- ]
- .into_iter()
- .fold(None, |acc, nulls| NullBuffer::union(acc.as_ref(), nulls))
+ ])
}
fn general_array_slice<O: OffsetSizeTrait>(
diff --git a/datafusion/functions/src/datetime/make_time.rs b/datafusion/functions/src/datetime/make_time.rs
index 718612fcad..d9e827ac23 100644
--- a/datafusion/functions/src/datetime/make_time.rs
+++ b/datafusion/functions/src/datetime/make_time.rs
@@ -142,10 +142,11 @@ impl ScalarUDFImpl for MakeTimeFunc {
let minutes = minutes.as_primitive::<Int32Type>();
let seconds = seconds.as_primitive::<Int32Type>();
- let nulls = NullBuffer::union(
- NullBuffer::union(hours.nulls(), minutes.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ hours.nulls(),
+ minutes.nulls(),
seconds.nulls(),
- );
+ ]);
let mut values = Vec::with_capacity(len);
for i in 0..len {
diff --git a/datafusion/functions/src/string/replace.rs b/datafusion/functions/src/string/replace.rs
index f588b8c3cc..b980baeb99 100644
--- a/datafusion/functions/src/string/replace.rs
+++ b/datafusion/functions/src/string/replace.rs
@@ -165,10 +165,11 @@ fn replace_view(args: &[ArrayRef]) -> Result<ArrayRef> {
let len = string_array.len();
let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(len, 0);
let mut buffer = String::new();
- let nulls = NullBuffer::union(
- NullBuffer::union(string_array.nulls(), from_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_array.nulls(),
+ from_array.nulls(),
to_array.nulls(),
- );
+ ]);
// Hoist the nulls.is_some() check out of the loop. LLVM does not always
// unswitch this loop on its own (the Utf8View body is large enough to
@@ -212,10 +213,11 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
let len = string_array.len();
let mut builder = GenericStringArrayBuilder::<T>::with_capacity(len, 0);
let mut buffer = String::new();
- let nulls = NullBuffer::union(
- NullBuffer::union(string_array.nulls(), from_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_array.nulls(),
+ from_array.nulls(),
to_array.nulls(),
- );
+ ]);
// Hoist the nulls.is_some() check out of the loop. LLVM unswitches this
// automatically today, but kept explicit so the no-nulls fast path is not
diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs
index cc0882cf30..951fc4c6cd 100644
--- a/datafusion/functions/src/unicode/substr.rs
+++ b/datafusion/functions/src/unicode/substr.rs
@@ -279,10 +279,11 @@ fn string_view_substr(
enable_ascii_fast_path(&string_view_array, start_array, count_array_opt);
// Combine null bitmaps from all inputs in bulk.
- let nulls = NullBuffer::union(
- NullBuffer::union(string_view_array.nulls(), start_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_view_array.nulls(),
+ start_array.nulls(),
count_array_opt.and_then(|a| a.nulls()),
- );
+ ]);
let mut views_buf = Vec::with_capacity(string_view_array.len());
@@ -363,10 +364,11 @@ fn generic_string_substr<T: OffsetSizeTrait>(
let mut has_out_of_line = false;
// Combine null bitmaps from all inputs in bulk.
- let nulls = NullBuffer::union(
- NullBuffer::union(string_array.nulls(), start_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_array.nulls(),
+ start_array.nulls(),
count_array_opt.and_then(|a| a.nulls()),
- );
+ ]);
for i in 0..string_array.len() {
if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
@@ -418,10 +420,11 @@ fn generic_string_substr_copy<T: OffsetSizeTrait>(
let is_ascii = enable_ascii_fast_path(&string_array, start_array, count_array_opt);
// Combine null bitmaps from all inputs in bulk.
- let nulls = NullBuffer::union(
- NullBuffer::union(string_array.nulls(), start_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_array.nulls(),
+ start_array.nulls(),
count_array_opt.and_then(|a| a.nulls()),
- );
+ ]);
let mut result_builder = StringViewBuilder::new();
diff --git a/datafusion/functions/src/unicode/substrindex.rs b/datafusion/functions/src/unicode/substrindex.rs
index 3aeff1b476..d122a34a9f 100644
--- a/datafusion/functions/src/unicode/substrindex.rs
+++ b/datafusion/functions/src/unicode/substrindex.rs
@@ -273,10 +273,11 @@ where
{
let num_rows = string_array.len();
// Output is null if and only if any input is null.
- let nulls = NullBuffer::union(
- NullBuffer::union(string_array.nulls(), delimiter_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_array.nulls(),
+ delimiter_array.nulls(),
count_array.nulls(),
- );
+ ]);
for i in 0..num_rows {
if nulls.as_ref().is_some_and(|n| n.is_null(i)) {
@@ -299,10 +300,11 @@ fn substr_index_view(
delimiter_array: &StringViewArray,
count_array: &PrimitiveArray<Int64Type>,
) -> Result<ArrayRef> {
- let nulls = NullBuffer::union(
- NullBuffer::union(string_array.nulls(), delimiter_array.nulls()).as_ref(),
+ let nulls = NullBuffer::union_many([
+ string_array.nulls(),
+ delimiter_array.nulls(),
count_array.nulls(),
- );
+ ]);
let views = string_array.views();
let mut views_buf = Vec::with_capacity(string_array.len());
let mut has_out_of_line = false;
diff --git a/datafusion/spark/src/function/array/repeat.rs b/datafusion/spark/src/function/array/repeat.rs
index c4e5e449c7..da9b19a768 100644
--- a/datafusion/spark/src/function/array/repeat.rs
+++ b/datafusion/spark/src/function/array/repeat.rs
@@ -100,7 +100,7 @@ fn spark_array_repeat(args: ScalarFunctionArgs) -> Result<ColumnarValue> {
let return_type = return_field.data_type().clone();
// A NULL element should be repeated into the array, not cause a NULL result.
- let null_mask = compute_null_mask(&arg_values[1..], number_rows)?;
+ let null_mask = compute_null_mask(&arg_values[1..]);
// If count is null then return NULL immediately
if matches!(null_mask, NullMaskResolution::ReturnNull) {
diff --git a/datafusion/spark/src/function/map/str_to_map.rs b/datafusion/spark/src/function/map/str_to_map.rs
index c603e775a6..d0f4cf03cd 100644
--- a/datafusion/spark/src/function/map/str_to_map.rs
+++ b/datafusion/spark/src/function/map/str_to_map.rs
@@ -185,14 +185,13 @@ fn str_to_map_impl<'a, V: StringArrayType<'a> + Copy>(
let num_rows = text_array.len();
// Precompute combined null buffer from all input arrays.
- // NullBuffer::union performs a bitmap-level AND, which is more efficient
- // than checking per-row nullability inline.
- let text_nulls = text_array.nulls().cloned();
- let pair_nulls = pair_delim_array.and_then(|a| a.nulls().cloned());
- let kv_nulls = kv_delim_array.and_then(|a| a.nulls().cloned());
- let combined_nulls = [text_nulls.as_ref(), pair_nulls.as_ref(), kv_nulls.as_ref()]
- .into_iter()
- .fold(None, |acc, nulls| NullBuffer::union(acc.as_ref(), nulls));
+ // NullBuffer::union_many performs a bitmap-level AND, which is more
+ // efficient than checking per-row nullability inline.
+ let combined_nulls = NullBuffer::union_many([
+ text_array.nulls(),
+ pair_delim_array.as_ref().and_then(|a| a.nulls()),
+ kv_delim_array.as_ref().and_then(|a| a.nulls()),
+ ]);
// Use field names matching map_type_from_key_value_types: "key" and "value"
let field_names = MapFieldNames {
diff --git a/datafusion/spark/src/function/null_utils.rs b/datafusion/spark/src/function/null_utils.rs
index b25dc07d0e..886b45d746 100644
--- a/datafusion/spark/src/function/null_utils.rs
+++ b/datafusion/spark/src/function/null_utils.rs
@@ -23,65 +23,51 @@ use datafusion_expr::ColumnarValue;
use std::sync::Arc;
pub(crate) enum NullMaskResolution {
- /// Return NULL as the result (e.g., scalar inputs with at least one NULL)
+ /// All inputs are scalars and at least one is NULL -> return NULL
ReturnNull,
- /// No null mask needed (e.g., all scalar inputs are non-NULL)
+ /// All inputs are non-NULL -> no null mask needed
NoMask,
/// Null mask to apply for arrays
Apply(NullBuffer),
}
-/// Compute NULL mask for the arguments using NullBuffer::union
-pub(crate) fn compute_null_mask(
- args: &[ColumnarValue],
- number_rows: usize,
-) -> Result<NullMaskResolution> {
- // Check if all arguments are scalars
- let all_scalars = args
- .iter()
- .all(|arg| matches!(arg, ColumnarValue::Scalar(_)));
+pub(crate) fn compute_null_mask(args: &[ColumnarValue]) -> NullMaskResolution {
+ let mut array_len = None;
+ let mut has_null_scalar = false;
- if all_scalars {
- // For scalars, check if any is NULL
- for arg in args {
- if let ColumnarValue::Scalar(scalar) = arg
- && scalar.is_null()
- {
- return Ok(NullMaskResolution::ReturnNull);
+ for arg in args {
+ match arg {
+ ColumnarValue::Array(array) => {
+ array_len.get_or_insert_with(|| array.len());
+ }
+ ColumnarValue::Scalar(scalar) => {
+ has_null_scalar |= scalar.is_null();
}
}
- // No NULLs in scalars
- Ok(NullMaskResolution::NoMask)
- } else {
- // For arrays, compute NULL mask for each row using NullBuffer::union
- let array_len = args
- .iter()
- .find_map(|arg| match arg {
- ColumnarValue::Array(array) => Some(array.len()),
- _ => None,
- })
- .unwrap_or(number_rows);
+ }
- // Convert all scalars to arrays for uniform processing
- let arrays: Result<Vec<_>> = args
- .iter()
- .map(|arg| match arg {
- ColumnarValue::Array(array) => Ok(Arc::clone(array)),
- ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(array_len),
- })
- .collect();
- let arrays = arrays?;
+ let Some(array_len) = array_len else {
+ // All arguments are scalars
+ return if has_null_scalar {
+ NullMaskResolution::ReturnNull
+ } else {
+ NullMaskResolution::NoMask
+ };
+ };
- // Use NullBuffer::union to combine all null buffers
- let combined_nulls = arrays
- .iter()
- .map(|arr| arr.nulls())
- .fold(None, |acc, nulls| NullBuffer::union(acc.as_ref(), nulls));
+ if has_null_scalar {
+ return NullMaskResolution::Apply(NullBuffer::new_null(array_len));
+ }
- match combined_nulls {
- Some(nulls) => Ok(NullMaskResolution::Apply(nulls)),
- None => Ok(NullMaskResolution::NoMask),
- }
+ let combined_nulls =
+ NullBuffer::union_many(args.iter().filter_map(|arg| match arg {
+ ColumnarValue::Array(array) => Some(array.nulls()),
+ ColumnarValue::Scalar(_) => None,
+ }));
+
+ match combined_nulls {
+ Some(nulls) => NullMaskResolution::Apply(nulls),
+ None => NullMaskResolution::NoMask,
}
}
diff --git a/datafusion/spark/src/function/string/concat.rs b/datafusion/spark/src/function/string/concat.rs
index 242c5fe471..57fd6cadd9 100644
--- a/datafusion/spark/src/function/string/concat.rs
+++ b/datafusion/spark/src/function/string/concat.rs
@@ -127,7 +127,7 @@ fn spark_concat(args: ScalarFunctionArgs) -> Result<ColumnarValue> {
}
// Step 1: Check for NULL mask in incoming args
- let null_mask = compute_null_mask(&arg_values, number_rows)?;
+ let null_mask = compute_null_mask(&arg_values);
// If all scalars and any is NULL, return NULL immediately
if matches!(null_mask, NullMaskResolution::ReturnNull) {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org