Fix set-returning functions in the sort key of ColumnarScan

Sorting a compressed chunk on a set-returning function failed with
"set-valued function called in context that cannot accept a set".
The function ended up in the sort key pushed down to the per-chunk
sorted scan, but a set-returning function can only be expanded above
the scan. Stop collecting per-chunk sort keys at the first
set-returning function so the sort happens above it.

Fixes #9927

Author: svenklemm

.unreleased/pr_9981

@@ -0,0 +1 @@
+Fixes: #9981 Fix set-returning functions in the sort key of ColumnarScan

tsl/src/nodes/columnar_scan/columnar_scan.c

@@ -3077,6 +3077,17 @@ build_sortinfo(PlannerInfo *root, const Chunk *chunk, RelOptInfo *chunk_rel,
 			break;
 		}
 
+		/*
+		 * We can't sort the ColumnarScan on a set-returning function. It is
+		 * expanded by a ProjectSet node above the scan, so stop here and let
+		 * the sort happen there. The leading sort keys collected before it are
+		 * still usable to sort the ColumnarScan.
+		 */
+		if (expression_returns_set((Node *) chunk_em_expr))
+		{
+			break;
+		}
+
 		sort_pathkeys = lappend(sort_pathkeys, pk);
 		sort_pathkey_exprs = lappend(sort_pathkey_exprs, chunk_em_expr);
 	}

tsl/test/expected/compress_sort_transform.out

@@ -381,3 +381,121 @@ group by 1
 reset max_parallel_workers_per_gather;
 reset work_mem;
 reset enable_hashagg;
+-- Sorting on a set-returning function.
+set max_parallel_workers_per_gather to 0;
+create table srf_sort(time timestamptz not null, dev int, seg int);
+select create_hypertable('srf_sort', 'time', chunk_time_interval => interval '1 day');
+   create_hypertable   
+-----------------------
+ (3,public,srf_sort,t)
+
+insert into srf_sort values
+    ('2025-01-01', 3, 1),
+    ('2025-01-02', 2, 1),
+    ('2025-01-03', 4, 2);
+alter table srf_sort set (timescaledb.compress, timescaledb.compress_segmentby='seg', timescaledb.compress_orderby='time');
+select count(compress_chunk(c)) from show_chunks('srf_sort') c;
+ count 
+-------
+     3
+
+-- Sorting by the set-returning function with a LIMIT.
+select generate_series(1, dev) g from srf_sort order by g limit 5;
+ g 
+---
+ 1
+ 1
+ 1
+ 2
+ 2
+
+explain (costs off) select generate_series(1, dev) g from srf_sort order by g limit 5;
+--- QUERY PLAN ---
+ Limit
+   ->  Sort
+         Sort Key: (generate_series(1, srf_sort.dev))
+         ->  ProjectSet
+               ->  Append
+                     ->  Custom Scan (ColumnarScan) on _hyper_3_7_chunk
+                           ->  Seq Scan on compress_hyper_4_10_chunk
+                     ->  Custom Scan (ColumnarScan) on _hyper_3_8_chunk
+                           ->  Seq Scan on compress_hyper_4_11_chunk
+                     ->  Custom Scan (ColumnarScan) on _hyper_3_9_chunk
+                           ->  Seq Scan on compress_hyper_4_12_chunk
+
+-- Without a LIMIT it produces the same data.
+select generate_series(1, dev) g from srf_sort order by g;
+ g 
+---
+ 1
+ 1
+ 1
+ 2
+ 2
+ 2
+ 3
+ 3
+ 4
+
+-- A set-returning function only in the output, sorting on a real column.
+explain (costs off) select generate_series(1, dev) g, time from srf_sort order by time limit 3;
+--- QUERY PLAN ---
+ Limit
+   ->  ProjectSet
+         ->  Custom Scan (ChunkAppend) on srf_sort
+               Order: srf_sort."time"
+               ->  Custom Scan (ColumnarScan) on _hyper_3_7_chunk
+                     ->  Sort
+                           Sort Key: compress_hyper_4_10_chunk._ts_meta_v2_first_time
+                           ->  Seq Scan on compress_hyper_4_10_chunk
+               ->  Custom Scan (ColumnarScan) on _hyper_3_8_chunk
+                     ->  Sort
+                           Sort Key: compress_hyper_4_11_chunk._ts_meta_v2_first_time
+                           ->  Seq Scan on compress_hyper_4_11_chunk
+               ->  Custom Scan (ColumnarScan) on _hyper_3_9_chunk
+                     ->  Sort
+                           Sort Key: compress_hyper_4_12_chunk._ts_meta_v2_first_time
+                           ->  Seq Scan on compress_hyper_4_12_chunk
+
+-- Sorting on a real column followed by a set-returning function.
+select dev, generate_series(1, dev) g from srf_sort order by dev, g limit 4;
+ dev | g 
+-----+---
+   2 | 1
+   2 | 2
+   3 | 1
+   3 | 2
+
+-- Sorting on a segment column followed by a set-returning function.
+explain (costs off) select seg, generate_series(1, dev) g from srf_sort order by seg, g limit 4;
+--- QUERY PLAN ---
+ Limit
+   ->  Incremental Sort
+         Sort Key: srf_sort.seg, (generate_series(1, srf_sort.dev))
+         Presorted Key: srf_sort.seg
+         ->  ProjectSet
+               ->  Merge Append
+                     Sort Key: srf_sort.seg
+                     ->  Sort
+                           Sort Key: _hyper_3_7_chunk.seg
+                           ->  Custom Scan (ColumnarScan) on _hyper_3_7_chunk
+                                 ->  Seq Scan on compress_hyper_4_10_chunk
+                     ->  Sort
+                           Sort Key: _hyper_3_8_chunk.seg
+                           ->  Custom Scan (ColumnarScan) on _hyper_3_8_chunk
+                                 ->  Seq Scan on compress_hyper_4_11_chunk
+                     ->  Sort
+                           Sort Key: _hyper_3_9_chunk.seg
+                           ->  Custom Scan (ColumnarScan) on _hyper_3_9_chunk
+                                 ->  Seq Scan on compress_hyper_4_12_chunk
+
+select seg, generate_series(1, dev) g from srf_sort order by seg, g limit 4;
+ seg | g 
+-----+---
+   1 | 1
+   1 | 1
+   1 | 2
+   1 | 2
+
+drop table srf_sort cascade;
+reset max_parallel_workers_per_gather;

