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 5 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.

⚠️ 15 benchmarks spent significant time in system calls

System calls cannot be consistently instrumented, so they are not included in the measure, which understates the real cost. Please switch to the Walltime instrument to accurately measure system calls.

Measurement and system calls

• cold[(16, 64)]
• density_sweep_dense_runs[0.001]
• pcodec_compress_f64
• case_when_fragmented[400]
• multipolygon_mixed_validity
• chunked_dict_primitive_into_canonical[f64, (1000, 1000, 10)]
• decompress[u32, (1000, 256)]
• take_varbin
• optimize_lookup_predicate[ids=256, shape=in_list]
• filter_run_end[len=4096_run=16_density=0.1]
• …and 5 more

⚠️ 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

⚡ 6 improved benchmarks
❌ 5 regressed benchmarks
✅ 2172 untouched benchmarks
⏩ 359 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.7 µs	40 µs	-28.15%
❌	WallTime	multiply_shapes_neon[(32768, PerRowPerRow)]	32.7 µs	38.2 µs	-14.47%
❌	WallTime	mul_i64_nonnull_neon	32.7 µs	38.2 µs	-14.33%
❌ ⚠️	Simulation	cold[(16, 64)]	311.4 µs	355.3 µs	-12.38%
❌	WallTime	bitpack_blocked_compress_avx2	6.7 µs	7.6 µs	-11.35%
⚡ ⚠️	Simulation	density_sweep_dense_runs[0.001]	47.8 µs	29.3 µs	+63.4%
⚡	Simulation	search_index_mixed_out_of_range_chunked	330.9 µs	225.6 µs	+46.66%
⚡	Simulation	search_index_below_min_chunked	332.5 µs	227.1 µs	+46.42%
⚡	Simulation	search_index_full_range_random_chunked	362.9 µs	257.7 µs	+40.81%
⚡	Simulation	search_index_above_max_chunked	459.6 µs	354.2 µs	+29.75%
⚡	Simulation	search_index_in_range_chunked	460.9 µs	355.7 µs	+29.58%

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 (0f919f0) with develop (731a231)}

Footnotes

1. 359 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

[github-actions]

This PR has been marked as stale because it has been open for 14 days with no activity. Please comment or remove the stale label if you wish to keep it active, otherwise it will be closed in 7 days