Fix crash in VectorAgg plan code

Due to a refactor in commit 947f7f4 that silenced a coverity warning,
a crash was introduced for certain aggregation query plans. However,
no such plan was covered in our tests, so the crash was never
triggered.

Add a fix and a test case that triggers crash when the fix is not
present.

Author: erimatnor

.unreleased/pr_7902

@@ -0,0 +1 @@
+Fixes: #7902 Fix crash in VectorAgg plan code

tsl/src/nodes/columnar_scan/columnar_scan.c

@@ -1038,6 +1038,7 @@ columnar_scan_plan_create(PlannerInfo *root, RelOptInfo *rel, CustomPath *best_p
 	VectorQualInfoHypercore vqih = {
 		.vqinfo = {
 			.rti = rel->relid,
+			.maxattno = hcinfo->num_columns,
 			.vector_attrs = columnar_scan_build_vector_attrs(hcinfo->columns, hcinfo->num_columns),
 		},
 		.hcinfo = hcinfo,

tsl/src/nodes/decompress_chunk/planner.c

@@ -884,6 +884,7 @@ find_vectorized_quals(DecompressionMapContext *context, DecompressChunkPath *pat
 					  List **vectorized, List **nonvectorized)
 {
 	VectorQualInfo vqi = {
+		.maxattno = path->info->chunk_rel->max_attr,
 		.vector_attrs = build_vector_attrs_array(context->uncompressed_attno_info, path->info),
 		.rti = path->info->chunk_rel->relid,
 	};

tsl/src/nodes/decompress_chunk/vector_quals.h

@@ -25,12 +25,18 @@ typedef struct VectorQualInfo
 	Index rti;
 
 	bool reverse;
+
 	/*
 	 * Arrays indexed by uncompressed attno indicating whether an
 	 * attribute/column is a vectorizable type and/or a segmentby attribute.
+	 *
+	 * Note: array lengths are maxattno + 1.
 	 */
 	bool *vector_attrs;
 	bool *segmentby_attrs;
+
+	/* Max attribute number found in arrays above */
+	AttrNumber maxattno;
 } VectorQualInfo;
 
 /*

tsl/src/nodes/vector_agg/plan.c

@@ -221,7 +221,9 @@ is_vector_var(const VectorQualInfo *vqinfo, Expr *expr)
 		return false;
 	}
 
-	return vqinfo->vector_attrs[var->varattno];
+	Assert(var->varattno <= vqinfo->maxattno);
+
+	return vqinfo->vector_attrs && vqinfo->vector_attrs[var->varattno];
 }
 
 /*
@@ -497,14 +499,26 @@ vectoragg_plan_possible(Plan *childplan, const List *rtable, VectorQualInfo *vqi
 	}
 
 	CustomScan *customscan = castNode(CustomScan, childplan);
-	RangeTblEntry *rte = rt_fetch(customscan->scan.scanrelid, rtable);
 
-	if (ts_is_hypercore_am(ts_get_rel_am(rte->relid)))
-		vectoragg_plan_tam(childplan, rtable, vqi);
-	else if (strcmp(customscan->methods->CustomName, "DecompressChunk") == 0)
+	if (strcmp(customscan->methods->CustomName, "DecompressChunk") == 0)
+	{
 		vectoragg_plan_decompress_chunk(childplan, vqi);
+		return true;
+	}
+
+	/* We're looking for a baserel scan */
+	if (customscan->scan.scanrelid > 0)
+	{
+		RangeTblEntry *rte = rt_fetch(customscan->scan.scanrelid, rtable);
 
-	return true;
+		if (rte && ts_is_hypercore_am(ts_get_rel_am(rte->relid)))
+		{
+			vectoragg_plan_tam(childplan, rtable, vqi);
+			return true;
+		}
+	}
+
+	return false;
 }
 
 /*
@@ -597,6 +611,7 @@ try_insert_vector_agg_node(Plan *plan, List *rtable)
 
 	Plan *childplan = agg->plan.lefttree;
 	VectorQualInfo vqi;
+	MemSet(&vqi, 0, sizeof(VectorQualInfo));
 
 	/*
 	 * Build supplementary info to determine whether we can vectorize the

tsl/src/nodes/vector_agg/plan_decompress_chunk.c

@@ -111,6 +111,7 @@ vectoragg_plan_decompress_chunk(Plan *childplan, VectorQualInfo *vqi)
 		}
 	}
 
+	vqi->maxattno = maxattno;
 	vqi->vector_attrs = (bool *) palloc0(sizeof(bool) * (maxattno + 1));
 	vqi->segmentby_attrs = (bool *) palloc0(sizeof(bool) * (maxattno + 1));
 

tsl/src/nodes/vector_agg/plan_tam.c

@@ -22,6 +22,7 @@ vectoragg_plan_tam(Plan *childplan, const List *rtable, VectorQualInfo *vqi)
 
 	*vqi = (VectorQualInfo){
 		.rti = customscan->scan.scanrelid,
+		.maxattno = hinfo->num_columns,
 		.vector_attrs = (bool *) palloc0(sizeof(bool) * (hinfo->num_columns + 1)),
 		.segmentby_attrs = (bool *) palloc0(sizeof(bool) * (hinfo->num_columns + 1)),
 		/*

tsl/test/expected/vectorized_aggregation.out

@@ -3602,3 +3602,42 @@ SELECT sum(float_value), int_value FROM testtable2 WHERE int_value = 1 GROUP BY
  3162 |         1
 (1 row)
 
+--
+-- Test handling of Agg on top of ChunkAppend. This reproduces a crash
+-- in the vector agg planning code introduced in commit 947f7f400.
+--
+CREATE TABLE testtable3 (time timestamptz, location_id int, device_id int, sensor_id int);
+SELECT create_hypertable('testtable3', 'time');
+NOTICE:  adding not-null constraint to column "time"
+    create_hypertable    
+-------------------------
+ (5,public,testtable3,t)
+(1 row)
+
+ALTER TABLE testtable3 SET (timescaledb.compress_orderby='time', timescaledb.compress_segmentby='location_id');
+INSERT INTO testtable3 SELECT t, ceil(random() * 20)::int, ceil(random() * 30)::int, ceil(random() * 20)::int FROM generate_series('2024-01-01'::timestamptz, '2024-01-10'::timestamptz, '1h') AS t;
+SELECT count(compress_chunk(ch)) FROM show_chunks('testtable3') ch;
+ count 
+-------
+     2
+(1 row)
+
+EXPLAIN (costs off) SELECT (date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour')::timestamp as time, TT.location_id as location_id, TT.device_id as device_id, 0 as sensor_id, date_trunc('day', current_timestamp) as discovered_date FROM testtable3 TT WHERE time >= date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour' GROUP BY TT.location_id, TT.device_id;
+                                                                                 QUERY PLAN                                                                                 
+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------
+ Group
+   Group Key: tt.location_id, tt.device_id
+   ->  Gather Merge
+         Workers Planned: 2
+         ->  Sort
+               Sort Key: tt.location_id, tt.device_id
+               ->  Partial HashAggregate
+                     Group Key: tt.location_id, tt.device_id
+                     ->  Parallel Custom Scan (ChunkAppend) on testtable3 tt
+                           Chunks excluded during startup: 1
+                           ->  Custom Scan (DecompressChunk) on _hyper_5_120_chunk tt_1
+                                 Vectorized Filter: ("time" >= (date_trunc('hour'::text, 'Tue Jan 09 00:00:00 2024 PST'::timestamp with time zone) - '@ 1 hour'::interval))
+                                 ->  Parallel Seq Scan on compress_hyper_6_122_chunk
+(13 rows)
+
+SELECT (date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour')::timestamp as time, TT.location_id as location_id, TT.device_id as device_id, 0 as sensor_id, date_trunc('day', current_timestamp) as discovered_date FROM testtable3 TT WHERE time >= date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour' GROUP BY TT.location_id, TT.device_id \g :TEST_OUTPUT_DIR/vectorized_aggregation_query_result.out

tsl/test/sql/vectorized_aggregation.sql

@@ -421,3 +421,19 @@ SELECT sum(float_value) FROM testtable2 GROUP BY tableoid ORDER BY 1 LIMIT 1;
 -- Postgres versions starting with 16 remove the grouping columns that are
 -- equated to a constant. Check that our planning code handles this well.
 SELECT sum(float_value), int_value FROM testtable2 WHERE int_value = 1 GROUP BY int_value;
+
+--
+-- Test handling of Agg on top of ChunkAppend. This reproduces a crash
+-- in the vector agg planning code introduced in commit 947f7f400.
+--
+CREATE TABLE testtable3 (time timestamptz, location_id int, device_id int, sensor_id int);
+SELECT create_hypertable('testtable3', 'time');
+ALTER TABLE testtable3 SET (timescaledb.compress_orderby='time', timescaledb.compress_segmentby='location_id');
+
+INSERT INTO testtable3 SELECT t, ceil(random() * 20)::int, ceil(random() * 30)::int, ceil(random() * 20)::int FROM generate_series('2024-01-01'::timestamptz, '2024-01-10'::timestamptz, '1h') AS t;
+
+SELECT count(compress_chunk(ch)) FROM show_chunks('testtable3') ch;
+
+EXPLAIN (costs off) SELECT (date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour')::timestamp as time, TT.location_id as location_id, TT.device_id as device_id, 0 as sensor_id, date_trunc('day', current_timestamp) as discovered_date FROM testtable3 TT WHERE time >= date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour' GROUP BY TT.location_id, TT.device_id;
+
+SELECT (date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour')::timestamp as time, TT.location_id as location_id, TT.device_id as device_id, 0 as sensor_id, date_trunc('day', current_timestamp) as discovered_date FROM testtable3 TT WHERE time >= date_trunc('hour', '2024-01-09'::timestamptz) - interval '1 hour' GROUP BY TT.location_id, TT.device_id \g :TEST_OUTPUT_DIR/vectorized_aggregation_query_result.out