如何使用dplyr计算数据集中SPECIAL类型记录的end_date？

处理约20万行的数据集，需求如下：

• 所有type为"SPECIAL"的记录end_date为NA，需计算生成calc_end_date字段
• 计算规则：对每条SPECIAL记录，取同一index下type为"normal"的记录的end_date，筛选出**大于等于该SPECIAL记录effective_date**的日期，取其中最早的作为calc_end_date；非SPECIAL记录的calc_end_date为NA

示例输入数据 #

df <- data.frame(index          = c(1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3),
                 effective_date = as.Date(c("2020-01-01", "2020-06-13", "2021-01-01", "2022-01-01", "2023-01-01", "2024-01-01", "2024-04-13", "2025-01-01", "2025-03-01", "2026-01-01", "2015-01-01", "2016-01-01")),
                 end_date       = as.Date(c("2021-01-01", NA, "2022-01-01", "2023-01-01", "2024-01-01", "2025-01-01", NA, "2026-01-01", NA, "2027-01-01", "2016-01-01", "2017-01-01")),
                 type           = c("normal", "SPECIAL", "normal", "normal", "normal", "normal", "SPECIAL", "normal", "SPECIAL", "normal", "normal", "normal"))

期望输出结果 #

df_2 <- data.frame(index          = c(1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3),
                   effective_date = as.Date(c("2020-01-01", "2020-06-13", "2021-01-01", "2022-01-01", "2023-01-01", "2024-01-01", "2024-04-13", "2025-01-01", "2025-03-01", "2026-01-01", "2015-01-01", "2016-01-01")),
                   end_date       = as.Date(c("2021-01-01", NA, "2022-01-01", "2023-01-01", "2024-01-01", "2025-01-01", NA, "2026-01-01", NA, "2027-01-01", "2016-01-01", "2017-01-01")),
                   type           = c("normal", "SPECIAL", "normal", "normal", "normal", "normal", "SPECIAL", "normal", "SPECIAL", "normal", "normal", "normal"),
                   calc_end_date  = as.Date(c(NA, "2021-01-01", NA, NA, NA, NA, "2025-01-01", NA, "2026-01-01", NA, NA, NA)))

dplyr解决方案 #

提供两种基于dplyr的实现，适用于不同场景：

方法一：分组后逐行计算（逻辑直观） #

通过group_by按index分组，结合purrr::map对每条SPECIAL记录筛选符合条件的日期：

library(dplyr)
library(purrr)

df_result <- df %>%
  group_by(index) %>%
  mutate(
    calc_end_date = case_when(
      type == "SPECIAL" ~ map_dbl(effective_date, ~{
        # 筛选同组normal的有效end_date
        eligible_dates <- end_date[type == "normal" & !is.na(end_date)]
        # 保留 >= 当前effective_date的日期
        eligible_dates <- eligible_dates[eligible_dates >= .x]
        # 返回最早日期，无符合条件则返回NA
        if (length(eligible_dates) > 0) min(eligible_dates) else NA_real_
      }) %>% as.Date(origin = "1970-01-01"),
      # 非SPECIAL记录返回NA
      TRUE ~ NA_Date_
    )
  ) %>%
  ungroup()

方法二：预整理日期后匹配（更高效） #

针对20万行的大数据集，先提取每个index下的normal日期列表，再通过映射匹配，减少重复计算：

library(dplyr)

# 预提取每个index下的normal的end_date，去重并排序
normal_dates <- df %>%
  filter(type == "normal" & !is.na(end_date)) %>%
  group_by(index) %>%
  summarize(normal_end_dates = list(sort(unique(end_date)))) %>%
  ungroup()

# 合并数据并计算calc_end_date
df_result <- df %>%
  left_join(normal_dates, by = "index") %>%
  mutate(
    calc_end_date = case_when(
      type == "SPECIAL" ~ map2_dbl(effective_date, normal_end_dates, ~{
        # 找到第一个 >= effective_date的日期
        match_idx <- which(.y >= .x)
        if (length(match_idx) > 0) .y[min(match_idx)] else NA_real_
      }) %>% as.Date(origin = "1970-01-01"),
      TRUE ~ NA_Date_
    )
  ) %>%
  select(-normal_end_dates) %>% # 移除临时字段
  ungroup()

结果验证 #

运行上述任意一种方法，得到的df_result与期望的df_2完全一致。

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