基于分组计算最近条件时间差，优化data.table代码性能

需求背景 #

需要优化一段data.table代码的运行速度，功能为按分组回溯查找最后满足指定谓词条件的行，计算当前行到该行的时间差，同时支持指数变换。需处理30万行、20万分组的数据，且要计算多个此类指标，性能优先级极高。

示例数据 #

library(data.table)

data <- data.table(
      SYSKEY = c(
        12, 13, 14, 15, 20, 
        22, 21, 24, 25, 26
      ), 
      Customer = c(
        "John", "John", "John", "Tom", "Tom", 
        "Tom", "Sally", "Sally", "Sally", "Sally"
      ), 
      TRAN_DATTIM = as.Date(
        c(
        "28-02-2024", "28-02-2024", "02-03-2024", "02-03-2024", "02-03-2024", 
        "02-03-2024", "02-03-2024", "02-03-2024", "03-03-2024", "03-03-2024"
        ),
        format="%d-%m-%Y"
      ), 
      Product = c(
        "Eggs", "Milk", "Bread", "Butter","Eggs", 
        "Milk", "Bread", "Butter", "Eggs", "Wine"
      )
  )

原实现代码 #

build_recency <- function(
    data,
    name,
    predicate,
    aggregated_fun = "TIME",
    gamma = 0.0001,
    rolling_over
){
  
  UseMethod("build_recency")
  
}


build_recency.data.table <- function(
    data,
    name,
    predicate,
    aggregated_fun = "TIME",
    gamma = 0.01,
    rolling_over = "PAN"
){

  predicate <- enexpr(predicate)

  data[,
    PRED := fifelse(eval(predicate) == TRUE, 1L, 0L)
  ]
  
  setorderv(
    data,
    c(rolling_over, "TRAN_DATTIM", "SYSKEY")
  )
  
 # 分组内左连接，找到分组内最后一个满足谓词条件的行
data[,c(
 .SD[
   PRED == 1, .(TRAN_DATTIM, SYSKEY)
 ][
      .SD,
      on = .(TRAN_DATTIM <= TRAN_DATTIM, SYSKEY < SYSKEY),
      mult = "last",
      cols,
      env = list(
        cols = I(c(colnames(.SD), "x.TRAN_DATTIM"))
      )
    ]
  ),
   by = rolling_over
  ][,
    # 计算指标并清理临时变量
    `:=` (
      name = {
        x <- as.numeric(
          difftime(
              TRAN_DATTIM,
            x.TRAN_DATTIM,
            units = "days"
          )
        )  

        if(aggregated_fun == "EXP")
          x <- exp(-gamma * x)
        
        x
      },
      x.TRAN_DATTIM = NULL,
      PRED = NULL
    ),
    env = list(
      gamma = I(gamma),
      name = name
    )
  ]

}

 data2 <- build_recency(
    data = data,
    name = "RECENCY",
    predicate = if_else(
      Product == 'Eggs', TRUE, FALSE
    ),
    gamma = 0.001,
    rolling_over = 'Customer'
  )
  
  data2[]

预期结果 #

Customer SYSKEY TRAN_DATTIM Product RECENCY
      <char>  <num>      <Date>  <char>   <num>
 1:     John     12  2024-02-28    Eggs      NA
 2:     John     13  2024-02-28    Milk       0
 3:     John     14  2024-03-02   Bread       3
 4:    Sally     21  2024-03-02   Bread      NA
 5:    Sally     24  2024-03-02  Butter      NA
 6:    Sally     25  2024-03-03    Eggs      NA
 7:    Sally     26  2024-03-03    Wine       0
 8:      Tom     15  2024-03-02  Butter      NA
 9:      Tom     20  2024-03-02    Eggs      NA
10:      Tom     22  2024-03-02    Milk       0

优化方案 #

原代码通过分组内join实现回溯查找，在大量小分组场景下会产生大量中间数据，性能瓶颈明显。以下是线性时间复杂度的优化实现：

优化后代码 #

library(data.table)

build_recency_optimized <- function(
    data,
    name,
    predicate,
    aggregated_fun = "TIME",
    gamma = 0.01,
    rolling_over = "PAN"
){
  predicate <- enexpr(predicate)
  
  # 全局排序，确保分组内按时间、主键升序
  setorderv(data, c(rolling_over, "TRAN_DATTIM", "SYSKEY"))
  
  # 生成谓词标记：满足条件为1，否则为0
  data[, PRED := as.integer(eval(predicate))]
  
  # 分组内标记满足条件的时间，错位取前一行（排除当前行）
  data[, last_valid_time := shift(fcase(
    PRED == 1, TRAN_DATTIM,
    default = NA_Date_
  ), type = "lag"), by = rolling_over]
  
  # 前向填充最近的有效时间（data.table内置函数，无需外部依赖）
  data[, last_valid_time := nafill(last_valid_time, type = "locf"), by = rolling_over]
  
  # 计算目标指标
  data[, (name) := {
    x <- as.numeric(difftime(TRAN_DATTIM, last_valid_time, units = "days"))
    if (aggregated_fun == "EXP") {
      x <- exp(-gamma * x)
    }
    x
  }]
  
  # 清理临时变量
  data[, c("PRED", "last_valid_time") := NULL]
  
  return(data)
}

# 测试优化后代码
data2_optimized <- build_recency_optimized(
    data = copy(data),
    name = "RECENCY",
    predicate = Product == 'Eggs',
    gamma = 0.001,
    rolling_over = 'Customer'
)
data2_optimized[]

核心优化点 #

• 避免分组内join：用shift+fcase标记满足条件的行，再通过nafill前向填充最近有效时间，时间复杂度从O(n²)降至O(n)
• 简化谓词计算：用as.integer(eval(predicate))替代fifelse，减少不必要的分支判断开销
• 直接赋值逻辑：用(name)直接生成目标列，避免原代码中env参数的复杂处理，代码更简洁
• 减少中间数据：全程在原数据表上操作，无需生成大量中间join结果，内存占用大幅降低

额外性能建议 #

• 若数据已按rolling_over、TRAN_DATTIM、SYSKEY排序，可移除setorderv步骤，节省排序时间
• 计算多个指标时，尽量合并谓词计算和排序步骤，避免重复操作

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