Merge pull request #3550 from citusdata/fix-generated-halfway

Fix create_distributed_table on a table using GENERATED ALWAYS AS

src/backend/distributed/commands/multi_copy.c

@@ -1290,8 +1290,8 @@ TypeForColumnName(Oid relationId, TupleDesc tupleDescriptor, char *columnName)
 
 
 /*
- * Walks a TupleDesc and returns an array of the types of each attribute. Will return
- * InvalidOid in the place of dropped attributes.
+ * Walks a TupleDesc and returns an array of the types of each attribute.
+ * Returns InvalidOid in the place of dropped or generated attributes.
  */
 static Oid *
 TypeArrayFromTupleDescriptor(TupleDesc tupleDescriptor)
@@ -1302,7 +1302,11 @@ TypeArrayFromTupleDescriptor(TupleDesc tupleDescriptor)
 	for (int columnIndex = 0; columnIndex < columnCount; columnIndex++)
 	{
 		Form_pg_attribute attr = TupleDescAttr(tupleDescriptor, columnIndex);
-		if (attr->attisdropped)
+		if (attr->attisdropped
+#if PG_VERSION_NUM >= 120000
+			|| attr->attgenerated == ATTRIBUTE_GENERATED_STORED
+#endif
+			)
 		{
 			typeArray[columnIndex] = InvalidOid;
 		}
@@ -1338,7 +1342,7 @@ ColumnCoercionPaths(TupleDesc destTupleDescriptor, TupleDesc inputTupleDescripto
 
 		if (inputTupleType == InvalidOid)
 		{
-			/* this was a dropped column and will not be in the incoming tuples */
+			/* TypeArrayFromTupleDescriptor decided to skip this column */
 			continue;
 		}
 
@@ -1378,9 +1382,9 @@ TypeOutputFunctions(uint32 columnCount, Oid *typeIdArray, bool binaryFormat)
 		bool typeVariableLength = false;
 		Oid outputFunctionId = InvalidOid;
 
-		/* If there are any dropped columns it'll show up as a NULL */
 		if (columnTypeId == InvalidOid)
 		{
+			/* TypeArrayFromTupleDescriptor decided to skip this column */
 			continue;
 		}
 		else if (binaryFormat)

src/test/regress/expected/pg12.out

@@ -21,10 +21,11 @@ insert into test_am values (1, 1);
 select create_distributed_table('test_am','id');
 ERROR:  cannot distribute relations using non-heap access methods
 -- Test generated columns
+-- val1 after val2 to test https://github.com/citusdata/citus/issues/3538
 create table gen1 (
 	id int,
-	val1 int,
-	val2 int GENERATED ALWAYS AS (val1 + 2) STORED
+	val2 int GENERATED ALWAYS AS (val1 + 2) STORED,
+	val1 int
 );
 create table gen2 (
 	id int,
@@ -46,16 +47,16 @@ DETAIL:  Distribution column must not use GENERATED ALWAYS AS (...) STORED.
 insert into gen1 (id, val1) values (2,4),(4,6),(6,2),(8,2);
 insert into gen2 (id, val1) values (2,4),(4,6),(6,2),(8,2);
 select * from gen1 order by 1,2,3;
- id | val1 | val2
+ id | val2 | val1
 ---------------------------------------------------------------------
-  1 |    4 |    6
-  2 |    4 |    6
-  3 |    6 |    8
-  4 |    6 |    8
-  5 |    2 |    4
-  6 |    2 |    4
-  7 |    2 |    4
-  8 |    2 |    4
+  1 |    6 |    4
+  2 |    6 |    4
+  3 |    8 |    6
+  4 |    8 |    6
+  5 |    4 |    2
+  6 |    4 |    2
+  7 |    4 |    2
+  8 |    4 |    2
 (8 rows)
 
 select * from gen2 order by 1,2,3;

src/test/regress/sql/pg12.sql

@@ -25,10 +25,11 @@ insert into test_am values (1, 1);
 select create_distributed_table('test_am','id');
 
 -- Test generated columns
+-- val1 after val2 to test https://github.com/citusdata/citus/issues/3538
 create table gen1 (
 	id int,
-	val1 int,
-	val2 int GENERATED ALWAYS AS (val1 + 2) STORED
+	val2 int GENERATED ALWAYS AS (val1 + 2) STORED,
+	val1 int
 );
 create table gen2 (
 	id int,