[BUGFIX] [OGR provider] Make filter by id(s) requests work again on OSM datasets (fixes #20308)

rouault · rouault · commit bf3e96cfb059 · 2018-11-01T14:32:37.000+01:00
diff --git a/src/providers/ogr/qgsogrfeatureiterator.cpp b/src/providers/ogr/qgsogrfeatureiterator.cpp
@@ -41,12 +41,16 @@
 
 QgsOgrFeatureIterator::QgsOgrFeatureIterator( QgsOgrFeatureSource *source, bool ownSource, const QgsFeatureRequest &request )
   : QgsAbstractFeatureIteratorFromSource<QgsOgrFeatureSource>( source, ownSource, request )
-  , mFilterFids( mRequest.filterFids() )
-  , mFilterFidsIt( mFilterFids.constBegin() )
   , mSharedDS( source->mSharedDS )
   , mFirstFieldIsFid( source->mFirstFieldIsFid )
   , mFieldsWithoutFid( source->mFieldsWithoutFid )
 {
+  for ( const auto &id :  mRequest.filterFids() )
+  {
+    mFilterFids.insert( id );
+  }
+  mFilterFidsIt = mFilterFids.begin();
+
   // Since connection timeout for OGR connections is problematic and can lead to crashes, disable for now.
   mRequest.setTimeout( -1 );
   if ( mSharedDS )
@@ -235,7 +239,47 @@ bool QgsOgrFeatureIterator::fetchFeatureWithId( QgsFeatureId id, QgsFeature &fea
 {
   feature.setValid( false );
   gdal::ogr_feature_unique_ptr fet;
-  fet.reset( OGR_L_GetFeature( mOgrLayer, FID_TO_NUMBER( id ) ) );
+
+#if GDAL_VERSION_NUM >= GDAL_COMPUTE_VERSION(2,2,0)
+  if ( !QgsOgrProviderUtils::canDriverShareSameDatasetAmongLayers( mSource->mDriverName ) )
+  {
+    OGRLayerH nextFeatureBelongingLayer;
+    bool found = false;
+    // First pass: try to read from the last feature, in the hope the dataset
+    // returns them in increasing feature id number (and as we use a std::set
+    // for mFilterFids, we get them in increasing number by the iterator)
+    // Second pass: reset before reading
+    for ( int passNumber = 0; passNumber < 2; passNumber++ )
+    {
+      while ( fet.reset( GDALDatasetGetNextFeature(
+                           mConn->ds, &nextFeatureBelongingLayer, nullptr, nullptr, nullptr ) ), fet )
+      {
+        if ( nextFeatureBelongingLayer == mOgrLayer )
+        {
+          if ( OGR_F_GetFID( fet.get() ) == FID_TO_NUMBER( id ) )
+          {
+            found = true;
+            break;
+          }
+        }
+      }
+      if ( found || passNumber == 1 )
+      {
+        break;
+      }
+      GDALDatasetResetReading( mConn->ds );
+    }
+
+    if ( !found )
+    {
+      return false;
+    }
+  }
+  else
+#endif
+  {
+    fet.reset( OGR_L_GetFeature( mOgrLayer, FID_TO_NUMBER( id ) ) );
+  }
 
   if ( !fet )
   {
@@ -281,10 +325,10 @@ bool QgsOgrFeatureIterator::fetchFeature( QgsFeature &feature )
   }
   else if ( mRequest.filterType() == QgsFeatureRequest::FilterFids )
   {
-    while ( mFilterFidsIt != mFilterFids.constEnd() )
+    while ( mFilterFidsIt != mFilterFids.end() )
     {
       QgsFeatureId nextId = *mFilterFidsIt;
-      mFilterFidsIt++;
+      ++mFilterFidsIt;
 
       if ( fetchFeatureWithId( nextId, feature ) )
         return true;
@@ -351,7 +395,7 @@ bool QgsOgrFeatureIterator::rewind()
 
   resetReading();
 
-  mFilterFidsIt = mFilterFids.constBegin();
+  mFilterFidsIt = mFilterFids.begin();
 
   return true;
 }
diff --git a/src/providers/ogr/qgsogrfeatureiterator.h b/src/providers/ogr/qgsogrfeatureiterator.h
@@ -22,6 +22,7 @@
 #include <ogr_api.h>
 
 #include <memory>
+#include <set>
 
 class QgsOgrFeatureIterator;
 class QgsOgrProvider;
@@ -86,8 +87,9 @@ class QgsOgrFeatureIterator : public QgsAbstractFeatureIteratorFromSource<QgsOgr
     bool mFetchGeometry = false;
 
     bool mExpressionCompiled = false;
-    QgsFeatureIds mFilterFids;
-    QgsFeatureIds::const_iterator mFilterFidsIt;
+    // use std::set to get sorted ids (needed for efficient QgsFeatureRequest::FilterFids requests on OSM datasource)
+    std::set<QgsFeatureId> mFilterFids;
+    std::set<QgsFeatureId>::iterator mFilterFidsIt;
 
     QgsRectangle mFilterRect;
     QgsCoordinateTransform mTransform;
diff --git a/tests/src/python/test_provider_ogr.py b/tests/src/python/test_provider_ogr.py
@@ -463,6 +463,32 @@ def testOSM(self):
         iter_multipolygons = vl_multipolygons.getFeatures(QgsFeatureRequest())
         self.assertTrue(iter_multipolygons.nextFeature(f))
 
+        # Test filter by id (#20308)
+        f = next(vl_multipolygons.getFeatures(QgsFeatureRequest().setFilterFid(8)))
+        self.assertTrue(f.isValid())
+        self.assertEqual(f.id(), 8)
+
+        f = next(vl_multipolygons.getFeatures(QgsFeatureRequest().setFilterFid(1)))
+        self.assertTrue(f.isValid())
+        self.assertEqual(f.id(), 1)
+
+        f = next(vl_multipolygons.getFeatures(QgsFeatureRequest().setFilterFid(5)))
+        self.assertTrue(f.isValid())
+        self.assertEqual(f.id(), 5)
+
+        # 6 doesn't exist
+        it = vl_multipolygons.getFeatures(QgsFeatureRequest().setFilterFids([1, 5, 6, 8]))
+        f = next(it)
+        self.assertTrue(f.isValid())
+        self.assertEqual(f.id(), 1)
+        f = next(it)
+        self.assertTrue(f.isValid())
+        self.assertEqual(f.id(), 5)
+        f = next(it)
+        self.assertTrue(f.isValid())
+        self.assertEqual(f.id(), 8)
+        del it
+
 
 if __name__ == '__main__':
     unittest.main()
