Implement provider side FilterFids iterators for OGR provider

nyalldawson · nyalldawson · commit 9dbd96961b55 · 2016-08-09T13:25:59.000+10:00
Makes some operations with OGR sources magnitudes faster, ie zoom to 20 selected features in a 4 million point dataset: before: 14 seconds of blocked gui after: instant (cherry-picked from 1f02fd4)
diff --git a/src/providers/ogr/qgsogrfeatureiterator.cpp b/src/providers/ogr/qgsogrfeatureiterator.cpp
@@ -43,6 +43,8 @@ QgsOgrFeatureIterator::QgsOgrFeatureIterator( QgsOgrFeatureSource* source, bool
     , mFetchGeometry( false )
     , mGeometrySimplifier( nullptr )
     , mExpressionCompiled( false )
+    , mFilterFids( mRequest.filterFids() )
+    , mFilterFidsIt( mFilterFids.constBegin() )
 {
   mConn = QgsOgrConnPool::instance()->acquireConnection( mSource->mProvider->dataSourceUri() );
   if ( !mConn->ds )
@@ -213,6 +215,22 @@ bool QgsOgrFeatureIterator::providerCanSimplify( QgsSimplifyMethod::MethodType m
   return false;
 }
 
+bool QgsOgrFeatureIterator::fetchFeatureWithId( QgsFeatureId id, QgsFeature& feature ) const
+{
+  feature.setValid( false );
+  OGRFeatureH fet = OGR_L_GetFeature( ogrLayer, FID_TO_NUMBER( id ) );
+  if ( !fet )
+  {
+    return false;
+  }
+
+  if ( readFeature( fet, feature ) )
+    OGR_F_Destroy( fet );
+
+  feature.setValid( true );
+  return true;
+}
+
 bool QgsOgrFeatureIterator::fetchFeature( QgsFeature& feature )
 {
   feature.setValid( false );
@@ -222,19 +240,22 @@ bool QgsOgrFeatureIterator::fetchFeature( QgsFeature& feature )
 
   if ( mRequest.filterType() == QgsFeatureRequest::FilterFid )
   {
-    OGRFeatureH fet = OGR_L_GetFeature( ogrLayer, FID_TO_NUMBER( mRequest.filterFid() ) );
-    if ( !fet )
+    bool result = fetchFeatureWithId( mRequest.filterFid(), feature );
+    close(); // the feature has been read or was not found: we have finished here
+    return result;
+  }
+  else if ( mRequest.filterType() == QgsFeatureRequest::FilterFids )
+  {
+    while ( mFilterFidsIt != mFilterFids.constEnd() )
     {
-      close();
-      return false;
-    }
-
-    if ( readFeature( fet, feature ) )
-      OGR_F_Destroy( fet );
+      QgsFeatureId nextId = *mFilterFidsIt;
+      mFilterFidsIt++;
 
-    feature.setValid( true );
-    close(); // the feature has been read: we have finished here
-    return true;
+      if ( fetchFeatureWithId( nextId, feature ) )
+        return true;
+    }
+    close();
+    return false;
   }
 
   OGRFeatureH fet;
@@ -267,6 +288,8 @@ bool QgsOgrFeatureIterator::rewind()
 
   OGR_L_ResetReading( ogrLayer );
 
+  mFilterFidsIt = mFilterFids.constBegin();
+
   return true;
 }
 
@@ -293,7 +316,7 @@ bool QgsOgrFeatureIterator::close()
 }
 
 
