如何带条件优化聚合函数？

Question

我有一个聚合函数，可以对数据组进行求和，然后根据一组条件创建一个标志并将其分配给该组。问题是有大量的组需要聚合，而每个组都非常小。

这意味着即使对于中等大小的数据集，执行聚合所需的时间也非常长，并且需要处理具有数百万行的数据集。

我在下面创建了一个可重现的示例，它说明了问题，并且具有与我的聚合函数类似的逻辑风格（我的有更多条件，但它们本质上相似）：

#install.packages("palmerpenguins")
library(data.table)
library(palmerpenguins)

# Create data
GROUPS <- 1:100
penguin_list <- lapply(GROUPS, \(x) data.table(group = x, penguins))
penguin_table <- rbindlist(penguin_list)

# Aggregation function
aggregatePenguinMass <- function(mass, sex, ratio = 2/3){
    data <- data.table(mass, sex)
    
    total <- sum(data[,mass], na.rm = TRUE)
    
    n_sex <- data[,.N, by = sex]
    n_male <- n_sex[sex == "male", N]
    n_female <- n_sex[sex == "female", N]
    
    if(n_female >= ratio * (n_male + n_female)){
        return(data.table(total = total,
                          flag = "F"))
    } else {
        return(data.table(total = total,
                          flag = "M"))
    }
}

# Perform aggregation and time
system.time(
penguin_table[, aggregatePenguinMass(body_mass_g, sex), by = .(group, species, year)]
)

#   user  system elapsed 
#   2.66    0.47    7.30 

如何更改此函数或执行聚合的方式以使其速度加快一个数量级？

---

基准

Unit: seconds
                                expr      min       lq     mean   median       uq      max neval
           agg(aggregatePenguinMass) 1.738215 1.763486 2.065458 1.778698 1.857921 4.550377    10
 agg(aggregatePenguinMass_RBarradas) 1.505988 1.534326 1.905495 1.541448 1.598832 5.086798    10

Answer 1

我认为您希望避免构建与组一样多的 data.table，而只使用一个大的 data.table。 我可以使用 dtyplr 重现您的聚合（因为我在 data.table 语法方面不太熟练）。 看起来像

library(dtplyr)

agg_dtyplr <- function(dt,mass, ratio = 2 / 3) {
  
  lazy_dt(dt) |> group_by(group, species, year)   |>
    summarise(
       total = sum({{mass}}, na.rm = TRUE),
    
       n_female = sum(1 * (sex == "female"),na.rm=TRUE),
       n = n())|> 
      mutate(
       flag = if_else(n_female >= ratio * n,
         "F", "M"
       )
     ) |> 
    select(group,species,year,total,flag) |> 
   as.data.table()
}

  agg_dtyplr(penguin_table,body_mass_g)