Go语言Partial Sort性能优化方案求助

我正在将Spotify的Annoy库移植到Go语言，项目名为goannoy，目前在排序环节（尤其是对标C++ std::partial_sort的部分排序功能）遭遇性能瓶颈。

我自行实现了两种部分排序方案：一种基于选择排序结合标准排序，另一种基于二叉堆，但在处理百万级规模的大数组时性能表现很差，推测核心原因是GC压力过大。现寻求**高性能、低内存分配（或零额外分配）**的Go语言partial sort实现方向或成熟库。

以下是我的实现代码：

排序用Pair结构体及优先队列实现 #

type Pair[T constraints.Ordered, S constraints.Ordered] struct {
    First  T
    Second S
}

func (p *Pair[T, S]) Less(other *Pair[T, S]) bool {
    return p.First < other.First ||
        (p.First == other.First && p.Second < other.Second)
}

type Pairs[T constraints.Ordered, S constraints.Ordered] []*Pair[T, S]

func (pq Pairs[_, _]) Len() int {
    return len(pq)
}

func (pq Pairs[_, _]) Less(i, j int) bool {
    return pq[i].First < pq[j].First ||
        (pq[i].First == pq[j].First && pq[i].Second < pq[j].Second)
}

func (pq Pairs[_, _]) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
}

func (pq *Pairs[TP, TV]) Push(x interface{}) {
    *pq = append(*pq, x.(*Pair[TP, TV]))
}

func (pq *Pairs[_, _]) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    *pq = old[:n-1]
    return item
}

func (pq *Pairs[TP, TV]) Top() *Pair[TP, TV] {
    if len(*pq) == 0 {
        return nil
    }

    return (*pq)[0]
}

第一种部分排序实现 #

func PartialSortSlice[TV interfaces.VectorType, TIX interfaces.IndexTypes](
    s []*Pair[TV, TIX],
    begin, middle, end int,
) {
    if begin >= end || middle <= begin || middle > end {
        return
    }

    // Find the N smallest elements
    N := middle - begin

    if end-begin > 20 && end-begin < 5000000 {
        SortPairs(s)
        return
    }

    for i := 0; i < N; i++ {
        minIndex := begin + i

        // Find the index of the smallest element in the unsorted part
        for j := begin + i + 1; j < end; j++ {
            if s[j].Less(s[minIndex]) {
                minIndex = j
            }
        }

        // Swap elements
        if minIndex != begin+i {
            s[begin+i], s[minIndex] = s[minIndex], s[begin+i]
        }
    }

    // Sort sub-range [begin, middle)
    if N > 15 {
        SortPairs(s[begin:middle])
    } else {
        for i := begin + 1; i < middle; i++ {
            for j := i; j > begin && s[j].Less(s[j-1]); j-- {
                s[j], s[j-1] = s[j-1], s[j]
            }
        }
    }
}

SortPairs实现 #

func SortPairs[TV interfaces.VectorType, TIX interfaces.IndexTypes](
    pairs []*Pair[TV, TIX],
) {
    sort.Slice(pairs, func(i, j int) bool {
        return pairs[i].Less(pairs[j])
    })
}

第二种基于二叉堆的部分排序实现 #

func PartialSortSlice2[TV interfaces.VectorType, TIX interfaces.IndexTypes](
    s []*Pair[TV, TIX],
    begin, middle, end int,
) {
    if begin >= end || middle <= begin || middle > end {
        return
    }

    // Find the N smallest elements using a binary heap
    N := middle - begin
    h := pmnh[TV, TIX]{indices: make([]int, N), data: s}
    for i := 0; i < N; i++ {
        h.indices[i] = i + begin
    }
    heap.Init(&h)
    for i := N; i < end-begin; i++ {
        if s[begin+i].Less(s[h.indices[0]]) {
            h.indices[0] = i + begin
            heap.Fix(&h, 0)
        }
    }

    // Swap elements
    for i := 0; i < N; i++ {
        s[begin+i], s[h.indices[i]-begin] = s[h.indices[i]-begin], s[begin+i]
    }

    // Sort sub-range [begin, middle) in place
    for i := begin + 1; i < middle; i++ {
        for j := i; j > begin && s[j].Less(s[j-1]); j-- {
            s[j], s[j-1] = s[j-1], s[j]
        }
    }
}

性能问题复现命令 #

git clone https://github.com/mariotoffia/goannoy.git
cd goannoy
go run cmd/precision/main.go -file -items 1000000 -prec 1000 -cpu-profile
cat results.txt

（使用-mem-profile参数可查看内存分配情况）

内容的提问来源于stack exchange，提问作者Mario Toffia