Use ordinals to force stable TopN heap sort

enginetest/queries/script_queries.go

@@ -81,6 +81,39 @@ type ScriptTestAssertion struct {
 // Unlike other engine tests, ScriptTests must be self-contained. No other tables are created outside the definition of
 // the tests.
 var ScriptTests = []ScriptTest{
+	{
+		Name: "topN stable output",
+		SetUpScript: []string{
+			"create table xy (x int primary key, y int)",
+			"insert into xy values (1,0),(2,0),(3,0),(4,0)",
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query:    "select * from xy order by y asc limit 1",
+				Expected: []sql.Row{{1, 0}},
+			},
+			{
+				Query:    "select * from xy order by y asc limit 1 offset 1",
+				Expected: []sql.Row{{2, 0}},
+			},
+			{
+				Query:    "select * from xy order by y asc limit 1 offset 2",
+				Expected: []sql.Row{{3, 0}},
+			},
+			{
+				Query:    "select * from xy order by y asc limit 1 offset 3",
+				Expected: []sql.Row{{4, 0}},
+			},
+			{
+				Query:    "(select * from xy order by y asc limit 1 offset 1) union (select * from xy order by y asc limit 1 offset 2)",
+				Expected: []sql.Row{{2, 0}, {3, 0}},
+			},
+			{
+				Query:    "with recursive cte as ((select * from xy order by y asc limit 1 offset 1) union (select * from xy order by y asc limit 1 offset 2)) select * from cte",
+				Expected: []sql.Row{{2, 0}, {3, 0}},
+			},
+		},
+	},
 	{
 		Name: "enums with default, case-sensitive collation (utf8mb4_0900_bin)",
 		SetUpScript: []string{

sql/plan/recursive_cte.go

@@ -128,7 +128,7 @@ func (r *RecursiveCte) RowIter(ctx *sql.Context, row sql.Row) (sql.RowIter, erro
 		if err != nil {
 			return nil, err
 		}
-		iter = newTopRowsIter(r.union.SortFields, limit, false, iter)
+		iter = newTopRowsIter(r.union.SortFields, limit, false, iter, len(r.union.Schema()))
 	} else if r.union.Limit != nil {
 		limit, err := getInt64Value(ctx, r.union.Limit)
 		if err != nil {

sql/plan/sort.go

@@ -21,6 +21,8 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/dolthub/go-mysql-server/sql/types"
+
 	"github.com/dolthub/go-mysql-server/sql"
 	"github.com/dolthub/go-mysql-server/sql/expression"
 )
@@ -320,7 +322,7 @@ func (n *TopN) RowIter(ctx *sql.Context, row sql.Row) (sql.RowIter, error) {
 	if err != nil {
 		return nil, err
 	}
-	return sql.NewSpanIter(span, newTopRowsIter(n.Fields, limit, n.CalcFoundRows, i)), nil
+	return sql.NewSpanIter(span, newTopRowsIter(n.Fields, limit, n.CalcFoundRows, i, len(n.Child.Schema()))), nil
 }
 
 func (n *TopN) String() string {
@@ -397,9 +399,9 @@ type topRowsIter struct {
 	idx           int
 }
 
-func newTopRowsIter(s sql.SortFields, limit int64, calcFoundRows bool, child sql.RowIter) *topRowsIter {
+func newTopRowsIter(s sql.SortFields, limit int64, calcFoundRows bool, child sql.RowIter, childSchemaLen int) *topRowsIter {
 	return &topRowsIter{
-		sortFields:    s,
+		sortFields:    append(s, sql.SortField{Column: expression.NewGetField(childSchemaLen, types.Int64, "order", false)}),
 		limit:         limit,
 		calcFoundRows: calcFoundRows,
 		childIter:     child,
@@ -421,7 +423,7 @@ func (i *topRowsIter) Next(ctx *sql.Context) (sql.Row, error) {
 	}
 	row := i.topRows[i.idx]
 	i.idx++
-	return row, nil
+	return row[:len(row)-1], nil
 }
 
 func (i *topRowsIter) Close(ctx *sql.Context) error {
@@ -453,6 +455,8 @@ func (i *topRowsIter) computeTopRows(ctx *sql.Context) error {
 		}
 		i.numFoundRows++
 
+		row = append(row, i.numFoundRows)
+
 		heap.Push(topRowsHeap, row)
 		if int64(topRowsHeap.Len()) > i.limit {
 			heap.Pop(topRowsHeap)

sql/plan/union.go

@@ -143,7 +143,7 @@ func (u *Union) RowIter(ctx *sql.Context, row sql.Row) (sql.RowIter, error) {
 		if err != nil {
 			return nil, err
 		}
-		iter = newTopRowsIter(u.SortFields, limit, false, iter)
+		iter = newTopRowsIter(u.SortFields, limit, false, iter, len(u.Schema()))
 	} else if u.Limit != nil {
 		limit, err := getInt64Value(ctx, u.Limit)
 		if err != nil {