如何泛化R语言data.frame子集筛选函数 支持自定义列判断常量/变量属性

需求说明 #

现有示例data.frame包含study、outcome、group、time共4列，需要筛选出满足「指定列取值为常量/变量」条件的study对应全部数据：
以校验列outcome、group为例，共4种匹配场景：

1. group为常量，outcome为变量
2. outcome为常量，group为变量
3. outcome和group均为变量
4. outcome和group均为常量
   原实现函数仅支持固定校验两列的场景，需要泛化为支持用户自定义输入校验列，且原场景下输出与原结果一致。

原示例数据与代码 #

# 示例数据
h = "
study outcome group time
a     1       1     0
a     2       1     1
b     1       1     0
b     1       2     0
c     2       1     0
c     3       2     1
d     1       1     0
d     1       1     0
e     1       1     0"
h = read.table(text=h,h=T)

# 原实现函数（仅支持固定校验两列）
foo <- function(dat, cond) {
  switch(cond, 
         `1` = dat %>% 
           group_by(study) %>%
           filter(n_distinct(group) == 1, n_distinct(outcome) > 1) %>%
           ungroup,
         `2` = dat %>% 
           group_by(study) %>%
           filter(n_distinct(group) > 1, n_distinct(outcome) == 1) %>%
           ungroup,
         `3` =  dat %>% 
           group_by(study) %>%
           filter(n_distinct(group) > 1, n_distinct(outcome) > 1) %>%
           ungroup,
         `4` = dat %>% 
           group_by(study) %>%
           filter(n_distinct(group) == 1, n_distinct(outcome) == 1) %>%
           ungroup )  } 

泛化后函数实现 #

依赖tidyverse套件实现，支持自定义校验列，兼容原函数调用方式：

library(tidyverse)

foo_general <- function(dat, check_cols = c("outcome", "group"), cond) {
  n_check <- length(check_cols)
  # 处理cond输入，兼容数值和逻辑向量两种传入方式
  if (is.numeric(cond)) {
    if (cond < 1 || cond > 2^n_check) {
      stop(sprintf("数值型cond取值范围应为1到%s（2的校验列数量次方）", 2^n_check))
    }
    # 二进制映射规则与原函数完全对齐，保证原场景输出一致
    cond_logic <- rev(as.integer(intToBits(cond-1))[1:n_check]) == 1
  } else if (is.logical(cond)) {
    if (length(cond) != n_check) {
      stop("逻辑型cond长度需与校验列数量保持一致")
    }
    cond_logic <- cond
  } else {
    stop("cond参数仅支持整数或逻辑向量类型")
  }
  
  # 动态构造每列的过滤条件
  filter_rules <- map2(check_cols, cond_logic, function(col, need_var) {
    if (need_var) {
      expr(n_distinct(!!sym(col)) > 1)
    } else {
      expr(n_distinct(!!sym(col)) == 1)
    }
  })
  
  # 执行过滤返回结果
  dat %>%
    group_by(study) %>%
    filter(!!!filter_rules) %>%
    ungroup()
}

验证测试 #

原场景兼容验证 #

调用方式与原函数完全一致，输出结果匹配：

# 对应原场景cond=1
foo_general(h, cond = 1)
# 对应原场景cond=2
foo_general(h, cond = 2)
# 对应原场景cond=3
foo_general(h, cond = 3)
# 对应原场景cond=4
foo_general(h, cond = 4)

自定义校验列示例 #

如需校验time、group两列，要求time为常量、group为变量，调用方式如下：

# 方式1：数值cond传入
foo_general(h, check_cols = c("time", "group"), cond = 1)
# 方式2：逻辑向量cond传入（TRUE表示要求该列为变量，FALSE表示要求为常量）
foo_general(h, check_cols = c("time", "group"), cond = c(FALSE, TRUE))

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