Add maximum search distance parameter to QgsSpatialIndex::nearestNeighbor

nyalldawson · nyalldawson · commit c8a4dff475ea · 2019-02-20T05:40:02.000+10:00
diff --git a/python/core/auto_generated/qgsspatialindex.sip.in b/python/core/auto_generated/qgsspatialindex.sip.in
@@ -147,16 +147,24 @@ Returns a list of features with a bounding box which intersects the specified ``
    when required.
 %End
 
-    QList<QgsFeatureId> nearestNeighbor( const QgsPointXY &point, int neighbors ) const;
+    QList<QgsFeatureId> nearestNeighbor( const QgsPointXY &point, int neighbors = 1, double maxDistance = 0 ) const;
 %Docstring
-Returns nearest neighbors to a ``point``. The number of neighbours returned is specified
-by the ``neighbours`` argument.
+Returns nearest neighbors to a ``point``. The number of neighbors returned is specified
+by the ``neighbors`` argument.
+
+If the ``maxDistance`` argument is greater than 0, then only features within the specified
+distance of ``point`` will be considered.
+
+Note that in some cases the number of returned features may differ from the requested
+number of ``neighbors``. E.g. if not enough features exist within the ``maxDistance`` of the
+search point. If multiple features are equidistant from the search ``point`` then the
+number of returned feature IDs may exceed ``neighbors``.
 
 .. warning::
 
    If this QgsSpatialIndex object was not constructed with the FlagStoreFeatureGeometries flag,
-   then the nearest neighbour test is performed based on the feature bounding boxes ONLY, so for non-point
-   geometry features this method is not guaranteed to return the actual closest neighbours.
+   then the nearest neighbor test is performed based on the feature bounding boxes ONLY, so for non-point
+   geometry features this method is not guaranteed to return the actual closest neighbors.
 %End
 
 
diff --git a/src/core/qgsspatialindex.cpp b/src/core/qgsspatialindex.cpp
@@ -88,29 +88,60 @@ class QgsSpatialIndexCopyVisitor : public SpatialIndex::IVisitor
 };
 
 ///@cond PRIVATE
-class QgsNearestNeighbourComparator : public INearestNeighborComparator
+class QgsNearestNeighborComparator : public INearestNeighborComparator
 {
   public:
 
-    QgsNearestNeighbourComparator( const QHash< QgsFeatureId, QgsGeometry > &geometries, const QgsPointXY &point )
+    QgsNearestNeighborComparator( const QHash< QgsFeatureId, QgsGeometry > *geometries, const QgsPointXY &point, double maxDistance )
       : mGeometries( geometries )
       , mPoint( QgsGeometry::fromPointXY( point ) )
+      , mMaxDistance( maxDistance )
     {
     }
 
-    QHash< QgsFeatureId, QgsGeometry > mGeometries;
+    const QHash< QgsFeatureId, QgsGeometry > *mGeometries = nullptr;
     QgsGeometry mPoint;
+    double mMaxDistance = 0;
+    QSet< QgsFeatureId > mFeaturesOutsideMaxDistance;
 
-    double getMinimumDistance( const IShape &, const IShape & ) override
+    double getMinimumDistance( const IShape &query, const IShape &entry ) override
     {
-      Q_ASSERT( false ); // not implemented!
-      return 0;
+      return query.getMinimumDistance( entry );
     }
 
-    double getMinimumDistance( const IShape &, const IData &data ) override
+    double getMinimumDistance( const IShape &query, const IData &data ) override
     {
-      QgsGeometry other = mGeometries.value( data.getIdentifier() );
-      return other.distance( mPoint );
+      // start with the default implementation, which gives distance to bounding box only
+      IShape *pS;
+      data.getShape( &pS );
+      double dist = query.getMinimumDistance( *pS );
+      delete pS;
+
+      // if doing exact distance search, AND either no max distance specified OR the
+      // distance to the bounding box is less than the max distance, calculate the exact
+      // distance now.
+      // (note: if bounding box is already greater than the distance away then max distance, there's no
+      // point doing this more expensive calculation, since we can't possibly use this feature anyway!)
+      if ( mGeometries && ( mMaxDistance <= 0.0 || dist <= mMaxDistance ) )
+      {
+        QgsGeometry other = mGeometries->value( data.getIdentifier() );
+        dist = other.distance( mPoint );
+      }
+
+      if ( mMaxDistance > 0 && dist > mMaxDistance )
+      {
+        // feature is outside of maximum distance threshold. Flag it,
+        // but "trick" libspatialindex into considering it as just outside
+        // our max distance region. This means if there's no other closer features (i.e.,
+        // within our actual maximum search distance), the libspatialindex
+        // nearest neighbor test will use this feature and not needlessly continue hunting
+        // through the remaining more distant features in the index.
+        // TODO: add proper API to libspatialindex to allow a maximum search distance in
+        // nearest neighbor tests
+        mFeaturesOutsideMaxDistance.insert( data.getIdentifier() );
+        return mMaxDistance + 0.00000001;
+      }
+      return dist;
     }
 };
 
