在两列中按范围对 data.table 进行二分搜索

Question

我正在尝试快速访问大型 data.table 的一个子集。数据有三列，都是数字（浮点数），几乎没有重复。两列是我要对其执行二进制搜索的数据，第三列包含我真正感兴趣的数字。本质上，我有 (x, y, z) 我想指定的数据x 范围和 y 范围以及 return 这些范围内的所有行。

# Generate some toy data of about the same size as the real data
DT <- data.table(x=runif(2000000), y=runif(2000000), z=runif(2000000))
head(DT)
#            x         y         z
# 1: 0.2675023 0.5725162 0.4162230
# 2: 0.1444540 0.8114941 0.1557195
# 3: 0.3607260 0.8159502 0.9705079
# 4: 0.3370213 0.9217284 0.5269885
# 5: 0.1085204 0.6312943 0.9676716
# 6: 0.1076674 0.1623447 0.1753712
ranges <- data.frame(x_min=runif(10000, max = 0.5), x_max=runif(10000, min = 0.5),
                     y_min=runif(10000, max = 0.5), y_max=runif(10000, min = 0.5))
head(ranges)
#        x_min     x_max      y_min     y_max
# 1 0.43817551 0.6720366 0.28052942 0.6309755
# 2 0.07469295 0.6744950 0.23170272 0.8431767
# 3 0.29520846 0.6991277 0.01882153 0.5162244
# 4 0.10500034 0.8977652 0.04806678 0.9528880
# 5 0.20168728 0.5655350 0.34401695 0.8241058
# 6 0.44158099 0.6739211 0.05359761 0.5832320

这是我正在尝试做的一个直观示例；我想要红色矩形内的所有点，其中矩形的边缘由 x 和 y 范围的最大值和最小值确定。但是，我有很多红色矩形，我将循环使用它们。

plot(DT$x, DT$y)
rect(xleft = ranges$x_min[1], xright = ranges$x_max[1],
     ybottom = ranges$y_min[1], ytop = ranges$y_max[1], border = "red")

目前，我正在使用的代码使用矢量扫描而不是二进制搜索（我认为），但它完全符合我的要求。

lapply(seq_len(nrow(ranges)), function(i){
  DT[x%between%c(ranges[i,]$x_min, ranges[i,]$x_max)&
       y%between%c(ranges[i,]$y_min, ranges[i,]$y_max)]
})

但是，根据 profvis，这仍然是过程中最慢的一步，鉴于我是 data.table 世界的新手，我想确保没有明显的东西我失踪了。据我所知，使用 data.table 键到运行二进制搜索而不是矢量扫描可以加快速度。但是，我一直无法弄清楚如何搜索范围而不是单个值。

This question asks something very similar but the best answer (from Matt) indicates that this wasn't doable easily in 2014 when the question was posted. He notes that this kind of problem really requires range join implementation and references a feature request 在 GitHub 页面上已经解决（打开几个月后）。

三年后，问题更新为我已经实现的新 %between% 功能，但我仍然不认为这对数据使用了二进制搜索。功能请求暗示理想的解决方案将采用 DT[J(id,DT(from,to)),...] 形式，这显然是使用 J() 语法来利用键。

%between% 语法实际上在幕后使用二进制搜索吗？如果没有，我怎样才能提供两个范围并仍然使用快速二进制搜索功能？

P.S。 dplyr 的 filter() 在数据集上的速度大约慢了 3 倍，所以就这样了。

Answer 1

我的理解是，rolling join 使用二进制搜索，但仅在最后一个 joining key 上，因此不可能同时对 4 个 key 执行 rolling join。此外，您的值本质上是非整数，因此无法使用二进制搜索精确定位 4 个角。

话虽如此，这里有一些选项可以加快子集化速度，非等连接是最快的，但我遇到了一些关于你的维度的内存限制问题：

m0 <- function()
    lapply(seq_len(nrow(ranges)), function(i){
        DT[x%between%c(ranges[i,]$x_min, ranges[i,]$x_max)&
                y%between%c(ranges[i,]$y_min, ranges[i,]$y_max)]
    })

m1 <- function()
    ranges[, DT[x %between% c(x_min, x_max) & y %between% c(y_min, y_max)], 1L:nrow(ranges)]

m2 <- function() {
    setkey(DT, x, y)
    setDT(ranges, key=c("x_min", "x_max", "y_min", "y_max"))
    DT[ranges, on=.(x>=x_min, x<=x_max, y>=y_min, y<=y_max), allow.cartesian=TRUE, .(x.x, x.y, x.z)]
}

m3 <- function() {
    setkey(DT3, x)[, rn := .I]
    ranges[, ixmin := DT3[.SD, on=.(x=x_min), roll=-Inf, rn]]
    ranges[, ixmax := DT3[.SD, on=.(x=x_max), roll=Inf, rn]]

    setkey(DT3, y)
    DT3[DT3[ranges, on=.(y>=y_min, y<=y_max),
        by=.EACHI, .(rn=rn[rn %between% c(ixmin, ixmax)])], on=.(rn),
        .(x, y, z)]
}

microbenchmark::microbenchmark(times=1L, m0(), m1(), m2(), m3())

时间安排：

Unit: milliseconds
 expr      min       lq     mean   median       uq      max neval
 m0() 782.6070 782.6070 782.6070 782.6070 782.6070 782.6070     1
 m1() 713.9469 713.9469 713.9469 713.9469 713.9469 713.9469     1
 m2() 272.6018 272.6018 272.6018 272.6018 272.6018 272.6018     1
 m3() 765.3667 765.3667 765.3667 765.3667 765.3667 765.3667     1

数据：

library(data.table)
set.seed(0L)
nr <- 2e4L
nrng <- 1e3L
dat <- data.table(x=runif(nr), y=runif(nr), z=runif(nr))
ranges <- data.frame(x_min=runif(nrng, max = 0.5), x_max=runif(nrng, min = 0.5),
    y_min=runif(nrng, max = 0.5), y_max=runif(nrng, min = 0.5))
dat[, rn := .I]

DT3 <- copy(dat)
DT <- copy(dat)

在两列中按范围对 data.table 进行二分搜索

Subsetting a data.table with a binary search by range in two columns

r

binary-search

data.table