Fix INSERT on hypertables using sub-selects with aggregates

When inserting into a hypertable using a sub-select clause that
involves an aggregate, the insert fails with the error "Aggref found
in non-Agg plan node". This happens because the target list from the
aggregate sub-select expects an aggregate parent node and we simply
reuse the target list when modifying the insert plan with new
CustomScan plan nodes.

This change creates a modified target list to use with the CustomScan
node that avoids this issue.

Author: erimatnor

src/chunk_dispatch_plan.c

@@ -5,6 +5,7 @@
 #include <nodes/readfuncs.h>
 #include <utils/rel.h>
 #include <catalog/pg_type.h>
+#include <rewrite/rewriteManip.h>
 
 #include "chunk_dispatch_plan.h"
 #include "chunk_dispatch_state.h"
@@ -27,6 +28,84 @@ static CustomScanMethods chunk_dispatch_plan_methods = {
 	.CreateCustomScanState = create_chunk_dispatch_state,
 };
 
+/*
+ * Build a target list for a CustomScan node.
+ *
+ * The CustomScan node needs to provide a modified target list in case of
+ * subplan nodes that have target lists with expressions that involve, e.g.,
+ * aggregates. Such target lists need a matching parent node type (e.g., a
+ * target list with an AggDef needs a Agg node parent). If we simply use the
+ * subplan's target list on a CustomScan parent node, there might be a mismatch.
+ *
+ * The code here is similar to ExecCheckPlanOutput() in nodeModifyTable.c in the
+ * PostgreSQL source code, but here we construct a modified target list for the
+ * CustomScan node instead of simply just checking validity.
+ *
+ */
+static List *
+build_customscan_targetlist(Relation rel, List *targetlist)
+{
+	TupleDesc	tupdesc = RelationGetDescr(rel);
+	List	   *new_targetlist = NIL;
+	ListCell   *lc;
+	int			attno = 0;
+
+	foreach(lc, targetlist)
+	{
+		TargetEntry *te = lfirst(lc);
+		Form_pg_attribute attr;
+		Expr	   *expr;
+
+		/* Ignore junk items */
+		if (te->resjunk)
+			continue;
+
+		if (attno >= tupdesc->natts)
+			ereport(ERROR,
+					(errcode(ERRCODE_DATATYPE_MISMATCH),
+					 errmsg("table row type and query-specified row type do not match"),
+					 errdetail("Query has too many columns.")));
+
+		attr = tupdesc->attrs[attno++];
+
+		if (attr->attisdropped)
+			expr = &makeConst(INT4OID,
+							  -1,
+							  InvalidOid,
+							  sizeof(int32),
+							  (Datum) 0,
+							  true,
+							  true)->xpr;
+		else
+
+			/*
+			 * According to nodes/primnodes.h, the INDEX_VAR Var type is
+			 * abused in CustomScan nodes to reference custom scan tuple
+			 * types.
+			 */
+			expr = &makeVar(INDEX_VAR,
+							attno,
+							exprType((Node *) te->expr),
+							exprTypmod((Node *) te->expr),
+							exprCollation((Node *) te->expr),
+							0)->xpr;
+
+		new_targetlist = lappend(new_targetlist,
+								 makeTargetEntry(expr,
+												 attno,
+												 NULL,
+												 te->resjunk));
+	}
+
+	if (attno != tupdesc->natts)
+		ereport(ERROR,
+				(errcode(ERRCODE_DATATYPE_MISMATCH),
+		  errmsg("table row type and query-specified row type do not match"),
+				 errdetail("Query has too few columns.")));
+
+	return new_targetlist;
+}
+
 /* Create a chunk dispatch plan node in the form of a CustomScan node. The
  * purpose of this plan node is to dispatch (route) tuples to the correct chunk
  * in a hypertable.
@@ -42,10 +121,11 @@ static CustomScanMethods chunk_dispatch_plan_methods = {
  * node.
  */
 CustomScan *
-chunk_dispatch_plan_create(Plan *subplan, Oid hypertable_relid, Query *parse)
+chunk_dispatch_plan_create(Plan *subplan, Index hypertable_rti, Oid hypertable_relid, Query *parse)
 {
 	CustomScan *cscan = makeNode(CustomScan);
 	ChunkDispatchInfo *info = chunk_dispatch_info_create(hypertable_relid, parse);
+	Relation	rel;
 
 	cscan->custom_private = list_make1(info);
 	cscan->methods = &chunk_dispatch_plan_methods;
@@ -60,11 +140,20 @@ chunk_dispatch_plan_create(Plan *subplan, Oid hypertable_relid, Query *parse)
 	cscan->scan.plan.plan_width = subplan->plan_width;
 
 	/*
-	 * Copy target list from parent table. This should work since hypertables
-	 * mandate that chunks have identical column definitions
+	 * Create a modified target list for the CustomScan based on the subplan's
+	 * original target list
+	 */
+	rel = relation_open(hypertable_relid, AccessShareLock);
+	cscan->scan.plan.targetlist = build_customscan_targetlist(rel, subplan->targetlist);
+	RelationClose(rel);
+
+	/*
+	 * We need to set a custom_scan_tlist for EXPLAIN (verbose). But this
+	 * target list needs to reference the proper rangetable entry so we must
+	 * replace all INDEX_VAR attnos with the hypertable's rangetable index.
 	 */
-	cscan->scan.plan.targetlist = subplan->targetlist;
-	cscan->custom_scan_tlist = NIL;
+	cscan->custom_scan_tlist = copyObject(cscan->scan.plan.targetlist);
+	ChangeVarNodes((Node *) cscan->custom_scan_tlist, INDEX_VAR, hypertable_rti, 0);
 
 	return cscan;
 }

src/chunk_dispatch_plan.h

@@ -6,6 +6,6 @@
 #include <nodes/parsenodes.h>
 #include <nodes/extensible.h>
 
-extern CustomScan *chunk_dispatch_plan_create(Plan *subplan, Oid hypertable_relid, Query *parse);
+extern CustomScan *chunk_dispatch_plan_create(Plan *subplan, Index hypertable_rti, Oid hypertable_relid, Query *parse);
 
 #endif   /* TIMESCALEDB_CHUNK_DISPATCH_PLAN_H */

src/planner.c

@@ -382,7 +382,8 @@ modifytable_plan_walker(Plan **planptr, void *pctx)
 			 */
 			forboth(lc_plan, mt->plans, lc_rel, mt->resultRelations)
 			{
-				RangeTblEntry *rte = rt_fetch(lfirst_int(lc_rel), ctx->rtable);
+				Index		rti = lfirst_int(lc_rel);
+				RangeTblEntry *rte = rt_fetch(rti, ctx->rtable);
 				Hypertable *ht = hypertable_cache_get_entry(ctx->hcache, rte->relid);
 
 				if (ht != NULL)
@@ -394,7 +395,7 @@ modifytable_plan_walker(Plan **planptr, void *pctx)
 					 * We replace the plan with our custom chunk dispatch
 					 * plan.
 					 */
-					*subplan_ptr = chunk_dispatch_plan_create(subplan, rte->relid, ctx->parse);
+					*subplan_ptr = chunk_dispatch_plan_create(subplan, rti, rte->relid, ctx->parse);
 					hypertable_found = true;
 				}
 			}