ARROW-13995: [R] Bindings for join node

cpp/src/arrow/compute/exec/options.h

@@ -155,9 +155,9 @@ class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {
         output_all(true),
         output_prefix_for_left(std::move(output_prefix_for_left)),
         output_prefix_for_right(std::move(output_prefix_for_right)) {
-    key_cmp.resize(left_keys.size());
-    for (size_t i = 0; i < left_keys.size(); ++i) {
-      key_cmp[i] = JoinKeyCmp::EQ;
+    this->key_cmp.resize(this->left_keys.size());
+    for (size_t i = 0; i < this->left_keys.size(); ++i) {
+      this->key_cmp[i] = JoinKeyCmp::EQ;
     }
   }
   HashJoinNodeOptions(
@@ -174,9 +174,9 @@ class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {
         right_output(std::move(right_output)),
         output_prefix_for_left(std::move(output_prefix_for_left)),
         output_prefix_for_right(std::move(output_prefix_for_right)) {
-    key_cmp.resize(left_keys.size());
-    for (size_t i = 0; i < left_keys.size(); ++i) {
-      key_cmp[i] = JoinKeyCmp::EQ;
+    this->key_cmp.resize(this->left_keys.size());
+    for (size_t i = 0; i < this->left_keys.size(); ++i) {
+      this->key_cmp[i] = JoinKeyCmp::EQ;
     }
   }
   HashJoinNodeOptions(

r/DESCRIPTION

@@ -92,6 +92,7 @@ Collate:
     'expression.R'
     'dplyr-functions.R'
     'dplyr-group-by.R'
+    'dplyr-join.R'
     'dplyr-mutate.R'
     'dplyr-select.R'
     'dplyr-summarize.R'

r/NAMESPACE

@@ -142,6 +142,7 @@ export(HivePartitioning)
 export(HivePartitioningFactory)
 export(InMemoryDataset)
 export(IpcFileFormat)
+export(JoinType)
 export(JsonParseOptions)
 export(JsonReadOptions)
 export(JsonTableReader)

r/R/arrow-package.R

@@ -36,7 +36,8 @@
       "select", "filter", "collect", "summarise", "group_by", "groups",
       "group_vars", "group_by_drop_default", "ungroup", "mutate", "transmute",
       "arrange", "rename", "pull", "relocate", "compute", "collapse",
-      "distinct"
+      "distinct", "left_join", "right_join", "inner_join", "full_join",
+      "semi_join", "anti_join"
     )
   )
   for (cl in c("Dataset", "ArrowTabular", "arrow_dplyr_query")) {

r/R/dplyr-collect.R

@@ -89,6 +89,14 @@ implicit_schema <- function(.data) {
 
   if (is.null(.data$aggregations)) {
     new_fields <- map(.data$selected_columns, ~ .$type(old_schm))
+    if (!is.null(.data$join) && !(.data$join$type %in% JoinType[1:4])) {
+      # Add cols from right side, except for semi/anti joins
+      right_cols <- .data$join$right_data$selected_columns
+      new_fields <- c(new_fields, map(
+        right_cols[setdiff(names(right_cols), .data$join$by)],
+        ~ .$type(.data$join$right_data$.data$schema)
+      ))
+    }
   } else {
     new_fields <- map(summarize_projection(.data), ~ .$type(old_schm))
     # * Put group_by_vars first (this can't be done by summarize,

r/R/dplyr-join.R

@@ -0,0 +1,126 @@
+# 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
+
+do_join <- function(x,
+                    y,
+                    by = NULL,
+                    copy = FALSE,
+                    suffix = c(".x", ".y"),
+                    ...,
+                    keep = FALSE,
+                    na_matches,
+                    join_type) {
+  # TODO: handle `copy` arg: ignore?
+  # TODO: handle `suffix` arg: Arrow does prefix
+  # TODO: handle `keep` arg: "Should the join keys from both ‘x’ and ‘y’ be preserved in the output?"
+  # TODO: handle `na_matches` arg
+  x <- as_adq(x)
+  y <- as_adq(y)
+  by <- handle_join_by(by, x, y)
+
+  x$join <- list(
+    type = JoinType[[join_type]],
+    right_data = y,
+    by = by
+  )
+  collapse.arrow_dplyr_query(x)
+}
+
+left_join.arrow_dplyr_query <- function(x,
+                                        y,
+                                        by = NULL,
+                                        copy = FALSE,
+                                        suffix = c(".x", ".y"),
+                                        ...,
+                                        keep = FALSE) {
+  do_join(x, y, by, copy, suffix, ..., keep = keep, join_type = "LEFT_OUTER")
+}
+left_join.Dataset <- left_join.ArrowTabular <- left_join.arrow_dplyr_query
+
+right_join.arrow_dplyr_query <- function(x,
+                                         y,
+                                         by = NULL,
+                                         copy = FALSE,
+                                         suffix = c(".x", ".y"),
+                                         ...,
+                                         keep = FALSE) {
+  do_join(x, y, by, copy, suffix, ..., keep = keep, join_type = "RIGHT_OUTER")
+}
+right_join.Dataset <- right_join.ArrowTabular <- right_join.arrow_dplyr_query
+
+inner_join.arrow_dplyr_query <- function(x,
+                                         y,
+                                         by = NULL,
+                                         copy = FALSE,
+                                         suffix = c(".x", ".y"),
+                                         ...,
+                                         keep = FALSE) {
+  do_join(x, y, by, copy, suffix, ..., keep = keep, join_type = "INNER")
+}
+inner_join.Dataset <- inner_join.ArrowTabular <- inner_join.arrow_dplyr_query
+
+full_join.arrow_dplyr_query <- function(x,
+                                        y,
+                                        by = NULL,
+                                        copy = FALSE,
+                                        suffix = c(".x", ".y"),
+                                        ...,
+                                        keep = FALSE) {
+  do_join(x, y, by, copy, suffix, ..., keep = keep, join_type = "FULL_OUTER")
+}
+full_join.Dataset <- full_join.ArrowTabular <- full_join.arrow_dplyr_query
+
+semi_join.arrow_dplyr_query <- function(x,
+                                        y,
+                                        by = NULL,
+                                        copy = FALSE,
+                                        suffix = c(".x", ".y"),
+                                        ...,
+                                        keep = FALSE) {
+  do_join(x, y, by, copy, suffix, ..., keep = keep, join_type = "LEFT_SEMI")
+}
+semi_join.Dataset <- semi_join.ArrowTabular <- semi_join.arrow_dplyr_query
+
+anti_join.arrow_dplyr_query <- function(x,
+                                        y,
+                                        by = NULL,
+                                        copy = FALSE,
+                                        suffix = c(".x", ".y"),
+                                        ...,
+                                        keep = FALSE) {
+  do_join(x, y, by, copy, suffix, ..., keep = keep, join_type = "LEFT_ANTI")
+}
+anti_join.Dataset <- anti_join.ArrowTabular <- anti_join.arrow_dplyr_query
+
+handle_join_by <- function(by, x, y) {
+  if (is.null(by)) {
+    return(set_names(intersect(names(x), names(y))))
+  }
+  stopifnot(is.character(by))
+  if (is.null(names(by))) {
+    by <- set_names(by)
+  }
+  # TODO: nicer messages?
+  stopifnot(
+    all(names(by) %in% names(x)),
+    all(by %in% names(y))
+  )
+  by
+}

r/R/enums.R

@@ -163,3 +163,16 @@ RoundMode <- enum("RoundMode",
   HALF_TO_EVEN = 8L,
   HALF_TO_ODD = 9L
 )
+
+#' @export
+#' @rdname enums
+JoinType <- enum("JoinType",
+  LEFT_SEMI = 0L,
+  RIGHT_SEMI = 1L,
+  LEFT_ANTI = 2L,
+  RIGHT_ANTI = 3L,
+  INNER = 4L,
+  LEFT_OUTER = 5L,
+  RIGHT_OUTER = 6L,
+  FULL_OUTER = 7L
+)

r/R/query-engine.R

@@ -143,7 +143,15 @@ ExecPlan <- R6Class("ExecPlan",
             )
           }
         }
-      } else {
+      } else if (!is.null(.data$join)) {
+        node <- node$Join(
+          type = .data$join$type,
+          right_node = self$Build(.data$join$right_data),
+          by = .data$join$by,
+          left_output = names(.data),
+          right_output = setdiff(names(.data$join$right_data), .data$join$by)
+        )
+      } else if (length(node$schema)) {
         # If any columns are derived, reordered, or renamed we need to Project
         # If there are aggregations, the projection was already handled above
         # We have to project at least once to eliminate some junk columns
@@ -206,6 +214,22 @@ ExecNode <- R6Class("ExecNode",
       self$preserve_sort(
         ExecNode_Aggregate(self, options, target_names, out_field_names, key_names)
       )
+    },
+    Join = function(type, right_node, by, left_output, right_output) {
+      self$preserve_sort(
+        ExecNode_Join(
+          self,
+          type,
+          right_node,
+          left_keys = names(by),
+          right_keys = by,
+          left_output = left_output,
+          right_output = right_output
+        )
+      )
     }
+  ),
+  active = list(
+    schema = function() ExecNode_output_schema(self)
   )
 )