Fix qual pushdown past grouping through simple CASE

Commit 44fb59fc6 taught the grouping-conflict walker to treat the arg
of a simple CASE as a direct operand of each WHEN comparison, but it
only checked the collation, on the assumption that the WHEN operator
is always the type-default "=" and thus matches the grouping eqop.
That assumption fails once the arg is relabeled to another type: the
WHEN then compares under that type's "=", which need not agree with
the grouping equality.  For instance, with a DISTINCT over a citext
column, a qual such as "CASE t::text WHEN 'A' THEN ..." was pushed
below the Unique, although the equivalent "t::text = 'A'" is correctly
kept above it.

Instead of special-casing the arg, have the walker bind a Var arg
while walking the WHEN conditions and resolve each CaseTestExpr to it,
so that the arg is checked exactly as each WHEN uses it: with the
opfamily and collation checks of a direct operand when the WHEN is a
comparison, and as a non-operand reference otherwise.  A non-Var arg
is walked once as a non-operand, as before.  The CaseTestExpr in an
ArrayCoerceExpr's elemexpr and a JsonConstructorExpr's coercion stand
for something else and are left alone.

Reported-by: Tender Wang <tndrwang@gmail.com>
Author: Richard Guo <guofenglinux@gmail.com>
Reviewed-by: Tender Wang <tndrwang@gmail.com>
Reviewed-by: Ewan Young <kdbase.hack@gmail.com>
Discussion: https://postgr.es/m/CAHewXNkvGTOgijRLjmudpg=wz0d-J8ChY2o7ksh3Be+Q_Bxwog@mail.gmail.com
Backpatch-through: 18

5 files changed, 171 insertions, 41 deletions

diff --git a/src/backend/optimizer/util/clauses.c b/src/backend/optimizer/util/clauses.c
index 2d145ae5fc9..6a6a54498b7 100644
--- a/src/backend/optimizer/util/clauses.c
+++ b/src/backend/optimizer/util/clauses.c
@@ -97,12 +97,15 @@ typedef struct
 /*
  * Walker context for expression_has_grouping_conflict.  get_eqop is a callback
  * that returns the equality operator used for grouping.  cb_context is opaque
- * to the walker and is forwarded to get_eqop unchanged.
+ * to the walker and is forwarded to get_eqop unchanged.  case_var is the Var
+ * that the CaseTestExprs of the simple CASE being walked stand for, or NULL if
+ * there is none.
  */
 typedef struct
 {
 	grouping_eqop_callback get_eqop;
 	void	   *cb_context;
+	Var		   *case_var;
 } grouping_walker_ctx;
 
 static bool contain_agg_clause_walker(Node *node, void *context);
@@ -5515,6 +5518,7 @@ expression_has_grouping_conflict(Node *expr,
 
 	ctx.get_eqop = get_eqop;
 	ctx.cb_context = context;
+	ctx.case_var = NULL;
 
 	return grouping_conflict_walker(expr, &ctx);
 }
@@ -5533,10 +5537,13 @@ expression_has_grouping_conflict(Node *expr,
  * member, and RowCompareExpr (one operator and collation per column).  A
  * simple CASE (CaseExpr with a non-NULL arg) is a comparison in disguise:
  * parse analysis builds each WHEN as "OpExpr(CaseTestExpr op val)", with the
- * CaseTestExpr standing in for the arg, so the arg is effectively an operand
- * of each WHEN's comparison.  Those WHEN operators are always the type-default
- * "=", matching the grouping eqop, so only a collation conflict is possible
- * there.
+ * CaseTestExpr standing in for the arg.  If the arg is a Var (after looking
+ * through RelabelType), it is bound in ctx->case_var while the WHEN
+ * conditions are walked and each CaseTestExpr is resolved to it, so the Var
+ * is checked exactly as each WHEN uses it.  Any other arg is walked once as
+ * a non-operand and its CaseTestExprs are ignored, as are those in an
+ * ArrayCoerceExpr's elemexpr and a JsonConstructorExpr's coercion, which
+ * stand for something else.
  */
 static bool
 grouping_conflict_walker(Node *node, grouping_walker_ctx *ctx)
@@ -5603,54 +5610,88 @@ grouping_conflict_walker(Node *node, grouping_walker_ctx *ctx)
 		}
 		return false;
 	}
