internal/postgres/search.go

// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package postgres

import (
	"context"
	"database/sql"
	"fmt"
	"math"
	"slices"
	"sort"
	"strings"
	"time"

	"github.com/lib/pq"
	"go.opencensus.io/plugin/ochttp"
	"go.opencensus.io/stats"
	"go.opencensus.io/stats/view"
	"go.opencensus.io/tag"
	"go.opencensus.io/trace"
	"golang.org/x/pkgsite/internal"
	"golang.org/x/pkgsite/internal/database"
	"golang.org/x/pkgsite/internal/dcensus"
	"golang.org/x/pkgsite/internal/derrors"
	"golang.org/x/pkgsite/internal/log"
	"golang.org/x/pkgsite/internal/postgres/search"
	"golang.org/x/pkgsite/internal/stdlib"
	"golang.org/x/pkgsite/internal/version"
)

var (
	// searchLatency holds observed latency in individual search queries.
	searchLatency = stats.Float64(
		"go-discovery/search/latency",
		"Latency of a search query.",
		stats.UnitMilliseconds,
	)
	// keySearchSource is a census tag for search query types.
	keySearchSource = tag.MustNewKey("search.source")
	// SearchLatencyDistribution aggregates search request latency by search
	// query type.
	SearchLatencyDistribution = &view.View{
		Name:        "go-discovery/search/latency",
		Measure:     searchLatency,
		Aggregation: ochttp.DefaultLatencyDistribution,
		Description: "Search latency, by result source query type.",
		TagKeys:     []tag.Key{keySearchSource},
	}
	// SearchResponseCount counts search responses by search query type.
	SearchResponseCount = &view.View{
		Name:        "go-discovery/search/count",
		Measure:     searchLatency,
		Aggregation: view.Count(),
		Description: "Search count, by result source query type.",
		TagKeys:     []tag.Key{keySearchSource},
	}
)

// searchResponse is used for internal bookkeeping when fanning-out search
// request to multiple different search queries.
type searchResponse struct {
	// source is a unique identifier for the search query type (e.g. 'deep',
	// 'popular'), to be used in logging and reporting.
	source string
	// results are partially filled out from only the search_documents table.
	results []*SearchResult
	// err indicates a technical failure of the search query, or that results are
	// not provably complete.
	err error
}

// searchEvent is used to log structured information about search events for
// later analysis. A 'search event' occurs when a searcher or count estimate
// returns.
type searchEvent struct {
	// Type is either the searcher name or 'estimate' (the count estimate).
	Type string
	// Latency is the duration that the operation took.
	Latency time.Duration
	// Err is the error returned by the operation, if any.
	Err error
}

// A searcher is used to execute a single search request.
type searcher func(db *DB, ctx context.Context, q string, limit int, opts SearchOptions) searchResponse

// The pkgSearchers used by Search.
var pkgSearchers = map[string]searcher{
	"popular": (*DB).popularSearch,
	"deep":    (*DB).deepSearch,
}

var symbolSearchers = map[string]searcher{
	"symbol": (*DB).symbolSearch,
}

type SearchOptions = internal.SearchOptions
type SearchScoringParams = internal.SearchScoringParams
type SearchResult = internal.SearchResult

// SearchSupport implements the DataSource interface, supporting all search
// types.
func (db *DB) SearchSupport() internal.SearchSupport { return internal.FullSearch }

// Search executes two search requests concurrently:
//   - a sequential scan of packages in descending order of popularity.
//   - all packages ("deep" search) using an inverted index to filter to search
//     terms.
//
// The sequential scan takes significantly less time when searching for very
// common terms (e.g. "errors", "cloud", or "kubernetes"), due to its ability
// to exit early once the requested page of search results is provably
// complete.
//
// Because 0 <= ts_rank() <= 1, we know that the highest score of any unscanned
// package is ln(e+N), where N is imported_by_count of the package we are
// currently considering.  Therefore if the lowest scoring result of popular
// search is greater than ln(e+N), we know that we haven't missed any results
// and can return the search result immediately, cancelling other searches.
//
// On the other hand, if the popular search is slow, it is likely that the
// search term is infrequent, and deep search will be fast due to our inverted
// gin index on search tokens.
//
// The gap in this optimization is search terms that are very frequent, but
// rarely relevant: "int" or "package", for example. In these cases we'll pay
// the penalty of a deep search that scans nearly every package.
func (db *DB) Search(ctx context.Context, q string, opts SearchOptions) (_ []*SearchResult, err error) {
	defer derrors.WrapStack(&err, "DB.Search(ctx, %q, %+v)", q, opts)
	if !opts.SearchSymbols {
		const (
			limitMultiplier1 = 3
			limitMultiplier2 = 5
		)
		// Limit search to more rows than the requested number of results, so
		// that it can find other packages in the modules it selects.
		srs, err := db.search(ctx, q, opts, limitMultiplier1*opts.MaxResults)
		if err != nil {
			return nil, err
		}
		if len(srs) >= opts.MaxResults || numRows(srs) <= limitMultiplier1*opts.MaxResults {
			return srs, nil
		}
		// Grouped search didn't find enough results, but there are more
		// rows that could potentially match. Try one more time, with a
		// larger limit.
		return db.search(ctx, q, opts, limitMultiplier2*opts.MaxResults)
	}
	return db.search(ctx, q, opts, opts.MaxResults)
}

