apache / apache/arrow

[C++][Python] Implement `search_sorted` kernel for all primitive types and run-end encoded arrays

Open
#49,677 0 comments 0 reactions 1 assignee Claimed by @Alex-PLACET View on GitHub
Component: Python Type: enhancement
Dominant language
C++
Stars
17.1k
Forks
4.3k
Avg merge
3d 13h
Merged PRs (30d)
88

Description

### Describe the enhancement requested

Implement a kernel that replicates the semantics of NumPy’s [`searchsorted`](https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html) function (without the `sorter` argument).
The kernel should support all primitive types as well as run-end encoded arrays.

## Example API

```python
>>> sorted_array = pa.array()
>>> to_search = pa.array()

>>> pc.search_sorted(sorted_array, to_search, side='left')

[
0,
1,
3,
5
]

>>> pc.search_sorted(sorted_array, to_search, side='right')

[
0,
3,
3,
5
]
```

### Explanation
- `50` < all → index `0`
- `200` equals first occurrence → index `1` (`side='left'`)
- `200` for `side='right'` → index `3`
- `250` between `200` and `300` → index `3`
- `400` > all values → index `5`

## Null Handling

### Nulls in the first (sorted) array

If nulls are **clustered first**:

```python
sorted_array = pa.array([null, 200, 300, 300])
to_search = pa.array()
```

Expected:
- `side='left'` → `[1, 1, 2, 4]`
- `side='right'` → `[1, 2, 2, 4]`

If nulls are **clustered last**:

```python
sorted_array = pa.array([200, 300, 300, null, null])
```

Expected:
- `side='left'` → `[0, 0, 1, 3]`
- `side='right'` → `[0, 1, 1, 3]`

### Nulls in the second (to_search) array

Two options:
1. Emit nulls for null search keys.
2. Match nulls within the null portion of the sorted array.

Example for (2):

```python
>>> sorted_array = pa.array([null, null, 200]) # nulls first
>>> to_search = pa.array([null, 100, 300])
>>> pc.search_sorted(sorted_array, to_search, side='left')

[
0,
2,
3
]
>>> pc.search_sorted(sorted_array, to_search, side='right')

[
1,
2,
3
]
>>> sorted_array = pa.array([200, null, null]) # nulls last
>>> pc.search_sorted(sorted_array, to_search, side='left')

[
1,
0,
1
]
>>> pc.search_sorted(sorted_array, to_search, side='left')

[
3,
0,
1
]
```

## Requirements

- Implement for all **primitive types** (`int*`, `float*`, `boolean`, etc.)
- Support **run-end encoded arrays**
- Handle both `side='left'` and `side='right'`
- Define consistent behavior for **null placement**
- Return a `UInt64Array` of insertion indices

## References

- [NumPy `searchsorted` documentation](https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html)

### Component(s)

Python, C++

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.