internal/postgres: improve vector search quality

Incorporate package popularity directly into the vector candidate
ranking CTE before computing the Reciprocal Rank Fusion to prevent
low-import package from diluting well-known packages.

Previously, vector search ranked candidates exclusively by semantic distance. This allowed low-quality or obscure packages to dominate the top vector results, pushing popular packages out of the aggregation window.

There is now this approach: 1. Fetch the top N nearest neighbors for the query using the index. 2. Rerank the candidates using a combined score that accounts for cosine similarity and popularity.

This ensures that the vector rank (rank_vec) provided better reflects package quality.

Change-Id: I1a5861bb5ba8268bf0c10fb129a2e0a3d7a8980b
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/815060
kokoro-CI: kokoro <noreply+kokoro@google.com>
Auto-Submit: Ethan Lee <ethanalee@google.com>
LUCI-TryBot-Result: golang-scoped@luci-project-accounts.iam.gserviceaccount.com <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>

internal/postgres/search.go

diff --git a/internal/postgres/search.go b/internal/postgres/search.go
index 56e0d03..6807d4e 100644
--- a/internal/postgres/search.go
+++ b/internal/postgres/search.go
@@ -315,6 +315,7 @@
 func buildVectorSearchQuery(opts SearchOptions, scoreExprStr string, q string, limit int) (string, []any) {
 	candidateLimit := max(100, opts.Offset+limit)
 	vectorWeight := sanitizeFloat(opts.ScoringParams.VectorWeight, 1.0)
+	popWeight := sanitizeFloat(opts.ScoringParams.PopularityWeight, 1.0)
 
 	query := fmt.Sprintf(`
 		WITH text_search AS (
@@ -327,11 +328,19 @@
 		),
 		vector_search AS (
 			SELECT package_path, version, module_path, commit_time, imported_by_count,
-				ROW_NUMBER() OVER (ORDER BY embedding <=> $2::halfvec, imported_by_count DESC) AS rank_vec
-			FROM search_documents
-			WHERE imported_by_count >= 1 AND embedding IS NOT NULL
-			ORDER BY embedding <=> $2::halfvec
-			LIMIT %d
+				ROW_NUMBER() OVER (
+					ORDER BY ((1.0 - (embedding <=> $2::halfvec)) * pow(ln(exp(1)+imported_by_count), %f) *
+					         CASE WHEN redistributable THEN 1 ELSE %f END *
+					         CASE WHEN COALESCE(has_go_mod, true) THEN 1 ELSE %f END) DESC,
+					         commit_time DESC, package_path
+				) AS rank_vec
+			FROM (
+				SELECT package_path, version, module_path, commit_time, imported_by_count, redistributable, has_go_mod, embedding
+				FROM search_documents
+				WHERE imported_by_count >= 1 AND embedding IS NOT NULL
+				ORDER BY embedding <=> $2::halfvec
+				LIMIT %d
+			) nn
 		),
 		combined AS (
 			SELECT
@@ -349,7 +358,7 @@
 		FROM combined
 		ORDER BY score DESC, commit_time DESC, package_path
 		LIMIT $3
-		OFFSET $4`, scoreExprStr, candidateLimit, candidateLimit, vectorWeight)
+		OFFSET $4`, scoreExprStr, candidateLimit, popWeight, nonRedistributablePenalty, noGoModPenalty, candidateLimit, vectorWeight)
 	args := []any{q, formatVector(opts.Vector), limit, opts.Offset}
 	return query, args
 }

internal/postgres/search_test.go

diff --git a/internal/postgres/search_test.go b/internal/postgres/search_test.go
index d0467cf..6a75610 100644
--- a/internal/postgres/search_test.go
+++ b/internal/postgres/search_test.go
@@ -1641,6 +1641,12 @@
 	if !strings.Contains(queryDefault, "LIMIT 100") {
 		t.Errorf("expected default candidate limit 100, query was:\n%s", queryDefault)
 	}
+	if !strings.Contains(queryDefault, fmt.Sprintf("CASE WHEN redistributable THEN 1 ELSE %f END", nonRedistributablePenalty)) {
+		t.Errorf("expected query to contain redistributable penalty (%f)", nonRedistributablePenalty)
+	}
+	if !strings.Contains(queryDefault, fmt.Sprintf("CASE WHEN COALESCE(has_go_mod, true) THEN 1 ELSE %f END", noGoModPenalty)) {
+		t.Errorf("expected query to contain no-go-mod penalty (%f)", noGoModPenalty)
+	}
 	if len(args) != 4 {
 		t.Errorf("expected 4 query args, got %d", len(args))
 	}