func (db *DB) search(ctx context.Context, q string, opts SearchOptions, limit int) (_ []*SearchResult, err error) {
	defer derrors.WrapStack(&err, "search(limit=%d)", limit)

	var searchers map[string]searcher
	if opts.SearchSymbols {
		searchers = symbolSearchers
	} else if len(opts.Vector) > 0 || opts.ScoringParams != (SearchScoringParams{}) {
		// Vector search and custom scoring parameters are only supported in deepSearch.
		// Bypass popularSearch to prevent fast text searches from prematurely cancelling deepSearch.
		searchers = map[string]searcher{"deep": (*DB).deepSearch}
	} else {
		searchers = pkgSearchers
	}
	resp, err := db.hedgedSearch(ctx, q, limit, opts, searchers, nil)
	if err != nil {
		return nil, err
	}
	// Filter out excluded paths.
	var results []*SearchResult
	for _, r := range resp.results {
		if !db.IsExcluded(ctx, r.PackagePath, "") {
			results = append(results, r)
		}
	}
	if opts.GroupResults && !opts.SearchSymbols {
		results = groupSearchResults(results)
	}
	if len(results) > opts.MaxResults {
		results = results[:opts.MaxResults]
	}
	return results, nil
}

// Penalties to search scores, applied as multipliers to the score.
const (
	// Module license is non-redistributable.
	nonRedistributablePenalty = 0.5
	// Module does not have a go.mod file.
	// Start this off gently (close to 1), but consider lowering
	// it as time goes by and more of the ecosystem converts to modules.
	noGoModPenalty = 0.8

	// Normalization adjusts the ts_rank score by some function of the document size.
	// See https://www.postgresql.org/docs/11/textsearch-controls.html#TEXTSEARCH-PARSING-QUERIES,
	// section 12.3.3.
	// Zero means no normalization.
	// If this value is changed, then the popular_search stored procedure must be changed as well.
	// See migrations/000005_change_b_weight.up.sql.
	normalization = 0
)

func sanitizeFloat(f float64, fallback float64) float64 {
	if math.IsNaN(f) || math.IsInf(f, 0) {
		return fallback
	}
	return f
}

// scoreExpr returns the SQL expression that computes search document relevance scores,
// customized according to the provided SearchScoringParams.
//
// It is the product of:
//   - The Postgres ts_rank score, based the relevance of the document to the query.
//   - The log of the module's popularity, estimated by the number of importing packages.
//
// TODO(golang/go#80242): s/packages/modules/
//
//	  The log factor contains exp(1) so that it is always >= 1. Taking the log
//	  of imported_by_count instead of using it directly makes the effect less
//	  dramatic: being 2x as popular only has an additive effect.
//	- A penalty factor for non-redistributable modules, since a lot of
//	  details cannot be displayed.
//
// The first argument to ts_rank is an array of weights for the four tsvector sections,
// in the order D, C, B, A.
// The weights below match the defaults except for B.
// TODO(golang/go#80242): s/imported_by_count/imported_by_module_count/ when all rows are populated.
func scoreExpr(p SearchScoringParams) string {
	w := p.TextWeights
	weights := fmt.Sprintf("{%f, %f, %f, %f}",
		sanitizeFloat(w[0], 0.1), sanitizeFloat(w[1], 0.2),
		sanitizeFloat(w[2], 1.0), sanitizeFloat(w[3], 1.0))
	popWeight := sanitizeFloat(p.PopularityWeight, 1.0)
	return fmt.Sprintf(`
		ts_rank('%s', tsv_search_tokens, websearch_to_tsquery($1), %d) *
		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
	`, weights, normalization, popWeight, nonRedistributablePenalty, noGoModPenalty)
}

