SERVER-15451 Adaptive initial shell in near query with 2d index

3 files changed, 29 insertions, 4 deletions

diff --git a/src/mongo/db/exec/geo_near.cpp b/src/mongo/db/exec/geo_near.cpp
index 452444c3c31..dc6b1fe3021 100644
--- a/src/mongo/db/exec/geo_near.cpp
+++ b/src/mongo/db/exec/geo_near.cpp
@@ -265,10 +265,14 @@ namespace mongo {
         return fullBounds;
     }
 
-    static double twoDBoundsIncrement(const GeoNearParams& nearParams) {
-        // TODO: Revisit and tune these
+    static double twoDBoundsIncrement(IndexDescriptor* twoDIndex, const GeoNearParams& nearParams) {
         if (FLAT == nearParams.nearQuery->centroid->crs) {
-            return 10;
+            GeoHashConverter::Parameters hashParams;
+            Status status = GeoHashConverter::parseParameters(twoDIndex->infoObj(), &hashParams);
+            invariant(status.isOK()); // The index status should always be valid
+
+            GeoHashConverter converter(hashParams);
+            return 5 * converter.sizeEdge(GeoHash(0u, 0u, hashParams.bits));
         }
         else {
             return kMaxEarthDistanceInMeters / 1000.0;
@@ -291,7 +295,7 @@ namespace mongo {
           _twoDIndex(twoDIndex),
           _fullBounds(twoDDistanceBounds(nearParams, twoDIndex)),
           _currBounds(_fullBounds.center(), -1, _fullBounds.getInner()),
-          _boundsIncrement(twoDBoundsIncrement(nearParams)) {
+          _boundsIncrement(twoDBoundsIncrement(twoDIndex, nearParams)) {
 
         getNearStats()->keyPattern = twoDIndex->keyPattern();
     }
diff --git a/src/mongo/db/geo/hash.cpp b/src/mongo/db/geo/hash.cpp
index baee21d2f4b..57996e0ba1e 100644
--- a/src/mongo/db/geo/hash.cpp
+++ b/src/mongo/db/geo/hash.cpp
@@ -690,6 +690,7 @@ namespace mongo {
         return distanceBetweenHashes(a, b);
     }
 
+    // TODO: Use ldexp(max - min, -level)
     double GeoHashConverter::sizeEdge(const GeoHash& a) const {
         if (!a.constrains())
             return _params.max - _params.min;
diff --git a/src/mongo/db/geo/hash_test.cpp b/src/mongo/db/geo/hash_test.cpp
index acb196faaca..dc9be31fe97 100644
--- a/src/mongo/db/geo/hash_test.cpp
+++ b/src/mongo/db/geo/hash_test.cpp
@@ -32,6 +32,7 @@
 
 #include <string>
 #include <sstream>
+#include <cmath>
 
 #include "mongo/db/geo/hash.h"
 #include "mongo/platform/random.h"
@@ -39,6 +40,7 @@
 #include "mongo/util/assert_util.h"
 
 using mongo::GeoHash;
+using mongo::GeoHashConverter;
 using std::string;
 using std::stringstream;
 
@@ -83,4 +85,22 @@ namespace {
         string a = makeRandomBitString(13);
         ASSERT_THROWS(makeHash(a), mongo::UserException);
     }
+
+    TEST(GeoHashConvertor, EdgeLength) {
+        const double kError = 10E-15;
+        GeoHashConverter::Parameters params;
+        params.max = 200.0;
+        params.min = 100.0;
+        params.bits = 32;
+        double numBuckets = (1024 * 1024 * 1024 * 4.0);
+        params.scaling = numBuckets / (params.max - params.min);
+
+        GeoHashConverter converter(params);
+
+        ASSERT_APPROX_EQUAL(100.0, converter.sizeEdge(GeoHash(0, 0, 0)), kError);
+        ASSERT_APPROX_EQUAL(50.0, converter.sizeEdge(GeoHash(0, 0, 1)), kError);
+        ASSERT_APPROX_EQUAL(25.0, converter.sizeEdge(GeoHash(0, 0, 2)), kError);
+        ASSERT_APPROX_EQUAL(ldexp(100.0, -26), converter.sizeEdge(GeoHash(0, 0, 26)), kError);
+        ASSERT_APPROX_EQUAL(ldexp(100.0, -32), converter.sizeEdge(GeoHash(0, 0, 32)), kError);
+    }
 }