+	else if (IsA(node, CaseTestExpr))
+	{
+		/*
+		 * A direct operand of a comparison is handled by
+		 * grouping_check_operand; any other use is a non-operand reference to
+		 * the Var it stands for, if any.
+		 */
+		return grouping_conflict_walker((Node *) ctx->case_var, ctx);
+	}
+	else if (IsA(node, ArrayCoerceExpr))
+	{
+		ArrayCoerceExpr *acexpr = (ArrayCoerceExpr *) node;
+		Var		   *save_case_var = ctx->case_var;
+		bool		result;
+
+		if (grouping_conflict_walker((Node *) acexpr->arg, ctx))
+			return true;
+
+		/* The CaseTestExpr in elemexpr is an array element, not case_var. */
+		ctx->case_var = NULL;
+		result = grouping_conflict_walker((Node *) acexpr->elemexpr, ctx);
+		ctx->case_var = save_case_var;
+		return result;
+	}
+	else if (IsA(node, JsonConstructorExpr))
+	{
+		JsonConstructorExpr *ctor = (JsonConstructorExpr *) node;
+		Var		   *save_case_var = ctx->case_var;
+		bool		result;
+
+		if (grouping_conflict_walker((Node *) ctor->args, ctx))
+			return true;
+		if (grouping_conflict_walker((Node *) ctor->func, ctx))
+			return true;
+
+		/* The CaseTestExpr in coercion is the JSON result, not case_var. */
+		ctx->case_var = NULL;
+		result = grouping_conflict_walker((Node *) ctor->coercion, ctx);
+		ctx->case_var = save_case_var;
+		return result;
+	}
 	else if (IsA(node, CaseExpr) && ((CaseExpr *) node)->arg != NULL)
 	{
 		CaseExpr   *cexpr = (CaseExpr *) node;
 		Node	   *arg = (Node *) cexpr->arg;
+		Var		   *save_case_var = ctx->case_var;
+		bool		result = false;
 
 		/* Look through RelabelType to find a direct Var arg. */
 		while (arg && IsA(arg, RelabelType))
 			arg = (Node *) ((RelabelType *) arg)->arg;
 
+		/*
+		 * A Var arg needs no walk of its own: each WHEN condition refers to
+		 * it through a CaseTestExpr, which is resolved to the Var and checked
+		 * as the WHEN uses it.  Any other arg is a non-operand reference in
+		 * its own right: walk it once here and ignore its CaseTestExprs.
+		 */
 		if (arg && IsA(arg, Var))
+			ctx->case_var = (Var *) arg;
+		else
 		{
-			Var		   *var = (Var *) arg;
-
-			/*
-			 * The arg is a grouping column compared by every WHEN.  For a
-			 * nondeterministic collation, reject if any WHEN applies a
-			 * different collation.
-			 */
-			if (OidIsValid(ctx->get_eqop(var, ctx->cb_context)) &&
-				OidIsValid(var->varcollid) &&
-				!get_collation_isdeterministic(var->varcollid))
+			if (grouping_conflict_walker(arg, ctx))
+				return true;
+			ctx->case_var = NULL;
+		}
+		foreach_node(CaseWhen, cw, cexpr->args)
+		{
+			if (grouping_conflict_walker((Node *) cw->expr, ctx))
 			{
-				foreach_node(CaseWhen, cw, cexpr->args)
-				{
-					Oid			collid = exprInputCollation((Node *) cw->expr);
-
-					if (OidIsValid(collid) && collid != var->varcollid)
-						return true;
-				}
+				result = true;
+				break;
 			}
 		}
-		else if (grouping_conflict_walker((Node *) cexpr->arg, ctx))
-		{
-			/* arg is a complex expression; walked as a non-operand */
+		ctx->case_var = save_case_var;
+		if (result)
 			return true;
-		}
 
-		/*
-		 * Walk the WHEN conditions, their results, and the default result as
-		 * non-operands.  The WHEN conditions hold a CaseTestExpr in place of
-		 * the arg, so they contribute no grouping operand of their own, but
-		 * the condition expression or the substitution result may reference
-		 * another grouping column.
-		 */
+		/* The results and the default result contain no CaseTestExpr. */
 		foreach_node(CaseWhen, cw, cexpr->args)
 		{
-			if (grouping_conflict_walker((Node *) cw->expr, ctx) ||
-				grouping_conflict_walker((Node *) cw->result, ctx))
+			if (grouping_conflict_walker((Node *) cw->result, ctx))
 				return true;
 		}
 		return grouping_conflict_walker((Node *) cexpr->defresult, ctx);
@@ -5683,12 +5724,13 @@ grouping_check_operands(Oid opno, Oid inputcollid, List *args,
  *		Handle one operand 'arg' of a comparison with operator 'opno' and
  *		collation 'inputcollid'.
  *
- * If 'arg' is a grouping column (after looking through RelabelType), verify
- * that comparison's operator has equality semantics compatible with the
- * grouping eqop and, for a nondeterministic collation, that it uses the same
- * collation; such a direct operand is then fully handled and is not recursed
- * into.  Any other operand is walked normally, so a grouping column buried
- * inside it is seen as a non-operand reference.
+ * If 'arg' is a grouping column (after looking through RelabelType, or through
+ * a CaseTestExpr to the Var it stands for), verify that comparison's operator
+ * has equality semantics compatible with the grouping eqop and, for a
+ * nondeterministic collation, that it uses the same collation; such a direct
+ * operand is then fully handled and is not recursed into.  Any other operand
+ * is walked normally, so a grouping column buried inside it is seen as a
+ * non-operand reference.
  */
 static bool
 grouping_check_operand(Node *arg, Oid opno, Oid inputcollid,
@@ -5699,6 +5741,9 @@ grouping_check_operand(Node *arg, Oid opno, Oid inputcollid,
 	while (node && IsA(node, RelabelType))
 		node = (Node *) ((RelabelType *) node)->arg;
 
+	if (node && IsA(node, CaseTestExpr))
+		node = (Node *) ctx->case_var;
+
 	if (node && IsA(node, Var))
 	{
 		Var		   *var = (Var *) node;
diff --git a/src/test/regress/expected/collate.icu.utf8.out b/src/test/regress/expected/collate.icu.utf8.out
index e700a5e1122..0fd5c1a2cb0 100644
--- a/src/test/regress/expected/collate.icu.utf8.out
+++ b/src/test/regress/expected/collate.icu.utf8.out
@@ -2344,6 +2344,41 @@ SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN 'abc' COLLATE cas
  abc |     2
 (1 row)
 
+-- Positive: the WHEN value's own CaseTestExpr (JSON RETURNING coercion) does
+-- not refer to the CASE arg
+EXPLAIN (COSTS OFF)
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN JSON_OBJECT('a': 'b' RETURNING text) THEN true ELSE false END);
+                                            QUERY PLAN                                            
+--------------------------------------------------------------------------------------------------
+ HashAggregate
+   Group Key: x
+   ->  Seq Scan on test3ci
+         Filter: CASE x WHEN JSON_OBJECT('a' : 'b'::text RETURNING text) THEN true ELSE false END
+(4 rows)
+
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN JSON_OBJECT('a': 'b' RETURNING text) THEN true ELSE false END);
+ x | count 
+---+-------
+(0 rows)
+
+-- Positive: likewise for the CaseTestExpr in an ArrayCoerceExpr's elemexpr
+CREATE DOMAIN nonempty_text AS text CHECK (VALUE <> '');
+EXPLAIN (COSTS OFF)
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN ('{abc}'::text[]::nonempty_text[])[1] THEN true ELSE false END);
+                                          QUERY PLAN                                          
+----------------------------------------------------------------------------------------------
+ HashAggregate
+   Group Key: x
+   ->  Seq Scan on test3ci
+         Filter: CASE x WHEN (('{abc}'::text[])::nonempty_text[])[1] THEN true ELSE false END
+(4 rows)
+
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN ('{abc}'::text[]::nonempty_text[])[1] THEN true ELSE false END);
+  x  | count 
+-----+-------
+ abc |     2
+(1 row)
+
 -- Negative: nested CASE with collation conflict
 EXPLAIN (COSTS OFF)
 SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE WHEN (CASE x WHEN 'abc' COLLATE case_sensitive THEN 1 ELSE 0 END) = 1 THEN true ELSE false END);
