Skip to content

chore: use NullBuffer::union for Spark concat #18087

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
comphead merged 3 commits into from
Oct 16, 2025
Merged

chore: use NullBuffer::union for Spark concat #18087

Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f5bf52b
chore: use `NullBuffer::union` for Spark `concat`
comphead Oct 15, 2025
e76800d
chore: use `NullBuffer::union` for Spark `concat`
comphead Oct 15, 2025
07b81fe
chore: use `NullBuffer::union` for Spark `concat`
comphead Oct 16, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 34 additions & 81 deletions datafusion/spark/src/function/string/concat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@
// specific language governing permissions and limitations
// under the License.

use arrow::array::{Array, ArrayBuilder};
use arrow::array::Array;
use arrow::buffer::NullBuffer;
use arrow::datatypes::DataType;
use datafusion_common::{Result, ScalarValue};
use datafusion_expr::{
Expand All @@ -31,6 +32,10 @@ use std::sync::Arc;
///
/// Concatenates multiple input strings into a single string.
/// Returns NULL if any input is NULL.
///
/// Differences with DataFusion concat:

@alamb alamb Oct 16, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️

/// - Support 0 arguments
/// - Return NULL if any input is NULL
#[derive(Debug, PartialEq, Eq, Hash)]
pub struct SparkConcat {
signature: Signature,
Expand Down Expand Up @@ -122,13 +127,13 @@ fn spark_concat(args: ScalarFunctionArgs) -> Result<ColumnarValue> {
apply_null_mask(result, null_mask)
}

/// Compute NULL mask for the arguments
/// Returns None if all scalars and any is NULL, or a Vector of
/// boolean representing the null mask for incoming arrays
/// Compute NULL mask for the arguments using NullBuffer::union
/// Returns None if all scalars and any is NULL, or an Option of NullBuffer
/// representing the combined null mask for incoming arrays
fn compute_null_mask(
args: &[ColumnarValue],
number_rows: usize,
) -> Result<Option<Vec<bool>>> {
) -> Result<Option<Option<NullBuffer>>> {

@Jefffrey Jefffrey Oct 16, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thoughts on defining a separate enum here instead of encoding meaning in a double Option?

enum NullMask {
    AllScalarsWithNull,
    AllScalarsNoNulls,
    Nulls(Option<NullBuffer>),
}

@comphead comphead Oct 16, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @Jefffrey I agree, originally I thought introducing another enum just for this function would be too much, but I think it would be better having enum

// Check if all arguments are scalars
let all_scalars = args
.iter()
Expand All @@ -145,9 +150,9 @@ fn compute_null_mask(
}
}
// No NULLs in scalars
Ok(Some(vec![]))
Ok(Some(None))
} else {
// For arrays, compute NULL mask for each row
// For arrays, compute NULL mask for each row using NullBuffer::union
let array_len = args
.iter()
.find_map(|arg| match arg {
Expand All @@ -166,95 +171,43 @@ fn compute_null_mask(
.collect();
let arrays = arrays?;

// Compute NULL mask
let mut null_mask = vec![false; array_len];
for array in &arrays {
for (i, null_flag) in null_mask.iter_mut().enumerate().take(array_len) {
if array.is_null(i) {
*null_flag = true;
}
}
}
// 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));

Ok(Some(null_mask))
Ok(Some(combined_nulls))
}
}

/// Apply NULL mask to the result
/// Apply NULL mask to the result using NullBuffer::union
fn apply_null_mask(
result: ColumnarValue,
null_mask: Option<Vec<bool>>,
null_mask: Option<Option<NullBuffer>>,
) -> Result<ColumnarValue> {
match (result, null_mask) {
// Scalar with NULL mask means return NULL
(ColumnarValue::Scalar(_), None) => {
Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None)))
}
// Scalar without NULL mask, return as-is
(scalar @ ColumnarValue::Scalar(_), Some(mask)) if mask.is_empty() => Ok(scalar),
// Array with NULL mask
(ColumnarValue::Array(array), Some(null_mask)) if !null_mask.is_empty() => {
let array_len = array.len();
let return_type = array.data_type();
(scalar @ ColumnarValue::Scalar(_), Some(None)) => Ok(scalar),
// Array with NULL mask - use NullBuffer::union to combine nulls
(ColumnarValue::Array(array), Some(Some(null_mask))) => {
// Combine the result's existing nulls with our computed null mask
let combined_nulls = NullBuffer::union(array.nulls(), Some(&null_mask));

let mut builder: Box<dyn ArrayBuilder> = match return_type {
DataType::Utf8 => {
let string_array = array
.as_any()
.downcast_ref::<arrow::array::StringArray>()
.unwrap();
let mut builder =
arrow::array::StringBuilder::with_capacity(array_len, 0);
for (i, &is_null) in null_mask.iter().enumerate().take(array_len) {
if is_null || string_array.is_null(i) {
builder.append_null();
} else {
builder.append_value(string_array.value(i));
}
}
Box::new(builder)
}
DataType::LargeUtf8 => {
let string_array = array
.as_any()
.downcast_ref::<arrow::array::LargeStringArray>()
.unwrap();
let mut builder =
arrow::array::LargeStringBuilder::with_capacity(array_len, 0);
for (i, &is_null) in null_mask.iter().enumerate().take(array_len) {
if is_null || string_array.is_null(i) {
builder.append_null();
} else {
builder.append_value(string_array.value(i));
}
}
Box::new(builder)
}
DataType::Utf8View => {
let string_array = array
.as_any()
.downcast_ref::<arrow::array::StringViewArray>()
.unwrap();
let mut builder =
arrow::array::StringViewBuilder::with_capacity(array_len);
for (i, &is_null) in null_mask.iter().enumerate().take(array_len) {
if is_null || string_array.is_null(i) {
builder.append_null();
} else {
builder.append_value(string_array.value(i));
}
}
Box::new(builder)
}
_ => {
return datafusion_common::exec_err!(
"Unsupported return type for concat: {:?}",
return_type
);
}
};
// Create new array with combined nulls
let new_array = array
.into_data()
.into_builder()
.nulls(combined_nulls)
.build()?;

Ok(ColumnarValue::Array(builder.finish()))
Ok(ColumnarValue::Array(Arc::new(arrow::array::make_array(
new_array,
))))
}
// Array without NULL mask, return as-is
(array @ ColumnarValue::Array(_), _) => Ok(array),
Expand Down