Fix set-returning functions in the sort key of ColumnarScan

.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

@@ -3060,6 +3060,17 @@ build_sortinfo(PlannerInfo *root, const Chunk *chunk, RelOptInfo *chunk_rel,
 		if (!ec->ec_has_volatile)
 		{
 			em_expr = ts_find_em_expr_for_rel(pk->pk_eclass, compression_info->chunk_rel);
+
+			/*
+			 * We can't sort the ColumnarScan on a set-returning function. It is
+			 * expanded by a ProjectSet node above the scan, so treat it as no
+			 * match and let the sort happen there. The leading sort keys
+			 * collected before it are still usable to sort the ColumnarScan.
+			 */
+			if (em_expr && expression_returns_set((Node *) em_expr))
+			{
+				em_expr = NULL;
+			}
 		}
 		chunk_em_exprs = lappend(chunk_em_exprs, 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;