[processing] Zonal Statistics create a new output instead of updating…

… an existing one Implements #29504
qgis · Sep 8, 2020 · c20ebcf · c20ebcf
1 parent 014de96
commit c20ebcf
Show file tree

Hide file tree

Showing 4 changed files with 233 additions and 178 deletions.
diff --git a/python/analysis/auto_generated/vector/qgszonalstatistics.sip.in b/python/analysis/auto_generated/vector/qgszonalstatistics.sip.in
@@ -42,6 +42,17 @@ A class that calculates raster statistics (count, sum, mean) for a polygon or mu
     typedef QFlags<QgsZonalStatistics::Statistic> Statistics;
 
 
+    enum Result
+    {
+      Success,
+      LayerTypeWrong,
+      LayerInvalid,
+      RasterInvalid,
+      RasterBandInvalid,
+      FailedToCreateField,
+      Cancelled
+    };
+
     QgsZonalStatistics( QgsVectorLayer *polygonLayer,
                         QgsRasterLayer *rasterLayer,
                         const QString &attributePrefix = QString(),
@@ -95,11 +106,27 @@ added to ``polygonLayer`` for each statistic calculated.
 .. versionadded:: 3.2
 %End
 
-    int calculateStatistics( QgsFeedback *feedback );
+
+    QgsZonalStatistics( QgsFeatureSource *source,
+                        QgsFeatureSink *sink,
+                        QgsRasterInterface *rasterInterface,
+                        const QgsCoordinateReferenceSystem &rasterCrs,
+                        const QMap<QgsZonalStatistics::Statistic, int> &statFieldIndexes,
+                        const QgsFields &fields,
+                        double rasterUnitsPerPixelX,
+                        double rasterUnitsPerPixelY,
+                        int rasterBand = 1,
+                        QgsZonalStatistics::Statistics stats = QgsZonalStatistics::Statistic::All );
 %Docstring
-Starts the calculation
 
-:return: 0 in case of success
+
+.. versionadded:: 3.16
+%End
+
+
+    QgsZonalStatistics::Result calculateStatistics( QgsFeedback *feedback );
+%Docstring
+Runs the calculation.
 %End
 
     static QString displayName( QgsZonalStatistics::Statistic statistic );

diff --git a/src/analysis/processing/qgsalgorithmzonalstatistics.cpp b/src/analysis/processing/qgsalgorithmzonalstatistics.cpp
@@ -89,14 +89,14 @@ void QgsZonalStatisticsAlgorithm::initAlgorithm( const QVariantMap & )
   addParameter( new QgsProcessingParameterRasterLayer( QStringLiteral( "INPUT_RASTER" ), QObject::tr( "Raster layer" ) ) );
   addParameter( new QgsProcessingParameterBand( QStringLiteral( "RASTER_BAND" ),
                 QObject::tr( "Raster band" ), 1, QStringLiteral( "INPUT_RASTER" ) ) );
-  addParameter( new QgsProcessingParameterVectorLayer( QStringLiteral( "INPUT_VECTOR" ), QObject::tr( "Vector layer containing zones" ),
+  addParameter( new QgsProcessingParameterFeatureSource( QStringLiteral( "INPUT_VECTOR" ), QObject::tr( "Feature source containing zones" ),
                 QList< int >() << QgsProcessing::TypeVectorPolygon ) );
   addParameter( new QgsProcessingParameterString( QStringLiteral( "COLUMN_PREFIX" ), QObject::tr( "Output column prefix" ), QStringLiteral( "_" ) ) );
 
   addParameter( new QgsProcessingParameterEnum( QStringLiteral( "STATISTICS" ), QObject::tr( "Statistics to calculate" ),
                 statChoices, true, QVariantList() << 0 << 1 << 2 ) );
 
-  addOutput( new QgsProcessingOutputVectorLayer( QStringLiteral( "INPUT_VECTOR" ), QObject::tr( "Zonal statistics" ), QgsProcessing::TypeVectorPolygon ) );
+  addParameter( new QgsProcessingParameterFeatureSink( QStringLiteral( "OUTPUT" ), QObject::tr( "Zonal statistics" ), QgsProcessing::TypeVectorPolygon, QVariant(), false, true, true ) );
 }
 
 bool QgsZonalStatisticsAlgorithm::prepareAlgorithm( const QVariantMap &parameters, QgsProcessingContext &context, QgsProcessingFeedback * )
@@ -129,24 +129,62 @@ bool QgsZonalStatisticsAlgorithm::prepareAlgorithm( const QVariantMap &parameter
 
 QVariantMap QgsZonalStatisticsAlgorithm::processAlgorithm( const QVariantMap &parameters, QgsProcessingContext &context, QgsProcessingFeedback *feedback )
 {
-  QgsVectorLayer *layer = parameterAsVectorLayer( parameters, QStringLiteral( "INPUT_VECTOR" ), context );
-  if ( !layer )
+  std::unique_ptr<QgsFeatureSource> source( parameterAsSource( parameters, QStringLiteral( "INPUT_VECTOR" ), context ) );
+  if ( !source )
     throw QgsProcessingException( QObject::tr( "Invalid zones layer" ) );
 
-  QgsZonalStatistics zs( layer,
+  QgsFields fields = source->fields();
+  QMap<QgsZonalStatistics::Statistic, int> statFieldIndexes;
+
+  for ( QgsZonalStatistics::Statistic stat :
+        {
+          QgsZonalStatistics::Count,
+          QgsZonalStatistics::Sum,
+          QgsZonalStatistics::Mean,
+          QgsZonalStatistics::Median,
+          QgsZonalStatistics::StDev,
+          QgsZonalStatistics::Min,
+          QgsZonalStatistics::Max,
+          QgsZonalStatistics::Range,
+          QgsZonalStatistics::Minority,
+          QgsZonalStatistics::Majority,
+          QgsZonalStatistics::Variety,
+          QgsZonalStatistics::Variance
+        } )
+  {
+    if ( mStats & stat )
+    {
+      QgsField field = QgsField( mPrefix + QgsZonalStatistics::shortName( stat ), QVariant::Double, QStringLiteral( "double precision" ) );
+      if ( fields.names().contains( field.name() ) )
+      {
+        throw QgsProcessingException( QObject::tr( "Field %1 already exists" ).arg( field.name() ) );
+      }
+      fields.append( field );
+      statFieldIndexes.insert( stat, fields.count() - 1 );
+    }
+  }
+
+  QString dest;
+  std::unique_ptr< QgsFeatureSink > sink( parameterAsSink( parameters, QStringLiteral( "OUTPUT" ), context, dest, fields, QgsWkbTypes::MultiPolygon, source->sourceCrs() ) );
+  if ( !sink )
+    throw QgsProcessingException( invalidSinkError( parameters, QStringLiteral( "OUTPUT" ) ) );
+
+  QgsZonalStatistics zs( source.get(),
+                         sink.get(),
                          mInterface.get(),
                          mCrs,
+                         statFieldIndexes,
+                         fields,
                          mPixelSizeX,
                          mPixelSizeY,
-                         mPrefix,
                          mBand,
-                         QgsZonalStatistics::Statistics( mStats )
+                         mStats
                        );
 
   zs.calculateStatistics( feedback );
 
   QVariantMap outputs;
-  outputs.insert( QStringLiteral( "INPUT_VECTOR" ), layer->id() );
+  outputs.insert( QStringLiteral( "OUTPUT" ), dest );
   return outputs;
 }