Uh oh!
There was an error while loading. Please reload this page.
GH-49677 [Python][C++][Compute] Add search sorted compute kernel - #49679
GH-49677 [Python][C++][Compute] Add search sorted compute kernel#49679Alex-PLACET wants to merge 50 commits into
Conversation
8e09ea3 to
4ec630eCompareUh oh!
There was an error while loading. Please reload this page.
pitrou
left a comment
There was a problem hiding this comment.
I haven't looked at the implementation yet, but have reviewed the tests and benchmarks (which are quite comprehensive, thank you!).
One missing item is support for chunked arrays. Besides that, see comments below :)
Uh oh!
There was an error while loading. Please reload this page.
| SearchSortedOptions(SearchSortedOptions::Left))); | ||
| ASSERT_OK_AND_ASSIGN(auto right, | ||
| SearchSorted(Datum(values), Datum(needles), | ||
| SearchSortedOptions(SearchSortedOptions::Right))); |
There was a problem hiding this comment.
Let's call ValidateFull on both results here too?
(also I'm curious, why not reuse CheckSimpleSearchSorted?)
| std::string scalar_needle_json; | ||
| uint64_t expected_scalar_left; | ||
| uint64_t expected_scalar_right; |
There was a problem hiding this comment.
Note that you could also generate the scalar needle tests automatically by calling GetScalar on the array needles and the expected results. This would make this easier to maintain later.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
| RunSearchSortedBenchmark(state, values, needles, side); | ||
| } | ||
| static void BM_SearchSortedInt64ScalarNeedle(benchmark::State& state, |
There was a problem hiding this comment.
If there is just one needle to search for, we're just benchmarking the function call overhead rather than any significant part of the sorted search kernel, right? Is it useful?
(benchmarks for other compute functions focus on array performance, not scalar performance)
Uh oh!
There was an error while loading. Please reload this page.
pitrou
commented
Apr 23, 2026
@Alex-PLACET Is this ready for review again? |
pitrou
commented
Apr 23, 2026
@Alex-PLACET It looks like the benchmarks don't compile anymore? |
pitrou
left a comment
There was a problem hiding this comment.
Now I also took a look at the implementation and some aspects of it seem overly-complicated (but I may miss something). See comments below.
| "[1, 3, 3, 5]", | ||
| "[0, 3, 4, 6]", |
There was a problem hiding this comment.
Can we make the values and needles arrays different lengths?
| auto values = ArrayFromJSON(int16(), "[]"); | ||
| auto needles = ArrayFromJSON(int16(), "[1, 2, 3]"); | ||
| ASSERT_OK_AND_ASSIGN(auto result, SearchSorted(Datum(values), Datum(needles))); |
There was a problem hiding this comment.
Why not reuse CheckSimpleSearchSorted? It would validate the results and also stress scalar needles.
There was a problem hiding this comment.
You are right, now it's the case
| auto values = ArrayFromJSON(int32(), "[null, 200, 300, 300]"); | ||
| auto needles = ArrayFromJSON(int32(), "[50, 200, 250, 400]"); | ||
| ASSERT_OK_AND_ASSIGN(auto left, |
There was a problem hiding this comment.
Again, let's re-use the testing helper CheckSimpleSearchSorted everywhere possible, otherwise we're forgetting some checks.
| {"Float32", float32(), "[1.0, 3.0, 3.0, 5.0]", "[0.0, 3.0, 4.0, 6.0]", | ||
| "[0, 1, 3, 4]", "[0, 3, 3, 4]"}, | ||
| {"Float64", float64(), "[1.0, 3.0, 3.0, 5.0]", "[0.0, 3.0, 4.0, 6.0]", | ||
| "[0, 1, 3, 4]", "[0, 3, 3, 4]"}, |
There was a problem hiding this comment.
Do we want to test NaNs here? They will probably need special treatment, whether in the values or the needles.
For the record, sorting an array with Arrow always put NaNs between nulls and non-null values.
e.g. either [1.0, 3.0, 3.0, 5.0, NaN, NaN, null] or [null, NaN, NaN, 1.0, 3.0, 3.0, 5.0].
There was a problem hiding this comment.
I completely missed this case, I added some tests for that:
FloatValuesWithTrailingNaNsAndNulls and FloatValuesWithTrailingNaNsAndNulls
| ASSERT_OK_AND_ASSIGN(auto left, | ||
| SearchSorted(Datum(values), Datum(needles), | ||
| SearchSortedOptions(SearchSortedOptions::Left))); | ||
| ASSERT_OK_AND_ASSIGN(auto right, | ||
| SearchSorted(Datum(values), Datum(needles), | ||
| SearchSortedOptions(SearchSortedOptions::Right))); |
There was a problem hiding this comment.
Either you write or reuse a helper here, or you need to call ValidateFull directly.
There was a problem hiding this comment.
And ideally we should also check scalar needles too.
| auto values_type = run_end_encoded(int32(), int32()); | ||
| ASSERT_OK_AND_ASSIGN(auto ree_values, | ||
| REEFromJSON(values_type, "[0, 0, 1, 1, 1, 4, 4, 9]")); | ||
| auto sliced = ree_values->Slice(2, 5); |
There was a problem hiding this comment.
Can we use a slice that does not happily falls on run boundaries? (for example Slice(1, 5))
| AssertArraysEqual(*ArrayFromJSON(uint64(), "[0, 0, 3, 3, 5]"), *result.make_array()); | ||
| } | ||
| TEST(SearchSorted, BinaryValues) { |
There was a problem hiding this comment.
This is already part of the smoke tests, right?
There was a problem hiding this comment.
It was a left over your are right, it's fixed now
Uh oh!
There was an error while loading. Please reload this page.
| const auto& needle_data = needles.array(); | ||
| if (needle_data->type->id() == Type::RUN_END_ENCODED) { | ||
| RunEndEncodedArray ree(needle_data); | ||
| return DispatchRunEndEncodedByRunEndType<Status>( |
There was a problem hiding this comment.
Similarly to what I pointed above about run-end-encoded values, you don't need complex indexing for run-end-encoded needles. It's actually even simpler as you can just search each physical needle and run-end-decode the (needle run ends, physical results).
| const int64_t leading_null_count = CountLeadingNulls( | ||
| values.length, [&](int64_t index) { return values.IsNull(index); }); | ||
| const int64_t trailing_null_count = | ||
| leading_null_count > 0 ? 0 : CountTrailingNulls(values.length, [&](int64_t index) { | ||
| return values.IsNull(index); | ||
| }); | ||
There was a problem hiding this comment.
This is... complicated. We know there are null_count nulls, and they are either clustered at the start or at the end. We just have to look up the first validity bit (values.IsNull(0)) and then we know how many leading/trailing nulls there are.
And we can do the exact same thing for chunked arrays.
There was a problem hiding this comment.
You are right, I simplified and refactored it
pitrou
left a comment
There was a problem hiding this comment.
Thanks for the update! This is getting better, though I think there are still opportunities for making the implementation simpler.
| ArrayVector{empty_chunk, null_chunk, empty_chunk, last_null_chunk}); | ||
| auto needles = ArrayFromJSON(int32(), "[1, 4, null]"); | ||
| CheckSearchSorted(Datum(values), Datum(needles), "[3, 3, null]", "[3, 3, null]"); |
There was a problem hiding this comment.
This one raises an interesting question: what is the intended null placement of the sorted values? At end or at start? If nulls are at end, we should return 0 (before the nulls), if nulls are at start, we should return 3 (after the nulls).
The default null placement for sorting in Arrow C++ is NullPlacement::AtEnd so I think this should return 0, not 3.
(we can make it configurable later if there's some demand)
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
| VISIT(UInt16Type) \ | ||
| VISIT(UInt32Type) \ | ||
| VISIT(UInt64Type) \ | ||
| VISIT(FloatType) \ |
There was a problem hiding this comment.
Thanks, but can you also add a test for it?
| RunEndEncodedValuesAccessor<ArrowType> non_null_values(values_accessor.array(), | ||
| non_null_values_range.offset, | ||
| non_null_values_range.length); |
There was a problem hiding this comment.
It seems that you are assuming that nulls in a REE array can be anywhere? But actually the nulls are stored in the REE values (a REE array does not have a top-level null bitmap).
So you don't need all the complication that seems to be related to computing physical offsets depending on the logical null range of the REE array. Just use the physical null range of the REE values.
You should therefore be able to simplify a lot of stuff away, and perhaps even all the REE Accessor classes. You might have to keep the Accessor / ChunkedAccessor distinction, however.
| /// This preserves the cheapest materialization strategy: write repeated | ||
| /// insertion indices directly into the preallocated result buffer, while still | ||
| /// sharing the same EmitInsertionIndices traversal used by the nullable path. | ||
| class PreallocatedInsertionIndexOutput { |
There was a problem hiding this comment.
If you like to have this optimized insertion strategy, you might want to use a TypedBufferBuilder<uint64_t> for both insertion strategies and a TypedBufferBuilder<bool> for the null bitmap (perhaps with a bool template parameter to say whether nulls should be emitted or not).
Something like (rough sketch):
template <boolkEmitNulls>
classInsertionIndexBuilder {
public:explicitInsertionIndexBuilder(MemoryPool* pool)
: indices_builder_(pool), bitmap_builder_pool_(pool) {}
Status Init(int64_t length) {
RETURN_NOT_OK(indices_builder_.Reserve(length));
ifconstexpr (kEmitNulls) {
RETURN_NOT_OK(bitmap_builder_.Reserve(length));
}
returnStatus::OK();
}
voidAppendIndices(uint64_t value, int64_t run_length) {
indices_builder_.UnsafeAppend(run_length, value);
ifconstexpr (kEmitNulls) {
bitmap_builder_.UnsafeAppend(run_length, /*value=*/true);
}
}
voidAppendNulls(int64_t run_length) {
ifconstexpr (kEmitNulls) {
bitmap_builder_.UnsafeAppend(run_length, /*value=*/false);
}
}
private:
TypedBufferBuilder<uint64_t> indices_builder_;
TypedBufferBuilder<uint64_t> bitmap_builder_;
};Uh oh!
There was an error while loading. Please reload this page.
| "[0, 0, 1, 3]", "[0, 1, 1, 3]"); | ||
| } | ||
| TEST(SearchSorted, FloatValuesWithTrailingNaNsAndNulls) { |
There was a problem hiding this comment.
Can we also have a leading NaNs and nulls test?
Uh oh!
There was an error while loading. Please reload this page.
pitrou
commented
Jun 1, 2026
I took a look at the compiled code size here and we're really generating too much code: 48M for the search_sorted kernels alone is almost 20% of the total size of compute kernels (256M for libarrow_compute): This is also confirmed by code sizes (the |
pitrou
commented
Jun 1, 2026
To alleviate this combinatorial explosion, it seems several axes of code generation can be decoupled:
Also, one important improvement is to use |
pitrou
left a comment
There was a problem hiding this comment.
Some more comments below. Most important though is the combinatorial explosion I mentioned in the comments.
| VISIT(UInt16Type) \ | ||
| VISIT(UInt32Type) \ | ||
| VISIT(UInt64Type) \ | ||
| VISIT(FloatType) \ |
There was a problem hiding this comment.
Thanks, but can you also add a test for it?
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
| SearchSortedOptions::Side side, uint64_t insertion_offset, | ||
| Output* output) { | ||
| auto emit_search_result = [&](const VisitedNeedle<ArrowType, EmitNulls>& needle, | ||
| int64_t run_length) -> Status { |
There was a problem hiding this comment.
run_length is always 1 in all call sites, right?
| template <typename ArrowType, typename ValuesAccessor> | ||
| Result<Datum> ComputeRunEndEncodedNeedleInsertionIndices( | ||
| const ValuesAccessor& sorted_values, const RunEndEncodedArray& needles, | ||
| SearchSortedOptions::Side side, uint64_t insertion_offset, ExecContext* ctx) { |
There was a problem hiding this comment.
I don't think this needs to be templated on ArrowType or ValuesAccessor. Running search_sorted on the needles should be an opaque operation. This function is mostly interested in:
- unpacking the physical needles from the REE needles
- running search_sorted on the physical needles
- re-packing the search_sorted results with the original run-ends
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Adds a new search_sorted compute function across C++ and Python, including options, docs, tests, and benchmarks.
Changes:
- Introduces C++
SearchSortedAPI + kernel registration and implementation, plus benchmarks. - Adds Python bindings for
SearchSortedOptionsand Python-level tests/docstrings forpc.search_sorted. - Updates C++ user docs to list
search_sortedand document its constraints/behavior.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| python/pyarrow/tests/test_compute.py | Adds Python tests validating pc.search_sorted behavior (left/right, nulls, REE, errors). |
| python/pyarrow/includes/libarrow.pxd | Exposes SearchSortedOptions and side enum to Cython. |
| python/pyarrow/compute.py | Exports SearchSortedOptions in the public Python compute namespace. |
| python/pyarrow/_compute_docstrings.py | Adds Python docstring examples for search_sorted. |
| python/pyarrow/_compute.pyx | Adds SearchSortedOptions Python option class + side parsing. |
| docs/source/cpp/compute.rst | Documents search_sorted in the C++ compute function table and notes constraints. |
| cpp/src/arrow/meson.build | Adds the new kernel source to Meson build. |
| cpp/src/arrow/compute/registry_internal.h | Declares RegisterVectorSearchSorted. |
| cpp/src/arrow/compute/kernels/vector_search_sorted_test.cc | Adds comprehensive C++ kernel tests. |
| cpp/src/arrow/compute/kernels/vector_search_sorted_benchmark.cc | Adds microbenchmarks for SearchSorted. |
| cpp/src/arrow/compute/kernels/vector_search_sorted.cc | Implements the search_sorted meta-function/kernel (arrays, chunked, REE, null handling). |
| cpp/src/arrow/compute/kernels/meson.build | Adds the new benchmark target to Meson. |
| cpp/src/arrow/compute/kernels/CMakeLists.txt | Adds the new C++ test and benchmark targets. |
| cpp/src/arrow/compute/initialize.cc | Registers the kernel at startup. |
| cpp/src/arrow/compute/api_vector.h | Adds SearchSortedOptions and the public SearchSorted API declaration/docs. |
| cpp/src/arrow/compute/api_vector.cc | Registers options type and wires SearchSorted to CallFunction. |
| cpp/src/arrow/CMakeLists.txt | Adds the new kernel source to CMake build. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Uh oh!
There was an error while loading. Please reload this page.
| /// The `values` datum must be a plain array or run-end encoded array sorted in | ||
| /// ascending order. `needles` may be a scalar, plain array, or run-end encoded | ||
| /// array whose logical value type matches `values`. |
Uh oh!
There was an error while loading. Please reload this page.
| const auto it = std::ranges::find_if( | ||
| values.chunks(), | ||
| [](const std::shared_ptr<Array>& chunk) { return chunk->length() != 0; }); |
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
pitrou
commented
Jun 10, 2026
@Alex-PLACET You don't have to answer all Copilot review comments by the way, if it's just hallucinating. |
pitrou
commented
Jun 10, 2026
And please ping when this is ready for review again (presumably not yet, given there are many CI failures) |
@pitrou can you launch the workflows please ? |
Alex-PLACET
commented
Jun 18, 2026
@raulcd Can you run the checks please |
Ensure ToPhysicalData updates the RunEndEncodedType to reflect the new physical values type when converting REE ArrayData. Add a RunEndEncoded date32 test and remove now-obsolete all-null test cases.
e89d93c to
72576b0CompareThere was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated no new comments.
Suppressed comments (2)
cpp/src/arrow/compute/kernels/vector_search_sorted.cc:1065
- ExecuteChunkedNeedles() calls ExecuteImpl() recursively for every needle chunk, which recomputes the non-null values range and repeats validation/type dispatch each iteration. This can add noticeable overhead when needles are highly chunked. Compute non_null_values_range once and call DispatchByType() directly for each chunk.
result_chunks.reserve(static_cast<size_t>(needles.num_chunks()));
for (const auto& chunk : needles.chunks()) {
ARROW_ASSIGN_OR_RAISE(auto chunk_result,
ExecuteImpl({values, Datum(chunk)}, &options, ctx));
result_chunks.push_back(chunk_result.make_array());
cpp/src/arrow/compute/kernels/meson.build:137
- Meson adds the new vector_search_sorted benchmark, but the corresponding test (vector_search_sorted_test.cc) isn’t included in vector_kernel_tests. As a result, Meson-based CI/builds won’t exercise the new kernel test coverage.
'vector_topk_benchmark',
'vector_replace_benchmark',
'vector_search_sorted_benchmark',
'vector_selection_benchmark',
]
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 17 out of 17 changed files in this pull request and generated no new comments.
Suppressed comments (2)
cpp/src/arrow/compute/kernels/vector_search_sorted.cc:638
- FindNonNullValuesRangeFromNullCount() infers null placement by only checking is_null_at(0) and using the total null_count. If nulls are not strictly clustered at one end (e.g. nulls on both ends, or a null in the middle), this will compute an incorrect non-null search window and can return wrong insertion indices instead of rejecting the input.
const bool has_leading_nulls = is_null_at(0);
return MakeNonNullValuesRangeFromNullPlacement(length, null_count, has_leading_nulls);
}
docs/source/cpp/compute.rst:1916
- This note lists supported needle input shapes, but it omits chunked arrays even though the C++ API and kernel accept chunked needles (and chunked run-end encoded needles). Updating this avoids underspecifying the public C++ behavior.
needles are matched only against the non-null portion. The second argument
may be a scalar, array, or run-end encoded array. Null needles yield null
outputs. Both arguments must have the same logical type. A scalar needle
pitrou
commented
Aug 27, 2026
I've noticed a bug where NaN values at the start (rather than at the end) at the haystack weren't found properly. I've fixed it along with a number of changes to reduce code size. We're down to ~250 kB of additional code on a release build, which is good (around a 2% size increase on |
pitrou
commented
Aug 27, 2026
I also added float16 support since we were close to adding it and it didn't make sense to sandbag it. |
Rationale for this change
Add the implemenation of the search sorted compute kernel based on the numpy function: https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html
What changes are included in this PR?
Implementation of the C++ kernel + Python API.
Tests in C++ and Python
Are these changes tested?
Yes
Are there any user-facing changes?
No breaking change
search_sortedkernel for all primitive types and run-end encoded arrays #49677