Correctly support joins using virtual fields (fix #14820)

nyalldawson · nyalldawson · commit a3768672c8fc · 2016-06-10T12:53:40.000+10:00
diff --git a/src/core/qgsvectorlayerfeatureiterator.cpp b/src/core/qgsvectorlayerfeatureiterator.cpp
@@ -93,11 +93,7 @@ QgsVectorLayerFeatureIterator::QgsVectorLayerFeatureIterator( QgsVectorLayerFeat
     , mFetchedFid( false )
     , mInterruptionChecker( nullptr )
 {
-  prepareExpressions();
-
-  // prepare joins: may add more attributes to fetch (in order to allow join)
-  if ( mSource->mJoinBuffer->containsJoins() )
-    prepareJoins();
+  prepareFields();
 
   mHasVirtualAttributes = !mFetchJoinInfo.isEmpty() || !mExpressionFieldInfo.isEmpty();
 
@@ -461,128 +457,132 @@ void QgsVectorLayerFeatureIterator::rewindEditBuffer()
   mFetchChangedGeomIt = mSource->mChangedGeometries.constBegin();
 }
 
+void QgsVectorLayerFeatureIterator::prepareJoin( int fieldIdx )
+{
+  if ( !mSource->mFields.exists( fieldIdx ) )
+    return;
 
+  if ( mSource->mFields.fieldOrigin( fieldIdx ) != QgsFields::OriginJoin )
+    return;
 
-void QgsVectorLayerFeatureIterator::prepareJoins()
-{
-  QgsAttributeList fetchAttributes = ( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes ) ? mRequest.subsetOfAttributes() : mSource->mFields.allAttributesList();
-  QgsAttributeList sourceJoinFields; // attributes that also need to be fetched from this layer in order to have joins working
+  int sourceLayerIndex;
+  const QgsVectorJoinInfo* joinInfo = mSource->mJoinBuffer->joinForFieldIndex( fieldIdx, mSource->mFields, sourceLayerIndex );
+  Q_ASSERT( joinInfo );
 
-  mFetchJoinInfo.clear();
+  QgsVectorLayer* joinLayer = qobject_cast<QgsVectorLayer*>( QgsMapLayerRegistry::instance()->mapLayer( joinInfo->joinLayerId ) );
+  Q_ASSERT( joinLayer );
 
-  for ( QgsAttributeList::const_iterator attIt = fetchAttributes.constBegin(); attIt != fetchAttributes.constEnd(); ++attIt )
+  if ( !mFetchJoinInfo.contains( joinInfo ) )
   {
-    if ( !mSource->mFields.exists( *attIt ) )
-      continue;
+    FetchJoinInfo info;
+    info.joinInfo = joinInfo;
+    info.joinLayer = joinLayer;
+    info.indexOffset = mSource->mJoinBuffer->joinedFieldsOffset( joinInfo, mSource->mFields );
 
-    if ( mSource->mFields.fieldOrigin( *attIt ) != QgsFields::OriginJoin )
-      continue;
+    if ( joinInfo->targetFieldName.isEmpty() )
+      info.targetField = joinInfo->targetFieldIndex;    //for compatibility with 1.x
+    else
+      info.targetField = mSource->mFields.indexFromName( joinInfo->targetFieldName );
 
-    int sourceLayerIndex;
-    const QgsVectorJoinInfo* joinInfo = mSource->mJoinBuffer->joinForFieldIndex( *attIt, mSource->mFields, sourceLayerIndex );
-    Q_ASSERT( joinInfo );
+    if ( joinInfo->joinFieldName.isEmpty() )
+      info.joinField = joinInfo->joinFieldIndex;      //for compatibility with 1.x
+    else
+      info.joinField = joinLayer->fields().indexFromName( joinInfo->joinFieldName );
 
-    QgsVectorLayer* joinLayer = qobject_cast<QgsVectorLayer*>( QgsMapLayerRegistry::instance()->mapLayer( joinInfo->joinLayerId ) );
-    Q_ASSERT( joinLayer );
+    // for joined fields, we always need to request the targetField from the provider too
+    if ( !mPreparedFields.contains( info.targetField ) && !mFieldsToPrepare.contains( info.targetField ) )
+      mFieldsToPrepare << info.targetField;
 
-    if ( !mFetchJoinInfo.contains( joinInfo ) )
-    {
-      FetchJoinInfo info;
-      info.joinInfo = joinInfo;
-      info.joinLayer = joinLayer;
-      info.indexOffset = mSource->mJoinBuffer->joinedFieldsOffset( joinInfo, mSource->mFields );
-
-      if ( joinInfo->targetFieldName.isEmpty() )
-        info.targetField = joinInfo->targetFieldIndex;    //for compatibility with 1.x
-      else
-        info.targetField = mSource->mFields.indexFromName( joinInfo->targetFieldName );
-
-      if ( joinInfo->joinFieldName.isEmpty() )
-        info.joinField = joinInfo->joinFieldIndex;      //for compatibility with 1.x
-      else
-        info.joinField = joinLayer->fields().indexFromName( joinInfo->joinFieldName );
-
-      // for joined fields, we always need to request the targetField from the provider too
-      if ( !fetchAttributes.contains( info.targetField ) )
-        sourceJoinFields << info.targetField;
-
-      mFetchJoinInfo.insert( joinInfo, info );
-    }
+    if ( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes && !mRequest.subsetOfAttributes().contains( info.targetField ) )
+      mRequest.setSubsetOfAttributes( mRequest.subsetOfAttributes() << info.targetField );
 
-    // store field source index - we'll need it when fetching from provider
-    mFetchJoinInfo[ joinInfo ].attributes.push_back( sourceLayerIndex );
+    mFetchJoinInfo.insert( joinInfo, info );
   }
 
-  // add sourceJoinFields if we're using a subset
-  if ( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes )
-    mRequest.setSubsetOfAttributes( mRequest.subsetOfAttributes() + sourceJoinFields );
+  // store field source index - we'll need it when fetching from provider
+  mFetchJoinInfo[ joinInfo ].attributes.push_back( sourceLayerIndex );
 }
 
-void QgsVectorLayerFeatureIterator::prepareExpressions()
+void QgsVectorLayerFeatureIterator::prepareExpression( int fieldIdx )
 {
-  const QList<QgsExpressionFieldBuffer::ExpressionField> exps = mSource->mExpressionFieldBuffer->expressions();
+  const QList<QgsExpressionFieldBuffer::ExpressionField>& exps = mSource->mExpressionFieldBuffer->expressions();
 
-  mExpressionContext.reset( new QgsExpressionContext() );
-  mExpressionContext->appendScope( QgsExpressionContextUtils::globalScope() );
-  mExpressionContext->appendScope( QgsExpressionContextUtils::projectScope() );
-  mExpressionContext->setFields( mSource->mFields );
+  int oi = mSource->mFields.fieldOriginIndex( fieldIdx );
+  QgsExpression* exp = new QgsExpression( exps[oi].cachedExpression );
 
-  QList< int > virtualFieldsToFetch;
-  for ( int i = 0; i < mSource->mFields.count(); i++ )
+  QgsDistanceArea da;
+  da.setSourceCrs( mSource->mCrsId );
+  da.setEllipsoidalMode( true );
+  da.setEllipsoid( QgsProject::instance()->readEntry( "Measure", "/Ellipsoid", GEO_NONE ) );
+  exp->setGeomCalculator( da );
+  exp->setDistanceUnits( QgsProject::instance()->distanceUnits() );
+  exp->setAreaUnits( QgsProject::instance()->areaUnits() );
+
+  exp->prepare( mExpressionContext.data() );
+  mExpressionFieldInfo.insert( fieldIdx, exp );
+
+  Q_FOREACH ( const QString& col, exp->referencedColumns() )
   {
-    if ( mSource->mFields.fieldOrigin( i ) == QgsFields::OriginExpression )
+    int dependantFieldIdx = mSource->mFields.fieldNameIndex( col );
+    if ( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes )
     {
-      // Only prepare if there is no subset defined or the subset contains this field
-      if ( !( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes )
-           || mRequest.subsetOfAttributes().contains( i ) )
-      {
-        virtualFieldsToFetch << i;
-      }
+      mRequest.setSubsetOfAttributes( mRequest.subsetOfAttributes() << dependantFieldIdx );
     }
+    // also need to fetch this dependant field
+    if ( !mPreparedFields.contains( dependantFieldIdx ) && !mFieldsToPrepare.contains( dependantFieldIdx ) )
+      mFieldsToPrepare << dependantFieldIdx;
   }
 
-  QList< int > virtualFieldsProcessed;
-  while ( !virtualFieldsToFetch.isEmpty() )
+  if ( exp->needsGeometry() )
   {
-    int fieldIdx = virtualFieldsToFetch.takeFirst();
-    if ( virtualFieldsProcessed.contains( fieldIdx ) )
-      continue;
+    mRequest.setFlags( mRequest.flags() & ~QgsFeatureRequest::NoGeometry );
+  }
+}
 
-    virtualFieldsProcessed << fieldIdx;
+void QgsVectorLayerFeatureIterator::prepareFields()
+{
+  mPreparedFields.clear();
+  mFieldsToPrepare.clear();
+  mFetchJoinInfo.clear();
 
+  mExpressionContext.reset( new QgsExpressionContext() );
+  mExpressionContext->appendScope( QgsExpressionContextUtils::globalScope() );
+  mExpressionContext->appendScope( QgsExpressionContextUtils::projectScope() );
+  mExpressionContext->setFields( mSource->mFields );
 
-    int oi = mSource->mFields.fieldOriginIndex( fieldIdx );
-    QgsExpression* exp = new QgsExpression( exps[oi].cachedExpression );
+  mFieldsToPrepare = ( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes ) ? mRequest.subsetOfAttributes() : mSource->mFields.allAttributesList();
 
-    QgsDistanceArea da;
-    da.setSourceCrs( mSource->mCrsId );
-    da.setEllipsoidalMode( true );
-    da.setEllipsoid( QgsProject::instance()->readEntry( "Measure", "/Ellipsoid", GEO_NONE ) );
-    exp->setGeomCalculator( da );
-    exp->setDistanceUnits( QgsProject::instance()->distanceUnits() );
-    exp->setAreaUnits( QgsProject::instance()->areaUnits() );
+  while ( !mFieldsToPrepare.isEmpty() )
+  {
+    int fieldIdx = mFieldsToPrepare.takeFirst();
+    if ( mPreparedFields.contains( fieldIdx ) )
+      continue;
 
-    exp->prepare( mExpressionContext.data() );
-    mExpressionFieldInfo.insert( fieldIdx, exp );
+    mPreparedFields << fieldIdx;
+    prepareField( fieldIdx );
+  }
+}
 
-    Q_FOREACH ( const QString& col, exp->referencedColumns() )
-    {
-      int dependantFieldIdx = mSource->mFields.fieldNameIndex( col );
-      if ( mRequest.flags() & QgsFeatureRequest::SubsetOfAttributes )
+void QgsVectorLayerFeatureIterator::prepareField( int fieldIdx )
+{
+  switch ( mSource->mFields.fieldOrigin( fieldIdx ) )
+  {
+    case QgsFields::OriginExpression:
+      prepareExpression( fieldIdx );
+      break;
+
+    case QgsFields::OriginJoin:
+      if ( mSource->mJoinBuffer->containsJoins() )
       {
-        mRequest.setSubsetOfAttributes( mRequest.subsetOfAttributes() << dependantFieldIdx );
+        prepareJoin( fieldIdx );
       }
-      // also need to fetch this dependant field
-      if ( mSource->mFields.fieldOrigin( dependantFieldIdx ) == QgsFields::OriginExpression )
-        virtualFieldsToFetch << dependantFieldIdx;
-    }
+      break;
 
-    if ( exp->needsGeometry() )
-    {
-      mRequest.setFlags( mRequest.flags() & ~QgsFeatureRequest::NoGeometry );
-    }
+    case QgsFields::OriginUnknown:
+    case QgsFields::OriginProvider:
+    case QgsFields::OriginEdit:
+      break;
   }
-
 }
 
 void QgsVectorLayerFeatureIterator::addJoinedAttributes( QgsFeature &f )
@@ -612,24 +612,48 @@ void QgsVectorLayerFeatureIterator::addVirtualAttributes( QgsFeature& f )
   attr.resize( mSource->mFields.count() );  // Provider attrs count + joined attrs count + expression attrs count
   f.setAttributes( attr );
 
+  // possible TODO - handle combinations of expression -> join -> expression -> join?
+  // but for now, write that off as too complex and an unlikely rare, unsupported use case
+
+  QList< int > fetchedVirtualAttributes;
+  //first, check through joins for any virtual fields we need
+  QMap<const QgsVectorJoinInfo*, FetchJoinInfo>::const_iterator joinIt = mFetchJoinInfo.constBegin();
+  for ( ; joinIt != mFetchJoinInfo.constEnd(); ++joinIt )
+  {
+    if ( mExpressionFieldInfo.contains( joinIt->targetField ) )
+    {
+      // have to calculate expression field before we can handle this join
+      addExpressionAttribute( f, joinIt->targetField );
+      fetchedVirtualAttributes << joinIt->targetField;
+    }
+  }
+
   if ( !mFetchJoinInfo.isEmpty() )
     addJoinedAttributes( f );
 
+  // add remaining expression fields
   if ( !mExpressionFieldInfo.isEmpty() )
   {
     QMap<int, QgsExpression*>::ConstIterator it = mExpressionFieldInfo.constBegin();
-
     for ( ; it != mExpressionFieldInfo.constEnd(); ++it )
     {
-      QgsExpression* exp = it.value();
-      mExpressionContext->setFeature( f );
-      QVariant val = exp->evaluate( mExpressionContext.data() );
-      mSource->mFields.at( it.key() ).convertCompatible( val );
-      f.setAttribute( it.key(), val );
+      if ( fetchedVirtualAttributes.contains( it.key() ) )
+        continue;
+
+      addExpressionAttribute( f, it.key() );
     }
   }
 }
 
+void QgsVectorLayerFeatureIterator::addExpressionAttribute( QgsFeature& f, int attrIndex )
+{
+  QgsExpression* exp = mExpressionFieldInfo.value( attrIndex );
+  mExpressionContext->setFeature( f );
+  QVariant val = exp->evaluate( mExpressionContext.data() );
+  mSource->mFields.at( attrIndex ).convertCompatible( val );
+  f.setAttribute( attrIndex, val );
+}
+
 bool QgsVectorLayerFeatureIterator::prepareSimplification( const QgsSimplifyMethod& simplifyMethod )
 {
   Q_UNUSED( simplifyMethod );
diff --git a/src/core/qgsvectorlayerfeatureiterator.h b/src/core/qgsvectorlayerfeatureiterator.h
@@ -97,10 +97,19 @@ class CORE_EXPORT QgsVectorLayerFeatureIterator : public QgsAbstractFeatureItera
 
     //! @note not available in Python bindings
     void rewindEditBuffer();
+
+    //! @note not available in Python bindings
+    void prepareJoin( int fieldIdx );
+
     //! @note not available in Python bindings
-    void prepareJoins();
+    void prepareExpression( int fieldIdx );
+
     //! @note not available in Python bindings
-    void prepareExpressions();
+    void prepareFields();
+
+    //! @note not available in Python bindings
+    void prepareField( int fieldIdx );
+
     //! @note not available in Python bindings
     bool fetchNextAddedFeature( QgsFeature& f );
     //! @note not available in Python bindings
@@ -127,6 +136,14 @@ class CORE_EXPORT QgsVectorLayerFeatureIterator : public QgsAbstractFeatureItera
      */
     void addVirtualAttributes( QgsFeature &f );
 
+    /** Adds an expression based attribute to a feature
+     * @param f feature
+     * @param attrIndex attribute index
+     * @note added in QGIS 2.14
+     * @note not available in Python bindings
+     */
+    void addExpressionAttribute( QgsFeature& f, int attrIndex );
+
     /** Update feature with uncommited attribute updates.
      * @note not available in Python bindings
      */
@@ -178,6 +195,9 @@ class CORE_EXPORT QgsVectorLayerFeatureIterator : public QgsAbstractFeatureItera
 
     QgsInterruptionChecker* mInterruptionChecker;
 
+    QList< int > mPreparedFields;
+    QList< int > mFieldsToPrepare;
+
     /**
      * Will always return true. We assume that ordering has been done on provider level already.
      *
diff --git a/tests/src/python/test_qgsfeatureiterator.py b/tests/src/python/test_qgsfeatureiterator.py

-Original file line number
+Diff line change
 import os
 -from qgis.core import QgsVectorLayer, QgsFeatureRequest, QgsFeature, QgsField, NULL
 +from qgis.core import QgsVectorLayer, QgsFeatureRequest, QgsFeature, QgsField, NULL, QgsMapLayerRegistry, QgsVectorJoinInfo
 from qgis.testing import start_app, unittest
 from qgis.PyQt.QtCore import QVariant
         fet = next(layer.getFeatures(QgsFeatureRequest().setSubsetOfAttributes(['exp2'], layer.fields())))
         self.assertEqual(fet['Class'], NULL)
 +    def test_JoinUsingExpression(self):
 +        """ test joining a layer using a virtual field """
 +        joinLayer = QgsVectorLayer(
 +            "Point?field=x:string&field=y:integer&field=z:integer",
 +            "joinlayer", "memory")
 +        pr = joinLayer.dataProvider()
 +        f1 = QgsFeature()
 +        f1.setAttributes(["foo", 246, 321])
 +        f2 = QgsFeature()
 +        f2.setAttributes(["bar", 456, 654])
 +        self.assertTrue(pr.addFeatures([f1, f2]))
++
 +        layer = QgsVectorLayer("Point?field=fldtxt:string&field=fldint:integer",
 +                               "addfeat", "memory")
 +        pr = layer.dataProvider()
 +        f = QgsFeature()
 +        f.setAttributes(["test", 123])
 +        self.assertTrue(pr.addFeatures([f]))
 +        layer.addExpressionField('"fldint"*2', QgsField('exp1', QVariant.LongLong))
++
 +        QgsMapLayerRegistry.instance().addMapLayers([layer, joinLayer])
++
 +        join = QgsVectorJoinInfo()
 +        join.targetFieldName = "exp1"
 +        join.joinLayerId = joinLayer.id()
 +        join.joinFieldName = "y"
 +        join.memoryCache = True
 +        layer.addJoin(join)
++
 +        f = QgsFeature()
 +        fi = layer.getFeatures()
 +        self.assertTrue(fi.nextFeature(f))
 +        attrs = f.attributes()
 +        self.assertEqual(attrs[0], "test")
 +        self.assertEqual(attrs[1], 123)
 +        self.assertEqual(attrs[2], "foo")
 +        self.assertEqual(attrs[3], 321)
 +        self.assertEqual(attrs[4], 246)
 +        self.assertFalse(fi.nextFeature(f))
++
 +        QgsMapLayerRegistry.instance().removeMapLayers([layer, joinLayer])
++
 +    def test_JoinUsingExpression2(self):
 +        """ test joining a layer using a virtual field (the other way!) """
 +        joinLayer = QgsVectorLayer(
 +            "Point?field=x:string&field=y:integer&field=z:integer",
 +            "joinlayer", "memory")
 +        pr = joinLayer.dataProvider()
 +        f1 = QgsFeature()
 +        f1.setAttributes(["foo", 246, 321])
 +        f2 = QgsFeature()
 +        f2.setAttributes(["bar", 456, 654])
 +        self.assertTrue(pr.addFeatures([f1, f2]))
 +        joinLayer.addExpressionField('"y"/2', QgsField('exp1', QVariant.LongLong))
++
 +        layer = QgsVectorLayer("Point?field=fldtxt:string&field=fldint:integer",
 +                               "addfeat", "memory")
 +        pr = layer.dataProvider()
 +        f = QgsFeature()
 +        f.setAttributes(["test", 123])
 +        self.assertTrue(pr.addFeatures([f]))
++
 +        QgsMapLayerRegistry.instance().addMapLayers([layer, joinLayer])
++
 +        join = QgsVectorJoinInfo()
 +        join.targetFieldName = "fldint"
 +        join.joinLayerId = joinLayer.id()
 +        join.joinFieldName = "exp1"
 +        join.memoryCache = True
 +        layer.addJoin(join)
++
 +        f = QgsFeature()
 +        fi = layer.getFeatures()
 +        self.assertTrue(fi.nextFeature(f))
 +        attrs = f.attributes()
 +        self.assertEqual(attrs[0], "test")
 +        self.assertEqual(attrs[1], 123)
 +        self.assertEqual(attrs[2], "foo")
 +        self.assertEqual(attrs[3], 246)
 +        self.assertEqual(attrs[4], 321)
 +        self.assertFalse(fi.nextFeature(f))
++
 +        QgsMapLayerRegistry.instance().removeMapLayers([layer, joinLayer])
 +        # try the other way too
++
 if __name__ == '__main__':
     unittest.main()