diff --git a/src/test/regress/expected/subselect.out b/src/test/regress/expected/subselect.out
index 501ffadc105..a4eb5a72912 100644
--- a/src/test/regress/expected/subselect.out
+++ b/src/test/regress/expected/subselect.out
@@ -1930,6 +1930,33 @@ WHERE a = ROW(1.0)::t_rec;
   1 | (1.00)
 (1 row)
 
+-- Simple CASE: the arg is compared by each WHEN, so the same rules apply.
+-- The relabeled arg is compared by oid's "=", not the grouping eqop.
+EXPLAIN (COSTS OFF)
+SELECT * FROM (SELECT DISTINCT id FROM pdt) s
+WHERE (CASE id::oid WHEN 1 THEN 1 ELSE 0 END) = 1;
+                            QUERY PLAN                            
+------------------------------------------------------------------
+ Subquery Scan on s
+   Filter: (CASE (s.id)::oid WHEN '1'::oid THEN 1 ELSE 0 END = 1)
+   ->  HashAggregate
+         Group Key: pdt.id
+         ->  Seq Scan on pdt
+(5 rows)
+
+-- Positive: compatible opfamily, safe to push past the grouping
+EXPLAIN (COSTS OFF)
+SELECT * FROM (SELECT DISTINCT id FROM pdt) s
+WHERE (CASE id WHEN 1 THEN 1 ELSE 0 END) = 1;
+                          QUERY PLAN                          
+--------------------------------------------------------------
+ Unique
+   ->  Sort
+         Sort Key: pdt.id
+         ->  Seq Scan on pdt
+               Filter: (CASE id WHEN 1 THEN 1 ELSE 0 END = 1)
+(5 rows)
+
 -- Set operations: any operation other than UNION ALL groups rows by equality,
 -- so the same opfamily-mismatch rules apply.
 CREATE TEMP TABLE u1 (a t_rec);
