如何在siuba中匹配列值并提取索引生成preceding_vel列？

我的目标是在指定frame_id下，查找preceding列的值在vehicle_id列中的匹配项，并提取对应的v_vel值到名为preceding_vel的新列中。我希望使用Python的siuba包来实现该需求，我的数据集如下：

import pandas as pd

df_mini_dict = {'vehicle_id': {884: 2, 885: 2, 886: 2, 14148: 44, 14149: 44, 14150: 44}, 
'frame_id': {884: 338, 885: 339, 886: 340, 14148: 338, 14149: 339, 14150: 340}, 
'preceding': {884: 44, 885: 44, 886: 44, 14148: 3355, 14149: 3355, 14150: 3355}, 
'v_vel': {884: 6.299857770322456, 885: 6.427411525504063, 886: 6.590098168958994, 14148: 7.22883474245701, 14149: 6.973590500351793, 14150: 6.727721962795176}}

df_mini = pd.DataFrame.from_dict(df_mini_dict)

我可以通过以下R代码实现该目标：

df_mini <- structure(list(vehicle_id = c(2L, 2L, 2L, 44L, 44L, 44L), 
                          frame_id = c(338L, 339L, 340L, 338L, 339L, 340L), 
                          preceding = c(44L, 44L, 44L, 3355L, 3355L, 3355L), 
                          v_vel = c(6.29985777032246, 6.42741152550406, 
                                    6.59009816895899, 7.22883474245701, 
                                    6.97359050035179, 6.72772196279518), 
                          preceding_vel = c(7.22883474245701, 6.97359050035179, 
                                            6.72772196279518, NA, NA, NA)), 
                     class = c("tbl_df", "tbl", "data.frame"), 
                     row.names = c(NA, -6L))

library(dplyr)

df_mini <- df_mini |> 
  dplyr::group_by(frame_id) |>  
  dplyr::mutate(preceding_vel = v_vel[match(preceding, vehicle_id)]) |> 
  dplyr::ungroup()

我本质上是想在siuba中实现dplyr的相同功能，但似乎需要用index()来替代match的作用。我尝试了以下代码但未成功：

def match(x, table):
  indicez = []
  for i in x:
    indicez.append(table.index(i))
  return indicez

from siuba import *

df_mini = (
    df_mini
    >> group_by(_.frame_id)  # grouping by frame id
    >> mutate(preceding_vel = _.v_vel[match(_.preceding, _.vehicle_id)])
)

报错：

TypeError: 'Symbolic' object is not iterable

请指导我如何正确定义match函数或使用其他方法来实现该目标，谢谢。

方法1：利用分组映射实现 #

siuba的分组操作可以结合pandas的字典映射逻辑，直接为每个分组构建vehicle_id到v_vel的映射关系，再通过preceding列匹配取值：

from siuba import *

df_mini = (
    df_mini
    >> group_by(_.frame_id)
    >> mutate(
        # 生成分组内vehicle_id与v_vel的映射字典
        vel_map = lambda d: dict(zip(d.vehicle_id, d.v_vel)),
        # 用preceding列匹配映射，无匹配返回NA
        preceding_vel = _.preceding.map(_.vel_map)
    )
    >> select(-_.vel_map)  # 删除临时映射列
    >> ungroup()
)

方法2：实现兼容siuba的向量化match函数 #

自定义函数报错是因为siuba传递的是Symbolic对象而非直接的Series，改用向量化的pandas操作实现match逻辑：

import pandas as pd
from siuba import *

def pd_match(x, table):
    # 构建table值到索引的映射，无匹配返回NaN
    pos_map = pd.Series(range(len(table)), index=table)
    return x.map(pos_map).fillna(-1).astype(int)

df_mini = (
    df_mini
    >> group_by(_.frame_id)
    >> mutate(
        # 获取匹配的索引位置
        match_idx = pd_match(_.preceding, _.vehicle_id),
        # 根据索引取v_vel，无匹配设为NA
        preceding_vel = lambda d: d.v_vel.iloc[d.match_idx].where(d.match_idx != -1)
    )
    >> select(-_.match_idx)  # 删除临时索引列
    >> ungroup()
)

验证结果 #

两种方法都会得到与R语言方案一致的输出：

内容的提问来源于stack exchange，提问作者umair durrani