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

时间:2018-01-17 10:37:00

标签: r data.table

我正在编写一个非常基本的函数来检测具有特定名称的列,然后返回除这些列之外的表。我不确定如何使这个函数类型稳定最简洁的方法,它接受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,但对于大型表我猜这将是计算上昂贵的。

愿意了解这种情况下的良好做法。

谢谢!

1 个答案:

答案 0 :(得分: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 - 就像子集一样,您也可以将您的功能调整为:

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

如果您想保留data.tabledata.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"