Fix a bug which caused queries with SRFs and function evalution to fail

src/backend/distributed/utils/citus_clauses.c

@@ -33,7 +33,7 @@ static Expr * citus_evaluate_expr(Expr *expr, Oid result_type, int32 result_typm
 								  MasterEvaluationContext *masterEvaluationContext);
 static bool CitusIsVolatileFunctionIdChecker(Oid func_id, void *context);
 static bool CitusIsMutableFunctionIdChecker(Oid func_id, void *context);
-static bool ShouldEvaluateExpressionType(NodeTag nodeTag);
+static bool ShouldEvaluateExpression(Expr *expression);
 static bool ShouldEvaluateFunctionWithMasterContext(MasterEvaluationContext *
 													evaluationContext);
 
@@ -117,7 +117,7 @@ PartiallyEvaluateExpression(Node *expression,
 											exprCollation(expression),
 											masterEvaluationContext);
 	}
-	else if (ShouldEvaluateExpressionType(nodeTag) &&
+	else if (ShouldEvaluateExpression((Expr *) expression) &&
 			 ShouldEvaluateFunctionWithMasterContext(masterEvaluationContext))
 	{
 		/* don't call citus_evaluate_expr on nodes it cannot handle */
@@ -171,15 +171,25 @@ ShouldEvaluateFunctionWithMasterContext(MasterEvaluationContext *evaluationConte
 
 
 /*
- * ShouldEvaluateExpressionType returns true if Citus should evaluate the
+ * ShouldEvaluateExpression returns true if Citus should evaluate the
  * input node on the coordinator.
  */
 static bool
-ShouldEvaluateExpressionType(NodeTag nodeTag)
+ShouldEvaluateExpression(Expr *expression)
 {
+	NodeTag nodeTag = nodeTag(expression);
+
 	switch (nodeTag)
 	{
 		case T_FuncExpr:
+		{
+			FuncExpr *funcExpr = (FuncExpr *) expression;
+
+			/* we cannot evaluate set returning functions */
+			bool isSetReturningFunction = funcExpr->funcretset;
+			return !isSetReturningFunction;
+		}
+
 		case T_OpExpr:
 		case T_DistinctExpr:
 		case T_NullIfExpr:
@@ -253,7 +263,7 @@ citus_evaluate_expr(Expr *expr, Oid result_type, int32 result_typmod,
 		else if (masterEvaluationContext->evaluationMode != EVALUATE_FUNCTIONS_PARAMS)
 		{
 			/* should only get here for node types we should evaluate */
-			Assert(ShouldEvaluateExpressionType(nodeTag(expr)));
+			Assert(ShouldEvaluateExpression(expr));
 
 			/* bail out, the caller doesn't want functions/expressions to be evaluated */
 			return expr;

src/test/regress/expected/multi_function_evaluation.out

@@ -1,6 +1,8 @@
 --
 -- MULTI_FUNCTION_EVALUATION
 --
+CREATE SCHEMA multi_function_evaluation;
+SET search_path TO multi_function_evaluation;
 SET citus.next_shard_id TO 1200000;
 -- many of the tests in this file is intended for testing non-fast-path
 -- router planner, so we're explicitly disabling it in this file.
@@ -138,5 +140,32 @@ SELECT * FROM example WHERE key = 44;
   44 | Tue Oct 10 00:00:00 2000 PDT
 (1 row)
 
-DROP FUNCTION stable_fn();
-DROP TABLE example;
+-- unnest is a set-returning function, which should not be evaluated
+UPDATE example SET value = stable_fn() + interval '2 hours' FROM UNNEST(ARRAY[44, 4]) AS k (key) WHERE example.key = k.key;
+NOTICE:  stable_fn called
+CONTEXT:  PL/pgSQL function stable_fn() line 3 at RAISE
+SELECT * FROM example WHERE key = 44;
+ key |            value
+---------------------------------------------------------------------
+  44 | Tue Oct 10 02:00:00 2000 PDT
+(1 row)
+
+-- create a table with multiple shards to trigger recursive planning
+CREATE TABLE table_1 (key int, value timestamptz);
+SELECT create_distributed_table('table_1', 'key');
+ create_distributed_table
+---------------------------------------------------------------------
+
+(1 row)
+
+-- the following query will have a read_intermediate_result call, but it should be skipped
+DELETE
+FROM table_1
+WHERE key >= (SELECT min(KEY) FROM table_1)
+AND value > now() - interval '1 hour';
+DROP SCHEMA multi_function_evaluation CASCADE;
+NOTICE:  drop cascades to 4 other objects
+DETAIL:  drop cascades to table example
+drop cascades to sequence example_value_seq
+drop cascades to function stable_fn()
+drop cascades to table table_1

src/test/regress/sql/multi_function_evaluation.sql

@@ -1,6 +1,8 @@
 --
 -- MULTI_FUNCTION_EVALUATION
 --
+CREATE SCHEMA multi_function_evaluation;
+SET search_path TO multi_function_evaluation;
 
 SET citus.next_shard_id TO 1200000;
 
@@ -129,6 +131,18 @@ $function$;
 INSERT INTO example VALUES (44, (ARRAY[stable_fn(),stable_fn()])[1]);
 SELECT * FROM example WHERE key = 44;
 
-DROP FUNCTION stable_fn();
+-- unnest is a set-returning function, which should not be evaluated
+UPDATE example SET value = stable_fn() + interval '2 hours' FROM UNNEST(ARRAY[44, 4]) AS k (key) WHERE example.key = k.key;
+SELECT * FROM example WHERE key = 44;
 
-DROP TABLE example;
+-- create a table with multiple shards to trigger recursive planning
+CREATE TABLE table_1 (key int, value timestamptz);
+SELECT create_distributed_table('table_1', 'key');
+
+-- the following query will have a read_intermediate_result call, but it should be skipped
+DELETE
+FROM table_1
+WHERE key >= (SELECT min(KEY) FROM table_1)
+AND value > now() - interval '1 hour';
+
+DROP SCHEMA multi_function_evaluation CASCADE;