Method Query

pkg/vector/pgvector/store.go:127–186

Query 在指定命名空间内执行向量检索。

(ctx context.Context, q vector.Query)

// Query 在指定命名空间内执行向量检索。
func (s *Store) Query(ctx context.Context, q vector.Query) ([]vector.Hit, error) {
	if len(q.Vector) != s.dim {
		return nil, fmt.Errorf("query vector dimension mismatch: got %d, want %d", len(q.Vector), s.dim)
	}

	topK := q.TopK
	if topK <= 0 {
		topK = 5
	}

	// pgvector 距离越小越相似, 我们将 score 定义为 1 - distance (cosine) 或负的 L2。
	var distanceExpr string
	switch s.metric {
	case "cosine":
		distanceExpr = "embedding <=> $1" // cosine distance
	case "l2":
		distanceExpr = "embedding <-> $1" // L2 distance
	default:
		distanceExpr = "embedding <=> $1"
	}

	ns := q.Namespace
	if ns == "" {
		ns = "default"
	}

	const baseTmpl = `
SELECT id, %s AS distance, metadata
FROM %s
WHERE namespace = $2
ORDER BY distance ASC
LIMIT $3;
`
	query := fmt.Sprintf(baseTmpl, distanceExpr, s.table)

	rows, err := s.pool.Query(ctx, query, float32SliceToPgVector(q.Vector), ns, topK)
	if err != nil {
		return nil, fmt.Errorf("query vectors: %w", err)
	}
	defer rows.Close()

	var hits []vector.Hit
	for rows.Next() {
		var (
			id       string
			distance float64
			meta     map[string]any
		)
		if err := rows.Scan(&id, &distance, &meta); err != nil {
			return nil, fmt.Errorf("scan row: %w", err)
		}
		score := distanceToScore(distance, s.metric)
		hits = append(hits, vector.Hit{
			ID:       id,
			Score:    score,
			Metadata: meta,
		})
	}