Optimize SearchSorted to reuse probes and avoid repeated context creation

[joseph-isaacs]

Summary

This change optimizes the SearchSorted implementations to reuse a single RepeatedArrayProbe and 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 SearchSortedArray adapter is introduced to support searching over arrays of any encoding using Scalar comparisons, with the same probe-reuse optimization.

Changes

SearchSortedPrimitiveArray Optimization

• Refactored SearchSortedPrimitiveArray from a tuple struct to a struct with named fields (reader, len, ctx, _ptype)
• Introduced a 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 cases
• Added reader() method to determine which access pattern to use based on array properties
• Added typed_value() method that returns Option<T> to distinguish nulls from zero values
• Simplified value() method to use typed_value() and map nulls to T::zero()
• Updated IndexOrd<Option<T>> implementation to use a single typed_value() call instead of separate validity and value checks
• All field accesses now use the struct fields instead of tuple indexing

New SearchSortedArray Implementation

• Created new scalar.rs module with SearchSortedArray struct for searching any array type using Scalar comparisons
• Implements IndexOrd<Scalar> with probe reuse across the entire search
• Includes tests demonstrating scalar search with and without nulls

Module Organization

• Exported SearchSortedArray from search_sorted/mod.rs
• Removed the generic IndexOrd<Scalar> implementation for ArrayRef (now handled by SearchSortedArray)

Patches Optimization

• Refactored search_index_chunked() to reuse a single probe and context for three reads of chunk_offsets, rather than creating a new pair per read
• Added #[allow(clippy::disallowed_methods)] annotation to document the intentional use of legacy_session()

Test Coverage

• Added search_sorted_reads_canonical_values_directly() test verifying the fast path for canonical arrays
• Added search_sorted_probes_when_values_are_not_addressable() test verifying probe-based access for nullable arrays
• Added search_sorted_scalar() and search_sorted_scalar_with_nulls() tests for the new SearchSortedArray
• Updated fuzz tests to pass execution context to assert_search_sorted()

API Changes

The SearchSortedPrimitiveArray struct 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 the SearchSorted trait.

A new public type SearchSortedArray is introduced for searching arrays with Scalar comparisons.

https://claude.ai/code/session_01UPw4v7SA2A3Ab5t9zwGbRV

[codspeed]

Merging this PR will regress 9 benchmarks

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚠️ 3 benchmarks measured no execution time

Nothing ran under measurement, usually because the compiler removed the code under test. These results are not comparable, so they count as unchanged.

Preventing compiler optimizations

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 9 improved benchmarks
❌ 9 regressed benchmarks
✅ 2162 untouched benchmarks
⏩ 385 skipped benchmarks¹

Warning

Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
❌	WallTime	mul_u64_nonnull_neon	28.5 µs	40.6 µs	-29.8%
❌	WallTime	decode_avx512[8192, (Inline, OneNullInEight)]	81.7 µs	98.2 µs	-16.79%
❌	WallTime	filtered_sink_i64_avx2[OneNullInEight]	21.9 µs	26.2 µs	-16.16%
❌	WallTime	mul_i64_nonnull_neon	32.7 µs	38.7 µs	-15.51%
❌	WallTime	multiply_shapes_neon[(32768, PerRowPerRow)]	32.8 µs	38.6 µs	-15.04%
❌	WallTime	decode_avx512[8192, (Inline, AllValid)]	96.6 µs	109.4 µs	-11.68%
❌	WallTime	decode_avx2[8192, (External, AllValid)]	166.4 µs	187.6 µs	-11.3%
❌	WallTime	dict_canonicalize_gt_u8_neon[1000000]	498.1 µs	560.7 µs	-11.17%
❌	WallTime	decode_avx2[8192, (Inline, AllValid)]	102.4 µs	114.3 µs	-10.39%
⚡	Simulation	take_fsl_f16_force_manual_range_copy[2048, 10]	62.7 µs	10.4 µs	×6
⚡	Simulation	take[duplicates/repeated/primitive/nonnull/chunks=16/indices=1000]	104.5 µs	24.7 µs	×4.2
⚡	Simulation	search_index_below_min_chunked	373.3 µs	283.2 µs	+31.83%
⚡	Simulation	search_index_mixed_out_of_range_chunked	396.7 µs	306.4 µs	+29.46%
⚡	Simulation	search_index_full_range_random_chunked	401.6 µs	310.9 µs	+29.17%
⚡	Simulation	from_vec_drop_arrow[16384]	40.4 µs	31.4 µs	+28.7%
⚡	Simulation	search_index_above_max_chunked	502.5 µs	410.3 µs	+22.47%
⚡	Simulation	search_index_in_range_chunked	503.7 µs	411.6 µs	+22.37%
⚡	WallTime	decode_avx2[8192, (Inline, OneNullInEight)]	97.9 µs	83.4 µs	+17.35%
⚠️	Simulation	fixed_16_advancing_ptr_safe[100]	< 1 ns	< 1 ns	N/A
⚠️	Simulation	preverify_advancing_ptr_unchecked[1000]	< 1 ns	< 1 ns	N/A
...	...	...	...	...	...

ℹ️ Only the first 20 benchmarks are displayed. Go to the app to view all benchmarks.

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 (e7a8cf1) with develop (035aeff)}

Footnotes

1. 385 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]

let's merge this, just one comment about ctx