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

Author: ajstorm

pkg/ccl/logictestccl/testdata/logic_test/plpgsql_srf

@@ -413,9 +413,7 @@ CREATE FUNCTION f() RETURNS TABLE (x INT, y INT) LANGUAGE PLpgSQL AS $$
   END
 $$;
 
-# TODO(144013): we should be able to use nosort here, but the ordering is
-# dropped.
-query T rowsort
+query T nosort
 SELECT f();
 ----
 (1,2)
@@ -432,6 +430,22 @@ SELECT * FROM f();
 statement ok
 DROP FUNCTION f;
 
+# Test that ordering is preserved with array_agg
+statement ok
+CREATE FUNCTION f_agg() RETURNS TABLE (x INT, y INT) LANGUAGE PLpgSQL AS $$
+  BEGIN
+    RETURN QUERY SELECT * FROM xy ORDER BY xy.x DESC;
+  END
+$$;
+
+query T
+SELECT array_agg(f) FROM f_agg() AS f;
+----
+{"(5,6)","(3,4)","(1,2)"}
+
+statement ok
+DROP FUNCTION f_agg;
+
 subtest end
 
 subtest record

pkg/sql/logictest/testdata/logic_test/udf_setof

@@ -325,3 +325,114 @@ 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;
+
+# Test with assignment casts using RETURNS SETOF RECORD and column definition list.
+# The function returns INT columns, but the column definition list specifies
+# FLOAT, which triggers assignment casts at query time.
+statement ok
+CREATE FUNCTION f_record_cast() RETURNS SETOF RECORD AS $$
+  SELECT x, y FROM test_ordering ORDER BY x DESC;
+$$ LANGUAGE SQL;
+
+query RR nosort
+SELECT * FROM f_record_cast() AS f(a FLOAT, b FLOAT);
+----
+4  40
+3  30
+2  20
+1  10
+
+statement ok
+DROP FUNCTION f_record_cast;
+
+# Test that ordering is preserved with array_agg
+statement ok
+CREATE FUNCTION f_out_for_agg(OUT a INT, OUT b INT) RETURNS SETOF RECORD AS $$
+  SELECT x, y FROM test_ordering ORDER BY x DESC;
+$$ LANGUAGE SQL;
+
+query T
+SELECT array_agg(f) FROM f_out_for_agg() AS f;
+----
+{"(4,40)","(3,30)","(2,20)","(1,10)"}
+
+statement ok
+DROP FUNCTION f_out_for_agg;
+
+# Test that ordering is preserved when a UDF returns a tuple and is used as a
+# data source. This exercises the expandRoutineTupleIntoCols code path.
+statement ok
+CREATE FUNCTION f_tuple() RETURNS SETOF RECORD AS $$
+  SELECT ROW(x, y) FROM test_ordering ORDER BY x DESC;
+$$ LANGUAGE SQL;
+
+query II nosort
+SELECT * FROM f_tuple() AS f(a INT, b INT);
+----
+4  40
+3  30
+2  20
+1  10
+
+statement ok
+DROP FUNCTION f_tuple;
+DROP TABLE test_ordering;
+
+subtest end

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 {
@@ -677,6 +678,7 @@ func (b *Builder) expandRoutineTupleIntoCols(stmtScope *scope) *scope {
 	}
 	tupleColID := stmtScope.cols[0].id
 	outScope := stmtScope.push()
+	outScope.copyOrdering(stmtScope)
 	colTyp := b.factory.Metadata().ColumnMeta(tupleColID).Type
 	for i := range colTyp.TupleContents() {
 		varExpr := b.factory.ConstructVariable(tupleColID)
@@ -721,6 +723,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]) {

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