Merge pull request #53 from philippgille/improve-query-sorting

Improve query sorting

CHANGELOG.md

@@ -8,7 +8,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 vNext
 -----
 
-In this release the main feature is query performance improvement (4x faster, 99% fewer memory allocations). There's also a new code example for semantic search across 5,000 arXiv papers.
+In this release the main feature is query performance improvement (5x faster, 99% fewer memory allocations). There's also a new code example for semantic search across 5,000 arXiv papers.
 
 ### Added
 
@@ -18,7 +18,7 @@ In this release the main feature is query performance improvement (4x faster, 99
 ### Improved
 
 - Changed the example link target to directory instead of `main.go` file (PR [#43](https://github.com/philippgille/chromem-go/pull/43))
-- Improved query performance (4x faster, 99% fewer memory allocations) (PR [#47](https://github.com/philippgille/chromem-go/pull/47))
+- Improved query performance (5x faster, 99% fewer memory allocations) (PR [#47](https://github.com/philippgille/chromem-go/pull/47), [#53](https://github.com/philippgille/chromem-go/pull/53))
 
 ### Fixed
 

README.md

@@ -11,7 +11,7 @@ Because `chromem-go` is embeddable it enables you to add retrieval augmented gen
 
 It's *not* a library to connect to Chroma and also not a reimplementation of it in Go. It's a database on its own.
 
-The focus is not scale (millions of documents) or number of features, but simplicity and performance for the most common use cases. On a mid-range 2020 Intel laptop CPU you can query 1,000 documents in 0.5 ms and 100,000 documents in 50-60 ms, both with just 44 memory allocations. See [Benchmarks](#benchmarks) for details.
+The focus is not scale (millions of documents) or number of features, but simplicity and performance for the most common use cases. On a mid-range 2020 Intel laptop CPU you can query 1,000 documents in 0.3 ms and 100,000 documents in 50-60 ms, both with just 41 memory allocations. See [Benchmarks](#benchmarks) for details.
 
 > ⚠️ The project is in beta, under heavy construction, and may introduce breaking changes in releases before `v1.0.0`. All changes are documented in the [`CHANGELOG`](./CHANGELOG.md).
 
@@ -197,18 +197,18 @@ goos: linux
 goarch: amd64
 pkg: github.com/philippgille/chromem-go
 cpu: 11th Gen Intel(R) Core(TM) i5-1135G7 @ 2.40GHz
-BenchmarkCollection_Query_NoContent_100-8          10000     109957 ns/op     6487 B/op       44 allocs/op
-BenchmarkCollection_Query_NoContent_1000-8          2043     541337 ns/op    35667 B/op       44 allocs/op
-BenchmarkCollection_Query_NoContent_5000-8           361    4230542 ns/op   166729 B/op       44 allocs/op
-BenchmarkCollection_Query_NoContent_25000-8           79   16343977 ns/op   813902 B/op       44 allocs/op
-BenchmarkCollection_Query_NoContent_100000-8          18   63417540 ns/op  3205961 B/op       44 allocs/op
-BenchmarkCollection_Query_100-8                    10861     110167 ns/op     6480 B/op       44 allocs/op
-BenchmarkCollection_Query_1000-8                    2146     540489 ns/op    35669 B/op       44 allocs/op
-BenchmarkCollection_Query_5000-8                     442    4913970 ns/op   166748 B/op       44 allocs/op
-BenchmarkCollection_Query_25000-8                     88   15213089 ns/op   813909 B/op       44 allocs/op
-BenchmarkCollection_Query_100000-8                    19   55963675 ns/op  3206027 B/op       44 allocs/op
+BenchmarkCollection_Query_NoContent_100-8          10000     106441 ns/op     6393 B/op       41 allocs/op
+BenchmarkCollection_Query_NoContent_1000-8          2278     494254 ns/op    35570 B/op       41 allocs/op
+BenchmarkCollection_Query_NoContent_5000-8           416    2767125 ns/op   166634 B/op       41 allocs/op
+BenchmarkCollection_Query_NoContent_25000-8           70   15165139 ns/op   813800 B/op       41 allocs/op
+BenchmarkCollection_Query_NoContent_100000-8          19   58823464 ns/op  3205865 B/op       41 allocs/op
+BenchmarkCollection_Query_100-8                    11269     105990 ns/op     6385 B/op       41 allocs/op
+BenchmarkCollection_Query_1000-8                    2364     494212 ns/op    35574 B/op       41 allocs/op
+BenchmarkCollection_Query_5000-8                     481    2750438 ns/op   166647 B/op       41 allocs/op
+BenchmarkCollection_Query_25000-8                     93   13143419 ns/op   813805 B/op       41 allocs/op
+BenchmarkCollection_Query_100000-8                    20   51727357 ns/op  3205871 B/op       41 allocs/op
 PASS
-ok   github.com/philippgille/chromem-go 27.887s
+ok   github.com/philippgille/chromem-go 26.187s
 ```
 
 ## Motivation

collection.go

@@ -1,13 +1,13 @@
 package chromem
 
 import (
+	"cmp"
 	"context"
 	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
 	"slices"
-	"sort"
 	"sync"
 )
 
@@ -324,29 +324,29 @@ func (c *Collection) QueryEmbedding(ctx context.Context, queryEmbedding []float3
 	}
 
 	// For the remaining documents, calculate cosine similarity.
-	docSim, err := calcDocSimilarity(ctx, queryEmbedding, filteredDocs)
+	docSims, err := calcDocSimilarity(ctx, queryEmbedding, filteredDocs)
 	if err != nil {
 		return nil, fmt.Errorf("couldn't calculate cosine similarity: %w", err)
 	}
 
 	// Sort by similarity
-	sort.Slice(docSim, func(i, j int) bool {
-		// The `less` function would usually use `<`, but we want to sort descending.
-		return docSim[i].similarity > docSim[j].similarity
+	slices.SortFunc(docSims, func(i, j docSim) int {
+		// i, j; for descending order
+		return cmp.Compare(j.similarity, i.similarity)
 	})
 
 	// Return the top nResults or len(docSim), whichever is smaller
-	if len(docSim) < nResults {
-		nResults = len(docSim)
+	if len(docSims) < nResults {
+		nResults = len(docSims)
 	}
 	res := make([]Result, 0, nResults)
 	for i := 0; i < nResults; i++ {
 		res = append(res, Result{
-			ID:         docSim[i].docID,
-			Metadata:   c.documents[docSim[i].docID].Metadata,
-			Embedding:  c.documents[docSim[i].docID].Embedding,
-			Content:    c.documents[docSim[i].docID].Content,
-			Similarity: docSim[i].similarity,
+			ID:         docSims[i].docID,
+			Metadata:   c.documents[docSims[i].docID].Metadata,
+			Embedding:  c.documents[docSims[i].docID].Embedding,
+			Content:    c.documents[docSims[i].docID].Content,
+			Similarity: docSims[i].similarity,
 		})
 	}