[processing] optimize random extract algorithm

qgis · Dec 12, 2019 · d167bab · d167bab
1 parent 19596a8
commit d167bab
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 9 deletions.
diff --git a/src/analysis/processing/qgsalgorithmrandomextract.cpp b/src/analysis/processing/qgsalgorithmrandomextract.cpp
@@ -110,22 +110,48 @@ QVariantMap QgsRandomExtractAlgorithm::processAlgorithm( const QVariantMap &para
   QVector< QgsFeatureId > fids( number );
   std::generate( fids.begin(), fids.end(), bind( fidsDistribution, mersenneTwister ) );
 
-  double step = number > 0 ? 100.0 / number : 1;
-  int i = 0;
-
-  QgsFeature f;
+  QHash< QgsFeatureId, int > idsCount;
   for ( QgsFeatureId id : fids )
   {
     if ( feedback->isCanceled() )
     {
       break;
     }
 
-    source->getFeatures( QgsFeatureRequest( id ), QgsProcessingFeatureSource::FlagSkipGeometryValidityChecks ).nextFeature( f );
+    if ( idsCount.contains( id ) )
+    {
+      idsCount.insert( id, idsCount.value( id ) + 1 );
+    }
+    else
+    {
+      idsCount.insert( id, 1 );
+    }
+  }
+
+  QgsFeatureIds ids = QSet< QgsFeatureId >::fromList( idsCount.keys() );
+  QgsFeatureIterator fit = source->getFeatures( QgsFeatureRequest().setFilterFids( ids ), QgsProcessingFeatureSource::FlagSkipGeometryValidityChecks );
+  sink->addFeatures( fit, QgsFeatureSink::FastInsert );
 
-    sink->addFeature( f, QgsFeatureSink::FastInsert );
-    feedback->setProgress( i * step );
-    i++;
+  // add duplicated features if any
+  if ( ids.size() != number )
+  {
+    QgsFeature f;
+    for ( auto it = idsCount.constBegin(); it != idsCount.constEnd(); ++it )
+    {
+      if ( feedback->isCanceled() )
+      {
+        break;
+      }
+
+      if ( it.value() > 1 )
+      {
+        source->getFeatures( QgsFeatureRequest( it.key() ), QgsProcessingFeatureSource::FlagSkipGeometryValidityChecks ).nextFeature( f );
+        for ( int j = 1; j < it.value(); j++ )
+        {
+          sink->addFeature( f, QgsFeatureSink::FastInsert );
+        }
+      }
+    }
   }
 
   QVariantMap outputs;

diff --git a/src/core/processing/qgsprocessingutils.cpp b/src/core/processing/qgsprocessingutils.cpp
@@ -1090,7 +1090,7 @@ bool QgsProcessingFeatureSink::addFeatures( QgsFeatureList &features, QgsFeature
 
 bool QgsProcessingFeatureSink::addFeatures( QgsFeatureIterator &iterator, QgsFeatureSink::Flags flags )
 {
-  bool result = !QgsProxyFeatureSink::addFeatures( iterator, flags );
+  bool result = QgsProxyFeatureSink::addFeatures( iterator, flags );
   if ( !result )
     mContext.feedback()->reportError( QObject::tr( "Features could not be written to %1" ).arg( mSinkName ) );
   return result;