RECORD: Add support for more expression types

- OpExpr
- NullIfExpr
- MinMaxExpr
- CoalesceExpr
- CaseExpr

Also fix case where ARRAY[(1,2), NULL] was rejected

src/backend/distributed/planner/distributed_planner.c

@@ -15,6 +15,7 @@
 
 #include "access/htup_details.h"
 #include "catalog/pg_class.h"
+#include "catalog/pg_proc.h"
 #include "catalog/pg_type.h"
 #include "distributed/citus_nodefuncs.h"
 #include "distributed/citus_nodes.h"
@@ -90,6 +91,7 @@ static PlannedStmt * FinalizeNonRouterPlan(PlannedStmt *localPlan,
 										   DistributedPlan *distributedPlan,
 										   CustomScan *customScan);
 static PlannedStmt * FinalizeRouterPlan(PlannedStmt *localPlan, CustomScan *customScan);
+static int32 BlessRecordExpressionList(List *exprs);
 static void CheckNodeIsDumpable(Node *node);
 static Node * CheckNodeCopyAndSerialization(Node *node);
 static void AdjustReadIntermediateResultCost(RangeTblEntry *rangeTableEntry,
@@ -1224,7 +1226,7 @@ BlessRecordExpression(Expr *expr)
 {
 	int32 typeMod = -1;
 
-	if (IsA(expr, FuncExpr))
+	if (IsA(expr, FuncExpr) || IsA(expr, OpExpr))
 	{
 		/*
 		 * Handle functions that return records on the target
@@ -1236,6 +1238,7 @@ BlessRecordExpression(Expr *expr)
 		/* get_expr_result_type blesses the tuple descriptor */
 		TypeFuncClass typeClass = get_expr_result_type((Node *) expr, &resultTypeId,
 													   &resultTupleDesc);
+
 		if (typeClass == TYPEFUNC_COMPOSITE)
 		{
 			typeMod = resultTupleDesc->tdtypmod;
@@ -1285,34 +1288,87 @@ BlessRecordExpression(Expr *expr)
 	{
 		/*
 		 * Handle row array expressions, e.g. SELECT ARRAY[(1,2)];
+		 * Postgres allows ARRAY[(1,2),(1,2,3)]. We do not.
 		 */
 		ArrayExpr *arrayExpr = (ArrayExpr *) expr;
-		ListCell *elemCell = NULL;
 
-		foreach(elemCell, arrayExpr->elements)
+		typeMod = BlessRecordExpressionList(arrayExpr->elements);
+	}
+	else if (IsA(expr, NullIfExpr))
+	{
+		NullIfExpr *nullIfExpr = (NullIfExpr *) expr;
+
+		typeMod = BlessRecordExpressionList(nullIfExpr->args);
+	}
+	else if (IsA(expr, MinMaxExpr))
+	{
+		MinMaxExpr *minMaxExpr = (MinMaxExpr *) expr;
+
+		typeMod = BlessRecordExpressionList(minMaxExpr->args);
+	}
+	else if (IsA(expr, CoalesceExpr))
+	{
+		CoalesceExpr *coalesceExpr = (CoalesceExpr *) expr;
+
+		typeMod = BlessRecordExpressionList(coalesceExpr->args);
+	}
+	else if (IsA(expr, CaseExpr))
+	{
+		CaseExpr *caseExpr = (CaseExpr *) expr;
+		List *results = NIL;
+		ListCell *whenCell = NULL;
+
+		foreach(whenCell, caseExpr->args)
 		{
-			Node *elemArg = (Node *) lfirst(elemCell);
+			CaseWhen *whenArg = (CaseWhen *) lfirst(whenCell);
-			int32 arrayTypeMod = BlessRecordExpression((Expr *) elemArg);
 
-			/*
+			results = lappend(results, whenArg->result);
-			 * Postgres allows ARRAY[(1,2),(1,2,3)]. We do not.
-			 * If multiple RECORD shapes exist in an ArrayExpr, bail out.
-			 */
-			if (typeMod == -1)
-			{
-				typeMod = arrayTypeMod;
-			}
-			else if (arrayTypeMod != typeMod)
-			{
-				return -1;
-			}
 		}
+
+		if (caseExpr->defresult != NULL)
+		{
+			results = lappend(results, caseExpr->defresult);
+		}
+
+		typeMod = BlessRecordExpressionList(results);
 	}
 
 	return typeMod;
 }
 
 
+/*
+ * BlessRecordExpressionList maps BlessRecordExpression over a list.
+ * Returns typmod of all expressions, or -1 if they are not all the same.
+ * Ignores expressions with a typmod of -1.
+ */
+static int32
+BlessRecordExpressionList(List *exprs)
+{
+	int32 finalTypeMod = -1;
+	ListCell *exprCell = NULL;
+	foreach(exprCell, exprs)
+	{
+		Node *exprArg = (Node *) lfirst(exprCell);
+		int32 exprTypeMod = BlessRecordExpression((Expr *) exprArg);
+
+		if (exprTypeMod == -1)
+		{
+			continue;
+		}
+		else if (finalTypeMod == -1)
+		{
+			finalTypeMod = exprTypeMod;
+		}
+		else if (finalTypeMod != exprTypeMod)
+		{
+			return -1;
+		}
+	}
+	return finalTypeMod;
+}
+
+
 /*
  * RemoteScanRangeTableEntry creates a range table entry from given column name
  * list to represent a remote scan.

src/test/regress/expected/row_types.out

@@ -34,6 +34,19 @@ SELECT create_distributed_function('record_returner(int)');
  
 (1 row)
 
+CREATE OR REPLACE FUNCTION identity_returner(x anyelement)
+RETURNS anyelement
+AS $$
+BEGIN
+	RETURN x;
+END;
+$$ language plpgsql;
+SELECT create_distributed_function('identity_returner(anyelement)');
+ create_distributed_function 
+-----------------------------
+ 
+(1 row)
+
 INSERT INTO test VALUES (1,2), (1,3), (2,2), (2,3);
 -- multi-shard queries support row types
 SELECT (x,y) FROM test ORDER BY x, y;
@@ -54,13 +67,13 @@ SELECT (x,y) FROM test GROUP BY x, y ORDER BY x, y;
  (2,3)
 (4 rows)
 
-SELECT ARRAY[(x,(y,x)),(y,(x,y))] FROM test ORDER BY x, y;
+SELECT ARRAY[NULL,(x,(y,x)),NULL,(y,(x,y))] FROM test ORDER BY x, y;
-               array               
+                    array                    
------------------------------------
+---------------------------------------------
- {"(1,\"(2,1)\")","(2,\"(1,2)\")"}
+ {NULL,"(1,\"(2,1)\")",NULL,"(2,\"(1,2)\")"}
- {"(1,\"(3,1)\")","(3,\"(1,3)\")"}
+ {NULL,"(1,\"(3,1)\")",NULL,"(3,\"(1,3)\")"}
- {"(2,\"(2,2)\")","(2,\"(2,2)\")"}
+ {NULL,"(2,\"(2,2)\")",NULL,"(2,\"(2,2)\")"}
- {"(2,\"(3,2)\")","(3,\"(2,3)\")"}
+ {NULL,"(2,\"(3,2)\")",NULL,"(3,\"(2,3)\")"}
 (4 rows)
 
 SELECT ARRAY[[(x,(y,x))],[(x,(x,y))]] FROM test ORDER BY x, y;
@@ -99,11 +112,72 @@ SELECT record_returner(x) FROM test ORDER BY x, y;
  (3,returned)
 (4 rows)
 
--- RECORD[] with varying shape unsupported
+SELECT NULLIF((x, y), (y, x)) FROM test ORDER BY x, y;
+ nullif 
+--------
+ (1,2)
+ (1,3)
+ 
+ (2,3)
+(4 rows)
+
+SELECT LEAST((x, y), (y, x)) FROM test ORDER BY x, y;
+ least 
+-------
+ (1,2)
+ (1,3)
+ (2,2)
+ (2,3)
+(4 rows)
+
+SELECT GREATEST((x, y), (y, x)) FROM test ORDER BY x, y;
+ greatest 
+----------
+ (2,1)
+ (3,1)
+ (2,2)
+ (3,2)
+(4 rows)
+
+SELECT COALESCE(NULL, (x, y), (y, x)) FROM test ORDER BY x, y;
+ coalesce 
+----------
+ (1,2)
+ (1,3)
+ (2,2)
+ (2,3)
+(4 rows)
+
+SELECT CASE x WHEN 2 THEN (x, y) ELSE (y, x) END FROM test ORDER BY x, y;
+  row  
+-------
+ (2,1)
+ (3,1)
+ (2,2)
+ (2,3)
+(4 rows)
+
+SELECT CASE x WHEN 2 THEN (x, y) END FROM test ORDER BY x, y;
+ case  
+-------
+ 
+ 
+ (2,2)
+ (2,3)
+(4 rows)
+
+-- varying shape unsupported
 SELECT ARRAY[(x,(y,x),y),(y,(x,y))] FROM test ORDER BY x, y;
 ERROR:  input of anonymous composite types is not implemented
 SELECT ARRAY[[(x,(y,x))],[((x,x),y)]] FROM test ORDER BY x, y;
 ERROR:  input of anonymous composite types is not implemented
+SELECT CASE x WHEN 2 THEN (x, y, x) ELSE (y, x) END FROM test ORDER BY x, y;
+ERROR:  input of anonymous composite types is not implemented
+-- RECORD from polymorphic types unsupported
+SELECT identity_returner((x, y)) FROM test ORDER BY x, y;
+ERROR:  input of anonymous composite types is not implemented
+SELECT array_agg((x, y)) FROM test;
+ERROR:  input of anonymous composite types is not implemented
 -- router queries support row types
 SELECT (x,y) FROM test WHERE x = 1 ORDER BY x, y;
   row  
@@ -119,11 +193,11 @@ SELECT (x,y) AS foo FROM test WHERE x = 1 ORDER BY x, y;
  (1,3)
 (2 rows)
 
-SELECT ARRAY[(x,(y,x)),(y,(x,y))] FROM test WHERE x = 1 ORDER BY x, y;
+SELECT ARRAY[NULL,(x,(y,x)),NULL,(y,(x,y))] FROM test WHERE x = 1 ORDER BY x, y;
-               array               
+                    array                    
------------------------------------
+---------------------------------------------
- {"(1,\"(2,1)\")","(2,\"(1,2)\")"}
+ {NULL,"(1,\"(2,1)\")",NULL,"(2,\"(1,2)\")"}
- {"(1,\"(3,1)\")","(3,\"(1,3)\")"}
+ {NULL,"(1,\"(3,1)\")",NULL,"(3,\"(1,3)\")"}
 (2 rows)
 
 SELECT ARRAY[[(x,(y,x))],[(x,(x,y))]] FROM test WHERE x = 1 ORDER BY x, y;
@@ -154,11 +228,60 @@ SELECT record_returner(x) FROM test WHERE x = 1 ORDER BY x, y;
  (2,returned)
 (2 rows)
 
--- RECORD[] with varying shape unsupported
+SELECT NULLIF((x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+ nullif 
+--------
+ (1,2)
+ (1,3)
+(2 rows)
+
+SELECT LEAST((x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+ least 
+-------
+ (1,2)
+ (1,3)
+(2 rows)
+
+SELECT GREATEST((x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+ greatest 
+----------
+ (2,1)
+ (3,1)
+(2 rows)
+
+SELECT COALESCE(NULL, (x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+ coalesce 
+----------
+ (1,2)
+ (1,3)
+(2 rows)
+
+SELECT CASE x WHEN 2 THEN (x, y) ELSE (y, x) END FROM test WHERE x = 1 ORDER BY x, y;
+  row  
+-------
+ (2,1)
+ (3,1)
+(2 rows)
+
+SELECT CASE x WHEN 2 THEN (x, y) END FROM test WHERE x = 1 ORDER BY x, y;
+ case 
+------
+ 
+ 
+(2 rows)
+
+-- varying shape unsupported
 SELECT ARRAY[(x,(y,x),y),(y,(x,y))] FROM test WHERE x = 1 ORDER BY x, y;
 ERROR:  input of anonymous composite types is not implemented
 SELECT ARRAY[[(x,(y,x))],[((x,x),y)]] FROM test WHERE x = 1 ORDER BY x, y;
 ERROR:  input of anonymous composite types is not implemented
+SELECT CASE x WHEN 2 THEN (x, y, x) ELSE (y, x) END FROM test WHERE x = 1 ORDER BY x, y;
+ERROR:  input of anonymous composite types is not implemented
+-- RECORD from polymorphic types unsupported
+SELECT identity_returner((x, y)) FROM test WHERE x = 1 ORDER BY x, y;
+ERROR:  input of anonymous composite types is not implemented
+SELECT array_agg((x, y)) FROM test WHERE x = 1;
+ERROR:  input of anonymous composite types is not implemented
 -- nested row expressions
 SELECT (x,(x,y)) AS foo FROM test WHERE x = 1 ORDER BY x, y;
      foo     
@@ -221,8 +344,5 @@ EXECUTE rec(1);
  (1,3)
 (2 rows)
 
+SET client_min_messages TO error;
 DROP SCHEMA row_types CASCADE;
-NOTICE:  drop cascades to 3 other objects
-DETAIL:  drop cascades to table test
-drop cascades to function table_returner(integer)
-drop cascades to function record_returner(integer)

src/test/regress/sql/row_types.sql

@@ -23,32 +23,61 @@ END;
 $$ language plpgsql;
 SELECT create_distributed_function('record_returner(int)');
 
+CREATE OR REPLACE FUNCTION identity_returner(x anyelement)
+RETURNS anyelement
+AS $$
+BEGIN
+	RETURN x;
+END;
+$$ language plpgsql;
+SELECT create_distributed_function('identity_returner(anyelement)');
+
 
 INSERT INTO test VALUES (1,2), (1,3), (2,2), (2,3);
 
 -- multi-shard queries support row types
 SELECT (x,y) FROM test ORDER BY x, y;
 SELECT (x,y) FROM test GROUP BY x, y ORDER BY x, y;
-SELECT ARRAY[(x,(y,x)),(y,(x,y))] FROM test ORDER BY x, y;
+SELECT ARRAY[NULL,(x,(y,x)),NULL,(y,(x,y))] FROM test ORDER BY x, y;
 SELECT ARRAY[[(x,(y,x))],[(x,(x,y))]] FROM test ORDER BY x, y;
 select distinct (x,y) AS foo, x, y FROM test ORDER BY x, y;
 SELECT table_returner(x) FROM test ORDER BY x, y;
 SELECT record_returner(x) FROM test ORDER BY x, y;
--- RECORD[] with varying shape unsupported
+SELECT NULLIF((x, y), (y, x)) FROM test ORDER BY x, y;
+SELECT LEAST((x, y), (y, x)) FROM test ORDER BY x, y;
+SELECT GREATEST((x, y), (y, x)) FROM test ORDER BY x, y;
+SELECT COALESCE(NULL, (x, y), (y, x)) FROM test ORDER BY x, y;
+SELECT CASE x WHEN 2 THEN (x, y) ELSE (y, x) END FROM test ORDER BY x, y;
+SELECT CASE x WHEN 2 THEN (x, y) END FROM test ORDER BY x, y;
+-- varying shape unsupported
 SELECT ARRAY[(x,(y,x),y),(y,(x,y))] FROM test ORDER BY x, y;
 SELECT ARRAY[[(x,(y,x))],[((x,x),y)]] FROM test ORDER BY x, y;
+SELECT CASE x WHEN 2 THEN (x, y, x) ELSE (y, x) END FROM test ORDER BY x, y;
+-- RECORD from polymorphic types unsupported
+SELECT identity_returner((x, y)) FROM test ORDER BY x, y;
+SELECT array_agg((x, y)) FROM test;
 
 -- router queries support row types
 SELECT (x,y) FROM test WHERE x = 1 ORDER BY x, y;
 SELECT (x,y) AS foo FROM test WHERE x = 1 ORDER BY x, y;
-SELECT ARRAY[(x,(y,x)),(y,(x,y))] FROM test WHERE x = 1 ORDER BY x, y;
+SELECT ARRAY[NULL,(x,(y,x)),NULL,(y,(x,y))] FROM test WHERE x = 1 ORDER BY x, y;
 SELECT ARRAY[[(x,(y,x))],[(x,(x,y))]] FROM test WHERE x = 1 ORDER BY x, y;
 select distinct (x,y) AS foo, x, y FROM test WHERE x = 1 ORDER BY x, y;
 SELECT table_returner(x) FROM test WHERE x = 1 ORDER BY x, y;
 SELECT record_returner(x) FROM test WHERE x = 1 ORDER BY x, y;
--- RECORD[] with varying shape unsupported
+SELECT NULLIF((x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+SELECT LEAST((x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+SELECT GREATEST((x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+SELECT COALESCE(NULL, (x, y), (y, x)) FROM test WHERE x = 1 ORDER BY x, y;
+SELECT CASE x WHEN 2 THEN (x, y) ELSE (y, x) END FROM test WHERE x = 1 ORDER BY x, y;
+SELECT CASE x WHEN 2 THEN (x, y) END FROM test WHERE x = 1 ORDER BY x, y;
+-- varying shape unsupported
 SELECT ARRAY[(x,(y,x),y),(y,(x,y))] FROM test WHERE x = 1 ORDER BY x, y;
 SELECT ARRAY[[(x,(y,x))],[((x,x),y)]] FROM test WHERE x = 1 ORDER BY x, y;
+SELECT CASE x WHEN 2 THEN (x, y, x) ELSE (y, x) END FROM test WHERE x = 1 ORDER BY x, y;
+-- RECORD from polymorphic types unsupported
+SELECT identity_returner((x, y)) FROM test WHERE x = 1 ORDER BY x, y;
+SELECT array_agg((x, y)) FROM test WHERE x = 1;
 
 -- nested row expressions
 SELECT (x,(x,y)) AS foo FROM test WHERE x = 1 ORDER BY x, y;
@@ -66,4 +95,5 @@ EXECUTE rec(1);
 EXECUTE rec(1);
 EXECUTE rec(1);
 
+SET client_min_messages TO error;
 DROP SCHEMA row_types CASCADE;