// hedgedSearch executes multiple search methods and returns the first
// available result.
// The optional guardTestResult func may be used to allow tests to control the
// order in which search results are returned.
func (db *DB) hedgedSearch(ctx context.Context, q string, limit int, opts SearchOptions, searchers map[string]searcher, guardTestResult func(string) func()) (_ *searchResponse, err error) {
	defer derrors.WrapStack(&err, "hedgedSearch(ctx, %q, %d, %+v)", q, limit, opts)

	searchStart := time.Now()
	responses := make(chan searchResponse, len(searchers))
	// cancel all unfinished searches when a result (or error) is returned. The
	// effectiveness of this depends on the database driver.
	searchCtx, cancel := context.WithCancel(ctx)
	defer cancel()

	// Fan out our search requests.
	for _, s := range searchers {
		go func() {
			start := time.Now()
			resp := s(db, searchCtx, q, limit, opts)
			log.Debug(ctx, searchEvent{
				Type:    resp.source,
				Latency: time.Since(start),
				Err:     resp.err,
			})
			if guardTestResult != nil {
				defer guardTestResult(resp.source)()
			}
			responses <- resp
		}()
	}
	// Note for future readers: in previous iterations of this code we kept
	// reading responses if the first one had an error, with the goal to minimize
	// error ratio. That didn't behave well if Postgres was overloaded.
	resp := <-responses
	if resp.err != nil {
		return nil, fmt.Errorf("%q search failed: %w", resp.source, resp.err)
	}
	// cancel proactively here: we've got the search result we need.
	cancel()
	// latency is only recorded for valid search results, as fast failures could
	// skew the latency distribution.
	// Note that this latency measurement might differ meaningfully from the
	// resp.Latency, if time was spent waiting for the result count estimate.
	stats.RecordWithTags(ctx,
		[]tag.Mutator{tag.Upsert(keySearchSource, resp.source)},
		dcensus.MDur(searchLatency, time.Since(searchStart)))
	// To avoid fighting with the query planner, our searches only hit the
	// search_documents table and we enrich after getting the results. In the
	// future, we may want to fully denormalize and put all search data in the
	// search_documents table.
	if _, ok := searchers["symbol"]; !ok {
		if err := db.addPackageDataToSearchResults(ctx, resp.results); err != nil {
			return nil, err
		}
	}
	return &resp, nil
}

const hllRegisterCount = 128

