如何制作选择data.frame子集的类型稳定函数

Question

我正在编写一个非常基本的函数，它检测具有特定名称的列，然后返回不包括这些列的表。我不知道什么是使这个函数类型稳定的最简洁的方法，其中它接受 data.frames 或 data.tables 并在它进入的同一个类中吐出结果表。

例如，我通常的工作流程是使用 data.table。

library(data.table)
dt <- data.table(names = sample(c("Ruby","Fire","Azure","Green"), 10, replace = T), age = 10:19, phone = 123456:123465)
df <- data.frame(names = sample(c("Ruby","Fire","Azure","Green"), 10, replace = T), age = 10:19, phone = 123456:123465)
detach("package:data.table")

removeAge <- function(db){
    ageCols <- grepl("age",names(db))
    db <- db[, !ageCols]
    return(db)
}

removeAge(df) # returns data.frame with age column removed
removeAge(dt) # returns vector of logical T,F,T

如何使我的示例中的removeAge 函数不知道输入表是data.frame 还是data.table？这个例子也将帮助我学习更复杂的功能。我假设一种方法是检查输入表类是否转换为 data.frame，但对于大型表，我猜这在计算上会很昂贵。

很想知道在这种情况下有什么好的做法。

谢谢！

Answer 1

最简单的选择是将函数更改为：

removeAge <- function(db){
  ageCols <- grepl('age', names(db))
  db <- as.data.frame(db)[, !ageCols]
  return(db)
}

现在使用removeAge(df) 或removeAge(dt) 给出预期结果：

> removeAge(df) 
   names  phone
1   Ruby 123456
2  Azure 123457
3   Ruby 123458
4   Ruby 123459
5   Fire 123460
6  Azure 123461
7  Green 123462
8   Ruby 123463
9  Green 123464
10 Azure 123465

> removeAge(dt)
   names  phone
1  Azure 123456
2   Fire 123457
3  Green 123458
4  Azure 123459
5   Fire 123460
6  Green 123461
7  Azure 123462
8  Green 123463
9   Fire 123464
10 Azure 123465

---

要使用data.table-like 子集，您还可以将您的函数调整为：

removeAge <- function(db){
  nonAgeCols <- setdiff(names(db), 'age')
  db <- setDT(db)[, ..nonAgeCols]
  return(db)
}

---

如果您想保留data.table 或data.frame 的类，那么您可以将函数更改为：

removeAge <- function(db) {
  if (any(class(db) == 'data.table')) {
    nonAgeCols <- setdiff(names(db), 'age')
    db <- setDT(db)[, ..nonAgeCols]
    return(db)
  } else {
    ageCols <- grepl("age",names(db))
    db <- db[, !ageCols]
    return(db)
  }
}

根据输入的类返回 data.table 或 data.frame：

> class(removeAge(df))
[1] "data.frame"
> class(removeAge(dt))
[1] "data.table" "data.frame"