tests/testthat/test-dplyr-select.R

In arrow: Integration to 'Apache' 'Arrow'

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

library(dplyr, warn.conflicts = FALSE)
library(stringr)

skip_if_not_available("acero")

tbl <- example_data

test_that("Empty select returns no columns", {
  compare_dplyr_binding(
    .input %>% select() %>% collect(),
    tbl
  )
})

test_that("Empty select still includes the group_by columns", {
  expect_message(
    compare_dplyr_binding(
      .input %>% group_by(chr) %>% select() %>% collect(),
      tbl
    ),
    "Adding missing grouping variables"
  )
})

test_that("Missing grouping columns are added to the beginning of the list", {
  expect_message(
    compare_dplyr_binding(
      .input %>% group_by(chr) %>% select(int) %>% collect(),
      tbl
    ),
    "Adding missing grouping variables"
  )
})

test_that("select/rename/rename_with", {
  compare_dplyr_binding(
    .input %>%
      select(string = chr, int) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      rename(string = chr) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      rename(strng = chr) %>%
      rename(other = strng) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      rename_with(
        ~ paste0(.x, "_suffix"),
        .cols = c("int", "chr")
      ) %>%
      collect(),
    tbl
  )
})

test_that("select/rename/rename_with using selection helpers", {
  compare_dplyr_binding(
    .input %>%
      select(everything()) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      select(any_of(c("int", "not_a_column", "lgl"))) %>%
      collect(),
    tbl
  )

  compare_dplyr_binding(
    .input %>%
      select(starts_with("d")) %>%
      collect(),
    tbl
  )

  compare_dplyr_binding(
    .input %>%
      select(where(is.numeric)) %>%
      collect(),
    tbl
  )

  compare_dplyr_binding(
    .input %>%
      rename_with(toupper) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      rename_with(toupper, .cols = c()) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      rename_with(
        ~ paste0(.x, "_suffix"),
        .cols = starts_with("d")
      ) %>%
      collect(),
    tbl
  )
})

test_that("filtering with rename", {
  compare_dplyr_binding(
    .input %>%
      filter(chr == "b") %>%
      select(string = chr, int) %>%
      collect(),
    tbl
  )
  compare_dplyr_binding(
    .input %>%
      select(string = chr, int) %>%
      filter(string == "b") %>%
      collect(),
    tbl
  )
})

test_that("relocate", {
  df <- tibble(a = 1, b = 1, c = 1, d = "a", e = "a", f = "a")
  compare_dplyr_binding(
    .input %>% relocate(f) %>% collect(),
    df,
  )
  compare_dplyr_binding(
    .input %>% relocate(a, .after = c) %>% collect(),
    df,
  )
  compare_dplyr_binding(
    .input %>% relocate(f, .before = b) %>% collect(),
    df,
  )
  compare_dplyr_binding(
    .input %>% relocate(a, .after = last_col()) %>% collect(),
    df,
  )
  compare_dplyr_binding(
    .input %>% relocate(ff = f) %>% collect(),
    df,
  )
})

test_that("relocate with selection helpers", {
  df <- tibble(a = 1, b = 1, c = 1, d = "a", e = "a", f = "a")
  compare_dplyr_binding(
    .input %>% relocate(any_of(c("a", "e", "i", "o", "u"))) %>% collect(),
    df
  )
  compare_dplyr_binding(
    .input %>% relocate(where(is.character)) %>% collect(),
    df
  )
  compare_dplyr_binding(
    .input %>% relocate(a, b, c, .after = where(is.character)) %>% collect(),
    df
  )
  compare_dplyr_binding(
    .input %>% relocate(d, e, f, .before = where(is.numeric)) %>% collect(),
    df
  )
  # works after other dplyr verbs
  compare_dplyr_binding(
    .input %>%
      mutate(c = as.character(c)) %>%
      relocate(d, e, f, .after = where(is.numeric)) %>%
      collect(),
    df
  )
})

test_that("multiple select/rename and group_by", {
  compare_dplyr_binding(
    .input %>%
      group_by(chr) %>%
      rename(string = chr, dub = dbl2) %>%
      rename(chr_actually = string) %>%
      collect(),
    tbl
  )

  compare_dplyr_binding(
    .input %>%
      group_by(chr) %>%
      select(string = chr, dub = dbl2) %>%
      rename(chr_actually = string) %>%
      collect(),
    tbl
  )

  compare_dplyr_binding(
    .input %>%
      group_by(chr) %>%
      rename(string = chr, dub = dbl2) %>%
      select(chr_actually = string) %>%
      collect(),
    tbl
  )
})