diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 96a64ce6b2c..ba0dcd3f15e 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -1559,7 +1559,6 @@ get_qual_for_range(PartitionKey key, PartitionBoundSpec *spec)
 	forboth(cell1, spec->lowerdatums, cell2, spec->upperdatums)
 	{
 		EState	   *estate;
-		MemoryContext oldcxt;
 		Expr	   *test_expr;
 		ExprState  *test_exprstate;
 		Datum		test_result;
@@ -1591,16 +1590,18 @@ get_qual_for_range(PartitionKey key, PartitionBoundSpec *spec)
 
 		/* Create the test expression */
 		estate = CreateExecutorState();
-		oldcxt = MemoryContextSwitchTo(estate->es_query_cxt);
-		test_expr = make_partition_op_expr(key, i, BTEqualStrategyNumber,
-										   (Expr *) lower_val,
-										   (Expr *) upper_val);
-		fix_opfuncids((Node *) test_expr);
-		test_exprstate = ExecInitExpr(test_expr, NULL);
-		test_result = ExecEvalExprSwitchContext(test_exprstate,
-												GetPerTupleExprContext(estate),
-												&isNull);
-		MemoryContextSwitchTo(oldcxt);
+		WITH_MEMORY_CONTEXT(estate->es_query_cxt)
+		{
+			test_expr = make_partition_op_expr(key, i, BTEqualStrategyNumber,
+											   (Expr *) lower_val,
+											   (Expr *) upper_val);
+			fix_opfuncids((Node *) test_expr);
+			test_exprstate = ExecInitExpr(test_expr, NULL);
+			test_result = ExecEvalExprSwitchContext(test_exprstate,
+													GetPerTupleExprContext(estate),
+													&isNull);
+		}
+		END_WITH_MEMORY_CONTEXT;
 		FreeExecutorState(estate);
 
 		/* If not equal, go generate the OR expressions */
diff --git a/src/include/utils/memutils.h b/src/include/utils/memutils.h
index c5533490668..02ec34e9f89 100644
--- a/src/include/utils/memutils.h
+++ b/src/include/utils/memutils.h
@@ -194,4 +194,27 @@ extern MemoryContext SlabContextCreate(MemoryContext parent,
 #define SLAB_DEFAULT_BLOCK_SIZE		(8 * 1024)
 #define SLAB_LARGE_BLOCK_SIZE		(8 * 1024 * 1024)
 
+/*
+ * Switch to a different memory context and for a limited scope.  The calling
+ * code should look like this:
+ *
+ * WITH_MEMORY_CONTEXT(x)
+ * {
+ *     ...;
+ * }
+ * END_WITH_MEMORY_CONTEXT;
+ *
+ * If the code in the enclosed scope needs to return early it should use
+ * RESTORE_MEMORY_CONTEXT_AND_RETURN(x) rather than return x.
+ */
+#define WITH_MEMORY_CONTEXT(context) { \
+	MemoryContext old_context = MemoryContextSwitchTo(context);
+#define END_WITH_MEMORY_CONTEXT MemoryContextSwitchTo(old_context); }
+#define RESTORE_MEMORY_CONTEXT_AND_RETURN(value) \
+	do \
+	{ \
+		MemoryContextSwitchTo(old_context); \
+		return (value); \
+	};
+
 #endif							/* MEMUTILS_H */