字典值的条件查找的极坐标列

Question

我想根据列之间的关系将一个 Polars DataFrame 中的键映射到另一个 Polars DF。这只是一个示例，完整的 DF1 和 DF2 要大得多（分别为 250 万行和 150 万行。

DF1 = pl.DataFrame({
'chr' : ["GL000008.2", "GL000008.2", "GL000008.2", "GL000008.2","GL000008.2", "GL000008.2"], 
'start': [14516,17380,17381,20177,22254,24357], 
'end': [14534,17399,17399,20195,22274,24377]
})

DF2 = pl.DataFrame({ 
'key' : [1,2,3,4,5,6],
'chrom' : ["GL000008.2", "GL000008.2", "GL000008.2", "GL000008.2","GL000008.2", "GL000008.2"], 
'start': [14516,15377,17376,20177,22254, 24357], 
'end': [14534,15403,17399,20195,22274,24377]})

我想要的是：

DF1 = pl.DataFrame({
'chr' : ["GL000008.2", "GL000008.2", "GL000008.2", "GL000008.2","GL000008.2", "GL000008.2"], 
'start': [14516,17380,17381,20177,22254,24357], 
'end': [14534,17399,17399,20195,22274,24377],
'key': [1,3,3,4,5,6]
})

当

key

与 chrom 匹配且 DF1 中的

chr

和

start

包含在

end

和

begin 中时，我想将 

DF2 中的 end

 分配给 

DF1 在 DF2。

我首先尝试迭代 DF1 的行，在 DF2 中查找匹配条目：

sz = len(DF1[:,0])

for i in range(sz):
    DF1[i,"key"] = DF2.filter(
        (pl.col("chrom") == DF1[i,"chr"])
        & (pl.col("begin") <= DF1[i,"start"])
        & (pl.col("end") >= DF1[i,"end"])
        ).select('key')[0,0]

通过 DF 进行行迭代非常慢。这大约需要10个小时。

我还尝试使用 np.array 而不是直接进入 df。那有点快，但仍然很慢。

我正在寻找一种使用本机 Polars 数据结构来实现此目的的方法。我没有加入密钥，因此

join

和

join_asof

策略似乎不起作用。

Answer 1

加入和过滤应该可以满足您的需求：

(
    df1.join(df2, left_on="chr", right_on="chrom")
    .filter(
        (pl.col("start") >= pl.col("start_right"))
        & (pl.col("end") <= pl.col("end_right"))
    )
    .drop(["start_right", "end_right"])
)

shape: (6, 4)
┌────────────┬───────┬───────┬─────┐
│ chr        ┆ start ┆ end   ┆ key │
│ ---        ┆ ---   ┆ ---   ┆ --- │
│ str        ┆ i64   ┆ i64   ┆ i64 │
╞════════════╪═══════╪═══════╪═════╡
│ GL000008.2 ┆ 14516 ┆ 14534 ┆ 1   │
│ GL000008.2 ┆ 17380 ┆ 17399 ┆ 3   │
│ GL000008.2 ┆ 17381 ┆ 17399 ┆ 3   │
│ GL000008.2 ┆ 20177 ┆ 20195 ┆ 4   │
│ GL000008.2 ┆ 22254 ┆ 22274 ┆ 5   │
│ GL000008.2 ┆ 24357 ┆ 24377 ┆ 6   │
└────────────┴───────┴───────┴─────┘

Answer 2

使用

join_asof

可能会提供高性能的解决方案：

(
    DF1
    .sort('start')
    .join_asof(
        DF2.sort('start'),
        by_left="chr",
        by_right="chrom",
        on="start",
        strategy="backward")
    .filter(
        pl.col('end') <= pl.col('end_right')
    )
)

shape: (6, 5)
┌────────────┬───────┬───────┬─────┬───────────┐
│ chr        ┆ start ┆ end   ┆ key ┆ end_right │
│ ---        ┆ ---   ┆ ---   ┆ --- ┆ ---       │
│ str        ┆ i64   ┆ i64   ┆ i64 ┆ i64       │
╞════════════╪═══════╪═══════╪═════╪═══════════╡
│ GL000008.2 ┆ 14516 ┆ 14534 ┆ 1   ┆ 14534     │
│ GL000008.2 ┆ 17380 ┆ 17399 ┆ 3   ┆ 17399     │
│ GL000008.2 ┆ 17381 ┆ 17399 ┆ 3   ┆ 17399     │
│ GL000008.2 ┆ 20177 ┆ 20195 ┆ 4   ┆ 20195     │
│ GL000008.2 ┆ 22254 ┆ 22274 ┆ 5   ┆ 22274     │
│ GL000008.2 ┆ 24357 ┆ 24377 ┆ 6   ┆ 24377     │
└────────────┴───────┴───────┴─────┴───────────┘

注意：这假设 DF2 中的开始结束间隔不重叠。