// buildVectorSearchQuery constructs a hybrid search SQL query combining full-text keyword search
// and semantic vector similarity using Reciprocal Rank Fusion (RRF).
//
// The query uses three Common Table Expressions (CTEs):
//  1. text_search: Retrieves top packages matching full-text query $1, ordered by ts_rank score (rank_text).
//  2. vector_search: Retrieves top packages closest to query embedding $2 using pgvector cosine distance (<=>)
//     for packages with imported_by_count >= 1 (rank_vec).
//  3. combined: Performs a FULL OUTER JOIN between text and vector candidate sets, calculating fused RRF score:
//     Score = 1 / (60 + rank_text) + VectorWeight * (1 / (60 + rank_vec))
//     (where 60 is the standard RRF smoothing constant k=60).
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 (
			SELECT package_path, version, module_path, commit_time, imported_by_count,
				ROW_NUMBER() OVER (ORDER BY (%s) DESC, commit_time DESC, package_path) AS rank_text,
				COUNT(*) OVER() AS total_text_matches
			FROM search_documents
			WHERE tsv_search_tokens @@ websearch_to_tsquery($1)
			LIMIT %d
		),
		vector_search AS (
			SELECT package_path, version, module_path, commit_time, imported_by_count,
				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
				COALESCE(t.package_path, v.package_path) AS package_path,
				COALESCE(t.version, v.version) AS version,
				COALESCE(t.module_path, v.module_path) AS module_path,
				COALESCE(t.commit_time, v.commit_time) AS commit_time,
				COALESCE(t.imported_by_count, v.imported_by_count) AS imported_by_count,
				(COALESCE(1.0 / (60 + t.rank_text), 0.0) + %f * COALESCE(1.0 / (60 + v.rank_vec), 0.0)) AS score,
				COALESCE(MAX(t.total_text_matches) OVER(), COUNT(*) OVER()) AS total
			FROM text_search t
			FULL OUTER JOIN vector_search v ON t.package_path = v.package_path
		)
		SELECT package_path, version, module_path, commit_time, imported_by_count, score, total
		FROM combined
		ORDER BY score DESC, commit_time DESC, package_path
		LIMIT $3
		OFFSET $4`, scoreExprStr, candidateLimit, popWeight, nonRedistributablePenalty, noGoModPenalty, candidateLimit, vectorWeight)
	args := []any{q, formatVector(opts.Vector), limit, opts.Offset}
	return query, args
}

// buildTextSearchQuery constructs a standard full-text search SQL query.
//
// The query:
//  1. Filters search_documents matching full-text query $1 via websearch_to_tsquery.
//  2. Computes relevance score via scoreExprStr, combining ts_rank section weights, popularity boost,
//     and non-redistributable / no-go-mod penalties.
//  3. Excludes low-relevance matches with score <= 0.1.
//  4. Orders results by score DESC, commit_time DESC, package_path, and computes windowed total match count.
func buildTextSearchQuery(opts SearchOptions, scoreExprStr string, q string, limit int) (string, []any) {
	query := fmt.Sprintf(`
		SELECT *, COUNT(*) OVER() AS total
		FROM (
			SELECT
				package_path,
				version,
				module_path,
				commit_time,
				imported_by_count,
				(%s) AS score
				FROM
					search_documents
				WHERE tsv_search_tokens @@ websearch_to_tsquery($1)
				ORDER BY
					score DESC,
					commit_time DESC,
					package_path
		) r
		WHERE r.score > 0.1
		LIMIT $2
		OFFSET $3`, scoreExprStr)
	args := []any{q, limit, opts.Offset}
	return query, args
}

// deepSearch searches all packages for the query. It is slower, but results
// are always valid. If opts.Vector is set, it performs a hybrid search using
// Reciprocal Rank Fusion (RRF) between text search and pgvector similarity.
func (db *DB) deepSearch(ctx context.Context, q string, limit int, opts SearchOptions) searchResponse {
	opts.ScoringParams = opts.ScoringParams.WithDefaults()
	scoreExprStr := scoreExpr(opts.ScoringParams)
	var query string
	var args []any
	if len(opts.Vector) > 0 {
		query, args = buildVectorSearchQuery(opts, scoreExprStr, q, limit)
	} else {
		query, args = buildTextSearchQuery(opts, scoreExprStr, q, limit)
	}

	var results []*SearchResult
	collect := func(rows *sql.Rows) error {
		var r SearchResult
		if err := rows.Scan(&r.PackagePath, &r.Version, &r.ModulePath, &r.CommitTime,
			&r.NumImportedBy, &r.Score, &r.NumResults); err != nil {
			return fmt.Errorf("rows.Scan(): %v", err)
		}
		results = append(results, &r)
		return nil
	}
	err := db.db.RunQuery(ctx, query, collect, args...)
	if err != nil {
		results = nil
	}
	for i, r := range results {
		r.Offset = opts.Offset + i
	}
	if len(results) > 0 && results[0].NumResults > uint64(opts.MaxResultCount) {
		for _, r := range results {
			r.NumResults = uint64(opts.MaxResultCount)
		}
	}
	return searchResponse{
		source:  "deep",
		results: results,
		err:     err,
	}
}

func (db *DB) popularSearch(ctx context.Context, searchQuery string, limit int, opts SearchOptions) searchResponse {
	query := `
		SELECT
			package_path,
			version,
			module_path,
			commit_time,
			imported_by_count,
			score
		FROM popular_search($1, $2, $3, $4, $5)`
	var results []*SearchResult
	collect := func(rows *sql.Rows) error {
		var r SearchResult
		if err := rows.Scan(&r.PackagePath, &r.Version, &r.ModulePath, &r.CommitTime,
			&r.NumImportedBy, &r.Score); err != nil {
			return fmt.Errorf("rows.Scan(): %v", err)
		}
		results = append(results, &r)
		return nil
	}
	err := db.db.RunQuery(ctx, query, collect, searchQuery, limit, opts.Offset, nonRedistributablePenalty, noGoModPenalty)
	if err != nil {
		results = nil
	}
	numResults := opts.MaxResultCount
	if opts.Offset+limit > opts.MaxResultCount || len(results) < limit {
		// It is practically impossible that len(results) < limit, because popular
		// search will never linearly scan everything before deep search completes,
		// but just to be slightly more theoretically correct, if our search
		// results are partial we know that we have exhausted all results.
		numResults = opts.Offset + len(results)
	}
	for i, r := range results {
		r.Offset = opts.Offset + i
		r.NumResults = uint64(numResults)
	}
	return searchResponse{
		source:  "popular",
		results: results,
		err:     err,
	}
}

// addPackageDataToSearchResults adds package information to SearchResults that is not stored
// in the search_documents table.
func (db *DB) addPackageDataToSearchResults(ctx context.Context, results []*SearchResult) (err error) {
	defer derrors.WrapStack(&err, "DB.addPackageDataToSearchResults(results)")
	if len(results) == 0 {
		return nil
	}
	var (
		keys []string
		// resultMap tracks PackagePath->SearchResult, to allow joining with the
		// returned package data.
		resultMap = make(map[string]*SearchResult)
	)
	for _, r := range results {
		resultMap[r.PackagePath] = r
		key := fmt.Sprintf("(%s, %s, %s)", pq.QuoteLiteral(r.PackagePath),
			pq.QuoteLiteral(r.Version), pq.QuoteLiteral(r.ModulePath))
		keys = append(keys, key)
	}
	query := fmt.Sprintf(`
		SELECT
			p.path,
			u.name,
			d.synopsis,
			u.license_types,
			u.redistributable
		FROM
			units u
		INNER JOIN
			paths p
		ON u.path_id = p.id
		INNER JOIN
			modules m
		ON u.module_id = m.id
		LEFT JOIN
			documentation d
		ON u.id = d.unit_id
		WHERE
			(p.path, m.version, m.module_path) IN (%s)`, strings.Join(keys, ","))
	collect := func(rows *sql.Rows) error {
		var (
			path, name, synopsis string
			licenseTypes         []string
			redist               bool
		)
		if err := rows.Scan(&path, &name, database.NullIsEmpty(&synopsis), pq.Array(&licenseTypes), &redist); err != nil {
			return fmt.Errorf("rows.Scan(): %v", err)
		}
		r, ok := resultMap[path]
		if !ok {
			return fmt.Errorf("BUG: unexpected package path: %q", path)
		}
		r.Name = name
		if redist || db.bypassLicenseCheck {
			r.Synopsis = synopsis
		}
		for _, l := range licenseTypes {
			if l != "" {
				r.Licenses = append(r.Licenses, l)
			}
		}
		r.Licenses = sortAndDedup(r.Licenses)
		return nil
	}
	return db.db.RunQuery(ctx, query, collect)
}

func sortAndDedup(s []string) []string {
	var r []string
	m := map[string]bool{}
	for _, x := range s {
		m[x] = true
	}
	for x := range m {
		r = append(r, x)
	}
	sort.Strings(r)
	return r
}

// groupSearchResults groups and re-orders the list of SearchResults by module
// and series path and returns a new list of SearchResults.
//
// The second and later packages from a module are grouped under the first package,
// and removed from the top-level list. The exception is the standard library: none
// of its packages are grouped under other packages.
//
// Higher tagged major versions of a module replace lower ones.
func groupSearchResults(rs []*SearchResult) []*SearchResult {
	bestInGroup := map[string]*SearchResult{} // series path to result with max major version
	// Since rs is sorted by score, the first package we see for a series is the
	// highest-ranked one. However, we may prefer to show a lower-ranked one from a
	// module in the same series with a higher major version.
	for _, r := range rs {
		group, rMajor := groupAndMajorVersion(r)
		b := bestInGroup[group]
		if b == nil {
			// First result (package) with this key; remember it.
			bestInGroup[group] = r
			r.OtherMajor = map[string]int{}
		} else {
			_, bMajor := groupAndMajorVersion(b)
			switch {
			case !version.IsPseudo(r.Version) && (rMajor > bMajor || version.IsPseudo(b.Version)):
				// r is tagged, and is either in a higher major version, or the current best
				// is not tagged. Either way, prefer r to b.
				bestInGroup[group] = r
				r.OtherMajor = b.OtherMajor
				r.OtherMajor[b.ModulePath] = bMajor
				r.Score = b.Score // inherit the lower major version's higher score
			case rMajor == bMajor:
				// r is another package in b's group; remember it there.
				b.SameModule = append(b.SameModule, r)

			default:
				// A package from a different major version (either lower, or
				// higher untagged). Remember the major version.
				b.OtherMajor[r.ModulePath] = rMajor
			}
		}
	}
	// Collect new results and re-sort by score.
	var results []*SearchResult
	for _, r := range bestInGroup {
		if len(r.OtherMajor) == 0 {
			r.OtherMajor = nil
		}
		results = append(results, r)
	}
	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})
	return results
}

func groupAndMajorVersion(r *SearchResult) (string, int) {
	// Packages in the standard library are not grouped.
	// We consider them all part of the same major version.
	if r.ModulePath == stdlib.ModulePath {
		return r.PackagePath, 1
	}
	series, major := internal.SeriesPathAndMajorVersion(r.ModulePath)
	if major == 1 && strings.HasPrefix(r.Version, "v0.") {
		major = 0
	}
	return series, major
}

// numRows counts the number of rows in a slice of SearchResults.
// Grouping will put some rows inside a SearchResult.
func numRows(rs []*SearchResult) int {
	n := 0
	for _, r := range rs {
		n += 1 + len(r.SameModule)
	}
	return n
}

var upsertSearchStatement = fmt.Sprintf(`
	INSERT INTO search_documents (
		package_path,
		package_path_id,
		version,
		module_path,
		module_path_id,
		unit_id,
		name,
		synopsis,
		license_types,
		redistributable,
		version_updated_at,
		commit_time,
		has_go_mod,
		-- TODO(https://golang.org/issue/44142): The path_tokens column is used
		-- to easily iterate on tsv_path_tokens, and can be removed once
		-- symbol search implementation is done.
		path_tokens,
		tsv_path_tokens,
		tsv_search_tokens,
		hll_register,
		hll_leading_zeros
	)
	SELECT
		p1.path,
		p1.id,
		m.version,
		m.module_path,
		p2.id,
		u.id AS unit_id,
		u.name,
		d.synopsis,
		u.license_types,
		u.redistributable,
		CURRENT_TIMESTAMP,
		m.commit_time,
		m.has_go_mod,
		$4,
		SETWEIGHT(TO_TSVECTOR('%s', replace($4, '_', '-')), 'A'),
		(
			SETWEIGHT(TO_TSVECTOR('path_tokens', $4), 'A') ||
			SETWEIGHT(TO_TSVECTOR($5), 'B') ||
			SETWEIGHT(TO_TSVECTOR($6), 'C') ||
			SETWEIGHT(TO_TSVECTOR($7), 'D')
		),
		hll_hash(p1.path) & (%d - 1),
		hll_zeros(hll_hash(p1.path))
	FROM units u
	INNER JOIN modules m ON u.module_id = m.id
	INNER JOIN paths p1 ON p1.id = u.path_id
	LEFT JOIN paths p2 ON p2.path = m.module_path
	LEFT JOIN documentation d ON u.id = d.unit_id
	WHERE
		p1.path = $1
		AND m.module_path = $2
		AND m.version = $3
	LIMIT 1 -- could be multiple build contexts
	ON CONFLICT (package_path_id)
	DO UPDATE SET
		package_path=excluded.package_path,
		version=excluded.version,
		module_path=excluded.module_path,
		unit_id=excluded.unit_id,
		name=excluded.name,
		synopsis=excluded.synopsis,
		license_types=excluded.license_types,
		redistributable=excluded.redistributable,
		commit_time=excluded.commit_time,
		has_go_mod=excluded.has_go_mod,
		path_tokens=excluded.path_tokens,
		tsv_path_tokens=excluded.tsv_path_tokens,
		tsv_search_tokens=excluded.tsv_search_tokens,
		-- the hll fields are functions of path, so they don't change
		version_updated_at=(
			CASE WHEN excluded.version = search_documents.version
			THEN search_documents.version_updated_at
			ELSE CURRENT_TIMESTAMP
			END)
	;`,
	search.SymbolTextSearchConfiguration,
	hllRegisterCount)

// upsertSearchDocuments adds search information for mod to the search_documents table.
// It assumes that all non-redistributable data has been removed from mod.
func upsertSearchDocuments(ctx context.Context, ddb *database.DB, mod *internal.Module) (err error) {
	defer derrors.WrapStack(&err, "upsertSearchDocuments(ctx, %q, %q)", mod.ModulePath, mod.Version)
	defer internal.RequestState(ctx, "updating search_documents table")()
	ctx, span := trace.StartSpan(ctx, "UpsertSearchDocuments")
	defer span.End()

	// Don't upsert a package if it is already present under a longer module
	// path. We need this because search_documents can have only one row per
	// import path, and we, like the go tool, prefer the package with the longer
	// module path. For example, if two packages have import path "a/b/c", one
	// in module "a" and the other in "a/b", we keep only the latter.
	lps, err := database.Collect1[string](ctx, ddb, `
		SELECT package_path
		FROM search_documents
		WHERE module_path LIKE $1 || '/%'
	`, mod.ModulePath)
	if err != nil {
		return err
	}
	longerPackages := map[string]bool{}
	for _, lp := range lps {
		longerPackages[lp] = true
	}

	for _, pkg := range mod.Packages() {
		if isInternalPackage(pkg.Path) {
			continue
		}
		if longerPackages[pkg.Path] {
			continue
		}
		args := UpsertSearchDocumentArgs{
			PackagePath: pkg.Path,
			ModulePath:  mod.ModulePath,
			Version:     mod.Version,
		}
		if len(pkg.Documentation) > 0 {
			// Use the synopsis of the first GOOS/GOARCH pair.
			args.Synopsis = pkg.Documentation[0].Synopsis
		}
		if pkg.Readme != nil {
			args.ReadmeFilePath = pkg.Readme.Filepath
			args.ReadmeContents = pkg.Readme.Contents
		}
		if err := UpsertSearchDocument(ctx, ddb, args); err != nil {
			return err
		}
	}
	return nil
}

type UpsertSearchDocumentArgs struct {
	PackagePath    string
	ModulePath     string
	Version        string
	Synopsis       string
	ReadmeFilePath string
	ReadmeContents string
}

// UpsertSearchDocument inserts a row in search_documents for the given package.
// The given module should have already been validated via a call to
// validateModule.
func UpsertSearchDocument(ctx context.Context, ddb *database.DB, args UpsertSearchDocumentArgs) (err error) {
	defer derrors.WrapStack(&err, "DB.UpsertSearchDocument(ctx, ddb, %q, %q)", args.PackagePath, args.ModulePath)

	// Only summarize the README if the package and module have the same path.
	// If this changes, fix DB.ReconcileSearch.
	if args.PackagePath != args.ModulePath {
		args.ReadmeFilePath = ""
		args.ReadmeContents = ""
	}
	pathTokens := strings.Join(GeneratePathTokens(args.PackagePath), " ")
	sectionB, sectionC, sectionD := SearchDocumentSections(args.Synopsis, args.ReadmeFilePath, args.ReadmeContents)
	_, err = ddb.Exec(ctx, upsertSearchStatement, args.PackagePath, args.ModulePath, args.Version, pathTokens, sectionB, sectionC, sectionD)
	return err
}

// PackageToEmbed contains the package path, synopsis, and top symbols needed to generate an embedding.
type PackageToEmbed struct {
	PackagePath   string
	PackagePathID int
	Synopsis      string
	TopSymbols    string // space-separated list of top symbol names for the package
}

// GetPackagesToEmbed fetches packages with imported_by_count >= 1 that do not have an embedding yet,
// returning the package information (synopsis and top symbols) required to generate their embeddings.
func (db *DB) GetPackagesToEmbed(ctx context.Context, limit int) (pkgs []PackageToEmbed, err error) {
	defer derrors.WrapStack(&err, "GetPackagesToEmbed(ctx, %d)", limit)

	// Top symbols are ordered by:
	// 1. Top-level symbols first (where symbol_name_id == parent_symbol_name_id) rather than methods/fields.
	// 2. Section priority: Types, Functions, Variables, Constants.
	// 3. Alphabetical order by symbol name.
	// TODO: add a test for this query.
	query := `
		SELECT sd.package_path, sd.package_path_id, COALESCE(sd.synopsis, ''),
		       (SELECT string_agg(s.name, ' ')
		        FROM (
		            SELECT s.name
		            FROM package_symbols ps
		            INNER JOIN symbol_names s ON ps.symbol_name_id = s.id
		            WHERE ps.package_path_id = sd.package_path_id
		            ORDER BY
		                (ps.symbol_name_id = ps.parent_symbol_name_id) DESC,
		                CASE ps.section
		                    WHEN 'Types' THEN 1
		                    WHEN 'Functions' THEN 2
		                    WHEN 'Variables' THEN 3
		                    WHEN 'Constants' THEN 4
		                    ELSE 5
		                END,
		                s.name
		            LIMIT 50
		        ) s
		       ) AS top_symbols
		FROM search_documents sd
		WHERE sd.imported_by_count >= 1 AND sd.embedding IS NULL
		LIMIT $1;`

	collect := func(rows *sql.Rows) error {
		var p PackageToEmbed
		if err := rows.Scan(&p.PackagePath, &p.PackagePathID, &p.Synopsis, database.NullIsEmpty(&p.TopSymbols)); err != nil {
			return err
		}
		pkgs = append(pkgs, p)
		return nil
	}

	if err := db.db.RunQuery(ctx, query, collect, limit); err != nil {
		return nil, err
	}
	return pkgs, nil
}

// UpdateSearchDocumentEmbedding updates the embedding vector for the package.
func (db *DB) UpdateSearchDocumentEmbedding(ctx context.Context, packagePath string, embedding []float32) (err error) {
	defer derrors.WrapStack(&err, "UpdateSearchDocumentEmbedding(ctx, %q)", packagePath)

	vecStr := formatVector(embedding)
	query := `UPDATE search_documents SET embedding = $1::halfvec WHERE package_path = $2;`
	_, err = db.db.Exec(ctx, query, vecStr, packagePath)
	return err
}

func formatVector(vec []float32) string {
	var sb strings.Builder
	sb.WriteByte('[')
	for i, v := range vec {
		if i > 0 {
			sb.WriteByte(',')
		}
		fmt.Fprintf(&sb, "%g", v)
	}
	sb.WriteByte(']')
	return sb.String()
}

// GetPackagesForSearchDocumentUpsert fetches search information for packages in search_documents
// whose update time is before the given time.
func (db *DB) GetPackagesForSearchDocumentUpsert(ctx context.Context, before time.Time, limit int) (argsList []UpsertSearchDocumentArgs, err error) {
	defer derrors.WrapStack(&err, "GetPackagesForSearchDocumentUpsert(ctx, %s, %d)", before, limit)

	query := `
		SELECT
			sd.package_path,
			sd.module_path,
			sd.version,
			sd.synopsis,
			sd.redistributable,
			r.file_path,
			r.contents
		FROM modules m
		INNER JOIN units u
		ON m.id = u.module_id
		INNER JOIN paths p
		ON p.id = u.path_id
		LEFT JOIN readmes r
		ON u.id = r.unit_id
		INNER JOIN search_documents sd
		ON sd.package_path = p.path
		    AND sd.module_path = m.module_path
		    AND sd.version = m.version
		WHERE sd.updated_at < $1
		LIMIT $2`

	collect := func(rows *sql.Rows) error {
		var (
			a      UpsertSearchDocumentArgs
			redist bool
		)
		if err := rows.Scan(&a.PackagePath, &a.ModulePath, &a.Version, database.NullIsEmpty(&a.Synopsis), &redist,
			database.NullIsEmpty(&a.ReadmeFilePath), database.NullIsEmpty(&a.ReadmeContents)); err != nil {
			return err
		}
		if !redist && !db.bypassLicenseCheck {
			a.Synopsis = ""
			a.ReadmeFilePath = ""
			a.ReadmeContents = ""
		}
		argsList = append(argsList, a)
		return nil
	}
	if err := db.db.RunQuery(ctx, query, collect, before, limit); err != nil {
		return nil, err
	}
	return argsList, nil
}

var (
	commonHostnames = map[string]bool{
		"bitbucket.org":         true,
		"code.cloudfoundry.org": true,
		"gitea.com":             true,
		"gitee.com":             true,
		"github.com":            true,
		"gitlab.com":            true,
		"go.etcd.io":            true,
		"go.googlesource.com":   true,
		"golang.org":            true,
		"google.golang.org":     true,
		"gopkg.in":              true,
	}
	commonHostParts = map[string]bool{
		"code":   true,
		"git":    true,
		"gitlab": true,
		"go":     true,
		"google": true,
		"www":    true,
	}
)

// GeneratePathTokens returns the subPaths and path token parts that will be
// indexed for search, which includes (1) the packagePath (2) all sub-paths of
// the packagePath (3) all parts for a path element that is delimited by a dash
// and (4) all parts of a path element that is delimited by a dot, except for
// the last element.
func GeneratePathTokens(packagePath string) []string {
	packagePath = strings.Trim(packagePath, "/")

	subPathSet := make(map[string]bool)
	parts := strings.Split(packagePath, "/")
	for i, part := range parts {
		dashParts := strings.Split(part, "-")
		if len(dashParts) > 1 {
			for _, p := range dashParts {
				subPathSet[p] = true
			}
		}
		for j := i + 2; j <= len(parts); j++ {
			p := strings.Join(parts[i:j], "/")
			p = strings.Trim(p, "/")
			subPathSet[p] = true
		}

		if i == 0 && commonHostnames[part] {
			continue
		}
		// Only index host names if they are not part of commonHostnames.
		subPathSet[part] = true
		dotParts := strings.Split(part, ".")
		if len(dotParts) > 1 {
			for _, p := range dotParts[:len(dotParts)-1] {
				if !commonHostParts[p] {
					// If the host is not in commonHostnames, we want to
					// index each element up to the extension. For example,
					// if the host is sigs.k8s.io, we want to index sigs
					// and k8s. Skip common host parts.
					subPathSet[p] = true
				}
			}
		}
	}

	var subPaths []string
	for sp := range subPathSet {
		if len(sp) > 0 {
			subPaths = append(subPaths, sp)
		}
	}
	sort.Strings(subPaths)
	return subPaths
}

// isInternalPackage reports whether the path represents an internal directory.
func isInternalPackage(path string) bool {
	return slices.Contains(strings.Split(path, "/"), "internal")
}