comphead commented on code in PR #25607:
URL: https://github.com/apache/datafusion/pull/25607#discussion_r4084774285
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
Review Comment:
Is the fused scalar loop actually beating `offsets + filter` for primitives?
Two independent out-of-tree harnesses both found the generic shape faster at
nearly every null density, by 1.15x to 4.5x (1M Int32, rows of 16, at 50% null:
~4.07 ms for this loop vs ~0.92 ms for per-row `count_set_bits_offset` plus
`filter` on a validity-stripped child). The per-element `if is_valid { push }`
is a branch per element and mispredicts badly at moderate density, whereas
arrow's `filter_native` picks between run copies and a branch-free gather via
`IterationStrategy::default_strategy`.
Important caveat that cuts the other way: with many **null parent rows**
this loop wins, because it skips whole rows. One harness measured 3.48 ms vs
6.97 ms at 50% null parents. So this is not a clear-cut delete.
Both measurements are out-of-tree, so please treat them as a prompt for an
in-tree number rather than a conclusion. If the fused loop does win on the
target shape, that data is worth putting in the PR body, along with the
row-length regime, because the three-path ladder is hard to justify by
inspection. If it does not, dropping `compact_primitive` also removes the
`downcast_primitive_array!` dispatch at line 168 and roughly 30 primitive types
x 2 offset widths of monomorphization.
Separately, line 187: `Vec::with_capacity(values_nulls.len() -
values_nulls.null_count())` counts valid children across the whole visible span
**including those under null parent rows**, which this loop never pushes.
`ScalarBuffer::from(Vec)` preserves the allocation, so the excess is retained
for the array's lifetime and shows up in `get_buffer_memory_size()`. Measured
100% over-allocation at 50% null parents. Slightly ironic given
`check_sliced_list_behavior` exists to catch exactly this class of
over-reservation.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
+ mask.append_n(start - mask.len(), false);
}
- } else if batch_start.is_none() {
- batch_start = Some(i);
+ mask.append_packed_range(
+ bit_start..bit_start + len,
+ values_nulls.validity(),
+ );
}
}
- // Flush any remaining batch after the loop
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, end)?;
- }
-
- offsets.push(offsets[row_index] + O::usize_as(kept));
+ offsets.push(O::usize_as(count));
}
+ let mask = mask
+ .map(|mut mask| {
+ mask.append_n(values.len() - mask.len(), false);
+ mask.finish()
+ })
+ .unwrap_or_else(|| values_nulls.inner().clone());
+ let output = match values.data_type() {
+ DataType::Utf8 => copy_string_spans(values.as_string::<i32>(), &mask,
count),
+ DataType::LargeUtf8 => copy_string_spans(values.as_string::<i64>(),
&mask, count),
+ _ => filter(values, &BooleanArray::new(mask, None))?,
Review Comment:
Worth stripping the child's validity buffer before calling `filter`.
`FilterPredicate::filter_nulls` (arrow-select 60.0.0 `filter.rs:511-532`)
early-returns only when `nulls.null_count() == 0`, which is never true here
since child nulls are the whole premise of `array_compact`. So it runs a full
`filter_bits` over the validity bitmap plus a popcount, discovers `null_count
== 0`, and throws the result away. That is guaranteed waste on every call,
because the mask is a subset of the validity.
Measured out-of-tree, 1M values: Utf8 goes 1.784 ms to 1.527 ms at 1% null
and 2.955 ms to 1.882 ms at 5% null. Int32 roughly doubles to triples.
The nice part is that this generalizes what `copy_string_spans` already does
by hand. Line 288 gets the win by passing `None` for nulls, but only for the
two types it reimplements. Boolean, Binary, StringView, Struct, List,
Dictionary, RunEndEncoded, and Map all still pay.
`values.to_data().into_builder().nulls(None)` is cheap (buffer `Arc` clones)
and is safe here since `mask` is a subset of `logical_nulls`.
Upstream follow-up worth filing: `filter_nulls` could pre-check with an
AND-NOT popcount (`count_set_bits(filter & !validity) == 0`, O(n/64) word ops)
before the O(n) `filter_bits`. That would benefit every arrow caller and remove
most of the motivation for `copy_string_spans` living in DataFusion.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
+ mask.append_n(start - mask.len(), false);
}
- } else if batch_start.is_none() {
- batch_start = Some(i);
+ mask.append_packed_range(
+ bit_start..bit_start + len,
+ values_nulls.validity(),
+ );
}
}
- // Flush any remaining batch after the loop
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, end)?;
- }
-
- offsets.push(offsets[row_index] + O::usize_as(kept));
+ offsets.push(O::usize_as(count));
}
+ let mask = mask
+ .map(|mut mask| {
+ mask.append_n(values.len() - mask.len(), false);
+ mask.finish()
+ })
+ .unwrap_or_else(|| values_nulls.inner().clone());
+ let output = match values.data_type() {
+ DataType::Utf8 => copy_string_spans(values.as_string::<i32>(), &mask,
count),
+ DataType::LargeUtf8 => copy_string_spans(values.as_string::<i64>(),
&mask, count),
+ _ => filter(values, &BooleanArray::new(mask, None))?,
+ };
+ Ok((OffsetBuffer::new(offsets.into()), output))
+}
- let new_values = make_array(mutable.freeze());
- Ok(Arc::new(GenericListArray::<O>::try_new(
- Arc::clone(field),
- OffsetBuffer::new(offsets.into()),
- new_values,
- list_nulls.cloned(),
- )?))
+/// Copy adjacent retained strings together. The mask selects only non-null
+/// strings, so the output does not need a validity bitmap.
+#[inline(always)]
Review Comment:
Neither inline attribute has a rationale comment, and they pull against each
other: this force-inlines a 35-line function into a function marked
`#[inline(never)]` at line 207.
Measured in-tree, Utf8 child, 1M values:
| variant | 10% null | 50% null |
| --- | --- | --- |
| as-is | 2.03-2.08 ms | 5.22-5.42 ms |
| minus `#[inline(always)]` | 1.73-1.83 ms | 5.10-5.27 ms |
| minus both | 1.73-1.83 ms | 5.10-5.21 ms |
So `#[inline(always)]` costs about 14% on the hot path, and
`#[inline(never)]` is measurement-neutral. I also checked the codegen theory
that would justify the latter: `downcast_primitive_array!` expands its fallback
arm once, not once per primitive type, so `compact_non_primitive` is
instantiated twice regardless of the attribute.
Suggest deleting both lines. If either is load-bearing, a one-line comment
with the measured effect would help, since otherwise the next person cannot
tell whether removing them regresses anything.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
+ mask.append_n(start - mask.len(), false);
}
- } else if batch_start.is_none() {
- batch_start = Some(i);
+ mask.append_packed_range(
+ bit_start..bit_start + len,
+ values_nulls.validity(),
+ );
}
}
- // Flush any remaining batch after the loop
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, end)?;
- }
-
- offsets.push(offsets[row_index] + O::usize_as(kept));
+ offsets.push(O::usize_as(count));
}
+ let mask = mask
+ .map(|mut mask| {
+ mask.append_n(values.len() - mask.len(), false);
+ mask.finish()
+ })
+ .unwrap_or_else(|| values_nulls.inner().clone());
+ let output = match values.data_type() {
+ DataType::Utf8 => copy_string_spans(values.as_string::<i32>(), &mask,
count),
+ DataType::LargeUtf8 => copy_string_spans(values.as_string::<i64>(),
&mask, count),
+ _ => filter(values, &BooleanArray::new(mask, None))?,
+ };
+ Ok((OffsetBuffer::new(offsets.into()), output))
+}
- let new_values = make_array(mutable.freeze());
- Ok(Arc::new(GenericListArray::<O>::try_new(
- Arc::clone(field),
- OffsetBuffer::new(offsets.into()),
- new_values,
- list_nulls.cloned(),
- )?))
+/// Copy adjacent retained strings together. The mask selects only non-null
+/// strings, so the output does not need a validity bitmap.
+#[inline(always)]
+fn copy_string_spans<S: OffsetSizeTrait>(
+ values: &GenericStringArray<S>,
+ mask: &BooleanBuffer,
+ count: usize,
+) -> ArrayRef {
+ let source_offsets = values.value_offsets();
+ let mut offsets = Vec::with_capacity(count + 1);
+ offsets.push(S::zero());
+ let mut length = S::zero();
+ for (start, end) in mask.set_slices() {
Review Comment:
Test coverage gap, and it is on the branch the optimization depends on.
`source_offsets[start + 1..=end]` is never exercised with more than one element
anywhere in the suite, so the multi-string span copy is untested.
I derived the mask for each test rather than guessing:
- `test_compact_child_types`, string child: after `child.slice(3, 12)` and
the outer `.slice(1, 5)`, the visible child validity is `[T,F,T,T,T,F,F,T,F,T]`
with row 2 a null parent, giving `mask = [T,F,T,F,F,F,F,T,F,T]` and
`set_slices() = [(0,1),(2,3),(7,8),(9,10)]`. Every span has length 1.
- `test_compact_nested_rows` uses a `List<Int32>` child, so it routes to
`filter`, not here.
- `test_compact_all_null_children` early-returns at the all-null check
before any dispatch.
- `test_compact_hidden_nulls_zero_copy` returns the input by `Arc::clone`.
- `check_sliced_list_behavior` uses a Float64 child.
- The SLT case `array_compact(['a', NULL, 'b', NULL, 'c'])` alternates, so
also max span 1.
A case with two or more adjacent non-null strings would cover both the
slice-wide offset `extend` and the multi-string `extend_from_slice`.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
Review Comment:
Name collision: `compact_list_values` already exists at
`datafusion/common/src/nested_struct.rs:384`, where it means "densify a sliced
list via `take`", which is a different operation from "remove null elements".
Same workspace, contradictory semantics, so grepping the name finds the wrong
one. `compact_list` vs `compact_list_values` also differ only by a suffix while
meaning "whole pipeline" vs "build the array".
This is also a single-call-site four-parameter wrapper whose only
contribution over the `downcast_primitive_array!` dispatch is the `try_new`
call. Inlining it into `compact_list` at line 156 would resolve both points at
once. If you prefer to keep it, `build_compacted_list` reads unambiguously.
Its doc comment at lines 159-161 describes `compact_primitive` and
`compact_non_primitive`, which have no doc comments of their own. Worth moving
each sentence onto the function it describes. The invariant that actually
licenses the `- first_offset` arithmetic in both is undocumented anywhere and
is the thing a future reader will need: **`values` is the child already
restricted to the visible span, so `values.len() == offsets.last() -
offsets.first()`, `values_nulls` is `values.logical_nulls()`, and child indices
are `offset - first_offset`.**
Naming nits in the same area:
- `count` (lines 216, 225, 237, 246) counts retained child elements and
doubles as the output length. The pre-PR code already had the right word: `let
kept = ...`.
- `length` (lines 264, 272) is a byte offset accumulator, not an element
count. Arrow calls the same thing `cur_offset`. `byte_len` would be clearer.
- `output` is used for both the `Vec<T::Native>` at line 187 and the
resulting `PrimitiveArray` at line 202. `kept_values` for the vec would avoid
the shadow.
- `mask`: this crate says `predicate` (`array_filter.rs:220`) or
`remove_mask` (`remove.rs:590`). `keep_mask` reconciles both and matches the
comment on line 217.
The `_primitive` / `_non_primitive` split itself is good and matches
`sort.rs:220` / `sort.rs:378`. No objection there.
One more, line 281: the SAFETY comment correctly discharges
`new_unchecked`'s obligations but does not say why the checked constructor is
not used, which is the part a reader wants. Something like "the checked
constructor would re-scan all output bytes for UTF-8 validity the input already
guarantees" would complete it.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
+ mask.append_n(start - mask.len(), false);
}
- } else if batch_start.is_none() {
- batch_start = Some(i);
+ mask.append_packed_range(
+ bit_start..bit_start + len,
+ values_nulls.validity(),
+ );
}
}
- // Flush any remaining batch after the loop
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, end)?;
- }
-
- offsets.push(offsets[row_index] + O::usize_as(kept));
+ offsets.push(O::usize_as(count));
}
+ let mask = mask
+ .map(|mut mask| {
+ mask.append_n(values.len() - mask.len(), false);
+ mask.finish()
+ })
+ .unwrap_or_else(|| values_nulls.inner().clone());
+ let output = match values.data_type() {
+ DataType::Utf8 => copy_string_spans(values.as_string::<i32>(), &mask,
count),
+ DataType::LargeUtf8 => copy_string_spans(values.as_string::<i64>(),
&mask, count),
+ _ => filter(values, &BooleanArray::new(mask, None))?,
+ };
+ Ok((OffsetBuffer::new(offsets.into()), output))
+}
- let new_values = make_array(mutable.freeze());
- Ok(Arc::new(GenericListArray::<O>::try_new(
- Arc::clone(field),
- OffsetBuffer::new(offsets.into()),
- new_values,
- list_nulls.cloned(),
- )?))
+/// Copy adjacent retained strings together. The mask selects only non-null
+/// strings, so the output does not need a validity bitmap.
+#[inline(always)]
+fn copy_string_spans<S: OffsetSizeTrait>(
+ values: &GenericStringArray<S>,
+ mask: &BooleanBuffer,
+ count: usize,
+) -> ArrayRef {
+ let source_offsets = values.value_offsets();
+ let mut offsets = Vec::with_capacity(count + 1);
+ offsets.push(S::zero());
+ let mut length = S::zero();
+ for (start, end) in mask.set_slices() {
+ let adjustment = length - source_offsets[start];
+ offsets.extend(
+ source_offsets[start + 1..=end]
+ .iter()
+ .map(|offset| *offset + adjustment),
+ );
+ length = source_offsets[end] + adjustment;
+ }
+ let mut bytes = Vec::with_capacity(length.as_usize());
+ for (start, end) in mask.set_slices() {
Review Comment:
This second walk of `mask.set_slices()` exists only to get an exact `bytes`
capacity. Folding the copy into the first pass removes a full
`BitSliceIterator` traversal and seven lines:
```rust
let mut bytes: Vec<u8> =
Vec::with_capacity(offset_span_len(values.offsets()));
for (start, end) in mask.set_slices() {
let adjustment = S::usize_as(bytes.len()) - source_offsets[start];
offsets.extend(
source_offsets[start + 1..=end]
.iter()
.map(|offset| *offset + adjustment),
);
bytes.extend_from_slice(
&values.value_data()
[source_offsets[start].as_usize()..source_offsets[end].as_usize()],
);
}
```
The `length` accumulator goes away, and this keeps `offset_span_len` in use,
which the suggestion on line 141 otherwise drops from the imports. Note
`values.value_data().len()` would be the wrong bound here, since it is the
entire backing buffer and can be arbitrarily larger than the visible span for a
sliced child.
Measured in-tree, combined with dropping the inline attributes: 2.03 to
1.55-1.68 ms at 10% null, 5.22 to 4.04-4.18 ms at 50% null.
Why this is worth doing rather than a nit. Routing Utf8/LargeUtf8 straight
through `filter` instead measures 3.75 ms / 4.79 ms. So at 10% null the
specialization earns its ~40 lines and its `unsafe` (1.85x), but **at 50% null
the current two-pass form is slower than just calling `filter`** (5.22 vs
4.79), because arrow drops to per-index copying below 0.8 selectivity
(`FILTER_SLICES_SELECTIVITY_THRESHOLD`) while this always uses spans. The
single-pass form fixes that (4.04 vs 4.79).
Two honest caveats. Capacity becomes an upper bound rather than exact, so
the output byte buffer retains capacity for dropped strings' bytes when null
slots carry payload. And a separate out-of-tree run saw a ~12% regression for
256-byte strings, where the copy is bandwidth-bound. Worth confirming on your
own workload.
Also worth saying: the two-pass shape matches arrow's own `filter_bytes`, so
it is idiomatic rather than sloppy, and this offset pass is actually better
than arrow's, because the slice-wide `extend` vectorizes where arrow's
`extend_offsets_slices` pushes one at a time.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
+ mask.append_n(start - mask.len(), false);
}
- } else if batch_start.is_none() {
- batch_start = Some(i);
+ mask.append_packed_range(
+ bit_start..bit_start + len,
+ values_nulls.validity(),
+ );
}
}
- // Flush any remaining batch after the loop
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, end)?;
- }
-
- offsets.push(offsets[row_index] + O::usize_as(kept));
+ offsets.push(O::usize_as(count));
}
+ let mask = mask
+ .map(|mut mask| {
+ mask.append_n(values.len() - mask.len(), false);
+ mask.finish()
+ })
+ .unwrap_or_else(|| values_nulls.inner().clone());
+ let output = match values.data_type() {
+ DataType::Utf8 => copy_string_spans(values.as_string::<i32>(), &mask,
count),
+ DataType::LargeUtf8 => copy_string_spans(values.as_string::<i64>(),
&mask, count),
+ _ => filter(values, &BooleanArray::new(mask, None))?,
+ };
+ Ok((OffsetBuffer::new(offsets.into()), output))
+}
- let new_values = make_array(mutable.freeze());
- Ok(Arc::new(GenericListArray::<O>::try_new(
- Arc::clone(field),
- OffsetBuffer::new(offsets.into()),
- new_values,
- list_nulls.cloned(),
- )?))
+/// Copy adjacent retained strings together. The mask selects only non-null
+/// strings, so the output does not need a validity bitmap.
+#[inline(always)]
+fn copy_string_spans<S: OffsetSizeTrait>(
+ values: &GenericStringArray<S>,
+ mask: &BooleanBuffer,
+ count: usize,
+) -> ArrayRef {
+ let source_offsets = values.value_offsets();
+ let mut offsets = Vec::with_capacity(count + 1);
+ offsets.push(S::zero());
+ let mut length = S::zero();
+ for (start, end) in mask.set_slices() {
+ let adjustment = length - source_offsets[start];
+ offsets.extend(
+ source_offsets[start + 1..=end]
+ .iter()
+ .map(|offset| *offset + adjustment),
+ );
+ length = source_offsets[end] + adjustment;
+ }
+ let mut bytes = Vec::with_capacity(length.as_usize());
+ for (start, end) in mask.set_slices() {
+ bytes.extend_from_slice(
+ &values.value_data()
+
[source_offsets[start].as_usize()..source_offsets[end].as_usize()],
+ );
+ }
+ // SAFETY: the mask selects only non-null strings. Each output string is
+ // copied unchanged from the input, so its UTF-8 remains valid. Rebasing
+ // offsets preserves their order, and the last offset equals bytes.len().
+ let output = unsafe {
+ GenericStringArray::<S>::new_unchecked(
+ OffsetBuffer::new(offsets.into()),
+ bytes.into(),
+ None,
+ )
+ };
+ Arc::new(output)
}
#[cfg(test)]
mod tests {
use super::*;
+ use arrow::array::{
+ AsArray, BooleanArray, DictionaryArray, FixedSizeListArray, Int32Array,
+ ListArray, MapArray, NullArray, RunArray, StringArray, StringViewArray,
+ StructArray, UInt32Array, new_empty_array,
+ };
+ use arrow::compute::take;
+ use arrow::datatypes::{Field, Int32Type, TimeUnit};
+ use datafusion_common::utils::list_values;
+
+ #[test]
+ fn test_compact_child_types() -> Result<()> {
+ let input = Int32Array::from(vec![
+ Some(1),
+ None,
+ Some(2), // Padding before the child slice.
+ Some(90), // Padding before the outer slice.
+ Some(10),
+ None,
+ Some(20),
+ Some(100),
+ Some(110), // Values belonging to a null list row.
+ None,
+ None, // A valid list with only null elements.
+ Some(30),
+ None,
+ Some(40),
+ Some(99), // Padding after the outer slice.
+ ]);
+ let mut children = vec![];
Review Comment:
Test dedup. The coverage here is genuinely good, so this is about runtime
and reading cost rather than lost coverage. Three spots:
**This matrix (lines 326-343).** 13 of the 14 listed types resolve to the
same `compact_primitive<T>` instantiation. Checked against arrow 60's
`downcast_primitive!` arms: Int8, UInt64, Int32, Int64, Float32, Float64,
Decimal32/64/128/256, Date32, Duration, and Timestamp all match primitive arms,
and only `Dictionary` falls through to `compact_non_primitive`.
`compact_primitive` has no width- or type-specific branch, so the only distinct
property exercised is `.with_data_type()` preservation. Four entries cover it:
`Int32` (plain), `Decimal256` (widest native plus precision/scale),
`Timestamp(Micro, tz)` (tz payload), `Dictionary` (different path).
**`test_compact_nested_rows` (line 451).** Mostly overlaps this test, which
already covers a `ListArray` child, an empty row, a null parent with a
non-empty hidden child range, and sliced outer and child. The one scenario
unique to it is an empty row sitting between two rows that both keep values,
where in `test_compact_child_types` the empty row abuts the null row. Worth
keeping that scenario and dropping the 2x2 `child_type` / `data_type` cast
matrix, since offset width is already swept by every other test and there is no
width-specific code on this path. Roughly 40 lines and 4 iterations down to ~15
lines and 1.
**`test_compact_all_null_children` line 498.** All five children hit the
same `values_nulls.null_count() == values_nulls.len()` early return before any
type dispatch, so the only per-type difference is the `logical_nulls()`
implementation. `Int32Array::new_null` and the all-null `StringArray` are both
"validity buffer, all bits unset". `NullArray`, `Dictionary`, and `RunArray`
each exercise a distinct `logical_nulls()`, so those three plus one plain array
suffice.
All three are optional. The one that is not optional is the missing
multi-element span case flagged on line 265.
##########
datafusion/functions-nested/src/array_filter.rs:
##########
@@ -204,27 +206,9 @@ impl HigherOrderUDFImpl for ArrayFilter {
/// Returns a list array with every non-null sublist emptied, preserving the
null buffer.
/// Used for the `x -> false` / `x -> null` scalar predicate short-circuit.
fn empty_filtered_list(list_array: &ArrayRef, field: FieldRef) ->
Result<ArrayRef> {
Review Comment:
With the extraction done, this is a seven-line single-call-site dispatch
wrapper that only downcasts and delegates. Giving `empty_list_values` the
`(&dyn Array, FieldRef) -> Result<ArrayRef>` signature with the match inside
would let this go away, and `array_compact.rs:153` could then call it with
`input_array`, which is already in scope. Net about ten lines, one helper
instead of two.
The `empty_list_values` extraction itself is a real win, three call sites
across two files. `pub(crate)` in `functions-nested/src/utils.rs` is also the
right home. I grepped for other hand-rolled "empty every valid list row" sites
(`new_zeroed`, `vec![0i32; n + 1]`, `from_repeated_length(0`) and found none,
so moving it to `datafusion-common` would widen the public API for no current
consumer.
Separate, and I would treat it as a follow-up issue rather than a change
request here. `filter_list_values` just below (line 218) and
`compact_non_primitive` in `array_compact.rs:208` are now the same primitive
one file apart, with the tuple order mirrored:
```rust
filter_list_values(...) -> Result<(ArrayRef, OffsetBuffer<O>)>
compact_non_primitive(...) -> Result<(OffsetBuffer<O>, ArrayRef)>
```
Both compute per-row kept counts by popcount over a keep bitmap, rebuild an
`OffsetBuffer`, call `arrow::compute::filter`, and carry a no-op short circuit.
They have already drifted: `array_compact` excludes children of null parent
rows, `array_filter` copies them into the values buffer where they are
invisible but retained. A shared `compact_list_with_mask(values, offsets,
parent_nulls, keep)` in `utils.rs` next to `empty_list_values` would remove the
drift by construction and give both the validity-strip and span-copy wins.
`general_array_distinct` (`set_ops.rs:555`), `array_except` (`except.rs:210`),
and `general_remove_with_scalar` (`remove.rs:560`) all build monotone `indices`
plus `take` and could use it too, so it plausibly serves five or six call sites.
Unifying is clearly out of scope for a perf PR on one function. Swapping one
tuple order so the two agree is not, and a linked issue would keep the
duplication tracked.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
+ mask.append_n(start - mask.len(), false);
}
- } else if batch_start.is_none() {
- batch_start = Some(i);
+ mask.append_packed_range(
+ bit_start..bit_start + len,
+ values_nulls.validity(),
+ );
}
}
- // Flush any remaining batch after the loop
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, end)?;
- }
-
- offsets.push(offsets[row_index] + O::usize_as(kept));
+ offsets.push(O::usize_as(count));
}
+ let mask = mask
+ .map(|mut mask| {
+ mask.append_n(values.len() - mask.len(), false);
+ mask.finish()
+ })
+ .unwrap_or_else(|| values_nulls.inner().clone());
+ let output = match values.data_type() {
+ DataType::Utf8 => copy_string_spans(values.as_string::<i32>(), &mask,
count),
Review Comment:
`Binary` and `LargeBinary` look excluded by accident rather than by design.
The body of `copy_string_spans` only touches `value_offsets()`, `value_data()`,
and `new_unchecked`, all of which exist on `GenericByteArray<T:
ByteArrayType>`. Generalizing the signature and adding two match arms here is
about three lines, and the SAFETY comment gets strictly weaker, since binary
has no encoding invariant to preserve.
Excluding `StringView` and `BinaryView` **is** principled, since those go
through `filter_byte_view` and have no offset chain to rebase, so there is
nothing to gain beyond the validity-strip on line 248.
Supporting evidence that it is an oversight: the type matrix at lines
325-409 includes `StringViewArray` but no `BinaryArray`, `LargeBinaryArray`,
`BinaryViewArray`, or `FixedSizeBinary`, so the gap is not exercised either way.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
+ .then(|| values.slice(start, visible_len));
+ let values = sliced_values.as_ref().unwrap_or(values);
// Use logical nulls so element types without a validity buffer
// (e.g. NullArray) are still treated as null.
- let Some(values_nulls) = values.logical_nulls() else {
- // Fast path: no validity buffer, no nulls to remove
- return Ok(Arc::new(list_array.clone()));
+ let Some(values_nulls) = values
+ .logical_nulls()
+ .filter(|nulls| nulls.null_count() != 0)
+ else {
+ return Ok(Arc::clone(input_array));
};
- if values_nulls.null_count() == 0 {
- // Fast path: validity buffer present but no nulls set
- return Ok(Arc::new(list_array.clone()));
+ if values_nulls.null_count() == values_nulls.len() {
+ return Ok(empty_list_values(list_array, Arc::clone(field)));
}
- let list_nulls = list_array.nulls();
+ compact_list_values(list_array, field, values.as_ref(), &values_nulls)
+}
+
+/// Copy primitive values and build list offsets in one pass. For other types,
+/// build offsets and a keep bitmap, then copy Utf8/LargeUtf8 strings in valid
+/// spans or use Arrow's filter kernel for the remaining types.
+fn compact_list_values<O: OffsetSizeTrait>(
+ list_array: &GenericListArray<O>,
+ field: &FieldRef,
+ values: &dyn Array,
+ values_nulls: &NullBuffer,
+) -> Result<ArrayRef> {
+ let (offsets, new_values) = downcast_primitive_array! {
+ values => Ok(compact_primitive(values, list_array, values_nulls)),
+ _ => compact_non_primitive(values, list_array, values_nulls),
+ }?;
+ Ok(Arc::new(GenericListArray::<O>::try_new(
+ Arc::clone(field),
+ offsets,
+ new_values,
+ list_array.nulls().cloned(),
+ )?))
+}
+
+fn compact_primitive<T: ArrowPrimitiveType, O: OffsetSizeTrait>(
+ values: &PrimitiveArray<T>,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> (OffsetBuffer<O>, ArrayRef) {
let list_offsets = list_array.offsets();
- let original_data = values.to_data();
- let (first_offset, visible_len) = offset_span(list_offsets);
- let capacity =
- visible_len - values_nulls.slice(first_offset,
visible_len).null_count();
- let mut offsets = Vec::<O>::with_capacity(list_array.len() + 1);
+ let first_offset = list_offsets[0].as_usize();
+ let mut output = Vec::with_capacity(values_nulls.len() -
values_nulls.null_count());
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
offsets.push(O::zero());
- let mut mutable = MutableArrayData::with_capacities(
- vec![&original_data],
- false,
- Capacities::Array(capacity),
- );
-
- for row_index in 0..list_array.len() {
- if list_nulls.is_some_and(|n| n.is_null(row_index)) {
- offsets.push(offsets[row_index]);
- continue;
+ for (row, window) in list_offsets.windows(2).enumerate() {
+ if list_array.is_valid(row) {
+ let start = window[0].as_usize() - first_offset;
+ let end = window[1].as_usize() - first_offset;
+ for i in start..end {
+ if values_nulls.is_valid(i) {
+ output.push(values.value(i));
+ }
+ }
}
+ offsets.push(O::usize_as(output.len()));
+ }
+ let output = PrimitiveArray::<T>::new(output.into(), None)
+ .with_data_type(values.data_type().clone());
+ (OffsetBuffer::new(offsets.into()), Arc::new(output))
+}
- let start = list_offsets[row_index].as_usize();
- let end = list_offsets[row_index + 1].as_usize();
- let row_null_count = values_nulls.slice(start, end -
start).null_count();
- let kept = (end - start) - row_null_count;
-
- // Batch consecutive non-null elements into single extend() calls
- // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this
- // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones.
- let mut batch_start: Option<usize> = None;
- for i in start..end {
- if values_nulls.is_null(i) {
- // Null breaks the current batch — flush it
- if let Some(bs) = batch_start {
- mutable.try_extend(0, bs, i)?;
- batch_start = None;
+#[inline(never)]
+fn compact_non_primitive<O: OffsetSizeTrait>(
+ values: &dyn Array,
+ list_array: &GenericListArray<O>,
+ values_nulls: &NullBuffer,
+) -> Result<(OffsetBuffer<O>, ArrayRef)> {
+ let first_offset = list_array.offsets()[0].as_usize();
+ let mut offsets = Vec::with_capacity(list_array.len() + 1);
+ offsets.push(O::zero());
+ let mut count = 0;
+ // Child validity is already the keep mask unless null parents hide values.
+ let mut mask =
+ (list_array.null_count() != 0).then(||
BooleanBufferBuilder::new(values.len()));
+ for (row, window) in list_array.offsets().windows(2).enumerate() {
+ let start = window[0].as_usize() - first_offset;
+ let len = window[1].as_usize() - window[0].as_usize();
+ if list_array.is_valid(row) {
+ let bit_start = values_nulls.offset() + start;
+ count += values_nulls.buffer().count_set_bits_offset(bit_start,
len);
+ if let Some(mask) = &mut mask {
+ // Fill the gap left by any preceding null parent rows.
+ if start > mask.len() {
Review Comment:
The builder only advances on valid rows, so its length silently falls behind
the child position, and two pieces of catch-up repair it: this leading backfill
and the trailing pad on line 241. List offsets are contiguous by construction,
so the gap a null row leaves is exactly `len`. Appending it when you see it
keeps the invariant `mask.len() == current child position` and removes both
fixups plus a branch from the hot valid-row path:
```rust
if list_array.is_valid(row) {
let bit_start = values_nulls.offset() + start;
count += values_nulls.buffer().count_set_bits_offset(bit_start, len);
if let Some(mask) = &mut mask {
mask.append_packed_range(bit_start..bit_start + len,
values_nulls.validity());
}
} else if let Some(mask) = &mut mask {
// Values under a null row are never selected.
mask.append_n(len, false);
}
```
Then line 239 collapses to `mask.map(|mut m| m.finish()).unwrap_or_else(||
values_nulls.inner().clone())`. Six lines fewer. Verified in-tree: all five
tests pass.
While here, line 220 re-fetches `list_array.offsets()` in the loop header
where `compact_primitive` binds it once at line 185. Free either way, but
inconsistent between two functions that are otherwise parallel.
##########
datafusion/functions-nested/src/utils.rs:
##########
@@ -640,11 +655,13 @@ pub(crate) mod tests {
let compact = arrow::compute::cast(&compact, &data_type)?;
// Middle slice, empty slice, empty row, and null row with child
data.
for (offset, len) in [(0, 3), (0, 0), (1, 1), (2, 1)] {
- let result = run(&input.slice(1 + offset, len))?;
+ let sliced = input.slice(1 + offset, len);
+ let result = run(&sliced)?;
let expected = run(&compact.slice(offset, len))?;
assert_eq!(result.as_ref(), expected.as_ref());
assert!(
- result.get_buffer_memory_size() < 1024,
+ Arc::ptr_eq(&result, &sliced)
Review Comment:
This loosens a shared assertion for all 11 callers (`array_add.rs:140`,
`array_compact.rs:537`, `array_normalize.rs:235`, `array_scale.rs:229`,
`array_subtract.rs:139`, `concat.rs:657` and `:664`, `extract.rs:1403` and
`:1413`, `remove.rs:1173`, `replace.rs:830`) to accommodate one of them.
The helper's contract is "the output must not reserve capacity based on the
full backing child". `array_compact` introduces a second legitimate outcome,
"the output is the input", and the escape hatch is now invisible at the other
ten call sites. Making it explicit, for example
`check_sliced_list_behavior(may_return_input: bool, run: ...)` with
`array_compact` the only `true`, keeps the invariant honest for everyone else.
It is a mechanical 11-line edit.
Worth noting the PR already has a dedicated positive test for the zero-copy
contract in `test_compact_hidden_nulls_zero_copy`, which is the right place for
it. The shared helper does not need to assert it loosely as well.
In practice the loosening is narrow, since `Arc::ptr_eq` needs the exact
same allocation and things like `remove.rs:566`'s
`Arc::new(list_array.clone())` would not satisfy it. So this is a clarity
point, not a hole.
##########
datafusion/functions-nested/src/array_compact.rs:
##########
@@ -112,101 +114,462 @@ fn array_compact_inner(arg: &[ArrayRef]) ->
Result<ArrayRef> {
let [input_array] = take_function_args("array_compact", arg)?;
match &input_array.data_type() {
- List(field) => {
- let array = as_list_array(input_array)?;
- compact_list::<i32>(array, field)
- }
- LargeList(field) => {
- let array = as_large_list_array(input_array)?;
- compact_list::<i64>(array, field)
- }
+ List(field) => compact_list::<i32>(input_array, field),
+ LargeList(field) => compact_list::<i64>(input_array, field),
Null => Ok(Arc::clone(input_array)),
array_type => exec_err!("array_compact does not support type
'{array_type}'."),
}
}
/// Remove null elements from each row of a list array.
+///
+/// Each row is a range in a shared child array. Compaction removes null child
+/// values and rebuilds the row offsets, preserving null list rows.
fn compact_list<O: OffsetSizeTrait>(
- list_array: &GenericListArray<O>,
- field: &Arc<arrow::datatypes::Field>,
+ input_array: &ArrayRef,
+ field: &FieldRef,
) -> Result<ArrayRef> {
+ let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
+ let visible_len = offset_span_len(list_array.offsets());
+ if visible_len == 0 || list_array.null_count() == list_array.len() {
+ return Ok(Arc::clone(input_array));
+ }
+ // Restrict the child to the visible values before computing logical nulls,
+ // which can be expensive.
let values = list_array.values();
+ let start = list_array.offsets()[0].as_usize();
+ let sliced_values = (start != 0 || visible_len != values.len())
Review Comment:
This is `datafusion_common::utils::list_values` reimplemented.
`sliced_list_values` (`common/src/utils/mod.rs:1275-1284`) has the identical
`start != 0 || len != values.len()` guard and the same `values.slice(first,
last - first)`. The test module in this very file imports `list_values` at line
304.
```rust
let list_array = as_generic_list_array::<O>(input_array.as_ref())?;
if list_array.null_count() == list_array.len() {
return Ok(Arc::clone(input_array));
}
// Restrict the child to the visible values before computing logical nulls,
// which can be expensive.
let values = list_values(input_array.as_ref())?;
```
Nine lines to six, and `visible_len`, `start`, and `sliced_values` all
disappear. Verified in-tree: all five tests pass, including the `(0, 0)`
empty-slice case in `test_sliced_capacity` and all of
`test_compact_hidden_nulls_zero_copy`. The `visible_len == 0` disjunct is
subsumed by the `nulls.null_count() != 0` filter three lines below, since an
empty child has null count 0.
Minor, line 133: the pre-PR code used `let (first_offset, visible_len) =
offset_span(list_offsets)`, one call returning both. This replaces it with
`offset_span_len` plus a hand-rolled `offsets()[0]`, and `offset_span_len` is
literally `offset_span(offsets).1`. `first_offset` is then re-derived a third
and fourth time at lines 186 and 213. No measurable cost, but it walks back an
existing helper use.
(Not for this PR: `common/src/nested_struct.rs:328-331` is a third copy of
the same predicate. Promoting `sliced_list_values` to `pub` and adopting it at
the hand-rolled sites in `position.rs`, `replace.rs`, `remove.rs`, and
`set_ops.rs` would be a decent follow-up.)
--
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]