diff --git a/src/test/regress/sql/collate.icu.utf8.sql b/src/test/regress/sql/collate.icu.utf8.sql
index a17e8ae27f9..b091f05c113 100644
--- a/src/test/regress/sql/collate.icu.utf8.sql
+++ b/src/test/regress/sql/collate.icu.utf8.sql
@@ -827,6 +827,18 @@ EXPLAIN (COSTS OFF)
 SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN 'abc' COLLATE case_insensitive THEN true ELSE false END);
 SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN 'abc' COLLATE case_insensitive THEN true ELSE false END);
 
+-- Positive: the WHEN value's own CaseTestExpr (JSON RETURNING coercion) does
+-- not refer to the CASE arg
+EXPLAIN (COSTS OFF)
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN JSON_OBJECT('a': 'b' RETURNING text) THEN true ELSE false END);
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN JSON_OBJECT('a': 'b' RETURNING text) THEN true ELSE false END);
+
+-- Positive: likewise for the CaseTestExpr in an ArrayCoerceExpr's elemexpr
+CREATE DOMAIN nonempty_text AS text CHECK (VALUE <> '');
+EXPLAIN (COSTS OFF)
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN ('{abc}'::text[]::nonempty_text[])[1] THEN true ELSE false END);
+SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE x WHEN ('{abc}'::text[]::nonempty_text[])[1] THEN true ELSE false END);
+
 -- Negative: nested CASE with collation conflict
 EXPLAIN (COSTS OFF)
 SELECT x, count(*) FROM test3ci GROUP BY x HAVING (CASE WHEN (CASE x WHEN 'abc' COLLATE case_sensitive THEN 1 ELSE 0 END) = 1 THEN true ELSE false END);
diff --git a/src/test/regress/sql/subselect.sql b/src/test/regress/sql/subselect.sql
index 10079f32ccd..f0f25fd569f 100644
--- a/src/test/regress/sql/subselect.sql
+++ b/src/test/regress/sql/subselect.sql
@@ -972,6 +972,17 @@ WHERE a = ROW(1.0)::t_rec;
 SELECT * FROM (SELECT DISTINCT ON (a) id, a FROM pdt ORDER BY a, id) s
 WHERE a = ROW(1.0)::t_rec;
 
+-- Simple CASE: the arg is compared by each WHEN, so the same rules apply.
+-- The relabeled arg is compared by oid's "=", not the grouping eqop.
+EXPLAIN (COSTS OFF)
+SELECT * FROM (SELECT DISTINCT id FROM pdt) s
+WHERE (CASE id::oid WHEN 1 THEN 1 ELSE 0 END) = 1;
+
+-- Positive: compatible opfamily, safe to push past the grouping
+EXPLAIN (COSTS OFF)
+SELECT * FROM (SELECT DISTINCT id FROM pdt) s
+WHERE (CASE id WHEN 1 THEN 1 ELSE 0 END) = 1;
+
 -- Set operations: any operation other than UNION ALL groups rows by equality,
 -- so the same opfamily-mismatch rules apply.
 CREATE TEMP TABLE u1 (a t_rec);