Optimize SearchSorted to reuse probes and avoid repeated context creation - #9904
Draft
joseph-isaacs wants to merge 4 commits into
Draft
joseph-isaacs wants to merge 4 commits into
joseph-isaacs wants to merge 4 commits into
Conversation
`SearchSortedPrimitiveArray` held a bare `&ArrayRef` and did a one-off `execute_scalar` per comparison, so every probe in a binary search rebuilt the state its predecessor had just thrown away. For a nullable array that meant resolving `Validity` from the encoding on each read, and `IndexOrd<Option<T>>` paid for it twice: once in its own `is_valid` call and again inside `execute_scalar`, which checks validity before dispatching. Hold a `RepeatedArrayProbe` instead. Validity is resolved on the first comparison and reused by the remaining ~log2(n), as is any state the encoding keeps. A null element reads back as a null scalar, so the separate `is_valid` call is redundant: `IndexOrd<Option<T>>` now decides from the one read, halving the probes on the nullable path. The array is no longer borrowed for the searcher's lifetime, since the probe owns a handle to it; only the execution context is. Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPw4v7SA2A3Ab5t9zwGbRV
…rray `Patches::search_index` already sidesteps the scalar path when its indices are a canonical primitive array, searching the buffer as a `&[T]`. `RunEnd::find_physical_index` has no such path, so a run-end array whose ends are plain `u32`s — which is what they are after a file read — still pays a scalar read per probe. Put the fast path in the searcher instead of in each caller, so every `SearchSortedPrimitiveArray` user gets it. The buffer can only be read as `[T]` when the array is canonical, is host-backed, and is non-nullable, since a null element leaves an arbitrary value in the buffer; everything else keeps the probe. The array is borrowed for the searcher's lifetime again, as the values are. This also settles the cost the previous commit adds on non-nullable arrays, where a retained read buys nothing while no encoding keeps state: those arrays no longer reach the probe at all. What still does — compressed ends, non-primitive patch indices — is where a retained probe pays off once those encodings keep state. Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPw4v7SA2A3Ab5t9zwGbRV
`IndexOrd<Scalar> for ArrayRef` built a whole execution context per comparison — `legacy_session().create_execution_ctx()` inside `index_cmp` — so a binary search over 65,536 elements created sixteen of them and threw each away after one read. A trait impl on `ArrayRef` has nowhere to keep anything, so this could not be fixed in place. Replace it with `SearchSortedArray`, which takes the caller's context and a `RepeatedArrayProbe`, matching `SearchSortedPrimitiveArray`. The searcher now has the same shape whether or not the element type is known, and the doc points at the typed one, which reads canonical values directly. The only caller was the fuzzer, which already had a context in scope. Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPw4v7SA2A3Ab5t9zwGbRV
…eads `Patches::search_index_chunked` read three chunk offsets out of the same array through `chunk_offset_at`, which builds an execution context and a one-off probe per call — three of each per lookup, for three reads of one small array. Read them through a single `RepeatedArrayProbe` and a single context. `chunk_offset_at` stays as the public single-read accessor. Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk> Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UPw4v7SA2A3Ab5t9zwGbRV
Merging this PR will regress 1 benchmark
|
| Mode | Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|---|
| ❌ | WallTime | filtered_owned_i64_avx512[OneNullInEight] |
22.8 µs | 26.8 µs | -14.81% |
| ⚡ | WallTime | arrow_checked_add_u32_neon[16384] |
20.4 µs | 13.5 µs | +51.34% |
| ⚡ | Simulation | allocate_drop_bytes[0] |
575.7 ns | 521.6 ns | +10.39% |
Tip
Investigate this regression by commenting @codspeedbot fix this regression on this PR, or directly use the CodSpeed MCP with your agent.
Comparing ji/patches-chunk-offset-probe (0ea6472) with develop (4af09a9)
Footnotes
-
218 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩
robert3005
reviewed
Sep 16, 2026
| // The three reads below are of the same array, so they share one probe and one context | ||
| // rather than building a pair per read as `Self::chunk_offset_at` does. | ||
| let mut probe = chunk_offsets.repeated_probe(); | ||
| let mut ctx = legacy_session().create_execution_ctx(); |
Contributor
There was a problem hiding this comment.
just pass the context to the method
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This change optimizes the
SearchSortedimplementations to reuse a singleRepeatedArrayProbeand execution context across all comparisons in a search, rather than creating a new probe and context for each element access. This significantly reduces overhead for searches over encoded arrays.Additionally, a new
SearchSortedArrayadapter is introduced to support searching over arrays of any encoding usingScalarcomparisons, with the same probe-reuse optimization.Changes
SearchSortedPrimitiveArrayOptimizationSearchSortedPrimitiveArrayfrom a tuple struct to a struct with named fields (reader,len,ctx,_ptype)Reader<'a, T>enum that distinguishes between two cases:Values(&'a [T]): Direct buffer access for canonical, non-nullable, host-backed arrays (zero-copy fast path)Probe(RefCell<RepeatedArrayProbe>): Probe-based access for all other casesreader()method to determine which access pattern to use based on array propertiestyped_value()method that returnsOption<T>to distinguish nulls from zero valuesvalue()method to usetyped_value()and map nulls toT::zero()IndexOrd<Option<T>>implementation to use a singletyped_value()call instead of separate validity and value checksNew
SearchSortedArrayImplementationscalar.rsmodule withSearchSortedArraystruct for searching any array type usingScalarcomparisonsIndexOrd<Scalar>with probe reuse across the entire searchModule Organization
SearchSortedArrayfromsearch_sorted/mod.rsIndexOrd<Scalar>implementation forArrayRef(now handled bySearchSortedArray)PatchesOptimizationsearch_index_chunked()to reuse a single probe and context for three reads ofchunk_offsets, rather than creating a new pair per read#[allow(clippy::disallowed_methods)]annotation to document the intentional use oflegacy_session()Test Coverage
search_sorted_reads_canonical_values_directly()test verifying the fast path for canonical arrayssearch_sorted_probes_when_values_are_not_addressable()test verifying probe-based access for nullable arrayssearch_sorted_scalar()andsearch_sorted_scalar_with_nulls()tests for the newSearchSortedArrayassert_search_sorted()API Changes
The
SearchSortedPrimitiveArraystruct layout changed from a tuple struct to a named-field struct. This is a breaking change for any code that directly constructs or pattern-matches on this type, though it is primarily used through theSearchSortedtrait.A new public type
SearchSortedArrayis introduced for searching arrays withScalarcomparisons.https://claude.ai/code/session_01UPw4v7SA2A3Ab5t9zwGbRV