в dplyr, как выбирать и фильтровать разные столбцы в зависимости от того, находится ли определенный столбец в фрейме данных?

Question

Мне нужно извлечь несколько строк и столбцов из фрейма данных:

library(dplyr)

foo <- structure(list(iso3c = c("SWZ", "SVN", "NZL", "JAM", "ESP", "LSO", 
"ATG", "GEO", "GIB", "BHS"), country = c("Eswatini", "Slovenia", 
"New Zealand", "Jamaica", "Spain", "Lesotho", "Antigua & Barbuda", 
"Georgia", "Gibraltar", "Bahamas"), confirmed = c(1, 141, 1522, 0, 148220, 4, 
19, 794, NA, 102), deaths = c(0, 0, 22, 0, 14792, 0, 2, 12, NA, 
11)), row.names = c(NA, -10L), class = c("tbl_df", "tbl", "data.frame"
))
compute_epidemic_curve_data <- function(df, country_code) {
  
  epidemic_curve_data <- df %>%
    filter(iso3c == country_code) %>%
    select(iso3c, 
           country,
           confirmed)
  
  return(epidemic_curve_data)
  
}

print(result <- compute_epidemic_curve_data(foo, "SVN"))

Однако данные могут поступать из разных источников, а это означает, что иногда фрейм данных будет иметь другую структуру. Обычно столбец iso3c называется id, столбец country называется admin_region, и присутствует дополнительный столбец с именем tests. Например:

bar <- structure(list(id = c("SWZ", "SVN", "NZL", "JAM", "ESP", "LSO", 
"ATG", "GEO", "GIB", "BHS"), admin_region = c("Eswatini", "Slovenia", 
"New Zealand", "Jamaica", "Spain", "Lesotho", "Antigua & Barbuda", 
"Georgia", "Gibraltar", "Bahamas"), confirmed = c(1, 141, 1522, 0, 148220, 4, 
19, 794, NA, 102), deaths = c(0, 0, 22, 0, 14792, 0, 2, 12, NA, 
11), tests = c(2, 282, 3044, 0, 296440, 8, 38, 1588, NA, 204)), row.names = c(NA, 
-10L), class = c("tbl_df", "tbl", "data.frame"))

Теперь, compute_epidemic_curve_data также должен возвращать tests, то есть он становится:

compute_epidemic_curve_data <- function(df, country_code) {
  
  epidemic_curve_data <- df %>%
    filter(id == country_code) %>%
    select(id, 
           admin_region,
           date, 
           confirmed,
           tests)
  
  return(epidemic_curve_data)
  
}

Вариант решения barbari c:

compute_epidemic_curve_data <- function(df, country_code) {
  
  if("id" %in% colnames(df))
    {
    epidemic_curve_data <- df %>%
      filter(id == country_code) %>%
      select(id, 
             admin_region,
             date, 
             confirmed,
             tests)
  }
  else 
  {
    epidemic_curve_data <- df %>%
      filter(iso3c == country_code) %>%
      select(iso3c, 
             country,
             date, 
             confirmed)    
  }
  
  return(epidemic_curve_data)
  
}

, но дублировать такой объем кода - плохая идея. Возможно ли, чтобы одна и та же функция обрабатывала два источника данных, в то же время уменьшая дублирование кода?

Answer 1

Идиоматический c способ динамического выбора между возможными именами столбцов в функции выбора tidyverse - использовать tidyselect::any_of:

compute_epidemic_curve_data <- function(df, country_code) 
{
  df <- if("iso3c" %in% names(df)) 
           filter(df, iso3c == country_code)
        else 
           filter(df, id == country_code)

  select(df, tidyselect::any_of(c("id", "iso3c","country", "confirmed", 
                                  "admin_region", "date", "tests")))
}

, что приводит к

print(result <- compute_epidemic_curve_data(foo, "SVN"))
#> # A tibble: 1 x 3
#>  iso3c country  confirmed
#>  <chr> <chr>        <dbl>
#> 1 SVN   Slovenia       141

print(result <- compute_epidemic_curve_data(bar, "SVN"))
#> # A tibble: 1 x 4
#>   id    confirmed admin_region tests
#>   <chr>     <dbl> <chr>        <dbl>
#> 1 SVN         141 Slovenia       282

Answer 2

Мы также можем использовать filter_at с matches

compute_epidemic_curve_data <- function(df, country_code) {

   df %>%
         filter_at(vars(matches('iso3c|id')), ~ . == country_code) %>% 
         #or with across
         #filter(across(matches('iso3c|id'), ~ . == country_code)) %>%
         select(matches('iso3c|id'), everything(), -deaths)
   
 }

-тестинг

compute_epidemic_curve_data(foo, "SVN")
# A tibble: 1 x 3
#  iso3c country  confirmed
#  <chr> <chr>        <dbl>
#1 SVN   Slovenia       141


compute_epidemic_curve_data(bar, "SVN")
# A tibble: 1 x 4
#  id    admin_region confirmed tests
#  <chr> <chr>            <dbl> <dbl>
#1 SVN   Slovenia           141   282

Answer 3

Это позволяет избежать некоторого дублирования, но вопрос о том, будет ли он более элегантным или читабельным, остается спорным. )