refactor(geo): Implement RTree SpatialIndex for faster spatial joins (facebookincubator#15492)

Summary:
Pull Request resolved: facebookincubator#15492

The previous flat spatial index required us to scan all bounding boxes
(minus the ones we could discard by binary search).  RTrees have a
hierarchical set of bounding boxes, allow us to skip many bounding
boxes with a single check.  This results in a faster spatial index

Results from the SpatialJoinBenchmark:

Strategy   | 1,1 Unf | 1,1 Cls | 50,5 Unf | 50,5 Cls | 25,5 L Unf | 25,5 L Cls | 200,50 Unf
-------------|---------|---------|----------|--------|------|------|--------
Flat         |  364.00 |  559.28 |     1.99 |   2.12 | 3.99 | 4.94 | 0.03027
Hilbert-32   |  383.03 |  545.45 |     4.50 |   2.62 | 8.78 | 5.76 | 0.13577
Hilbert-64   |  388.38 |  534.55 |     4.44 |   2.76 | 8.63 | 5.69 | 0.13694
Hilbert-128  |  364.52 |  520.60 |     4.31 |   2.76 | 8.42 | 5.38 | 0.12061
Hilbert-256  |  351.38 |  498.73 |     4.08 |   2.64 | 8.04 | 5.92 | 0.13350

Entries are in iterations per second, larger is better.

Legend:
Hilbert-N: Hilbert sorted RTree with branch size N
N,M: Probe/build size, in thousands of rows.  This 50,5 means 50k probe rows, 5k build
Unf: Uniform distribution
Cls: Clustered distribution
L: Left join (otherwise inner)

Reviewed By: kgpai

Differential Revision: D86127954

fbshipit-source-id: 4b5473215db47318a729dc26c574e75ace1e55c3

Author: jagill

velox/exec/SpatialIndex.cpp

@@ -18,54 +18,157 @@
 #include <vector>
 
 #include "velox/common/base/Exceptions.h"
+#include "velox/exec/HilbertIndex.h"
 #include "velox/exec/SpatialIndex.h"
 
 namespace facebook::velox::exec {
 
-SpatialIndex::SpatialIndex(Envelope bounds, std::vector<Envelope> envelopes)
-    : bounds_(std::move(bounds)) {
-  std::ranges::sort(envelopes, {}, &Envelope::minX);
+std::vector<size_t> RTreeLevel::query(
+    const Envelope& queryEnv,
+    const std::vector<size_t>& branchIndices) const {
+  std::vector<size_t> result;
 
-  minXs_.reserve(envelopes.size());
-  minYs_.reserve(envelopes.size());
-  maxXs_.reserve(envelopes.size());
-  maxYs_.reserve(envelopes.size());
-  rowIndices_.reserve(envelopes.size());
+  for (size_t branchIdx : branchIndices) {
+    size_t startIdx = branchIdx * branchSize_;
+    size_t endIdx = std::min(startIdx + branchSize_, minXs_.size());
+    for (size_t idx = startIdx; idx < endIdx; ++idx) {
+      bool intersects = (queryEnv.maxX >= minXs_[idx]) &&
+          (queryEnv.maxY >= minYs_[idx]) && (queryEnv.minX <= maxXs_[idx]) &&
+          (queryEnv.minY <= maxYs_[idx]);
+      if (intersects) {
+        result.push_back(idx);
+      }
+    }
+  }
 
+  return result;
+}
+
+namespace {
+std::pair<RTreeLevel, std::vector<Envelope>> buildLevel(
+    uint32_t branchSize,
+    const std::vector<Envelope>& envelopes) {
+  std::vector<float> minXs;
+  minXs.reserve(envelopes.size());
+  std::vector<float> minYs;
+  minYs.reserve(envelopes.size());
+  std::vector<float> maxXs;
+  maxXs.reserve(envelopes.size());
+  std::vector<float> maxYs;
+  maxYs.reserve(envelopes.size());
+
+  std::vector<Envelope> parentEnvelopes;
+  parentEnvelopes.reserve((envelopes.size() + branchSize - 1) / branchSize);
+  Envelope currentBounds = Envelope::empty();
+
+  uint32_t idx = 0;
   for (const auto& env : envelopes) {
-    minXs_.push_back(env.minX);
-    minYs_.push_back(env.minY);
-    maxXs_.push_back(env.maxX);
-    maxYs_.push_back(env.maxY);
-    rowIndices_.push_back(env.rowIndex);
+    ++idx;
+    currentBounds.maxX = std::max(currentBounds.maxX, env.maxX);
+    currentBounds.maxY = std::max(currentBounds.maxY, env.maxY);
+    currentBounds.minX = std::min(currentBounds.minX, env.minX);
+    currentBounds.minY = std::min(currentBounds.minY, env.minY);
+    if (idx % branchSize == 0) {
+      parentEnvelopes.push_back(currentBounds);
+      currentBounds = Envelope::empty();
+    }
+
+    minXs.push_back(env.minX);
+    minYs.push_back(env.minY);
+    maxXs.push_back(env.maxX);
+    maxYs.push_back(env.maxY);
+  }
+
+  if (!currentBounds.isEmpty()) {
+    parentEnvelopes.push_back(currentBounds);
   }
+
+  return {
+      RTreeLevel(
+          branchSize,
+          std::move(minXs),
+          std::move(minYs),
+          std::move(maxXs),
+          std::move(maxYs)),
+      std::move(parentEnvelopes)};
 }
+} // namespace
 
-std::vector<int32_t> SpatialIndex::query(const Envelope& queryEnv) const {
-  std::vector<int32_t> result;
-  if (!Envelope::intersects(queryEnv, bounds_)) {
-    return result;
+SpatialIndex::SpatialIndex(
+    Envelope bounds,
+    std::vector<Envelope> envelopes,
+    uint32_t branchSize)
+    : branchSize_(branchSize), bounds_(std::move(bounds)) {
+  VELOX_CHECK_GT(branchSize_, 1);
+
+  if (!bounds_.isEmpty()) {
+    HilbertIndex hilbert(
+        bounds_.minX, bounds_.minY, bounds_.maxX, bounds_.maxY);
+
+    std::sort(
+        envelopes.begin(), envelopes.end(), [&](const auto& a, const auto& b) {
+          return hilbert.indexOf(a.minX, a.minY) <
+              hilbert.indexOf(b.minX, b.minY);
+        });
   }
 
-  // Find the last minX that is <= queryEnv.maxX . These first envelopes
-  // are the only ones that can intersect the query envelope.
-  // `it` is _one past_ the last element, so we iterate up to it - 1.
-  auto it = std::upper_bound(minXs_.begin(), minXs_.end(), queryEnv.maxX);
-  if (it == minXs_.begin()) {
+  rowIndices_.reserve(envelopes.size());
+  for (const auto& env : envelopes) {
+    VELOX_CHECK(env.minX >= bounds_.minX);
+    VELOX_CHECK(env.minY >= bounds_.minY);
+    VELOX_CHECK(env.maxX <= bounds_.maxX);
+    VELOX_CHECK(env.maxY <= bounds_.maxY);
+    rowIndices_.push_back(env.rowIndex);
+  }
+
+  if (envelopes.size() > 0) {
+    size_t numLevels =
+        std::ceil(std::log(envelopes.size()) / std::log(branchSize_));
+    levels_.reserve(numLevels);
+  }
+
+  while (envelopes.size() > branchSize_) {
+    auto [level, parentEnvelopes] = buildLevel(branchSize_, envelopes);
+    levels_.push_back(std::move(level));
+    envelopes = std::move(parentEnvelopes);
+  }
+
+  if (envelopes.size() > 1 || levels_.empty()) {
+    levels_.push_back(buildLevel(branchSize_, envelopes).first);
+  }
+
+  VELOX_CHECK_GT(branchSize_ + 1, levels_.back().size());
+}
+
+std::vector<vector_size_t> SpatialIndex::query(const Envelope& queryEnv) const {
+  std::vector<vector_size_t> result;
+  if (!Envelope::intersects(queryEnv, bounds_)) {
     return result;
   }
 
-  auto lastIdx = std::distance(minXs_.begin(), it);
-  VELOX_CHECK_GT(lastIdx, 0);
+  size_t thisLevel = levels_.size() - 1;
+  VELOX_CHECK_GT(levels_[thisLevel].size(), 0);
+  VELOX_CHECK_GT(branchSize_ + 1, levels_[thisLevel].size());
 
-  for (size_t idx = 0; idx < lastIdx; ++idx) {
-    bool intersects = (queryEnv.maxY >= minYs_[idx]) &&
-        (queryEnv.minX <= maxXs_[idx]) && (queryEnv.minY <= maxYs_[idx]);
-    if (intersects) {
-      result.push_back(rowIndices_[idx]);
+  // The top level should have only one branch.
+  std::vector<size_t> childIndices = {0};
+  for (; thisLevel > 0; --thisLevel) {
+    // Avoiding thisLevel = 0 due to int underflow
+    childIndices = levels_[thisLevel].query(queryEnv, childIndices);
+    // If we have no matches, return.
+    if (childIndices.empty()) {
+      return result;
     }
   }
 
+  // We're at level 0 now.  The indices index into rowIndices.
+  VELOX_DCHECK_EQ(thisLevel, 0);
+  childIndices = levels_[thisLevel].query(queryEnv, childIndices);
+  result.reserve(childIndices.size());
+  for (auto idx : childIndices) {
+    result.push_back(rowIndices_[idx]);
+  }
+
   return result;
 }
 

velox/exec/SpatialIndex.h

@@ -20,6 +20,7 @@
 #include <cstdint>
 #include <limits>
 #include <vector>
+#include "velox/vector/TypeAliases.h"
 
 namespace facebook::velox::exec {
 
@@ -73,11 +74,11 @@ namespace facebook::velox::exec {
 /// nextUp to maxX/Ys, the float32 envelope intersects in all cases that the
 /// float64 envelope would (but not necessarily the converse).
 struct Envelope {
-  float minX;
-  float minY;
-  float maxX;
-  float maxY;
-  int32_t rowIndex = -1;
+  float minX{std::numeric_limits<float>::infinity()};
+  float minY{std::numeric_limits<float>::infinity()};
+  float maxX{-std::numeric_limits<float>::infinity()};
+  float maxY{-std::numeric_limits<float>::infinity()};
+  vector_size_t rowIndex = -1;
 
   /// Returns true if the intersection of two envelopes is not empty.
   static inline bool intersects(const Envelope& left, const Envelope& right) {
@@ -114,7 +115,7 @@ struct Envelope {
       double minY,
       double maxX,
       double maxY,
-      int32_t rowIndex = -1) {
+      vector_size_t rowIndex = -1) {
     return Envelope{
         .minX = std::nextafterf(
             static_cast<float>(minX), -std::numeric_limits<float>::infinity()),
@@ -136,6 +137,52 @@ struct Envelope {
   }
 };
 
+/// A single level of an R-tree.  It is a set of envelopes that can be linearly
+/// scanned for envelope intersection.
+class RTreeLevel {
+ public:
+  RTreeLevel(const RTreeLevel&) = delete;
+  RTreeLevel& operator=(const RTreeLevel&) = delete;
+
+  RTreeLevel() = default;
+  RTreeLevel(RTreeLevel&&) = default;
+  RTreeLevel& operator=(RTreeLevel&&) = default;
+  ~RTreeLevel() = default;
+
+  explicit RTreeLevel(
+      size_t branchSize,
+      std::vector<float> minXs,
+      std::vector<float> minYs,
+      std::vector<float> maxXs,
+      std::vector<float> maxYs)
+      : branchSize_{branchSize},
+        minXs_(std::move(minXs)),
+        minYs_(std::move(minYs)),
+        maxXs_(std::move(maxXs)),
+        maxYs_(std::move(maxYs)) {}
+
+  /// Returns the internal indices of all envelopes that probeEnv intersects.
+  /// Order of the returned indices is an implementation detail and cannot be
+  /// relied upon.
+  /// This does not do a short-circuit bounds check: the caller should do that
+  /// first.
+  std::vector<size_t> query(
+      const Envelope& queryEnv,
+      const std::vector<size_t>& branchIndices) const;
+
+  size_t size() const {
+    return minXs_.size();
+  }
+
+ private:
+  size_t branchSize_{};
+  Envelope bounds_;
+  std::vector<float> minXs_{};
+  std::vector<float> minYs_{};
+  std::vector<float> maxXs_{};
+  std::vector<float> maxYs_{};
+};
+
 /// A spatial index for a set of geometries. The index only cares about the
 /// envelopes of the geometries, and an index into the geometries (not stored in
 /// SpatialIndex).
@@ -154,29 +201,31 @@ class SpatialIndex {
   SpatialIndex& operator=(SpatialIndex&&) = default;
   ~SpatialIndex() = default;
 
+  static const uint32_t kDefaultRTreeBranchSize = 32;
+
   /// Constructs a spatial index from envelopes contained with `bounds`.
   /// `bounds` must contain all envelopes in `envelopes`, otherwise the
-  /// an assertio will fail.  Envelopes should not contains
-  /// and NaN coordinates.
-  explicit SpatialIndex(Envelope bounds, std::vector<Envelope> envelopes);
+  /// an assertio will fail.  Envelopes should not contain NaN coordinates.
+  explicit SpatialIndex(
+      Envelope bounds,
+      std::vector<Envelope> envelopes,
+      uint32_t branchSize = kDefaultRTreeBranchSize);
 
   /// Returns the row indices of all envelopes that probeEnv intersects.
   /// Order of the returned indices is an implementation detail and cannot be
   /// relied upon.
-  std::vector<int32_t> query(const Envelope& queryEnv) const;
+  std::vector<vector_size_t> query(const Envelope& queryEnv) const;
 
   /// Returns the envelope of the all envelopes in the index.
   /// The returned envelope will have index = -1.
   Envelope bounds() const;
 
  private:
-  Envelope bounds_ = Envelope::empty();
+  uint32_t branchSize_ = kDefaultRTreeBranchSize;
 
-  std::vector<double> minXs_{};
-  std::vector<double> minYs_{};
-  std::vector<double> maxXs_{};
-  std::vector<double> maxYs_{};
-  std::vector<int32_t> rowIndices_{};
+  Envelope bounds_ = Envelope::empty();
+  std::vector<RTreeLevel> levels_{};
+  std::vector<vector_size_t> rowIndices_{};
 };
 
 } // namespace facebook::velox::exec