Add subtransactions for JsonExpr execution

Author: Nikita Glukhov

src/backend/executor/execExprInterp.c

@@ -57,6 +57,8 @@
 #include "postgres.h"
 
 #include "access/tuptoaster.h"
+#include "access/xact.h"
+#include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "commands/sequence.h"
 #include "executor/execExpr.h"
@@ -77,6 +79,7 @@
 #include "utils/jsonb.h"
 #include "utils/jsonpath.h"
 #include "utils/lsyscache.h"
+#include "utils/resowner.h"
 #include "utils/timestamp.h"
 #include "utils/typcache.h"
 #include "utils/xml.h"
@@ -4448,6 +4451,12 @@ ExecEvalJsonExpr(ExprState *state, ExprEvalStep *op, ExprContext *econtext,
 	return res;
 }
 
+bool
+ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr)
+{
+	return jsexpr->on_error.btype != JSON_BEHAVIOR_ERROR;
+}
+
 /* ----------------------------------------------------------------
  *		ExecEvalJson
  * ----------------------------------------------------------------
@@ -4487,20 +4496,43 @@ ExecEvalJson(ExprState *state, ExprEvalStep *op, ExprContext *econtext)
 		var->evaluated = false;
 	}
 
-	if (jexpr->on_error.btype == JSON_BEHAVIOR_ERROR)
+	if (!ExecEvalJsonNeedsSubTransaction(jexpr))
 	{
 		/* No need to use PG_TRY/PG_CATCH with subtransactions. */
 		res = ExecEvalJsonExpr(state, op, econtext, jexpr, path, item,
 							   op->resnull);
 	}
 	else
 	{
+		/*
+		 * We should catch exceptions of category ERRCODE_DATA_EXCEPTION and
+		 * execute corresponding ON ERROR behavior.
+		 */
 		MemoryContext oldcontext = CurrentMemoryContext;
+		ResourceOwner oldowner = CurrentResourceOwner;
+		ExprContext *newecontext;
+
+		BeginInternalSubTransaction(NULL);
+		/* Want to execute expressions inside function's memory context */
+		MemoryContextSwitchTo(oldcontext);
+		/*
+		 * We need to execute expressions with a new econtext
+		 * that belongs to the current subtransaction; if we try to use
+		 * the outer econtext then ExprContext shutdown callbacks will be
+		 * called at the wrong times.
+		 */
+		newecontext = CreateExprContext(econtext->ecxt_estate);
 
 		PG_TRY();
 		{
-			res = ExecEvalJsonExpr(state, op, econtext, jexpr, path, item,
+			res = ExecEvalJsonExpr(state, op, newecontext, jexpr, path, item,
 								   op->resnull);
+
+			/* Commit the inner transaction, return to outer xact context */
+			ReleaseCurrentSubTransaction();
+			MemoryContextSwitchTo(oldcontext);
+			CurrentResourceOwner = oldowner;
+			FreeExprContext(newecontext, true);
 		}
 		PG_CATCH();
 		{
@@ -4511,6 +4543,12 @@ ExecEvalJson(ExprState *state, ExprEvalStep *op, ExprContext *econtext)
 			edata = CopyErrorData();
 			FlushErrorState();
 
+			/* Abort the inner transaction */
+			RollbackAndReleaseCurrentSubTransaction();
+			MemoryContextSwitchTo(oldcontext);
+			CurrentResourceOwner = oldowner;
+			FreeExprContext(newecontext, false);
+
 			if (ERRCODE_TO_CATEGORY(edata->sqlerrcode) != ERRCODE_DATA_EXCEPTION)
 				ReThrowError(edata);
 

src/backend/optimizer/util/clauses.c

@@ -28,6 +28,7 @@
 #include "catalog/pg_type.h"
 #include "executor/executor.h"
 #include "executor/functions.h"
+#include "executor/execExpr.h"
 #include "funcapi.h"
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
@@ -1288,6 +1289,16 @@ max_parallel_hazard_walker(Node *node, max_parallel_hazard_context *context)
 								 context, 0);
 	}
 
+	/* JsonExpr is parallel-unsafe if subtransactions can be used. */
+	else if (IsA(node, JsonExpr))
+	{
+		JsonExpr  *jsexpr = (JsonExpr *) node;
+
+		if (ExecEvalJsonNeedsSubTransaction(jsexpr))
+			context->max_hazard = PROPARALLEL_UNSAFE;
+			return true;
+	}
+
 	/* Recurse to check arguments */
 	return expression_tree_walker(node,
 								  max_parallel_hazard_walker,

src/include/executor/execExpr.h

@@ -790,6 +790,7 @@ extern Datum ExecPrepareJsonItemCoercion(struct JsonbValue *item,
 							JsonReturning *returning,
 							struct JsonCoercionsState *coercions,
 							struct JsonCoercionState **pjcstate);
+extern bool ExecEvalJsonNeedsSubTransaction(JsonExpr *jsexpr);
 
 extern void ExecAggInitGroup(AggState *aggstate, AggStatePerTrans pertrans, AggStatePerGroup pergroup);
 extern Datum ExecAggTransReparent(AggState *aggstate, AggStatePerTrans pertrans,

src/test/regress/expected/jsonb_sqljson.out

@@ -903,3 +903,40 @@ SELECT JSON_QUERY(jsonb '{"a": 123}', '$' || '.' || 'a' WITH WRAPPER);
 SELECT JSON_QUERY(jsonb '{"a": 123}', 'error' || ' ' || 'error');
 ERROR:  bad jsonpath representation
 DETAIL:  syntax error, unexpected IDENT_P at or near " "
+-- Test parallel JSON_VALUE()
+CREATE TABLE test_parallel_jsonb_value AS
+SELECT i::text::jsonb AS js
+FROM generate_series(1, 1000000) i;
+-- Should be non-parallel due to subtransactions
+EXPLAIN (COSTS OFF)
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric)) FROM test_parallel_jsonb_value;
+                 QUERY PLAN                  
+---------------------------------------------
+ Aggregate
+   ->  Seq Scan on test_parallel_jsonb_value
+(2 rows)
+
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric)) FROM test_parallel_jsonb_value;
+     sum      
+--------------
+ 500000500000
+(1 row)
+
+-- Should be parallel
+EXPLAIN (COSTS OFF)
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric ERROR ON ERROR)) FROM test_parallel_jsonb_value;
+                            QUERY PLAN                            
+------------------------------------------------------------------
+ Finalize Aggregate
+   ->  Gather
+         Workers Planned: 2
+         ->  Partial Aggregate
+               ->  Parallel Seq Scan on test_parallel_jsonb_value
+(5 rows)
+
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric ERROR ON ERROR)) FROM test_parallel_jsonb_value;
+     sum      
+--------------
+ 500000500000
+(1 row)
+

src/test/regress/sql/jsonb_sqljson.sql

@@ -273,3 +273,19 @@ SELECT JSON_QUERY(jsonb '{"a": 123}', '$' || '.' || 'a');
 SELECT JSON_QUERY(jsonb '{"a": 123}', '$' || '.' || 'a' WITH WRAPPER);
 -- Should fail (invalid path)
 SELECT JSON_QUERY(jsonb '{"a": 123}', 'error' || ' ' || 'error');
+
+-- Test parallel JSON_VALUE()
+CREATE TABLE test_parallel_jsonb_value AS
+SELECT i::text::jsonb AS js
+FROM generate_series(1, 1000000) i;
+
+-- Should be non-parallel due to subtransactions
+EXPLAIN (COSTS OFF)
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric)) FROM test_parallel_jsonb_value;
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric)) FROM test_parallel_jsonb_value;
+
+-- Should be parallel
+EXPLAIN (COSTS OFF)
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric ERROR ON ERROR)) FROM test_parallel_jsonb_value;
+SELECT sum(JSON_VALUE(js, '$' RETURNING numeric ERROR ON ERROR)) FROM test_parallel_jsonb_value;
+