Skip to content

perf: optimize spark_array_compact (no-op fast path for null-free arrays)#4934

Open
andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:perf-optimize-array-expr
Open

perf: optimize spark_array_compact (no-op fast path for null-free arrays)#4934
andygrove wants to merge 2 commits into
apache:mainfrom
andygrove:perf-optimize-array-expr

Conversation

@andygrove

Copy link
Copy Markdown
Member

Which issue does this PR close?

N/A

Rationale for this change

Optimize existing expression.

What changes are included in this PR?

array_compact only 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 to compact_list that returns the input array 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 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 CometArrayExpressionSuite continue to cover end-to-end behavior.

Benchmark (criterion), baseline main vs this branch, 8192-row list columns:

array_compact: no nulls short:              316 µs -> 287 ns   (~99.9% faster)
array_compact: no nulls long:              2.46 ms -> 254 ns   (~99.99% faster)
array_compact: no element nulls, null rows: 296 µs -> 258 ns   (~99.9% faster)
array_compact: sparse nulls:                489 µs -> 488 µs   (within noise)
array_compact: dense nulls:                 258 µs -> 257 µs   (within noise)

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.

…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 mbutrovich left a comment

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.

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();

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.

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

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.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants