From 833807906235a0c52d5be0403d326fc0b0179414 Mon Sep 17 00:00:00 2001
From: Adam Storm <storm@cockroachlabs.com>
Date: Thu, 20 Nov 2025 17:21:59 -0500
Subject: [PATCH] optbuilder: preserve ORDER BY in set-returning UDFs with OUT
 parameters

When set-returning UDFs with OUT parameters are called directly in a
SELECT list (e.g., SELECT f()), CockroachDB wraps multiple result
columns into a single tuple column. During this transformation, two
functions created new scopes without preserving ordering information,
causing ORDER BY clauses in the UDF body to be ignored.

This commit fixes the issue by calling copyOrdering() in two places:
1. combineRoutineColsIntoTuple - when wrapping columns into a tuple
2. maybeAddRoutineAssignmentCasts - when adding type casts

The copyOrdering() method not only copies the ordering metadata but
also adds the columns referenced by the ordering to extraCols, ensuring
they remain available for the optimizer to enforce the ordering.

Fixes #144013

Release note (bug fix): Fixed a bug where ORDER BY clauses in
set-returning SQL user-defined functions with OUT parameters were
ignored when the function was called directly in a SELECT list
(e.g., SELECT f()). The ordering is now properly preserved and
enforced.

Epic: None
---
 .../logictest/testdata/logic_test/udf_setof   | 60 +++++++++++++++++++
 pkg/sql/opt/optbuilder/routine.go             |  2 +
 pkg/sql/opt/optbuilder/testdata/udf           |  8 +--
 3 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/pkg/sql/logictest/testdata/logic_test/udf_setof b/pkg/sql/logictest/testdata/logic_test/udf_setof
index e50a1ff5a643..f4c7c682cc99 100644
--- a/pkg/sql/logictest/testdata/logic_test/udf_setof
+++ b/pkg/sql/logictest/testdata/logic_test/udf_setof
@@ -325,3 +325,63 @@ statement ok
 DROP FUNCTION f145414;
 
 subtest end
+
+# Test that ORDER BY is preserved in SQL UDFs with OUT parameters
+subtest out_params_ordering
+
+statement ok
+CREATE TABLE test_ordering (x INT PRIMARY KEY, y INT);
+
+statement ok
+INSERT INTO test_ordering VALUES (1, 10), (2, 20), (3, 30), (4, 40);
+
+statement ok
+CREATE FUNCTION f_out_desc(OUT a INT, OUT b INT) RETURNS SETOF RECORD AS $$
+  SELECT x, y FROM test_ordering ORDER BY x DESC;
+$$ LANGUAGE SQL;
+
+query T nosort
+SELECT f_out_desc();
+----
+(4,40)
+(3,30)
+(2,20)
+(1,10)
+
+query II nosort
+SELECT * FROM f_out_desc();
+----
+4  40
+3  30
+2  20
+1  10
+
+statement ok
+DROP FUNCTION f_out_desc;
+
+statement ok
+CREATE FUNCTION f_out_asc(OUT a INT, OUT b INT) RETURNS SETOF RECORD AS $$
+  SELECT x, y FROM test_ordering ORDER BY x ASC;
+$$ LANGUAGE SQL;
+
+query T nosort
+SELECT f_out_asc();
+----
+(1,10)
+(2,20)
+(3,30)
+(4,40)
+
+query II nosort
+SELECT * FROM f_out_asc();
+----
+1  10
+2  20
+3  30
+4  40
+
+statement ok
+DROP FUNCTION f_out_asc;
+DROP TABLE test_ordering;
+
+subtest end
diff --git a/pkg/sql/opt/optbuilder/routine.go b/pkg/sql/opt/optbuilder/routine.go
index b249afbc0a94..34908cfc9db9 100644
--- a/pkg/sql/opt/optbuilder/routine.go
+++ b/pkg/sql/opt/optbuilder/routine.go
@@ -653,6 +653,7 @@ func (b *Builder) finalizeRoutineReturnType(
 // into a single tuple column.
 func (b *Builder) combineRoutineColsIntoTuple(stmtScope *scope) *scope {
 	outScope := stmtScope.push()
+	outScope.copyOrdering(stmtScope)
 	elems := make(memo.ScalarListExpr, len(stmtScope.cols))
 	typContents := make([]*types.T, len(stmtScope.cols))
 	for i := range stmtScope.cols {
@@ -721,6 +722,7 @@ func (b *Builder) maybeAddRoutineAssignmentCasts(
 		return stmtScope
 	}
 	outScope := stmtScope.push()
+	outScope.copyOrdering(stmtScope)
 	for i, col := range stmtScope.cols {
 		scalar := b.factory.ConstructVariable(col.id)
 		if !col.typ.Identical(desiredTypes[i]) {
diff --git a/pkg/sql/opt/optbuilder/testdata/udf b/pkg/sql/opt/optbuilder/testdata/udf
index 5dc565ccd9a8..68ac7c01690b 100644
--- a/pkg/sql/opt/optbuilder/testdata/udf
+++ b/pkg/sql/opt/optbuilder/testdata/udf
@@ -865,9 +865,9 @@ project
            ├── params: i:1
            └── body
                 └── project
-                     ├── columns: column8:8
+                     ├── columns: column8:8 b:3
                      ├── project
-                     │    ├── columns: column7:7
+                     │    ├── columns: column7:7 b:3
                      │    ├── limit
                      │    │    ├── columns: a:2!null b:3 c:4!null
                      │    │    ├── internal-ordering: -3
@@ -911,9 +911,9 @@ project
            ├── params: i:1
            └── body
                 └── project
-                     ├── columns: column8:8
+                     ├── columns: column8:8 b:3
                      ├── project
-                     │    ├── columns: column7:7
+                     │    ├── columns: column7:7 b:3
                      │    ├── limit
                      │    │    ├── columns: a:2!null b:3 c:4!null
                      │    │    ├── internal-ordering: -3