如何优化R包中VhgPreprocessTaxa函数的str_extract性能？

问题背景 #

我正在开发一款R包，其中VhgPreprocessTaxa函数负责处理病毒分类数据：从形如taxid:2065037|Betatorquevirus|Anelloviridae的ViralRefSeq_taxonomy列中提取指定分类阶元（如Family），结合ICTV数据库补充分类信息，缺失值标记为unclassified。

当前性能瓶颈：处理百万级数据集时，即便已通过多线程拆分并行优化，str_extract仍是核心耗时点——单核心处理10万条数据需7.05秒，7核心处理百万条数据仍需22秒。经profvis定位，str_extract是主要性能拖累，求进一步优化方案。

原核心代码 #

分块处理函数 #

process_chunk <- function(chunk, ictv_formatted, taxa_rank) {
  taxon_filter <- paste(unique(ictv_formatted$name), collapse = "|")

  chunk_processed <- chunk %>%
    mutate(
      ViralRefSeq_taxonomy = str_remove_all(.data$ViralRefSeq_taxonomy, "taxid:\\d+\\||\\w+\\s\\w+\\|"),
      name = str_extract(.data$ViralRefSeq_taxonomy, taxon_filter),
      ViralRefSeq_taxonomy = str_extract(.data$ViralRefSeq_taxonomy, paste0("\\w+", taxa_rank))
    ) %>%
    left_join(ictv_formatted, join_by("name" == "name")) %>%
    mutate(
      ViralRefSeq_taxonomy = case_when(
        is.na(.data$ViralRefSeq_taxonomy) & is.na(.data$Phylum) ~ "unclassified",
        is.na(.data$ViralRefSeq_taxonomy) ~ paste("unclassified", .data$Phylum),
        .default = .data$ViralRefSeq_taxonomy
      )
    ) %>% select(-c(.data$name:.data$level)) %>%
    mutate(
      ViralRefSeq_taxonomy = if_else(.data$ViralRefSeq_taxonomy == "unclassified unclassified", "unclassified", .data$ViralRefSeq_taxonomy),
      ViralRefSeq_taxonomy = if_else(.data$ViralRefSeq_taxonomy == "unclassified NA", "unclassified", .data$ViralRefSeq_taxonomy)
    )

  return(chunk_processed)
}

主函数 #

VhgPreprocessTaxa <- function(file,taxa_rank, num_cores = 1) {
  if (is_file_empty(file)) {
    return(invisible(NULL))
  }

  if (!all(grepl("^taxid:", file$ViralRefSeq_taxonomy))) {
    message("The 'ViralRefSeq_taxonomy' column is expected to start with 'taxid:' followed by taxa ranks separated by '|'.\n",
            "However, Your column has a different structure, possibly indicating that the data has already been processed by 'VhgPreprocessTaxa'.\n",
            "Skipping Taxonomy processing ...")
    return(file)
  }

  taxa_rank <- taxonomy_rank_hierarchy(taxa_rank)
  ictv_formatted <- format_ICTV(taxa_rank)

  if (num_cores < 1) {
    stop("Number of cores must be at least 1")
  }
  num_cores <- min(num_cores, detectCores() - 1)

  num_chunks <- num_cores
  chunks <- split(file, rep(1:num_chunks, length.out = nrow(file)))

  processed_chunks <- mclapply(chunks, process_chunk, ictv_formatted = ictv_formatted, taxa_rank = taxa_rank, mc.cores = num_cores)
  file_processed <- bind_rows(processed_chunks)

  return(file_processed)
}

优化方案 #

1. 替换stringr为底层高效实现+预编译正则 #

stringr的封装会带来额外开销，直接用stringi（stringr的底层依赖）并预编译正则，避免每个分块重复生成正则表达式；同时针对|分隔的固定格式，用strsplit拆分后匹配，比正则提取更快。

2. 用哈希表替代left_join #

将ICTV数据转换为命名向量（哈希表），直接映射分类信息，避免left_join的耗时操作。

3. 合并字符串操作，减少中间变量 #

减少对ViralRefSeq_taxonomy列的重复修改，合并缺失值处理逻辑。

优化后的代码 #

分块处理函数

