Polars中search_sorted功能的相关技术问题咨询

在NumPy/Pandas中，searchsorted是用于批量二分搜索的工具，常用于时序数据场景，示例如下：

import numpy as np

np.searchsorted(['a', 'a', 'b', 'c'], ['a', 'b', 'c']) # 返回 [0, 2, 3]
np.searchsorted(['a', 'a', 'b', 'c'], ['a', 'b', 'c'], side='right') # 返回 [2, 3, 4]

针对你提出的Polars相关问题，解答如下：

1. 是否可以以向量化方式在Polars中使用search_sorted？ #

不可以。当前Polars的Series.search_sorted()仅支持单个标量输入，直接传入列表/Series会抛出PanicException: dtype List not implemented错误，无法直接实现批量向量化搜索。

2. 是否可以为Polars的search_sorted指定side=right参数？ #

不可以。当前Polars的search_sorted()未提供公开的side参数，且其默认行为与NumPy的side='left'/side='right'均不一致（例如你测试中搜索0在[0,0,1,2]中返回1，与NumPy两种模式的结果都不同）。

3. Polars的search_sorted是否支持非数值数据？ #

不可以。传入Utf8类型的Series调用search_sorted()会抛出PanicException: not implemented for Utf8，仅支持数值类型的标量搜索。

4. 替代方案 #

实现向量化+多side参数+非数值支持 #

可以借助NumPy的searchsorted实现对齐预期的功能，以下是几种实用方式：

方式1：处理单个Series的批量搜索

import polars as pl
import numpy as np

# 非数值类型示例
s = pl.Series(['a', 'a', 'b', 'c'])
targets = pl.Series(['a', 'b', 'c'])

result_left = pl.Series(np.searchsorted(s.to_numpy(), targets.to_numpy(), side='left'))
result_right = pl.Series(np.searchsorted(s.to_numpy(), targets.to_numpy(), side='right'))

print(result_left)  # shape: (3,), dtype: int64, values: [0, 2, 3]
print(result_right) # shape: (3,), dtype: int64, values: [2, 3, 4]

方式2：DataFrame列级表达式（推荐）

使用map_elements或自定义UDF实现列间的批量搜索：

df = pl.DataFrame({
    'a': [0, 0, 1, 2],
    'b': [0, 1, 2, 3],
})

# 方法A：map_elements
df = df.with_columns(
    pl.col('b').map_elements(
        lambda x: np.searchsorted(df['a'].to_numpy(), x, side='left'),
        return_dtype=pl.Int64
    ).alias('c_left'),
    pl.col('b').map_elements(
        lambda x: np.searchsorted(df['a'].to_numpy(), x, side='right'),
        return_dtype=pl.Int64
    ).alias('c_right')
)

# 方法B：自定义UDF（更符合表达式风格）
from polars import udf

@udf(return_dtype=pl.Int64)
def search_sorted_left(arr: list, target: int) -> int:
    return np.searchsorted(arr, target, side='left')

@udf(return_dtype=pl.Int64)
def search_sorted_right(arr: list, target: int) -> int:
    return np.searchsorted(arr, target, side='right')

df = df.with_columns(
    search_sorted_left(pl.col('a'), pl.col('b')).alias('c_left_udf'),
    search_sorted_right(pl.col('a'), pl.col('b')).alias('c_right_udf')
)

print(df)

输出结果：

shape: (4, 6)
┌─────┬─────┬────────┬─────────┬────────────┬─────────────┐
│ a   ┆ b   ┆ c_left ┆ c_right ┆ c_left_udf ┆ c_right_udf │
│ --- ┆ --- ┆ ---    ┆ ---     ┆ ---        ┆ ---         │
│ i64 ┆ i64 ┆ i64    ┆ i64     ┆ i64        ┆ i64         │
╞═════╪═════╪════════╪═════════╪════════════╪═════════════╡
│ 0   ┆ 0   ┆ 0      ┆ 2       ┆ 0          ┆ 2           │
│ 0   ┆ 1   ┆ 2      ┆ 3       ┆ 2          ┆ 3           │
│ 1   ┆ 2   ┆ 3      ┆ 4       ┆ 3          ┆ 4           │
│ 2   ┆ 3   ┆ 4      ┆ 4       ┆ 4          ┆ 4           │
└─────┴─────┴────────┴─────────┴────────────┴─────────────┘

注意事项 #

• UDF和map_elements的性能在超大数据量下略低于Polars原生表达式，但足以覆盖大多数场景；
• Polars仍在快速迭代，后续可关注官方是否会完善原生的search_sorted功能，支持向量化、多side参数和非数值类型。

内容的提问来源于stack exchange，提问作者T.H Rice