Merge #158162

158162: optbuilder: preserve ORDER BY in set-returning UDFs with OUT parameters r=drewkimball a=ajstorm

Fixes #144013

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.

## Solution

This PR fixes the issue by calling `copyOrdering()` in two places in `pkg/sql/opt/optbuilder/routine.go`:

1. **`combineRoutineColsIntoTuple`** - preserves ordering when wrapping columns into a tuple
2. **`maybeAddRoutineAssignmentCasts`** - preserves ordering 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.

## Testing

Added a new regression test `out_params_ordering` in `pkg/sql/logictest/testdata/logic_test/udf_setof` that verifies ORDER BY is preserved for both ascending and descending order with OUT parameters.

## Release note

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.

Co-authored-by: Adam Storm <storm@cockroachlabs.com>

Author: craig[bot]

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

@@ -658,6 +658,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 {
@@ -682,6 +683,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)
@@ -726,6 +728,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