Added an option to limit the number of results

qgis · Dec 21, 2020 · e10e42d · e10e42d
1 parent 7380f8d
commit e10e42d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 3 deletions.
diff --git a/resources/function_help/json/array_filter b/resources/function_help/json/array_filter
@@ -5,9 +5,11 @@
   "description": "Returns an array with only the items for which the expression evaluates to true.",
   "arguments": [
     {"arg":"array","description":"an array"},
-    {"arg":"expression","description":"an expression to evaluate on each item. The variable `@element` will be replaced by the current value."}
+    {"arg":"expression","description":"an expression to evaluate on each item. The variable `@element` will be replaced by the current value."},
+    {"arg":"limit","optional":true,"default":"0","description":"maximum number of elements to be returned. Use 0 to return all values."}
   ],
   "examples": [
-    { "expression": "array_filter(array(1,2,3),@element &lt; 3)", "returns":"[ 1, 2 ]"}
+    { "expression": "array_filter(array(1,2,3),@element &lt; 3)", "returns":"[ 1, 2 ]"},
+    { "expression": "array_filter(array(1,2,3),@element &lt; 3, 1)", "returns":"[ 1 ]"}
   ]
 }
diff --git a/src/core/expression/qgsexpressionfunction.cpp b/src/core/expression/qgsexpressionfunction.cpp
@@ -7133,7 +7133,8 @@ bool QgsArrayForeachExpressionFunction::prepare( const QgsExpressionNodeFunction
 QgsArrayFilterExpressionFunction::QgsArrayFilterExpressionFunction()
   : QgsExpressionFunction( QStringLiteral( "array_filter" ), QgsExpressionFunction::ParameterList()
                            << QgsExpressionFunction::Parameter( QStringLiteral( "array" ) )
-                           << QgsExpressionFunction::Parameter( QStringLiteral( "expression" ) ),
+                           << QgsExpressionFunction::Parameter( QStringLiteral( "expression" ) )
+                           << QgsExpressionFunction::Parameter( QStringLiteral( "limit" ), true, 0 ),
                            QStringLiteral( "Arrays" ) )
 {
 
@@ -7177,11 +7178,31 @@ QVariant QgsArrayFilterExpressionFunction::run( QgsExpressionNode::NodeList *arg
   QgsExpressionContextScope *subScope = new QgsExpressionContextScope();
   subContext->appendScope( subScope );
 
+  int limit = 0;
+  if ( args->count() >= 3 )
+  {
+    const QVariant limitVar = args->at( 2 )->eval( parent, context );
+
+    if ( QgsExpressionUtils::isIntSafe( limitVar ) )
+    {
+      limit = limitVar.toInt();
+    }
+    else
+    {
+      return result;
+    }
+  }
+
   for ( const QVariant &value : array )
   {
     subScope->addVariable( QgsExpressionContextScope::StaticVariable( QStringLiteral( "element" ), value, true ) );
     if ( args->at( 1 )->eval( parent, subContext ).toBool() )
+    {
       result << value;
+
+      if ( limit > 0 && limit == result.size() )
+        break;
+    }
   }
 
   if ( context )

diff --git a/tests/src/core/testqgsexpression.cpp b/tests/src/core/testqgsexpression.cpp
@@ -3366,6 +3366,8 @@ class TestQgsExpression: public QObject
 
       QVariantList filterExpected = QVariantList() << QStringLiteral( "A: a" ) << QStringLiteral( "A: d" );
       QCOMPARE( QgsExpression( "array_filter(array:=array('A: a', 'B: b', 'C: c', 'A: d'), expression:=substr(@element, 1, 2) = 'A:')" ).evaluate( &context ), QVariant( filterExpected ) );
+      QVariantList filterExpectedLimit = QVariantList() << QStringLiteral( "A: a" );
+      QCOMPARE( QgsExpression( "array_filter(array:=array('A: a', 'B: b', 'C: c', 'A: d'), expression:=substr(@element, 1, 2) = 'A:', limit:=1)" ).evaluate( &context ), QVariant( filterExpectedLimit ) );
 
       QCOMPARE( QgsExpression( "array_intersect(array('1', '2', '3', '4'), array('4', '0', '2', '5'))" ).evaluate( &context ), QVariant( true ) );
       QCOMPARE( QgsExpression( "array_intersect(array('1', '2', '3', '4'), array('0', '5'))" ).evaluate( &context ), QVariant( false ) );
@@ -3460,6 +3462,8 @@ class TestQgsExpression: public QObject
 
       QVariantList filterExpected = QVariantList() << 1 << 2;
       QCOMPARE( QgsExpression( "array_filter(array(1, 2, 4), @element < 3)" ).evaluate( &context ), QVariant( filterExpected ) );
+      QVariantList filterExpectedLimit = QVariantList() << 1;
+      QCOMPARE( QgsExpression( "array_filter(array(1, 2, 4), @element < 3, 1)" ).evaluate( &context ), QVariant( filterExpectedLimit ) );
 
       QgsExpression badArray( QStringLiteral( "array_get('not an array', 0)" ) );
       QCOMPARE( badArray.evaluate( &context ), QVariant() );