SERVER-26492 terminate density estimation after exceeding $maxDistance

1 files changed, 84 insertions, 6 deletions

diff --git a/src/mongo/db/exec/geo_near.cpp b/src/mongo/db/exec/geo_near.cpp
index 56b10994c3c..2acb29c6b02 100644
--- a/src/mongo/db/exec/geo_near.cpp
+++ b/src/mongo/db/exec/geo_near.cpp
@@ -267,8 +267,13 @@ class GeoNear2DStage::DensityEstimator {
 public:
     DensityEstimator(PlanStage::Children* children,
                      const IndexDescriptor* twoDindex,
-                     const GeoNearParams* nearParams)
-        : _children(children), _twoDIndex(twoDindex), _nearParams(nearParams), _currentLevel(0) {
+                     const GeoNearParams* nearParams,
+                     const R2Annulus& fullBounds)
+        : _children(children),
+          _twoDIndex(twoDindex),
+          _nearParams(nearParams),
+          _fullBounds(fullBounds),
+          _currentLevel(0) {
         GeoHashConverter::Parameters hashParams;
         Status status = GeoHashConverter::parseParameters(_twoDIndex->infoObj(), &hashParams);
         // The index status should always be valid.
@@ -295,7 +300,8 @@ private:
     PlanStage::Children* _children;     // Points to PlanStage::_children in the NearStage.
     const IndexDescriptor* _twoDIndex;  // Not owned here.
     const GeoNearParams* _nearParams;   // Not owned here.
-    IndexScan* _indexScan = nullptr;    // Owned in PlanStage::_children.
+    const R2Annulus& _fullBounds;
+    IndexScan* _indexScan = nullptr;  // Owned in PlanStage::_children.
     unique_ptr<GeoHashConverter> _converter;
     GeoHash _centroidCell;
     unsigned _currentLevel;
@@ -363,6 +369,39 @@ PlanStage::StageState GeoNear2DStage::DensityEstimator::work(OperationContext* t
 
     if (state == PlanStage::IS_EOF) {
         // We ran through the neighbors but found nothing.
+        //
+        // Before going to the next-coarsest level, check whether our search area contains the
+        // entire search annulus, since we don't want to spend time doing density estimation over
+        // areas that are much larger than the requested $maxDistance.
+        //
+        // The search area consists of four cells with side length S. Within its cell, the closest
+        // vertex to the search point must be the vertex shared with the other three cells. If the
+        // search point lies in the upper left cell, this means that it must lie in the lower right
+        // quadrant of that cell. Furthermore, this lower-right quadrant has a side-length of S/2.
+        //
+        //   +-----------+-----------+
+        //   |           |           |
+        //   |       S/2 |           |
+        //   +     +-----+           |
+        //   |     | o   |           |
+        //   |     |     |           |
+        //   +-----+-----+-----------+
+        //   |           |           |
+        //   |           |           |
+        //   |           |           |
+        //   |           |           |
+        //   |           |           |
+        //   +-----------+-----------+
+        //         S
+        //
+        // As long as the outer radius of the search annulus is less than S/2, it must be entirely
+        // contained within these four cells.
+        if (_fullBounds.getOuter() < (0.5 * _converter->sizeEdge(_currentLevel))) {
+            // We're covering the entire search annulus. Return EOF to indicate we're done.
+            *estimatedDistance = 0.5 * _converter->sizeEdge(_currentLevel);
+            return PlanStage::IS_EOF;
+        }
+
         if (_currentLevel > 0u) {
             // Advance to the next level and search again.
             _currentLevel--;
@@ -395,7 +434,8 @@ PlanStage::StageState GeoNear2DStage::initialize(OperationContext* txn,
                                                  Collection* collection,
                                                  WorkingSetID* out) {
     if (!_densityEstimator) {
-        _densityEstimator.reset(new DensityEstimator(&_children, _twoDIndex, &_nearParams));
+        _densityEstimator.reset(
+            new DensityEstimator(&_children, _twoDIndex, &_nearParams, _fullBounds));
     }
 
     double estimatedDistance;
@@ -805,11 +845,13 @@ public:
     DensityEstimator(PlanStage::Children* children,
                      const IndexDescriptor* s2Index,
                      const GeoNearParams* nearParams,
-                     const S2IndexingParams& indexParams)
+                     const S2IndexingParams& indexParams,
+                     const R2Annulus& fullBounds)
         : _children(children),
           _s2Index(s2Index),
           _nearParams(nearParams),
           _indexParams(indexParams),
+          _fullBounds(fullBounds),
           _currentLevel(0) {
         // cellId.AppendVertexNeighbors(level, output) requires level < finest,
         // so we use the minimum of max_level - 1 and the user specified finest
@@ -832,6 +874,7 @@ private:
     const IndexDescriptor* _s2Index;   // Not owned here.
     const GeoNearParams* _nearParams;  // Not owned here.
     const S2IndexingParams _indexParams;
+    const R2Annulus& _fullBounds;
     int _currentLevel;
     IndexScan* _indexScan = nullptr;  // Owned in PlanStage::_children.
 };
@@ -885,6 +928,41 @@ PlanStage::StageState GeoNear2DSphereStage::DensityEstimator::work(OperationCont
 
     if (state == PlanStage::IS_EOF) {
         // We ran through the neighbors but found nothing.
+        //
+        // Before going to the next-coarsest level, check whether our search area contains the
+        // entire search annulus, since we don't want to spend time doing density estimation over
+        // areas that are much larger than the requested $maxDistance.
+        //
+        // The search area consists of four cells at level L. Within its cell, the closest vertex to
+        // the search point must be the vertex shared with the other three cells. If the search
+        // point lies in the upper left cell, this means that it must lie in the lower right
+        // sub-cell at level L+1.
+        //
+        //   +-----------+-----------+
+        //   |           |           |
+        //   |        S  |           |
+        //   +     +-----+           |
+        //   |     | o   |           |
+        //   |     |     |           |
+        //   +-----+-----+-----------+
+        //   |           |           |
+        //   |           |           |
+        //   |           |           |
+        //   |           |           |
+        //   |           |           |
+        //   +-----------+-----------+
+        //
+        // In the diagram above, S is the width of the cell at level L+1. We can determine a lower
+        // bound for the width any cell at this level, i.e. S > minWidth(L+1). As long as the outer
+        // radius of the search annulus is less than minWidth(L+1), it must be entirely contained
+        // within these four level L cells.
+        if (_fullBounds.getOuter() <
+            (S2::kMinWidth.GetValue(_currentLevel + 1) * kRadiusOfEarthInMeters)) {
+            // We're covering the entire search annulus. Return EOF to indicate we're done.
+            *estimatedDistance = S2::kMinWidth.GetValue(_currentLevel + 1) * kRadiusOfEarthInMeters;
+            return PlanStage::IS_EOF;
+        }
+
         if (_currentLevel > 0) {
             // Advance to the next level and search again.
             _currentLevel--;
@@ -919,7 +997,7 @@ PlanStage::StageState GeoNear2DSphereStage::initialize(OperationContext* txn,
                                                        WorkingSetID* out) {
     if (!_densityEstimator) {
         _densityEstimator.reset(
-            new DensityEstimator(&_children, _s2Index, &_nearParams, _indexParams));
+            new DensityEstimator(&_children, _s2Index, &_nearParams, _indexParams, _fullBounds));
     }
 
     double estimatedDistance;