-void QgsOgrFeatureIterator::getFeatureAttribute( OGRFeatureH ogrFet, QgsFeature & f, int attindex )
+void QgsOgrFeatureIterator::getFeatureAttribute( OGRFeatureH ogrFet, QgsFeature & f, int attindex ) const
 {
   OGRFieldDefnH fldDef = OGR_F_GetFieldDefnRef( ogrFet, attindex );
 
@@ -351,7 +374,7 @@ void QgsOgrFeatureIterator::getFeatureAttribute( OGRFeatureH ogrFet, QgsFeature
 }
 
 
-bool QgsOgrFeatureIterator::readFeature( OGRFeatureH fet, QgsFeature& feature )
+bool QgsOgrFeatureIterator::readFeature( OGRFeatureH fet, QgsFeature& feature ) const
 {
   feature.setFeatureId( OGR_F_GetFID( fet ) );
   feature.initAttributes( mSource->mFields.count() );
diff --git a/src/providers/ogr/qgsogrfeatureiterator.h b/src/providers/ogr/qgsogrfeatureiterator.h
@@ -70,10 +70,10 @@ class QgsOgrFeatureIterator : public QgsAbstractFeatureIteratorFromSource<QgsOgr
     //! fetch next feature filter expression
     bool nextFeatureFilterExpression( QgsFeature& f ) override;
 
-    bool readFeature( OGRFeatureH fet, QgsFeature& feature );
+    bool readFeature( OGRFeatureH fet, QgsFeature& feature ) const;
 
     //! Get an attribute associated with a feature
-    void getFeatureAttribute( OGRFeatureH ogrFet, QgsFeature & f, int attindex );
+    void getFeatureAttribute( OGRFeatureH ogrFet, QgsFeature & f, int attindex ) const;
 
     bool mFeatureFetched;
 
@@ -90,6 +90,10 @@ class QgsOgrFeatureIterator : public QgsAbstractFeatureIteratorFromSource<QgsOgr
     QgsOgrAbstractGeometrySimplifier* mGeometrySimplifier;
 
     bool mExpressionCompiled;
+    QgsFeatureIds mFilterFids;
+    QgsFeatureIds::const_iterator mFilterFidsIt;
+
+    bool fetchFeatureWithId( QgsFeatureId id, QgsFeature& feature ) const;
 
     //! returns whether the iterator supports simplify geometries on provider side
     virtual bool providerCanSimplify( QgsSimplifyMethod::MethodType methodType ) const override;
diff --git a/tests/src/python/providertestbase.py b/tests/src/python/providertestbase.py
@@ -352,10 +352,36 @@ def testGetFeaturesFidsTests(self):
         expected = set([fids[1], fids[3], fids[4]])
         assert result == expected, 'Expected {} and got {} when testing for feature IDs filter'.format(expected, result)
 
+        #providers should ignore non-existant fids
+        result = set([f.id() for f in self.provider.getFeatures(QgsFeatureRequest().setFilterFids([-101, fids[1], -102, fids[3], -103, fids[4], -104]))])
+        expected = set([fids[1], fids[3], fids[4]])
+        assert result == expected, 'Expected {} and got {} when testing for feature IDs filter'.format(expected, result)
+
         result = set([f.id() for f in self.provider.getFeatures(QgsFeatureRequest().setFilterFids([]))])
         expected = set([])
         assert result == expected, 'Expected {} and got {} when testing for feature IDs filter'.format(expected, result)
 
+        # Rewind mid-way
+        request = QgsFeatureRequest().setFilterFids([fids[1], fids[3], fids[4]])
+        feature_it = self.provider.getFeatures(request)
+        feature = QgsFeature()
+        feature.setValid(True)
+        self.assertTrue(feature_it.nextFeature(feature))
+        self.assertIn(feature.id(), [fids[1], fids[3], fids[4]])
+        first_feature = feature
+        self.assertTrue(feature.isValid())
+        # rewind
+        self.assertTrue(feature_it.rewind())
+        self.assertTrue(feature_it.nextFeature(feature))
+        self.assertEqual(feature.id(), first_feature.id())
+        self.assertTrue(feature.isValid())
+        # grab all features
+        self.assertTrue(feature_it.nextFeature(feature))
+        self.assertTrue(feature_it.nextFeature(feature))
+        # none left
+        self.assertFalse(feature_it.nextFeature(feature))
+        self.assertFalse(feature.isValid())
+
     def testGetFeaturesFilterRectTests(self):
         extent = QgsRectangle(-70, 67, -60, 80)
         request = QgsFeatureRequest().setFilterRect(extent)
diff --git a/tests/src/python/test_qgsfeatureiterator.py b/tests/src/python/test_qgsfeatureiterator.py
@@ -81,23 +81,20 @@ def test_FilterFids(self):
 
         ids = [feat.id() for feat in pointLayer.getFeatures(QgsFeatureRequest().setFilterFids([7, 8, 12, 30]))]
         expectedIds = [7, 8, 12]
-        myMessage = '\nExpected: {0} features\nGot: {1} features'.format(repr(expectedIds), repr(ids))
-        assert ids == expectedIds, myMessage
+        self.assertEquals(set(ids), set(expectedIds))
 
         pointLayer.startEditing()
         self.addFeatures(pointLayer)
 
         ids = [feat.id() for feat in pointLayer.getFeatures(QgsFeatureRequest().setFilterFids([-4, 7, 8, 12, 30]))]
         expectedIds = [-4, 7, 8, 12]
-        myMessage = '\nExpected: {0} features\nGot: {1} features'.format(repr(expectedIds), repr(ids))
-        assert ids == expectedIds, myMessage
+        self.assertEquals(set(ids), set(expectedIds))
 
         pointLayer.rollBack()
 
         ids = [feat.id() for feat in pointLayer.getFeatures(QgsFeatureRequest().setFilterFids([-2, 7, 8, 12, 30]))]
         expectedIds = [7, 8, 12]
-        myMessage = '\nExpected: {0} features\nGot: {1} features'.format(repr(expectedIds), repr(ids))
-        assert ids == expectedIds, myMessage
+        self.assertEquals(set(ids), set(expectedIds))
 
     def addFeatures(self, vl):
         feat = QgsFeature()
