perf(query): use quickselect instead of sorting while pagination

types/select.go

@@ -0,0 +1,111 @@
+/*
+ * Copyright 2023 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package types
+
+// Below functions are taken from go's sort library zsortinterface.go
+// https://go.dev/src/sort/zsortinterface.go
+func insertionSort(data byValue, a, b int) {
+	for i := a + 1; i < b; i++ {
+		for j := i; j > a && data.Less(j, j-1); j-- {
+			data.Swap(j, j-1)
+		}
+	}
+}
+
+func order2(data byValue, a, b int) (int, int) {
+	if data.Less(b, a) {
+		return b, a
+	}
+	return a, b
+}
+
+func median(data byValue, a, b, c int) int {
+	a, b = order2(data, a, b)
+	b, _ = order2(data, b, c)
+	_, b = order2(data, a, b)
+	return b
+}
+
+func medianAdjacent(data byValue, a int) int {
+	return median(data, a-1, a, a+1)
+}
+
+// [shortestNinther,∞): uses the Tukey ninther method.
+func choosePivot(data byValue, a, b int) (pivot int) {
+	const (
+		shortestNinther = 50
+		maxSwaps        = 4 * 3
+	)
+
+	l := b - a
+
+	var (
+		i = a + l/4*1
+		j = a + l/4*2
+		k = a + l/4*3
+	)
+
+	if l >= 8 {
+		if l >= shortestNinther {
+			// Tukey ninther method, the idea came from Rust's implementation.
+			i = medianAdjacent(data, i)
+			j = medianAdjacent(data, j)
+			k = medianAdjacent(data, k)
+		}
+		// Find the median among i, j, k and stores it into j.
+		j = median(data, i, j, k)
+	}
+
+	return j
+}
+
+func partition(data byValue, a, b, pivot int) int {
+	partitionIndex := a
+	data.Swap(pivot, b)
+	for i := a; i < b; i++ {
+		if data.Less(i, b) {
+			data.Swap(i, partitionIndex)
+			partitionIndex++
+		}
+	}
+	data.Swap(partitionIndex, b)
+	return partitionIndex
+}
+
+func quickSelect(data byValue, low, high, k int) {
+	var pivotIndex int
+
+	for {
+		if low >= high {
+			return
+		} else if high-low <= 8 {
+			insertionSort(data, low, high+1)
+			return
+		}
+
+		pivotIndex = choosePivot(data, low, high)
+		pivotIndex = partition(data, low, high, pivotIndex)
+
+		if k < pivotIndex {
+			high = pivotIndex - 1
+		} else if k > pivotIndex {
+			low = pivotIndex + 1
+		} else {
+			return
+		}
+	}
+}

types/sort.go

@@ -50,6 +50,11 @@ func (s sortBase) Swap(i, j int) {
 
 type byValue struct{ sortBase }
 
+func (s byValue) isNil(i int) bool {
+	first := s.values[i]
+	return len(first) == 0 || first[0].Value == nil
+}
+
 // Less compares two elements
 func (s byValue) Less(i, j int) bool {
 	first, second := s.values[i], s.values[j]
@@ -96,6 +101,52 @@ func IsSortable(tid TypeID) bool {
 	}
 }
 
+// SortTopN finds and places the first n elements in 0-N
+func SortTopN(v [][]Val, ul *[]uint64, desc []bool, lang string, n int) error {
+	if len(v) == 0 || len(v[0]) == 0 {
+		return nil
+	}
+
+	for _, val := range v[0] {
+		if !IsSortable(val.Tid) {
+			return errors.Errorf("Value of type: %v isn't sortable", val.Tid.Name())
+		}
+	}
+
+	var cl *collate.Collator
+	if lang != "" {
+		// Collator is nil if we are unable to parse the language.
+		// We default to bytewise comparison in that case.
+		if langTag, err := language.Parse(lang); err == nil {
+			cl = collate.New(langTag)
+		}
+	}
+
+	b := sortBase{v, desc, ul, nil, cl}
+	toBeSorted := byValue{b}
+
+	nul := 0
+	for i := 0; i < len(*ul); i++ {
+		if toBeSorted.isNil(i) {
+			continue
+		}
+		if i != nul {
+			toBeSorted.Swap(i, nul)
+		}
+		nul += 1
+	}
+
+	if nul > n {
+		b1 := sortBase{v[:nul], desc, ul, nil, cl}
+		toBeSorted1 := byValue{b1}
+		quickSelect(toBeSorted1, 0, nul-1, n)
+	}
+	toBeSorted.values = toBeSorted.values[:n]
+	sort.Sort(toBeSorted)
+
+	return nil
+}
+
 // SortWithFacet sorts the given array in-place and considers the given facets to calculate
 // the proper ordering.
 func SortWithFacet(v [][]Val, ul *[]uint64, l []*pb.Facets, desc []bool, lang string) error {

types/sort_test.go

@@ -17,7 +17,10 @@
 package types
 
 import (
+	"fmt"
+	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -53,6 +56,91 @@ func getUIDList(n int) *pb.List {
 	return &pb.List{Uids: data}
 }
 
+const charset = "abcdefghijklmnopqrstuvwxyz"
+
+func StringWithCharset(length int) string {
+	var seededRand *rand.Rand = rand.New(rand.NewSource(time.Now().UnixNano()))
+	b := make([]byte, length)
+	for i := range b {
+		b[i] = charset[seededRand.Intn(len(charset))]
+	}
+	return string(b)
+}
+
+func TestQuickSelect(t *testing.T) {
+	n := 10000
+	k := 10
+	getList := func() [][]Val {
+		strs := make([]string, n)
+		for i := 0; i < n; i++ {
+			strs[i] = fmt.Sprintf("%d", rand.Intn(100000))
+		}
+
+		list := make([][]Val, len(strs))
+		for i, s := range strs {
+			va := Val{StringID, []byte(s)}
+			v, _ := Convert(va, IntID)
+			list[i] = []Val{v}
+		}
+
+		return list
+	}
+
+	ul := getUIDList(n)
+	list := getList()
+	require.NoError(t, SortTopN(list, &ul.Uids, []bool{false}, "", k))
+
+	for i := 0; i < k; i++ {
+		for j := k; j < n; j++ {
+			require.Equal(t, list[i][0].Value.(int64) <= list[j][0].Value.(int64), true)
+		}
+	}
+
+}
+
+func BenchmarkSortQuickSort(b *testing.B) {
+	n := 1000000
+	getList := func() [][]Val {
+		strs := make([]string, n)
+		for i := 0; i < n; i++ {
+			strs[i] = StringWithCharset(10)
+		}
+
+		list := make([][]Val, len(strs))
+		for i, s := range strs {
+			va := Val{StringID, []byte(s)}
+			v, _ := Convert(va, StringID)
+			list[i] = []Val{v}
+		}
+
+		return list
+	}
+
+	b.Run(fmt.Sprintf("Normal Sort Ratio %d ", n), func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			ul := getUIDList(n)
+			list := getList()
+			b.ResetTimer()
+			err := Sort(list, &ul.Uids, []bool{false}, "")
+			b.StopTimer()
+			require.NoError(b, err)
+		}
+	})
+
+	for j := 1; j < n; j += n / 6 {
+		b.Run(fmt.Sprintf("QuickSort Sort Ratio %d %d", n, j), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				ul := getUIDList(n)
+				list := getList()
+				b.ResetTimer()
+				err := SortTopN(list, &ul.Uids, []bool{false}, "", j)
+				b.StopTimer()
+				require.NoError(b, err)
+			}
+		})
+	}
+}
+
 func TestSortStrings(t *testing.T) {
 	list := getInput(t, StringID, []string{"bb", "aaa", "aa", "bab"})
 	ul := getUIDList(4)

worker/sort.go

@@ -453,12 +453,20 @@ func multiSort(ctx context.Context, r *sortresult, ts *pb.SortMessage) error {
 			x.AssertTrue(idx >= 0)
 			vals[j] = sortVals[idx]
 		}
-		//nolint:gosec
-		if err := types.Sort(vals, &ul.Uids, desc, ""); err != nil {
-			return err
-		}
-		// Paginate
+
 		start, end := x.PageRange(int(ts.Count), int(r.multiSortOffsets[i]), len(ul.Uids))
+		if end < len(ul.Uids)/2 {
+			//nolint:gosec
+			if err := types.SortTopN(vals, &ul.Uids, desc, "", end); err != nil {
+				return err
+			}
+		} else {
+			//nolint:gosec
+			if err := types.Sort(vals, &ul.Uids, desc, ""); err != nil {
+				return err
+			}
+		}
+
 		ul.Uids = ul.Uids[start:end]
 		r.reply.UidMatrix[i] = ul
 	}