@@ -442,7 +473,7 @@ QList<QgsFeatureId> QgsSpatialIndex::intersects( const QgsRectangle &rect ) cons
   return list;
 }
 
-QList<QgsFeatureId> QgsSpatialIndex::nearestNeighbor( const QgsPointXY &point, int neighbors ) const
+QList<QgsFeatureId> QgsSpatialIndex::nearestNeighbor( const QgsPointXY &point, const int neighbors, const double maxDistance ) const
 {
   QList<QgsFeatureId> list;
   QgisVisitor visitor( list );
@@ -451,15 +482,18 @@ QList<QgsFeatureId> QgsSpatialIndex::nearestNeighbor( const QgsPointXY &point, i
   Point p( pt, 2 );
 
   QMutexLocker locker( &d->mMutex );
-  if ( d->mFlags & QgsSpatialIndex::FlagStoreFeatureGeometries )
-  {
-    QgsNearestNeighbourComparator nnc( d->mGeometries, point );
-    d->mRTree->nearestNeighborQuery( neighbors, p, visitor, nnc );
-  }
-  else
+  QgsNearestNeighborComparator nnc( d->mFlags & QgsSpatialIndex::FlagStoreFeatureGeometries ? &d->mGeometries : nullptr,
+                                    point, maxDistance );
+  d->mRTree->nearestNeighborQuery( neighbors, p, visitor, nnc );
+
+  if ( maxDistance > 0 )
   {
-    // simple bounding box search - meaningless for non-points
-    d->mRTree->nearestNeighborQuery( neighbors, p, visitor );
+    // trim features outside of max distance
+    list.erase( std::remove_if( list.begin(), list.end(),
+                                [&nnc]( QgsFeatureId id )
+    {
+      return nnc.mFeaturesOutsideMaxDistance.contains( id );
+    } ), list.end() );
   }
 
   return list;
diff --git a/src/core/qgsspatialindex.h b/src/core/qgsspatialindex.h
@@ -74,7 +74,7 @@ class CORE_EXPORT QgsSpatialIndex : public QgsFeatureSink
     //! Flags controlling index behavior
     enum Flag
     {
-      FlagStoreFeatureGeometries = 1 << 0, //!< Indicates that the spatial index should also store feature geometries. This requires more memory, but can speed up operations by avoiding additional requests to data providers to fetch matching feature geometries. Additionally, it is required for non-bounding box nearest neighbour searches.
+      FlagStoreFeatureGeometries = 1 << 0, //!< Indicates that the spatial index should also store feature geometries. This requires more memory, but can speed up operations by avoiding additional requests to data providers to fetch matching feature geometries. Additionally, it is required for non-bounding box nearest neighbor searches.
     };
     Q_DECLARE_FLAGS( Flags, Flag )
 
@@ -175,14 +175,22 @@ class CORE_EXPORT QgsSpatialIndex : public QgsFeatureSink
     QList<QgsFeatureId> intersects( const QgsRectangle &rectangle ) const;
 
     /**
-     * Returns nearest neighbors to a \a point. The number of neighbours returned is specified
-     * by the \a neighbours argument.
+     * Returns nearest neighbors to a \a point. The number of neighbors returned is specified
+     * by the \a neighbors argument.
+     *
+     * If the \a maxDistance argument is greater than 0, then only features within the specified
+     * distance of \a point will be considered.
+     *
+     * Note that in some cases the number of returned features may differ from the requested
+     * number of \a neighbors. E.g. if not enough features exist within the \a maxDistance of the
+     * search point. If multiple features are equidistant from the search \a point then the
+     * number of returned feature IDs may exceed \a neighbors.
      *
      * \warning If this QgsSpatialIndex object was not constructed with the FlagStoreFeatureGeometries flag,
-     * then the nearest neighbour test is performed based on the feature bounding boxes ONLY, so for non-point
-     * geometry features this method is not guaranteed to return the actual closest neighbours.
+     * then the nearest neighbor test is performed based on the feature bounding boxes ONLY, so for non-point
+     * geometry features this method is not guaranteed to return the actual closest neighbors.
      */
-    QList<QgsFeatureId> nearestNeighbor( const QgsPointXY &point, int neighbors ) const;
+    QList<QgsFeatureId> nearestNeighbor( const QgsPointXY &point, int neighbors = 1, double maxDistance = 0 ) const;
 
 #ifndef SIP_RUN
 
diff --git a/tests/src/core/testqgsspatialindex.cpp b/tests/src/core/testqgsspatialindex.cpp
@@ -315,15 +315,33 @@ class TestQgsSpatialIndex : public QObject
       f1.setGeometry( QgsGeometry::fromWkt( QStringLiteral( "LineString(1 1, 3 1, 3 3)" ) ) );
       QgsFeature f2( 2 );
       f2.setGeometry( QgsGeometry::fromWkt( QStringLiteral( "LineString(0 1, 0 3)" ) ) );
+      QgsFeature f3( 3 );
+      f3.setGeometry( QgsGeometry::fromWkt( QStringLiteral( "LineString(0 4, 1 5, 3 3)" ) ) );
       i.addFeature( f1 );
       i2.addFeature( f1 );
       i.addFeature( f2 );
       i2.addFeature( f2 );
+      i.addFeature( f3 );
+      i2.addFeature( f3 );
 
       // i does not store feature geometries, so nearest neighbour search uses bounding box only
-      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 3 ), 1 ), QList< QgsFeatureId >() << 1 );
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 1 ), QList< QgsFeatureId >() << 1 );
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2 ), QList< QgsFeatureId >() << 1 << 3 );
       // i2 does store feature geometries, so nearest neighbour is exact
-      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 3 ), 1 ), QList< QgsFeatureId >() << 2 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 1 ), QList< QgsFeatureId >() << 2 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2 ), QList< QgsFeatureId >() << 2 << 3 );
+
+      // with maximum distance
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 1, 0.5 ), QList< QgsFeatureId >() << 1 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 1, 0.5 ), QList< QgsFeatureId >() );
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2, 0.5 ), QList< QgsFeatureId >() << 1 << 3 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2, 0.5 ), QList< QgsFeatureId >() );
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 1, 1.1 ), QList< QgsFeatureId >() << 1 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 1, 1.1 ), QList< QgsFeatureId >() << 2 );
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2, 1.1 ), QList< QgsFeatureId >() << 1 << 3 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2, 1.1 ), QList< QgsFeatureId >() << 2 );
+      QCOMPARE( i.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2, 2 ), QList< QgsFeatureId >() << 1 << 3 );
+      QCOMPARE( i2.nearestNeighbor( QgsPointXY( 1, 2.9 ), 2, 2 ), QList< QgsFeatureId >() << 2 << 3 );
 
     }
 
