Fix heap sorting

query/raw_result.go

@@ -113,9 +113,6 @@ func NewRawResultHeap(
 	}
 }
 
-// Data returns the underlying array backing the heap.
-func (h RawResultHeap) Data() []RawResult { return h.dv }
-
 // Min returns the "smallest" heap element according to the `lessThan` function.
 func (h RawResultHeap) Min() RawResult { return h.dv[0] }
 
@@ -149,6 +146,20 @@ func (h *RawResultHeap) Pop() RawResult {
 	return val
 }
 
+// SortInPlace sorts the heap in place and returns the sorted data,
+// with the smallest element at the end of the returned array.
+// NB: The heap becomes invalid after this is called.
+func (h *RawResultHeap) SortInPlace() []RawResult {
+	numElems := len(h.dv)
+	for len(h.dv) > 0 {
+		h.Pop()
+	}
+	res := h.dv[:numElems]
+	h.dv = nil
+	h.lessThanFn = nil
+	return res
+}
+
 func (h RawResultHeap) shiftUp(i int) {
 	for {
 		parent := (i - 1) / 2

query/raw_result_test.go

@@ -0,0 +1,36 @@
+package query
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestRawResultHeapSortInPlace(t *testing.T) {
+	input := []RawResult{
+		{
+			Data:  "foo",
+			DocID: 145,
+		},
+		{
+			Data:  "bar",
+			DocID: 34,
+		},
+		{
+			Data:  "baz",
+			DocID: 69,
+		},
+		{
+			Data:  "cat",
+			DocID: 254,
+		},
+	}
+	lessThanFn := func(v1, v2 RawResult) bool { return v1.DocID < v2.DocID }
+	h := NewRawResultHeap(0, lessThanFn)
+	for _, r := range input {
+		h.Push(r)
+	}
+	sortedResults := h.SortInPlace()
+	expected := []RawResult{input[3], input[0], input[2], input[1]}
+	require.Equal(t, expected, sortedResults)
+}

storage/segment_query.go

@@ -340,10 +340,7 @@ func collectTopNRawResultDocIDOrderByValues(
 	// Sort the result heap in place, and when done the items are sorted from left to
 	// in the right order based on the query sorting criteria (i.e., if the sort order
 	// is ascending, the leftmost item is the smallest item).
-	for results.Len() > 0 {
-		results.Pop()
-	}
-	return results.Data(), nil
+	return results.SortInPlace(), nil
 }
 
 // collectTopNRawResults collects the top N raw results from the raw doc source field