Merge pull request #3606 from tidyverse/b-3587-set-operations-grouped

reconstruct groups metadata for set operations. closes #3587
tidyverse · May 28, 2018 · b5e98e2 · b5e98e2
2 parents fce9e4c + 29ae8c8
commit b5e98e2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/src/set.cpp b/src/set.cpp
@@ -16,6 +16,7 @@
 #include <dplyr/DataFrameJoinVisitors.h>
 
 #include <dplyr/train.h>
+#include <dplyr/GroupedDataFrame.h>
 
 using namespace Rcpp;
 using namespace dplyr;
@@ -241,6 +242,14 @@ dplyr::BoolResult equal_data_frame(DataFrame x, DataFrame y, bool ignore_col_ord
   return yes();
 }
 
+DataFrame reconstruct_metadata(DataFrame out, const DataFrame& x) {
+  if (is<GroupedDataFrame>(x)) {
+    out = GroupedDataFrame(out, x).data();
+  }
+  // nothing to do for rowwise and natural data frames
+  return out ;
+}
+
 // [[Rcpp::export]]
 DataFrame union_data_frame(DataFrame x, DataFrame y) {
   BoolResult compat = compatible_data_frame(x, y, true, true);
@@ -256,7 +265,7 @@ DataFrame union_data_frame(DataFrame x, DataFrame y) {
   train_insert(set, x.nrows());
   train_insert_right(set, y.nrows());
 
-  return visitors.subset(set, get_class(x));
+  return reconstruct_metadata(visitors.subset(set, get_class(x)), x);
 }
 
 // [[Rcpp::export]]
@@ -283,7 +292,7 @@ DataFrame intersect_data_frame(DataFrame x, DataFrame y) {
     }
   }
 
-  return visitors.subset(indices, get_class(x));
+  return reconstruct_metadata(visitors.subset(indices, get_class(x)), x);
 }
 
 // [[Rcpp::export]]
@@ -310,5 +319,5 @@ DataFrame setdiff_data_frame(DataFrame x, DataFrame y) {
     }
   }
 
-  return visitors.subset(indices, get_class(x));
+  return reconstruct_metadata(visitors.subset(indices, get_class(x)), x);
 }
diff --git a/tests/testthat/test-sets.R b/tests/testthat/test-sets.R
@@ -68,3 +68,12 @@ test_that("intersect does not unnecessarily coerce (#1722)", {
   res <- intersect(df, df)
   expect_is(res$a, "integer")
 })
+
+test_that("set operations reconstruct grouping metadata (#3587)", {
+  df1 <- tibble(x = 1:4, g = rep(1:2, each = 2)) %>% group_by(g)
+  df2 <- tibble(x = 3:6, g = rep(2:3, each = 2))
+
+  expect_equal( setdiff(df1, df2) %>% group_rows(), list(0:1))
+  expect_equal( intersect(df1, df2) %>% group_rows(), list(0:1))
+  expect_equal( union(df1, df2) %>% group_rows(), list(4:5, 2:3, 0:1))
+})