process_chunk <- function(chunk, ictv_map, taxa_rank_regex, taxon_regex) {
  # 拆分固定格式字符串，替代多次正则删除/提取
  chunk_split <- strsplit(chunk$ViralRefSeq_taxonomy, "\\|", fixed = TRUE)
  
  # 匹配ICTV分类名称
  chunk$name <- sapply(chunk_split, function(x) {
    match_idx <- stringi::stri_detect_regex(x, taxon_regex)
    if (any(match_idx)) x[match_idx][1] else NA_character_
  })
  
  # 提取目标分类阶元
  chunk$ViralRefSeq_taxonomy <- sapply(chunk_split, function(x) {
    match_idx <- stringi::stri_detect_regex(x, taxa_rank_regex)
    if (any(match_idx)) x[match_idx][1] else NA_character_
  })
  
  # 哈希表快速映射Phylum
  chunk$Phylum <- ictv_map[chunk$name]
  
  # 统一处理缺失值
  chunk$ViralRefSeq_taxonomy <- case_when(
    is.na(chunk$ViralRefSeq_taxonomy) & is.na(chunk$Phylum) ~ "unclassified",
    is.na(chunk$ViralRefSeq_taxonomy) ~ paste("unclassified", chunk$Phylum, sep = " "),
    .default = chunk$ViralRefSeq_taxonomy
  )
  
  # 清理异常值
  chunk$ViralRefSeq_taxonomy <- ifelse(chunk$ViralRefSeq_taxonomy %in% c("unclassified unclassified", "unclassified NA"), 
                                       "unclassified", chunk$ViralRefSeq_taxonomy)
  
  # 移除临时列
  chunk <- select(chunk, -c(name, Phylum))
  
  return(chunk)
}

主函数

VhgPreprocessTaxa <- function(file, taxa_rank, num_cores = 1) {
  if (is_file_empty(file)) {
    return(invisible(NULL))
  }
  
  if (!all(grepl("^taxid:", file$ViralRefSeq_taxonomy))) {
    message("The 'ViralRefSeq_taxonomy' column is expected to start with 'taxid:' followed by taxa ranks separated by '|'.\n",
            "However, Your column has a different structure, possibly indicating that the data has already been processed by 'VhgPreprocessTaxa'.\n",
            "Skipping Taxonomy processing ...")
    return(file)
  }
  
  taxa_rank <- taxonomy_rank_hierarchy(taxa_rank)
  ictv_formatted <- format_ICTV(taxa_rank)
  
  # 预编译正则，避免重复生成
  taxa_rank_regex <- stringi::stri_regex_compile(paste0("\\w+", taxa_rank))
  taxon_regex <- stringi::stri_regex_compile(paste(unique(ictv_formatted$name), collapse = "|"))
  
  # 构建哈希表，用于快速匹配Phylum
  ictv_map <- setNames(ictv_formatted$Phylum, ictv_formatted$name)
  
  # 验证核心数
  if (num_cores < 1) {
    stop("Number of cores must be at least 1")
  }
  num_cores <- min(num_cores, detectCores() - 1)
  
  # 拆分数据
  num_chunks <- num_cores
  chunks <- split(file, rep(1:num_chunks, length.out = nrow(file)))
  
  # 并行处理
  processed_chunks <- mclapply(chunks, process_chunk, 
                               ictv_map = ictv_map, 
                               taxa_rank_regex = taxa_rank_regex, 
                               taxon_regex = taxon_regex, 
                               mc.cores = num_cores)
  
  # 合并结果
  file_processed <- bind_rows(processed_chunks)
  
  return(file_processed)
}

额外优化建议 #

• 改用data.table：data.table的字符串操作和行处理速度远快于dplyr，百万级数据场景下能大幅降低耗时。
• 预计算唯一值：如果数据中有大量重复的分类字符串，先处理唯一值再映射回原数据，避免重复处理相同内容：

  # 提取唯一分类字符串处理，再映射回原数据
  unique_taxa <- unique(file$ViralRefSeq_taxonomy)
  processed_unique <- process_chunk(data.frame(ViralRefSeq_taxonomy = unique_taxa), ictv_map, taxa_rank_regex, taxon_regex)
  file_processed <- file %>% left_join(processed_unique, by = "ViralRefSeq_taxonomy") %>% 
    select(-ViralRefSeq_taxonomy.x) %>% rename(ViralRefSeq_taxonomy = ViralRefSeq_taxonomy.y)
  

• 更换并行框架：用furrr（基于future）替代mclapply，Windows系统兼容性更好，且能更灵活控制并行粒度。

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