How to define a keyed (grouped) join in data table? - r

Question

I would like to overload the [.data.table operator so that grouped joins can be achieved. I will use dplyr to illustrate what I mean by a grouped join. Take for example the function below that wraps around dplyr::inner_join:
inner_join_grp <- function(x, y, by = NULL, copy = FALSE, suffix = c(".x", ".y"),
...) {
stopifnot(identical(group_vars(x), group_vars(y)))
grp <- group_vars(x)
by <- c(by, grp)
return(inner_join(x, y, by, copy, suffix, ...))
}
All this function does is to check that input data.frames x and y have the same groups, and if they do, to add the grouping variables to the join key specified by by = ....
An application would be like this:
n_grp <- 5
n_grp2 <- 3
set.seed(0)
tmp_df <-
data.frame(grp = rep(letters[1:n_grp2], each = n_grp),
grp2 = rep(1:n_grp, times = n_grp2),
x = runif(n_grp * n_grp2))
tmp_df_2 <-
data.frame(grp = rep(letters[1:n_grp2], each = n_grp),
grp2 = rep(1:n_grp, times = n_grp2),
y = runif(n_grp * n_grp2))
Using my function:
tmp_df %>%
group_by(grp) %>%
inner_join_grp(tmp_df_2 %>%
group_by(grp),
by = "grp2")
Produces the expected fifteen row table:
# A tibble: 15 x 4
# Groups: grp [?]
grp grp2 x y
<fct> <int> <dbl> <dbl>
1 a 1 0.897 0.770
2 a 2 0.266 0.498
3 a 3 0.372 0.718
4 a 4 0.573 0.992
5 a 5 0.908 0.380
...
whereas using dplyr::inner_join produces a 45 row table since the join is only on grp2:
# A tibble: 45 x 5
# Groups: grp.x [?]
grp.x grp2 x grp.y y
<fct> <int> <dbl> <fct> <dbl>
1 a 1 0.897 a 0.770
2 a 1 0.897 b 0.777
3 a 1 0.897 c 0.267
...
My question is whether is possible to achieve something similar using [, when the input data tables are keyed. I know I can overload merge.data.table (see below for a rough example) in the same way but I would prefer to pass an additional argument to [ to achieve a keyed join.
merge_dt_keyed <-
function (x, y, by = NULL, by.x = NULL, by.y = NULL, all = FALSE,
all.x = all, all.y = all, sort = TRUE,
suffixes = c(".x", ".y"),
allow.cartesian = getOption("datatable.allow.cartesian"),
...){
stopifnot(identical(key(x), key(y)))
new_key <- key(x)
by.x <- c(new_key, by, by.x)
by.y <- c(new_key, by, by.y)
return(merge(x, y, by.x = by.x, by.y = by.y, all = FALSE,
all.x = all, all.y = all, sort = TRUE,
suffixes = c(".x", ".y"),
allow.cartesian = getOption("datatable.allow.cartesian"),
...))
}

Answer 1

Something along these lines should work:
`[.data.table` = function(x, i, ...) {
args = match.call()
if ('on' %in% names(args)) {
args[['on']] = union(args[['on']], intersect(key(x), key(i)))
}
args[[1]] = data.table:::`[.data.table`
eval(args)
}