perf: optimize spark_array_compact (no-op fast path for null-free arrays)#4934
perf: optimize spark_array_compact (no-op fast path for null-free arrays)#4934andygrove wants to merge 2 commits into
spark_array_compact (no-op fast path for null-free arrays)#4934Conversation
…ree arrays) array_compact only removes null elements. When the values buffer has no null elements there is nothing to remove, so the result is bit-identical to the input. Add a fast path that returns the input unchanged in that case, reusing its buffers zero-copy and skipping the per-element MutableArrayData::extend loop. The null-containing path is left byte-for-byte unchanged so shapes with sparse or dense element nulls do not regress. Add Rust unit tests (the expression previously had only SQL/Scala coverage) and a criterion benchmark covering no-null, null-row, sparse-null, and dense-null shapes.
mbutrovich
left a comment
There was a problem hiding this comment.
First pass, thanks @andygrove!
| // logical_nulls() (not nulls()) is used so a NullArray, whose elements are | ||
| // all logically null, is correctly reported as fully null. NullArray::nulls() | ||
| // returns None, which would make is_null() report false for every element. | ||
| let value_nulls = values.logical_nulls(); |
There was a problem hiding this comment.
native/spark-expr/src/array_funcs/array_compact.rs:124,132
let value_nulls = values.logical_nulls();
...
if value_nulls.as_ref().map(|n| n.null_count()).unwrap_or(0) == 0 {
return Ok(Arc::new(list_array.clone()));
}For the common element types (primitive, string, struct, list) this is already cheap. Array::logical_nulls() is self.nulls().cloned() (arrow-rs arrow-array/src/array/mod.rs:243-245), and since a NullBuffer wraps an Arc-backed buffer that clone is an O(1) refcount bump, not a bitmap copy, while NullBuffer::null_count() is a precomputed O(1) read. The default logical_null_count() (mod.rs:324-328) is the same clone, so for these types switching changes nothing.
The case worth guarding against is a values child whose type overrides logical_nulls() to compute a fresh buffer: DictionaryArray (dictionary_array.rs:741), RunArray (run_array.rs:717), and NullArray (null_array.rs:117). Each of these also overrides logical_null_count() (dictionary_array.rs:762, run_array.rs:721, null_array.rs:125) to count without materializing, so for those types logical_nulls() is an O(n) allocation on exactly the fast path this PR exists to speed up while logical_null_count() is not. (UnionArray overrides logical_nulls() at union_array.rs:794 but has no logical_null_count() override, so it allocates either way and does not benefit.) Using logical_null_count() makes the fast path allocation-free for the types that have the override and is never worse than the current code.
Suggested change: gate the fast path on the count, and only materialize the buffer on the slow path where it is actually consumed.
let values = list_array.values();
// Fast path: array_compact only removes null elements. When the values buffer
// has no logical null elements there is nothing to remove, so every row is
// returned unchanged and the result is bit-identical to the input. logical_null_count
// (not null_count) is used so a NullArray, whose elements are all logically null,
// is counted correctly; NullArray::nulls() is None.
if values.logical_null_count() == 0 {
return Ok(Arc::new(list_array.clone()));
}
let value_nulls = values.logical_nulls();This keeps the slow path byte-for-byte identical and removes the allocation from the common case.
|
|
||
| let values = list_array.values(); | ||
|
|
||
| // logical_nulls() (not nulls()) is used so a NullArray, whose elements are |
There was a problem hiding this comment.
The comment justifies logical_nulls over is_null/nulls specifically because a NullArray values child reports nulls() == None yet is logically all-null. array_compact(array(null, null)) produces a List<Null> and is exactly this shape. None of the three new unit tests exercise it, so the branch the comment protects has no direct coverage and a future refactor to nulls() would pass the suite while silently breaking. This also interacts with finding 1: it is the case that proves logical_null_count must be used rather than null_count.
Suggested change: add a test building a List<Null> (via ListBuilder::new(NullBuilder::new()) or an explicit NullArray values child) and assert all-null rows compact to empty rows.
Which issue does this PR close?
N/A
Rationale for this change
Optimize existing expression.
What changes are included in this PR?
array_compactonly removes null elements. When the values buffer contains no null elements there is nothing to remove, so the output is bit-identical to the input. This adds a fast path tocompact_listthat returns the input array unchanged in that case, reusing its buffers zero-copy and skipping the per-elementMutableArrayData::extendloop.The null-containing path is left byte-for-byte unchanged, so arrays with sparse or dense element nulls take exactly the same code as before and do not regress. This deliberately avoids the per-row run-batching approach that regressed the dense-null shape in a previous attempt (#4886).
Also adds Rust unit tests (the expression previously had only SQL file / Scala coverage) and a criterion benchmark covering no-null, null-row, sparse-null, and dense-null shapes.
How are these changes tested?
New Rust unit tests assert the fast path is bit-identical to the input (including null and empty rows) and that the null-removal path is unchanged; existing SQL file tests and
CometArrayExpressionSuitecontinue to cover end-to-end behavior.Benchmark (criterion), baseline
mainvs this branch, 8192-row list columns:The null-free shapes hit the zero-copy fast path; the sparse/dense-null shapes run the unchanged loop and are within noise across repeated samples.