diff --git a/datafusion/functions-nested/benches/array_remove.rs b/datafusion/functions-nested/benches/array_remove.rs index a494d322392a8..6fd153548a63b 100644 --- a/datafusion/functions-nested/benches/array_remove.rs +++ b/datafusion/functions-nested/benches/array_remove.rs @@ -389,7 +389,7 @@ fn create_f64_list_array( if rng.random::() < null_density { None } else { - Some(rng.random_range(0.0..array_size as f64)) + Some(rng.random_range(0..array_size as i64) as f64) } }) .collect::(); diff --git a/datafusion/functions-nested/src/remove.rs b/datafusion/functions-nested/src/remove.rs index 730b4b39ac1ca..c428009425483 100644 --- a/datafusion/functions-nested/src/remove.rs +++ b/datafusion/functions-nested/src/remove.rs @@ -20,10 +20,10 @@ use crate::utils; use crate::utils::make_scalar_function; use arrow::array::{ - Array, ArrayRef, Capacities, GenericListArray, MutableArrayData, NullBufferBuilder, - OffsetSizeTrait, cast::AsArray, make_array, + Array, ArrayRef, Capacities, GenericListArray, MutableArrayData, OffsetSizeTrait, + cast::AsArray, make_array, }; -use arrow::buffer::OffsetBuffer; +use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{DataType, FieldRef}; use datafusion_common::cast::as_int64_array; use datafusion_common::utils::ListCoercion; @@ -386,12 +386,13 @@ fn general_remove( false, Capacities::Array(original_data.len()), ); - let mut valid = NullBufferBuilder::new(list_array.len()); + + // Pre-compute combined null bitmap + let nulls = NullBuffer::union(list_array.nulls(), element_array.nulls()); for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() { - if list_array.is_null(row_index) || element_array.is_null(row_index) { + if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row_index)) { offsets.push(offsets[row_index]); - valid.append_null(); continue; } @@ -414,7 +415,6 @@ fn general_remove( if num_to_remove == 0 { mutable.extend(0, start, end); offsets.push(offsets[row_index] + OffsetSize::usize_as(end - start)); - valid.append_non_null(); continue; } @@ -445,7 +445,6 @@ fn general_remove( } offsets.push(offsets[row_index] + OffsetSize::usize_as(copied)); - valid.append_non_null(); } let new_values = make_array(mutable.freeze()); @@ -453,7 +452,7 @@ fn general_remove( Arc::clone(list_field), OffsetBuffer::new(offsets.into()), new_values, - valid.finish(), + nulls, )?)) }