如何将R/Dplyr中自定义lead值计算逻辑封装为通用函数？

需求说明 #

现有R/dplyr代码按id分组后，为目标列计算t+5位置的lead值；若该值不存在，则依次 fallback 到t+4、t+3…t+1位置的可用值，取最晚的那个：

data <- data %>%
  group_by(id) %>%
  mutate(
    lead_value = coalesce(lead(VALUE, n = 5), 
                           lead(VALUE, n = 4),
                           lead(VALUE, n = 3),
                           lead(VALUE, n = 2),
                           lead(VALUE, n = 1))
)

需要将该逻辑封装为可传入目标列名和偏移量n的通用函数，同时支持多分组维度场景。

示例数据 #

单分组示例 #

data <- data.frame(
  yyyymm = c(202401, 202402, 202403, 202404, 202405, 202406, 202407, 202408, 202409, 202410, 202411, 202412,
             202401, 202402, 202403, 202404, 202405, 202406, 202407),
  id = c(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2),
  value = c(1,2,3,4,5,6,7,8,9,10,11,12,1,2,3,4,5,6,7),
  lead = c(6,7,8,9,10,11,12,12,12,12,12,NA,6,7,7,7,7,7,NA)
)

多分组示例 #

data2 <- data.frame(
  yyyymm = c(202401, 202402, 202403, 202404, 202405, 202406, 202407, 202408, 202409, 202410, 202411, 202412,
             202401, 202402, 202403, 202404, 202405, 202406, 202407,202401, 202402, 202403, 202404, 202405, 202406, 202407, 202408, 202409, 202410, 202411, 202412),
  id = c(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2,3,3,3,3,3,3,3,3,3,3,3,3),
  id2 = c(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,1,1,1,1,2,2,2,2,2,2,2,2),
  value = c(1,2,3,4,5,6,7,8,9,10,11,12,1,2,3,4,5,6,7,1,2,3,4,5,6,7,8,9,10,11,12),
  lead = c(6,7,8,9,10,11,12,12,12,12,12,NA,6,7,7,7,7,7,NA,4,4,4,NA,10,11,12,12,12,12,12,NA)
)

通用函数实现 #

利用dplyr的tidyeval特性，动态生成从n到1的lead调用，再通过coalesce取第一个非NA值：

library(dplyr)
library(purrr)
library(rlang)

lead_fallback <- function(data, col, n) {
  # 生成从n到1的lead表达式列表
  lead_exprs <- map(n:1, ~expr(lead({{col}}, n = .x)))
  # 组合成coalesce调用
  coalesce_expr <- expr(coalesce(!!!lead_exprs))
  # 执行mutate生成结果列
  data %>% mutate(lead_value = !!coalesce_expr)
}

使用示例 #

单分组场景 #

# 按id分组后计算lead值
result1 <- data %>%
  group_by(id) %>%
  lead_fallback(col = value, n = 5)

# 验证结果与期望列是否一致
all.equal(result1$lead_value, result1$lead)
# [1] TRUE

多分组场景 #

# 按id和id2分组后计算lead值
result2 <- data2 %>%
  group_by(id, id2) %>%
  lead_fallback(col = value, n = 5)

# 验证结果与期望列是否一致
all.equal(result2$lead_value, result2$lead)
# [1] TRUE

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