From 3a92f567b9204e6b2859865c5684a43366af07c1 Mon Sep 17 00:00:00 2001
From: Fei Han <hanfei19910905@gmail.com>
Date: Tue, 12 Apr 2016 12:09:25 +0800
Subject: [PATCH 1/2] executor: support limit sort

---
 executor/builder.go           |  1 +
 executor/executor.go          | 17 +++++++++++++++++
 executor/executor_test.go     | 26 ++++++++++++++++++++++----
 optimizer/plan/planbuilder.go |  4 ++++
 optimizer/plan/plans.go       |  2 ++
 5 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/executor/builder.go b/executor/builder.go
index 84e16e7b06..23785d173a 100644
--- a/executor/builder.go
+++ b/executor/builder.go
@@ -332,6 +332,7 @@ func (b *executorBuilder) buildSort(v *plan.Sort) Executor {
 		Src:     src,
 		ByItems: v.ByItems,
 		ctx:     b.ctx,
+		Limit:   v.ExecLimit,
 	}
 	return e
 }
diff --git a/executor/executor.go b/executor/executor.go
index 28551b8d0d..4b76b2e9fd 100644
--- a/executor/executor.go
+++ b/executor/executor.go
@@ -856,6 +856,7 @@ type SortExec struct {
 	ByItems []*ast.ByItem
 	Rows    []*orderByRow
 	ctx     context.Context
+	Limit   *plan.Limit
 	Idx     int
 	fetched bool
 	err     error
@@ -902,9 +903,18 @@ func (e *SortExec) Less(i, j int) bool {
 	return false
 }
 
+// SortBufferSize represents the total extra row count that sort can use.
+var SortBufferSize = 500
+
 // Next implements Executor Next interface.
 func (e *SortExec) Next() (*Row, error) {
 	if !e.fetched {
+		offset := -1
+		totalCount := -1
+		if e.Limit != nil {
+			offset = int(e.Limit.Offset)
+			totalCount = offset + int(e.Limit.Count)
+		}
 		for {
 			srcRow, err := e.Src.Next()
 			if err != nil {
@@ -924,8 +934,15 @@ func (e *SortExec) Next() (*Row, error) {
 				}
 			}
 			e.Rows = append(e.Rows, orderRow)
+			if totalCount != -1 && e.Len() >= totalCount+SortBufferSize {
+				sort.Sort(e)
+				e.Rows = e.Rows[:totalCount]
+			}
 		}
 		sort.Sort(e)
+		if offset > 0 {
+			e.Rows = e.Rows[offset:totalCount]
+		}
 		e.fetched = true
 	}
 	if e.err != nil {
diff --git a/executor/executor_test.go b/executor/executor_test.go
index de340f1641..48297b5892 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -15,10 +15,6 @@ package executor_test
 
 import (
 	"fmt"
-	"strings"
-	"testing"
-	"time"
-
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb"
 	"github.com/pingcap/tidb/domain"
@@ -29,6 +25,9 @@ import (
 	"github.com/pingcap/tidb/util/testkit"
 	"github.com/pingcap/tidb/util/testleak"
 	"github.com/pingcap/tidb/util/types"
+	"strings"
+	"testing"
+	"time"
 )
 
 func TestT(t *testing.T) {
@@ -594,6 +593,25 @@ func (s *testSuite) TestSelectOrderBy(c *C) {
 	r.Check(testkit.Rows(rowStr))
 	tk.MustExec("commit")
 
+	// Test limit + order by
+	tk.MustExec("begin")
+	executor.SortBufferSize = 10
+	for i := 3; i <= 10; i += 1 {
+		tk.MustExec(fmt.Sprintf("insert INTO select_order_test VALUES (%d, \"zz\");", i))
+	}
+	tk.MustExec("insert INTO select_order_test VALUES (10086, \"hi\");")
+	for i := 11; i <= 20; i += 1 {
+		tk.MustExec(fmt.Sprintf("insert INTO select_order_test VALUES (%d, \"hh\");", i))
+	}
+	for i := 21; i <= 30; i += 1 {
+		tk.MustExec(fmt.Sprintf("insert INTO select_order_test VALUES (%d, \"zz\");", i))
+	}
+	tk.MustExec("insert INTO select_order_test VALUES (1501, \"aa\");")
+	r = tk.MustQuery("select * from select_order_test order by name, id limit 1 offset 3;")
+	rowStr = fmt.Sprintf("%v %v", 11, []byte("hh"))
+	r.Check(testkit.Rows(rowStr))
+	tk.MustExec("commit")
+	executor.SortBufferSize = 500
 	tk.MustExec("drop table select_order_test")
 }
 
diff --git a/optimizer/plan/planbuilder.go b/optimizer/plan/planbuilder.go
index 2872053a70..bb0552e344 100644
--- a/optimizer/plan/planbuilder.go
+++ b/optimizer/plan/planbuilder.go
@@ -486,6 +486,10 @@ func (b *planBuilder) buildLimit(src Plan, limit *ast.Limit) Plan {
 		Offset: limit.Offset,
 		Count:  limit.Count,
 	}
+	if s, ok := src.(*Sort); ok {
+		s.ExecLimit = li
+		return s
+	}
 	li.SetSrc(src)
 	li.SetFields(src.Fields())
 	return li
diff --git a/optimizer/plan/plans.go b/optimizer/plan/plans.go
index 761ebee2b9..d7af39fc50 100644
--- a/optimizer/plan/plans.go
+++ b/optimizer/plan/plans.go
@@ -273,6 +273,8 @@ type Sort struct {
 	planWithSrc
 
 	ByItems []*ast.ByItem
+
+	ExecLimit *Limit
 }
 
 // Accept implements Plan Accept interface.

From 19310f1511208d931a5f0f275dc8afb42a64cc20 Mon Sep 17 00:00:00 2001
From: Fei Han <hanfei19910905@gmail.com>
Date: Tue, 12 Apr 2016 13:27:07 +0800
Subject: [PATCH 2/2] change format

---
 executor/executor_test.go | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/executor/executor_test.go b/executor/executor_test.go
index 48297b5892..3a1e0f8a25 100644
--- a/executor/executor_test.go
+++ b/executor/executor_test.go
@@ -15,6 +15,10 @@ package executor_test
 
 import (
 	"fmt"
+	"strings"
+	"testing"
+	"time"
+
 	. "github.com/pingcap/check"
 	"github.com/pingcap/tidb"
 	"github.com/pingcap/tidb/domain"
@@ -25,9 +29,6 @@ import (
 	"github.com/pingcap/tidb/util/testkit"
 	"github.com/pingcap/tidb/util/testleak"
 	"github.com/pingcap/tidb/util/types"
-	"strings"
-	"testing"
-	"time"
 )
 
 func TestT(t *testing.T) {