tsl/test/sql/compress_sort_transform.sql

@@ -96,3 +96,35 @@ reset max_parallel_workers_per_gather;
 
 reset work_mem;
 reset enable_hashagg;
+
+-- Sorting on a set-returning function.
+set max_parallel_workers_per_gather to 0;
+
+create table srf_sort(time timestamptz not null, dev int, seg int);
+select create_hypertable('srf_sort', 'time', chunk_time_interval => interval '1 day');
+insert into srf_sort values
+    ('2025-01-01', 3, 1),
+    ('2025-01-02', 2, 1),
+    ('2025-01-03', 4, 2);
+alter table srf_sort set (timescaledb.compress, timescaledb.compress_segmentby='seg', timescaledb.compress_orderby='time');
+select count(compress_chunk(c)) from show_chunks('srf_sort') c;
+
+-- Sorting by the set-returning function with a LIMIT.
+select generate_series(1, dev) g from srf_sort order by g limit 5;
+explain (costs off) select generate_series(1, dev) g from srf_sort order by g limit 5;
+
+-- Without a LIMIT it produces the same data.
+select generate_series(1, dev) g from srf_sort order by g;
+
+-- A set-returning function only in the output, sorting on a real column.
+explain (costs off) select generate_series(1, dev) g, time from srf_sort order by time limit 3;
+
+-- Sorting on a real column followed by a set-returning function.
+select dev, generate_series(1, dev) g from srf_sort order by dev, g limit 4;
+
+-- Sorting on a segment column followed by a set-returning function.
+explain (costs off) select seg, generate_series(1, dev) g from srf_sort order by seg, g limit 4;
+select seg, generate_series(1, dev) g from srf_sort order by seg, g limit 4;
+
+drop table srf_sort cascade;
+reset max_parallel_workers_per_gather;