Repository navigation
perf: Optimize repeated counting of set bits - #26507
neilconway wants to merge 6 commits into
Conversation
…arks Cover Float64 lists of 4, 32, and 256 elements per row, with the same total number of elements, both without NULL elements and with one NULL element in 10% of rows.
array_contains: lists of 4, 16, and 128 Int32 elements and a fixed-size list of 16, with 10% NULL elements and a value that never occurs, so every row checks whether its list contains a NULL element. map_from_entries: lists of 4, 16, and 64 entries with unique keys, where 10% of rows contain one NULL entry.
…ulls Make SetBitCounter public in datafusion_common::utils so other crates can count the set bits in each row of a list in one pass, and add count_range for the set bits within one row, which prepare_list_filter now uses for the elements stored under NULL rows. Add rows_have_nulls, which returns a bitmap of the list rows that contain a NULL element. It counts from the list's first element, which can be far into the validity bitmap when the list is sliced, builds the bitmap with BooleanBuffer::collect_bool, and reuses each row's end count as the next row's start.
array_remove, inner_product, cosine_distance, and array_normalize sliced
a bitmap or the values array for each list row to count its set bits or
NULLs. Each slice clones a reference-counted buffer and sets up a new
bit count.
array_remove now counts each row's matches with SetBitCounter, which
reads each 64-bit word of the bitmap once over the whole list, and
slices a row's bits only when it has matches to remove. The vector
functions get the rows that contain a NULL from rows_have_nulls, through
a new rows_with_null_elements helper, and read each row's values as a
plain slice of the values buffer.
Test inner_product, cosine_distance, and array_remove with a scalar
needle through check_sliced_list_behavior, and pad its input by a length
that is not a multiple of 64, so the list's elements start partway
through a 64-bit word.
x86 (AMD EPYC-Milan, Hetzner ccx33), the branch against main, fastest
of 3 alternating rounds; no other commit on the branch changes these
functions:
no_nulls null_elements
inner_product/width_4 -88.8% -81.2%
inner_product/width_32 -88.9% -84.0%
inner_product/width_256 -79.5% -76.7%
cosine_distance/width_4 -78.5% -73.2%
cosine_distance/width_32 -57.1% -57.5%
cosine_distance/width_256 -17.7% -16.3%
array_normalize/width_4 -80.9% -71.0%
array_normalize/width_32 -77.2% -73.6%
array_normalize/width_256 -65.2% -64.2%
array_remove with a scalar needle, lists of 10 elements:
array_remove_int64 -11.7%
array_remove_n_int64 -14.8%
array_remove_all_int64 -9.8%
array_remove_strings -3.0%
array_remove_boolean -12.7%
array_remove_fixed_size_binary -11.6%
With lists of 100 and 500 elements, which do more work per row,
array_remove moved -2.1% to +0.8%. The nested-list cases, whose list
needle takes a path this does not change, moved -1.9% to +2.7%.
…and map_from_entries Both sliced the element validity bitmap for each list row to check whether the row contains a NULL. Use rows_have_nulls, which counts the valid elements of each row in one pass over the bitmap. Add tests that a sliced list gives the same rows as the whole list. x86 (AMD EPYC-Milan, Hetzner ccx33), the branch against main, fastest of 3 alternating rounds; no other commit on the branch changes these functions: array_contains/list/width_4 -67.2% array_contains/list/width_16 -53.6% array_contains/list/width_128 -10.2% array_contains/fixed_size_list/width_16 -48.6% map_from_entries/width_4 -8.2% map_from_entries/width_16 -3.1% map_from_entries/width_64 -2.8%
For lists of non-primitive values, array_compact counted each valid row's non-NULL elements with count_set_bits_offset, which has a fixed cost per call that dominates for short rows, in a loop that otherwise repeated prepare_list_filter. Call prepare_list_filter instead, which counts with SetBitCounter in one pass over the element validity and reuses that validity as the keep bitmap until a NULL row stores elements. The old array_compact loop did no work for NULL rows, but prepare_list_filter counted each NULL row's set bits and appended its cleared bits to the keep bitmap one row at a time. Calling it as it was made array_compact nearly three times slower when most rows were NULL and stored elements. prepare_list_filter now handles the elements between two valid rows together: when a valid row starts after the previous one ended, it counts the set bits in between once and clears them with one append, and it clears any elements after the last valid row at the end. A NULL row only adds its offset. A valid row that directly follows the previous valid row still costs one count, so inputs whose NULL rows store no elements do the same work as before. array_filter also builds its keep bitmap with prepare_list_filter, so it shares this handling. x86 (AMD EPYC-Milan, Hetzner ccx33), the branch against main, fastest of 3 alternating rounds; no other commit on the branch changes array_compact: dictionary/mixed_nulls/2048 -41.4% dictionary/mixed_nulls/1048576 -40.5% random/utf8/random_sparse -44.4% random/utf8/random -40.1% random/utf8/random_half -39.3% random/utf8/random_dense -65.4% random/utf8_view/random_sparse -56.2% random/utf8_view/random -40.1% random/utf8_view/random_half -47.1% random/utf8_view/random_dense -68.9% types/utf8/mixed -35.7% nulls/utf8/clustered -49.0% types/boolean/mixed -36.0% types/nested/sparse -31.0% types/struct_nested/sparse -28.6% types/fixed_nested/sparse -18.7% types/map/sparse -21.0% types/run/mixed -23.1% strings/utf8_long -19.9% shape/utf8/width_128 -4.2% large_list/utf8 -36.5% child_slice/utf8 -36.2% large/utf8/random_tenth -17.1% large/utf8/random_half -17.2% large/utf8_long/random_half -0.6% binary/random_half -26.9% large_binary/random_half -24.0% strings/null_payload -68.4% With NULL rows that store elements: parent_nulls/utf8/some +0.5% parent_nulls/utf8/mostly +42.5% parent_nulls/nested/mostly +9.5% strings/null_payload_null_rows +0.6% Lists whose rows are mostly NULL but still store elements remain slower than before, because SetBitCounter counts every 64-bit word of the element validity, including the words under NULL rows, which the old loop never read. NULL rows that store elements are rare. Cases with primitive values, which take a separate path, and cases that return before reaching this code moved -2.0% to +5.3%. array_filter's cases with NULL rows moved -5.6% to +8.5%, and its cases without NULL rows, which run none of the code this changes, moved -5.7% to +12.1%.
github-actions
Bot
added
common
|
nice |
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Thanks for the thorough benchmarks. The prepare_list_filter run-of-NULL-rows rewrite reads correct to me for leading, interior and trailing hidden rows. Suggestions below, none blocking.
Small items:
SetBitCounter::count_rangerunscount_before(start)for every row, but for contiguous rowsstartequals the previousend, so the same mask and popcount is repeated.rows_have_nullsalready carries the previous count by hand. Caching the last(end, count)inside the counter would give both call sites the same behavior. Unmeasured, so treat it as optional.prepare_list_filterbuildsSetBitCounter::new(keep.bit_chunks().iter_padded())in both branches. It can be hoisted above theif.map_from_entries.rs:.and_then(|nulls| Some(..).filter(..))is a.map(..).filter(..). The inversion!&has_nullsalso allocates a second bitmap and thennull_countscans it again.
| } | ||
|
|
||
| impl<I: Iterator<Item = u64>> SetBitCounter<I> { | ||
| pub fn new(mut words: I) -> Self { |
There was a problem hiding this comment.
Now that this is public, two things about new.
- It has no rustdoc, and the zero-padding of the final word is an invariant the caller must uphold.
bit_chunks().iter()withoutiter_padded()compiles, and the counter then reads the tail word as0viaunwrap_or(0), so counts for positions in the last partial word are silently low. A constructor taking&BooleanBuffer(returningSetBitCounter<impl Iterator<Item = u64> + '_>) would remove that footgun and replace the six repeatedSetBitCounter::new(x.bit_chunks().iter_padded())call sites. - The monotonic-position precondition is only a
debug_assert!(end >= self.word_start), which does not catch a regression inside the current word. By reading, in a release build a call withend < word_startwrapsend - self.word_startand thewhileloop does not terminate in practice. Please document it on the type, or use a cheapassert!. I did not run this.
|
|
||
| /// Returns a bitmap whose bit `i` is set if row `i` of `list` contains a NULL | ||
| /// element, or `None` if the elements have no validity bitmap. | ||
| pub(crate) fn rows_with_null_elements<O: OffsetSizeTrait>( |
There was a problem hiding this comment.
This body is the same as generic_list_row_has_nulls in datafusion/spark/src/function/array/array_contains.rs (list.values().nulls() then rows_have_nulls(nulls, offsets[0], offsets[1..])). A GenericListArray<O> variant next to rows_have_nulls in datafusion-common would serve both crates and drop both wrappers.
Also, nulls() returns Some for an all-valid bitmap, in which case this scans every word for nothing. Filtering on nulls.null_count() != 0 (O(1)) restores the old per-row null_count() != 0 short cut. The new benches only use None for no_nulls, so that case is not measured.
There was a problem hiding this comment.
The new benches only use
Noneforno_nulls, so that case is not measured.
With None, rows_with_null_elements returns before calling rows_have_nulls, so the no_nulls cases in vector_math never run SetBitCounter. As I read the diff, their speedup comes from replacing the per-row values.slice(start, len) with &values.values()[start..end]. Could float_lists also build an all-valid NullBuffer as a third case? That would measure the path this comment is about, and the counter's cost when nothing is NULL.
| let slice1 = values1.slice(start1, len1); | ||
| let slice2 = values2.slice(start2, len2); | ||
| if slice1.null_count() != 0 || slice2.null_count() != 0 { | ||
| let has_null = |
There was a problem hiding this comment.
The has_null closure is rebuilt per row and does two bitmap lookups per row. The same code is in cosine_distance.rs. OR-ing the two bitmaps once before the loop (Option<BooleanBuffer> combined with &a | &b) gives one lookup per row and removes the duplicated closure in both files.
| use std::sync::Arc; | ||
|
|
||
| #[test] | ||
| fn test_sliced_list() -> datafusion_common::Result<()> { |
There was a problem hiding this comment.
Two points on the new tests.
- Sibling UDFs name this test
test_sliced_capacity(array_normalize,array_add,array_scale,remove). Please use the same name here and incosine_distance.rs. - Both new tests pass the same array as both arguments, so an indexing mistake in the second bitmap (
null_rows2) cannot fail them.inner_product.sltandcosine_distance.sltonly cover NULL elements with single constant rows. A multi-rowVALUEScase with a NULL element only in the second argument, and one only in the first, would cover the per-row indexing and belongs in the.sltfiles.
|
Thank you for opening this pull request! Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). Details |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #26507 +/- ##
=======================================
Coverage 82.92% 82.92%
=======================================
Files 1155 1155
Lines 458204 458274 +70
Branches 458204 458274 +70
=======================================
+ Hits 379965 380041 +76
+ Misses 55205 55192 -13
- Partials 23034 23041 +7 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
mbutrovich
left a comment
mbutrovich
left a comment
There was a problem hiding this comment.
Thanks @neilconway. The benchmarks for every call site make this easy to follow. I ran vector_math and array_contains on aarch64 (Apple M5 Max, default target CPU) against the base commit, and every case is faster, from 4.5% (array_normalize/width_256/no_nulls) to 80% (inner_product/width_32/null_elements).
sort_rows_with_nulls in array_sort still slices the element validity and counts it for each row (sort.rs). Should it use SetBitCounter too, or does the per-row sort cost hide the difference there?
|
|
||
| /// Returns a bitmap whose bit `i` is set if row `i` of `list` contains a NULL | ||
| /// element, or `None` if the elements have no validity bitmap. | ||
| pub(crate) fn rows_with_null_elements<O: OffsetSizeTrait>( |
There was a problem hiding this comment.
The new benches only use
Noneforno_nulls, so that case is not measured.
With None, rows_with_null_elements returns before calling rows_have_nulls, so the no_nulls cases in vector_math never run SetBitCounter. As I read the diff, their speedup comes from replacing the per-row values.slice(start, len) with &values.values()[start..end]. Could float_lists also build an all-valid NullBuffer as a third case? That would measure the path this comment is about, and the counter's cost when nothing is NULL.
| /// assert_eq!(counter.count_before(1), 1); | ||
| /// assert_eq!(counter.count_range(1, 5), 2); | ||
| /// ``` | ||
| pub struct SetBitCounter<I> { |
There was a problem hiding this comment.
Should this move to arrow-rs eventually? As far as I can tell, arrow-rs has no counterpart. Its closest APIs answer one range per call. Buffer::count_set_bits_offset builds a new UnalignedBitChunk on every call, which is the overhead this type avoids. BooleanBuffer::find_nth_set_bit_position answers the inverse question (where the n-th set bit is) and also rescans from start on every call. A counter like this one, or a kernel that takes an OffsetBuffer and returns per-row counts, would fit next to them in arrow-buffer. arrow-rs can ship an additive API like that in a minor release.
I'm happy to land it here first. Could you open an arrow-rs issue for the upstream API, and a DataFusion issue to switch to it once it lands? Because SetBitCounter becomes public in datafusion-common, switching later means a deprecation cycle. That makes the constructor question in @comphead's thread worth settling with the upstream API in mind. arrow-rs exposes its bitmap iterators as &self methods on BooleanBuffer (bit_chunks, set_indices, set_slices), so a constructor that takes &BooleanBuffer would also match how this would look upstream.
4 participants
Which issue does this PR close?
array_filter#25754Rationale for this change
There are several places where we want to count the set bits within a certain range of indices in a bitmap: for example, to count how many non-null elements are in a given row of a list column. Previously, this was typically done by slicing the bitmap and then calling Arrow's
count_set_bits_offset. This has a lot of per-call overhead, which can be avoided when we want to count bits repeatedly on the same bitmap (e.g., for each row in the scenario above). #25754 introducedSetBitCounteras a small abstraction to do this more efficiently for thearray_filterUDF. This PR refines that API, makes it public, and uses it in several other places in DataFusion: Spark'sarray_containsandmap_from_entriesUDFs, and core'sarray_compact,array_remove,inner_product,cosine_distance, andarray_normalizeUDFs.Benchmarks: (x86, AMD EPYC-Milan)
vector_math
array_remove (scalar needle)
array_contains (Spark)
map_from_entries (Spark)
array_compact (non-primitive values)
array_filter (lists with NULL rows)
What changes are included in this PR?
array_containsandmap_from_entriesand core'sinner_product,cosine_distance, andarray_normalizeSetBitCountertodatafusion-common, make it public, and generalize the APISetBitCounterAPI inarray_remove,inner_product,cosine_distance, andarray_normalizecore UDFs, and inarray_containsandmap_from_entriesSpark UDFsarray_compact's result for non-primitive values withprepare_list_filter, whicharray_filteralready uses. This is slower when NULL list rows still store elements, which is rare but can occur (e.g., the result ofnullif). For such inputs,SetBitCounterreads the validity bits for "hidden" rows, which the old loop skipped. To keep that cost down,prepare_list_filternow handles each run of NULL rows between two valid rows with one count and one append, instead of one per NULL row. The remaining cost shows in extreme cases:parent_nulls/utf8/mostly(99% of rows NULL, 256 elements each) is 42% slower andparent_nulls/nested/mostlyis 9.5% slower. A counter that skips the words under NULL rows avoided this, but it slowedarray_filteron lists without NULL rows by 34% on x86, so I didn't pursue it.What is the testing strategy for this PR?
Existing tests pass; new tests added.
Are there any user-facing changes?
New public
SetBitCounterAPI.