ARROW-15622: [R] Implement union_all and union for arrow_dplyr_query

r/DESCRIPTION

@@ -109,6 +109,7 @@ Collate:
     'dplyr-mutate.R'
     'dplyr-select.R'
     'dplyr-summarize.R'
+    'dplyr-union.R'
     'record-batch.R'
     'table.R'
     'dplyr.R'

r/R/arrow-package.R

@@ -41,7 +41,7 @@
       "group_vars", "group_by_drop_default", "ungroup", "mutate", "transmute",
       "arrange", "rename", "pull", "relocate", "compute", "collapse",
       "distinct", "left_join", "right_join", "inner_join", "full_join",
-      "semi_join", "anti_join", "count", "tally", "rename_with"
+      "semi_join", "anti_join", "count", "tally", "rename_with", "union", "union_all"
     )
   )
   for (cl in c("Dataset", "ArrowTabular", "RecordBatchReader", "arrow_dplyr_query")) {

r/R/dplyr-union.R

@@ -0,0 +1,37 @@
+# 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.
+
+# The following S3 methods are registered on load if dplyr is present
+
+union.arrow_dplyr_query <- function(x, y, ...) {
+  x <- as_adq(x)
+  y <- as_adq(y)
+
+  distinct(union_all(x, y))
+}
+
+union.Dataset <- union.ArrowTabular <- union.RecordBatchReader <- union.arrow_dplyr_query
+
+union_all.arrow_dplyr_query <- function(x, y, ...) {
+  x <- as_adq(x)
+  y <- as_adq(y)
+
+  x$union_all <- list(right_data = y)
+  collapse.arrow_dplyr_query(x)
+}
+
+union_all.Dataset <- union_all.ArrowTabular <- union_all.RecordBatchReader <- union_all.arrow_dplyr_query

r/R/query-engine.R

@@ -138,6 +138,10 @@ ExecPlan <- R6Class("ExecPlan",
             right_suffix = .data$join$suffix[[2]]
           )
         }
+
+        if (!is.null(.data$union_all)) {
+          node <- node$UnionAll(self$Build(.data$union_all$right_data))
+        }
       }
 
       # Apply sorting: this is currently not an ExecNode itself, it is a
@@ -271,6 +275,9 @@ ExecNode <- R6Class("ExecNode",
           output_suffix_for_right = right_suffix
         )
       )
+    },
+    UnionAll = function(right_node) {
+      self$preserve_sort(ExecNode_Union(self, right_node))
     }
   ),
   active = list(

r/src/compute-exec.cpp

@@ -309,6 +309,13 @@ std::shared_ptr<compute::ExecNode> ExecNode_Join(
           std::move(output_suffix_for_left), std::move(output_suffix_for_right)});
 }
 
+// [[arrow::export]]
+std::shared_ptr<compute::ExecNode> ExecNode_Union(
+    const std::shared_ptr<compute::ExecNode>& input,
+    const std::shared_ptr<compute::ExecNode>& right_data) {
+  return MakeExecNodeOrStop("union", input->plan(), {input.get(), right_data.get()}, {});
+}
+
 // [[arrow::export]]
 std::shared_ptr<compute::ExecNode> ExecNode_SourceNode(
     const std::shared_ptr<compute::ExecPlan>& plan,

r/tests/testthat/test-dplyr-union.R

@@ -0,0 +1,74 @@
+# 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.
+
+skip_if(on_old_windows())
+
+library(dplyr, warn.conflicts = FALSE)
+
+withr::local_options(list(arrow.summarise.sort = FALSE))
+
+test_that("union_all", {
+  compare_dplyr_binding(
+    .input %>%
+      union_all(example_data) %>%
+      collect(),
+    example_data
+  )
+
+  test_table <- arrow_table(x = 1:10)
+
+  # Union with empty table produces same dataset
+  expect_equal(
+    test_table %>%
+      union_all(test_table$Slice(0, 0)) %>%
+      compute(),
+    test_table
+  )
+
+  expect_error(
+    test_table %>%
+      union_all(arrow_table(y = 1:10)) %>%
+      collect(),
+    regex = "input schemas must all match"
+  )
+})
+
+test_that("union", {
+  compare_dplyr_binding(
+    .input %>%
+      dplyr::union(example_data) %>%
+      collect(),
+    example_data
+  )
+
+  test_table <- arrow_table(x = 1:10)
+
+  # Union with empty table produces same dataset
+  expect_equal(
+    test_table %>%
+      dplyr::union(test_table$Slice(0, 0)) %>%
+      compute(),
+    test_table
+  )
+
+  expect_error(
+    test_table %>%
+      dplyr::union(arrow_table(y = 1:10)) %>%
+      collect(),
+    regex = "input schemas must all match"
+  )
+})