add test for 2d+geoNear+within, add fix to let within work in s2+geoNear

mongodb · Dec 11, 2012 · f17fbd3 · f17fbd3
1 parent 7fcaffc
commit f17fbd3
Show file tree

Hide file tree

Showing 7 changed files with 96 additions and 21 deletions.
diff --git a/jstests/geo_nearwithin.js b/jstests/geo_nearwithin.js
@@ -0,0 +1,27 @@
+// Test geoNear + $within.
+t = db.geo_nearwithin
+t.drop();
+
+points = 10
+for (var x = -points; x < points; x += 1) {
+    for (var y = -points; y < points; y += 1) {
+        t.insert({geo: [x, y]})
+    }
+}
+
+t.ensureIndex({ geo : "2d" })
+
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[0, 0], 1]}}}})
+assert.eq(resNear.results.length, 5)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[0, 0], 0]}}}})
+assert.eq(resNear.results.length, 1)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[1, 0], 0.5]}}}})
+assert.eq(resNear.results.length, 1)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[1, 0], 1.5]}}}})
+assert.eq(resNear.results.length, 9)
+
+// We want everything distance >1 from us but <1.5
+// These points are (-+1, -+1)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {$and: [{geo: {$within: {$center: [[0, 0], 1.5]}}},
+                                                                             {geo: {$not: {$within: {$center: [[0,0], 1]}}}}]}})
+assert.eq(resNear.results.length, 4)
diff --git a/jstests/geo_s2nearwithin.js b/jstests/geo_s2nearwithin.js
@@ -0,0 +1,24 @@
+// Test geoNear + $within.
+t = db.geo_s2nearwithin
+t.drop();
+
+points = 10
+for (var x = -points; x < points; x += 1) {
+    for (var y = -points; y < points; y += 1) {
+        t.insert({geo: [x, y]})
+    }
+}
+
+t.ensureIndex({ geo : "2dsphere" })
+
+// So, this is kind of messed-up.  The math for $within assumes a plane, but
+// 2dsphere is spherical.  And yet...it still works, though it's inconsistent.
+// TODO(hk): clarify this all.  :)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[0, 0], 1]}}}})
+assert.eq(resNear.results.length, 5)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[0, 0], 0]}}}})
+assert.eq(resNear.results.length, 1)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[1, 0], 0.5]}}}})
+assert.eq(resNear.results.length, 1)
+resNear = db.runCommand({geoNear : t.getName(), near: [0, 0], query: {geo: {$within: {$center: [[1, 0], 1.5]}}}})
+assert.eq(resNear.results.length, 9)
diff --git a/src/mongo/db/geo/s2index.cpp b/src/mongo/db/geo/s2index.cpp
@@ -236,13 +236,22 @@ namespace mongo {
             if (numWanted < 0) numWanted *= -1;
             if (0 == numWanted) numWanted = INT_MAX;
 
+            BSONObjBuilder geoFieldsToNuke;
+            for (size_t i = 0; i < _fields.size(); ++i) {
+                const IndexedField &field = _fields[i];
+                if (IndexedField::GEO != field.type) { continue; }
+                geoFieldsToNuke.append(field.name, "");
+            }
+            // false means we want to filter OUT geoFieldsToNuke, not filter to include only that.
+            BSONObj filteredQuery = query.filterFieldsUndotted(geoFieldsToNuke.obj(), false);
+
             if (isNear) {
-                S2NearCursor *cursor = new S2NearCursor(keyPattern(), getDetails(), query, regions,
+                S2NearCursor *cursor = new S2NearCursor(keyPattern(), getDetails(), filteredQuery, regions,
                                                         _params, numWanted, maxDistance);
                 return shared_ptr<Cursor>(cursor);
             } else {
                 // Default to intersect.
-                S2Cursor *cursor = new S2Cursor(keyPattern(), getDetails(), query, regions, _params,
+                S2Cursor *cursor = new S2Cursor(keyPattern(), getDetails(), filteredQuery, regions, _params,
                                                 numWanted);
                 return shared_ptr<Cursor>(cursor);
             }
@@ -407,8 +416,9 @@ namespace mongo {
         vector<QueryGeometry> regions;
         regions.push_back(queryGeo);
 
-        scoped_ptr<S2NearCursor> cursor(new S2NearCursor(idxType->keyPattern(), idxType->getDetails(), query, regions,
+        scoped_ptr<S2NearCursor> cursor(new S2NearCursor(idxType->keyPattern(), idxType->getDetails(),
-                    idxType->getParams(), numWanted, maxDistance));
+                                                         query, regions, idxType->getParams(),
+                                                         numWanted, maxDistance));
 
         double totalDistance = 0;
         int results = 0;

diff --git a/src/mongo/db/geo/s2nearcursor.cpp b/src/mongo/db/geo/s2nearcursor.cpp
@@ -35,7 +35,8 @@ namespace mongo {
         }
         // false means we want to filter OUT geoFieldsToNuke, not filter to include only that.
         _filteredQuery = query.filterFieldsUndotted(geoFieldsToNuke.obj(), false);
-        _matcher.reset(new CoveredIndexMatcher(_filteredQuery, keyPattern));
+        // We match on the whole query, since it might have $within.
+        _matcher.reset(new CoveredIndexMatcher(query, keyPattern));
 
         // More indexing machinery.
         BSONObjBuilder specBuilder;

diff --git a/src/mongo/db/geo/s2nearcursor.h b/src/mongo/db/geo/s2nearcursor.h
@@ -82,7 +82,13 @@ namespace mongo {
 
         // Need this to make a FieldRangeSet.
         const IndexDetails *_details;
-        // The query with the geo stuff taken out.  We use this with a matcher.
+
+        // How we need/use the query:
+        // Matcher: Can have geo fields in it, but only with $within.
+        //          This only really happens (right now) from geoNear command.
+        //          We assume the caller takes care of this in the right way.
+        // FRS:     No geo fields allowed!
+        // So, on that note: the query with the geo stuff taken out, used by makeFRSObject().
         BSONObj _filteredQuery;
         // What geo regions are we looking for?
         vector<QueryGeometry> _fields;

diff --git a/src/mongo/db/matcher.cpp b/src/mongo/db/matcher.cpp
@@ -387,15 +387,15 @@ namespace mongo {
                 uassert(16519, "Malformed $box: " + obj.toString(), coordIt.more());
                 BSONElement maxE = coordIt.next();
                 uassert(16520, "Malformed $box: " + obj.toString(), minE.isABSONObj());
-                _geo.push_back(GeoMatcher::makeBox(e.fieldName(), minE.Obj(), maxE.Obj()));
+                _geo.push_back(GeoMatcher::makeBox(e.fieldName(), minE.Obj(), maxE.Obj(), isNot));
             } else if (str::equals(elt.fieldName(), "$center")) {
                 BSONElement center = coordIt.next();
                 uassert(16521, "Malformed $center: " + obj.toString(), center.isABSONObj());
                 uassert(16522, "Malformed $center: " + obj.toString(), coordIt.more());
                 BSONElement radius = coordIt.next();
                 uassert(16523, "Malformed $center: " + obj.toString(), radius.isNumber());
                 _geo.push_back(
-                        GeoMatcher::makeCircle(e.fieldName(), center.Obj(), radius.number()));
+                        GeoMatcher::makeCircle(e.fieldName(), center.Obj(), radius.number(), isNot));
             } else if (str::equals(elt.fieldName(), "$polygon")) {
                 while (coordIt.more()) {
                     BSONElement coord = coordIt.next();
@@ -408,7 +408,7 @@ namespace mongo {
                     BSONElement y = numIt.next();
                     uassert(16528, "Malformed $polygon: " + obj.toString(), y.isNumber());
                 }
-                _geo.push_back(GeoMatcher::makePolygon(e.fieldName(), elt.Obj()));
+                _geo.push_back(GeoMatcher::makePolygon(e.fieldName(), elt.Obj(), isNot));
             } else {
                 uasserted(16529, "Couldn't pull any geometry out of $within query: " + obj.toString());
             }

diff --git a/src/mongo/db/matcher.h b/src/mongo/db/matcher.h
@@ -52,8 +52,9 @@ namespace mongo {
 
     class GeoMatcher {
     private:
-        GeoMatcher(const string& field) : _isBox(false), _isCircle(false), _isPolygon(false),
+        GeoMatcher(const string& field, bool isNot) : _isBox(false), _isCircle(false),
-                                          _fieldName(field) {}
+                                                      _isPolygon(false), _fieldName(field),
+                                                      _isNot(isNot) {}
         bool _isBox;
         Box _box;
 
@@ -65,27 +66,31 @@ namespace mongo {
         Polygon _polygon;
 
         string _fieldName;
+        bool _isNot;
     public:
         const string& getFieldName() const { return _fieldName; }
 
-        static GeoMatcher makeBox(const string& field, const BSONObj &min, const BSONObj &max) {
+        static GeoMatcher makeBox(const string& field, const BSONObj &min, const BSONObj &max,
-            GeoMatcher m(field);
+                                  bool isNot) {
+            GeoMatcher m(field, isNot);
             m._isBox = true;
             uassert(16511, "Malformed coord: " + min.toString(), pointFrom(min, &m._box._min));
             uassert(16512, "Malformed coord: " + max.toString(), pointFrom(max, &m._box._max));
             return m;
         }
 
-        static GeoMatcher makeCircle(const string& field, const BSONObj &center, double rad) {
+        static GeoMatcher makeCircle(const string& field, const BSONObj &center, double rad,
-            GeoMatcher m(field);
+                                     bool isNot) {
+            GeoMatcher m(field, isNot);
             m._isCircle = true;
             uassert(16513, "Malformed coord: " + center.toString(), pointFrom(center, &m._center));
             m._radius = rad;
             return m;
         }
 
-        static GeoMatcher makePolygon(const string& field, const BSONObj &poly) {
+        static GeoMatcher makePolygon(const string& field, const BSONObj &poly,
-            GeoMatcher m(field);
+                                      bool isNot) {
+            GeoMatcher m(field, isNot);
             vector<Point> points;
 
             m._isPolygon = true;
@@ -102,15 +107,17 @@ namespace mongo {
         }
 
         bool containsPoint(Point p) const {
+            bool ret;
             if (_isBox) {
-                return _box.inside(p, 0);
+                ret = _box.inside(p, 0);
             } else if (_isCircle) {
-                return distance(p, _center) <= _radius;
+                ret = distance(p, _center) <= _radius;
             } else if (_isPolygon) {
-                return _polygon.contains(p);
+                ret = _polygon.contains(p);
             } else {
-                return false;
+                ret = false;
             }
+            return _isNot ? !ret : ret;
         }
 
         string toString() const {