snapshot of https://github.com/blevesearch/bleve.git (master, 2026-09-15; blobs > 131072B omitted)
Some checks failed
Coverage / coverage (push) Failing after 6s
Tests / test (1.24.x, ubuntu-latest) (push) Failing after 42s
Tests / test (1.26.x, ubuntu-latest) (push) Failing after 38s
Tests / test (1.25.x, ubuntu-latest) (push) Failing after 35s
Tests / test (1.25.x, windows-latest) (push) Has been cancelled
Tests / test (1.26.x, macos-latest) (push) Has been cancelled
Tests / test (1.26.x, windows-latest) (push) Has been cancelled
Tests / test (1.24.x, macos-latest) (push) Has been cancelled
Tests / test (1.24.x, windows-latest) (push) Has been cancelled
Tests / test (1.25.x, macos-latest) (push) Has been cancelled
Some checks failed
Coverage / coverage (push) Failing after 6s
Tests / test (1.24.x, ubuntu-latest) (push) Failing after 42s
Tests / test (1.26.x, ubuntu-latest) (push) Failing after 38s
Tests / test (1.25.x, ubuntu-latest) (push) Failing after 35s
Tests / test (1.25.x, windows-latest) (push) Has been cancelled
Tests / test (1.26.x, macos-latest) (push) Has been cancelled
Tests / test (1.26.x, windows-latest) (push) Has been cancelled
Tests / test (1.24.x, macos-latest) (push) Has been cancelled
Tests / test (1.24.x, windows-latest) (push) Has been cancelled
Tests / test (1.25.x, macos-latest) (push) Has been cancelled
This commit is contained in:
26
fusion/fusion.go
Normal file
26
fusion/fusion.go
Normal file
@@ -0,0 +1,26 @@
|
||||
// Copyright (c) 2025 Couchbase, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package fusion
|
||||
|
||||
import (
|
||||
"github.com/blevesearch/bleve/v2/search"
|
||||
)
|
||||
|
||||
type FusionResult struct {
|
||||
Hits search.DocumentMatchCollection
|
||||
Total uint64
|
||||
MaxScore float64
|
||||
}
|
||||
143
fusion/rrf.go
Normal file
143
fusion/rrf.go
Normal file
@@ -0,0 +1,143 @@
|
||||
// Copyright (c) 2025 Couchbase, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package fusion
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/blevesearch/bleve/v2/search"
|
||||
)
|
||||
|
||||
// formatRRFMessage builds the explanation string for a single component of the
|
||||
// Reciprocal Rank Fusion calculation.
|
||||
func formatRRFMessage(weight float64, rank int, rankConstant int) string {
|
||||
return fmt.Sprintf("rrf score (weight=%.3f, rank=%d, rank_constant=%d), normalized score of", weight, rank, rankConstant)
|
||||
}
|
||||
|
||||
// ReciprocalRankFusion applies Reciprocal Rank Fusion across the primary FTS
|
||||
// results and each KNN sub-query. Ranks are limited to `windowSize` per source,
|
||||
// weighted, and combined into a single fused score, with optional explanation
|
||||
// details.
|
||||
func ReciprocalRankFusion(hits search.DocumentMatchCollection, weights []float64, rankConstant int, windowSize int, numKNNQueries int, explain bool) *FusionResult {
|
||||
nHits := len(hits)
|
||||
if nHits == 0 || windowSize == 0 {
|
||||
return &FusionResult{
|
||||
Hits: search.DocumentMatchCollection{},
|
||||
Total: 0,
|
||||
MaxScore: 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
limit := min(nHits, windowSize)
|
||||
|
||||
// precompute rank+scores to prevent additional division ops later
|
||||
rankReciprocals := make([]float64, limit)
|
||||
for i := range rankReciprocals {
|
||||
rankReciprocals[i] = 1.0 / float64(rankConstant+i+1)
|
||||
}
|
||||
|
||||
// init explanations if required
|
||||
var fusionExpl map[*search.DocumentMatch][]*search.Explanation
|
||||
if explain {
|
||||
fusionExpl = make(map[*search.DocumentMatch][]*search.Explanation, nHits)
|
||||
}
|
||||
|
||||
// The code here mainly deals with obtaining rank/score for fts hits.
|
||||
// First sort hits by score
|
||||
sortDocMatchesByScore(hits)
|
||||
|
||||
// Calculate fts rank+scores
|
||||
ftsWeight := weights[0]
|
||||
for i := 0; i < nHits; i++ {
|
||||
if i < windowSize {
|
||||
hit := hits[i]
|
||||
|
||||
// No fts scores from this hit onwards, break loop
|
||||
if hit.Score == 0.0 {
|
||||
break
|
||||
}
|
||||
|
||||
contrib := ftsWeight * rankReciprocals[i]
|
||||
hit.Score = contrib
|
||||
|
||||
if explain {
|
||||
expl := getFusionExplAt(
|
||||
hit,
|
||||
0,
|
||||
contrib,
|
||||
formatRRFMessage(ftsWeight, i+1, rankConstant),
|
||||
)
|
||||
fusionExpl[hit] = append(fusionExpl[hit], expl)
|
||||
}
|
||||
} else {
|
||||
// These FTS hits are not counted in the results, so set to 0
|
||||
hits[i].Score = 0.0
|
||||
}
|
||||
}
|
||||
|
||||
// Code from here is to calculate knn ranks and scores
|
||||
// iterate over each knn query and calculate knn rank+scores
|
||||
for queryIdx := 0; queryIdx < numKNNQueries; queryIdx++ {
|
||||
knnWeight := weights[queryIdx+1]
|
||||
// Sorts hits in decreasing order of hit.ScoreBreakdown[i]
|
||||
sortDocMatchesByBreakdown(hits, queryIdx)
|
||||
|
||||
for i := 0; i < nHits; i++ {
|
||||
// break if score breakdown doesn't exist (sort function puts these hits at the end)
|
||||
// or if we go past the windowSize
|
||||
_, scoreBreakdownExists := scoreBreakdownForQuery(hits[i], queryIdx)
|
||||
if i >= windowSize || !scoreBreakdownExists {
|
||||
break
|
||||
}
|
||||
|
||||
hit := hits[i]
|
||||
contrib := knnWeight * rankReciprocals[i]
|
||||
hit.Score += contrib
|
||||
|
||||
if explain {
|
||||
expl := getFusionExplAt(
|
||||
hit,
|
||||
queryIdx+1,
|
||||
contrib,
|
||||
formatRRFMessage(knnWeight, i+1, rankConstant),
|
||||
)
|
||||
fusionExpl[hit] = append(fusionExpl[hit], expl)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var maxScore float64
|
||||
for _, hit := range hits {
|
||||
if explain {
|
||||
finalizeFusionExpl(hit, fusionExpl[hit])
|
||||
}
|
||||
hit.ScoreBreakdown = nil
|
||||
|
||||
if hit.Score > maxScore {
|
||||
maxScore = hit.Score
|
||||
}
|
||||
}
|
||||
|
||||
sortDocMatchesByScore(hits)
|
||||
if nHits > windowSize {
|
||||
hits = hits[:windowSize]
|
||||
}
|
||||
return &FusionResult{
|
||||
Hits: hits,
|
||||
Total: uint64(len(hits)),
|
||||
MaxScore: maxScore,
|
||||
}
|
||||
}
|
||||
212
fusion/rrf_test.go
Normal file
212
fusion/rrf_test.go
Normal file
@@ -0,0 +1,212 @@
|
||||
// Copyright (c) 2025 Couchbase, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package fusion
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/blevesearch/bleve/v2/search"
|
||||
)
|
||||
|
||||
const epsilon float64 = 1e-3
|
||||
|
||||
func nearlyEqual(a float64, b float64, epsilon float64) bool {
|
||||
return math.Abs(a-b) < epsilon
|
||||
}
|
||||
|
||||
func compareFusionResults(a, b FusionResult) bool {
|
||||
if a.Total != b.Total || !nearlyEqual(a.MaxScore, b.MaxScore, epsilon) || len(a.Hits) != len(b.Hits) {
|
||||
return false
|
||||
}
|
||||
for i := range a.Hits {
|
||||
if a.Hits[i].ID != b.Hits[i].ID || !nearlyEqual(a.Hits[i].Score, b.Hits[i].Score, epsilon) {
|
||||
return false
|
||||
}
|
||||
|
||||
if a.Hits[i].ScoreBreakdown != nil || b.Hits[i].ScoreBreakdown != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func TestReciprocalRankFusion(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
hits search.DocumentMatchCollection
|
||||
weights []float64
|
||||
rank_constant int
|
||||
window_size int
|
||||
numKNNQueries int
|
||||
want FusionResult
|
||||
}{
|
||||
{
|
||||
name: "empty hits",
|
||||
hits: search.DocumentMatchCollection{},
|
||||
weights: []float64{0.5, 0.5},
|
||||
rank_constant: 60,
|
||||
window_size: 10,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{},
|
||||
Total: 0,
|
||||
MaxScore: 0.0,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "single knn query",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.8}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.7}},
|
||||
},
|
||||
weights: []float64{0.4, 0.6},
|
||||
rank_constant: 1,
|
||||
window_size: 3,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "b", Score: 0.433},
|
||||
{ID: "a", Score: 0.4},
|
||||
{ID: "c", Score: 0.25},
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.433,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "multiple knn queries",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.8, 1: 0.6}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9, 1: 0.5}},
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.7, 1: 0.7}},
|
||||
},
|
||||
weights: []float64{0.3, 0.4, 0.3},
|
||||
rank_constant: 1,
|
||||
window_size: 3,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.383},
|
||||
{ID: "b", Score: 0.375},
|
||||
{ID: "c", Score: 0.325},
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.383,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "window size smaller than hits",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.7}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.8}},
|
||||
},
|
||||
weights: []float64{0.4, 0.6},
|
||||
rank_constant: 1,
|
||||
window_size: 2,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "b", Score: 0.433},
|
||||
{ID: "a", Score: 0.2},
|
||||
},
|
||||
Total: 2,
|
||||
MaxScore: 0.433,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "documents with partial scores missing KNN scores",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.8}}, // has FTS and KNN query 0, missing KNN query 1
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{1: 0.7}}, // has FTS and KNN query 1, missing KNN query 0
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.6, 1: 0.9}}, // has all scores
|
||||
{ID: "d", Score: 0.6, ScoreBreakdown: map[int]float64{}}, // has only FTS, missing all KNN scores
|
||||
},
|
||||
weights: []float64{0.3, 0.4, 0.3}, // FTS, KNN query 0, KNN query 1
|
||||
rank_constant: 1,
|
||||
window_size: 4,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "c", Score: 0.358}, // FTS rank 3, KNN0 rank 2, KNN1 rank 1: 0.3/4 + 0.4/3 + 0.3/2 = 0.075 + 0.133 + 0.15 = 0.358
|
||||
{ID: "a", Score: 0.35}, // FTS rank 1, KNN0 rank 1, no KNN1: 0.3/2 + 0.4/2 + 0 = 0.15 + 0.2 + 0 = 0.35
|
||||
{ID: "b", Score: 0.2}, // FTS rank 2, no KNN0, KNN1 rank 2: 0.3/3 + 0 + 0.3/3 = 0.1 + 0 + 0.1 = 0.2
|
||||
{ID: "d", Score: 0.06}, // FTS rank 4, no KNN0, no KNN1: 0.3/5 + 0 + 0 = 0.06
|
||||
},
|
||||
Total: 4,
|
||||
MaxScore: 0.358,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "documents with only KNN scores",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.0, ScoreBreakdown: map[int]float64{0: 0.9}}, // no FTS rank (Score 0.0), only KNN query 0
|
||||
{ID: "b", Score: 0.0, ScoreBreakdown: map[int]float64{1: 0.8}}, // no FTS rank (Score 0.0), only KNN query 1
|
||||
{ID: "c", Score: 0.0, ScoreBreakdown: map[int]float64{0: 0.7, 1: 0.6}}, // no FTS rank (Score 0.0), both KNN queries
|
||||
},
|
||||
weights: []float64{0.5, 0.3, 0.2}, // FTS, KNN query 0, KNN query 1
|
||||
rank_constant: 1,
|
||||
window_size: 3,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "c", Score: 0.167}, // no FTS rank, KNN0 rank 2, KNN1 rank 2: 0 + 0.3/3 + 0.2/3 = 0 + 0.1 + 0.067 = 0.167
|
||||
{ID: "a", Score: 0.15}, // no FTS rank, KNN0 rank 1, no KNN1: 0 + 0.3/2 + 0 = 0 + 0.15 + 0 = 0.15
|
||||
{ID: "b", Score: 0.1}, // no FTS rank, no KNN0, KNN1 rank 1: 0 + 0 + 0.2/2 = 0 + 0 + 0.1 = 0.1
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.167,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "mixed scenario with gaps in KNN queries",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.8, ScoreBreakdown: map[int]float64{1: 0.9}}, // has FTS and KNN query 1, missing KNN query 0
|
||||
{ID: "b", Score: 0.6, ScoreBreakdown: map[int]float64{0: 0.8}}, // has FTS and KNN query 0, missing KNN query 1
|
||||
{ID: "c", Score: 0.0, ScoreBreakdown: map[int]float64{0: 0.7}}, // no FTS rank (Score 0.0), only KNN query 0
|
||||
{ID: "d", Score: 0.4, ScoreBreakdown: map[int]float64{}}, // only FTS, no KNN scores
|
||||
},
|
||||
weights: []float64{0.4, 0.3, 0.3}, // FTS, KNN query 0, KNN query 1
|
||||
rank_constant: 1,
|
||||
window_size: 4,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.35}, // FTS rank 1, no KNN0, KNN1 rank 1: 0.4/2 + 0 + 0.3/2 = 0.2 + 0 + 0.15 = 0.35
|
||||
{ID: "b", Score: 0.283}, // FTS rank 2, KNN0 rank 1, no KNN1: 0.4/3 + 0.3/2 + 0 = 0.133 + 0.15 + 0 = 0.283
|
||||
{ID: "d", Score: 0.1}, // FTS rank 3, no KNN0, no KNN1: 0.4/4 + 0 + 0 = 0.1
|
||||
{ID: "c", Score: 0.1}, // no FTS rank, KNN0 rank 2, no KNN1: 0 + 0.3/3 + 0 = 0 + 0.1 + 0 = 0.1
|
||||
},
|
||||
Total: 4,
|
||||
MaxScore: 0.35,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
for i, hit := range tt.hits {
|
||||
hit.HitNumber = uint64(i)
|
||||
}
|
||||
|
||||
if got := ReciprocalRankFusion(tt.hits, tt.weights, tt.rank_constant, tt.window_size, tt.numKNNQueries, false); !compareFusionResults(*got, tt.want) {
|
||||
t.Errorf("ReciprocalRankFusion() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
164
fusion/rsf.go
Normal file
164
fusion/rsf.go
Normal file
@@ -0,0 +1,164 @@
|
||||
// Copyright (c) 2025 Couchbase, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package fusion
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/blevesearch/bleve/v2/search"
|
||||
)
|
||||
|
||||
// formatRSFMessage builds the explanation string associated with a single
|
||||
// component of the Relative Score Fusion calculation.
|
||||
func formatRSFMessage(weight float64, normalizedScore float64, minScore float64, maxScore float64) string {
|
||||
return fmt.Sprintf("rsf score (weight=%.3f, normalized=%.6f, min=%.6f, max=%.6f), normalized score of",
|
||||
weight, normalizedScore, minScore, maxScore)
|
||||
}
|
||||
|
||||
// RelativeScoreFusion normalizes the best-scoring documents from the primary
|
||||
// FTS query and each KNN query, scales those normalized values by the supplied
|
||||
// weights, and combines them into a single fused score. Only the top
|
||||
// `windowSize` documents per source are considered, and explanations are
|
||||
// materialized lazily when requested.
|
||||
func RelativeScoreFusion(hits search.DocumentMatchCollection, weights []float64, windowSize int, numKNNQueries int, explain bool) *FusionResult {
|
||||
nHits := len(hits)
|
||||
if nHits == 0 || windowSize == 0 {
|
||||
return &FusionResult{
|
||||
Hits: search.DocumentMatchCollection{},
|
||||
Total: 0,
|
||||
MaxScore: 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
// init explanations if required
|
||||
var fusionExpl map[*search.DocumentMatch][]*search.Explanation
|
||||
if explain {
|
||||
fusionExpl = make(map[*search.DocumentMatch][]*search.Explanation, nHits)
|
||||
}
|
||||
|
||||
// Code here for calculating fts results
|
||||
// Sort by fts scores
|
||||
sortDocMatchesByScore(hits)
|
||||
|
||||
// ftsLimit holds the total number of fts hits to consider for rsf
|
||||
ftsLimit := 0
|
||||
for _, hit := range hits {
|
||||
if hit.Score == 0.0 {
|
||||
break
|
||||
}
|
||||
ftsLimit++
|
||||
}
|
||||
ftsLimit = min(ftsLimit, windowSize)
|
||||
|
||||
// calculate fts scores
|
||||
if ftsLimit > 0 {
|
||||
max := hits[0].Score
|
||||
min := hits[ftsLimit-1].Score
|
||||
denom := max - min
|
||||
weight := weights[0]
|
||||
|
||||
for i := 0; i < ftsLimit; i++ {
|
||||
hit := hits[i]
|
||||
norm := 1.0
|
||||
if denom > 0 {
|
||||
norm = (hit.Score - min) / denom
|
||||
}
|
||||
contrib := weight * norm
|
||||
if explain {
|
||||
expl := getFusionExplAt(
|
||||
hit,
|
||||
0,
|
||||
norm,
|
||||
formatRSFMessage(weight, norm, min, max),
|
||||
)
|
||||
fusionExpl[hit] = append(fusionExpl[hit], expl)
|
||||
}
|
||||
hit.Score = contrib
|
||||
}
|
||||
for i := ftsLimit; i < nHits; i++ {
|
||||
// These FTS hits are not counted in the results, so set to 0
|
||||
hits[i].Score = 0.0
|
||||
}
|
||||
}
|
||||
|
||||
// Code from here is for calculating knn scores
|
||||
for queryIdx := 0; queryIdx < numKNNQueries; queryIdx++ {
|
||||
sortDocMatchesByBreakdown(hits, queryIdx)
|
||||
|
||||
// knnLimit holds the total number of knn hits retrieved for a specific knn query
|
||||
knnLimit := 0
|
||||
for _, hit := range hits {
|
||||
if _, ok := scoreBreakdownForQuery(hit, queryIdx); !ok {
|
||||
break
|
||||
}
|
||||
knnLimit++
|
||||
}
|
||||
knnLimit = min(knnLimit, windowSize)
|
||||
|
||||
// if limit is 0, skip calculating
|
||||
if knnLimit == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
max, _ := scoreBreakdownForQuery(hits[0], queryIdx)
|
||||
min, _ := scoreBreakdownForQuery(hits[knnLimit-1], queryIdx)
|
||||
denom := max - min
|
||||
weight := weights[queryIdx+1]
|
||||
|
||||
for i := 0; i < knnLimit; i++ {
|
||||
hit := hits[i]
|
||||
score, _ := scoreBreakdownForQuery(hit, queryIdx)
|
||||
norm := 1.0
|
||||
if denom > 0 {
|
||||
norm = (score - min) / denom
|
||||
}
|
||||
contrib := weight * norm
|
||||
if explain {
|
||||
expl := getFusionExplAt(
|
||||
hit,
|
||||
queryIdx+1,
|
||||
norm,
|
||||
formatRSFMessage(weight, norm, min, max),
|
||||
)
|
||||
fusionExpl[hit] = append(fusionExpl[hit], expl)
|
||||
}
|
||||
hit.Score += contrib
|
||||
}
|
||||
}
|
||||
|
||||
// Finalize scores
|
||||
var maxScore float64
|
||||
for _, hit := range hits {
|
||||
if explain {
|
||||
finalizeFusionExpl(hit, fusionExpl[hit])
|
||||
}
|
||||
if hit.Score > maxScore {
|
||||
maxScore = hit.Score
|
||||
}
|
||||
hit.ScoreBreakdown = nil
|
||||
}
|
||||
|
||||
sortDocMatchesByScore(hits)
|
||||
|
||||
if nHits > windowSize {
|
||||
hits = hits[:windowSize]
|
||||
}
|
||||
|
||||
return &FusionResult{
|
||||
Hits: hits,
|
||||
Total: uint64(len(hits)),
|
||||
MaxScore: maxScore,
|
||||
}
|
||||
}
|
||||
209
fusion/rsf_test.go
Normal file
209
fusion/rsf_test.go
Normal file
@@ -0,0 +1,209 @@
|
||||
// Copyright (c) 2025 Couchbase, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package fusion
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/blevesearch/bleve/v2/search"
|
||||
)
|
||||
|
||||
func TestRelativeScoreFusion(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
hits search.DocumentMatchCollection
|
||||
weights []float64
|
||||
windowSize int
|
||||
numKNNQueries int
|
||||
want FusionResult
|
||||
}{
|
||||
{
|
||||
name: "empty hits",
|
||||
hits: search.DocumentMatchCollection{},
|
||||
weights: []float64{0.5, 0.5},
|
||||
windowSize: 10,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{},
|
||||
Total: 0,
|
||||
MaxScore: 0.0,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "single knn query",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.8}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.7}},
|
||||
},
|
||||
weights: []float64{0.4, 0.6},
|
||||
windowSize: 3,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "b", Score: 0.8}, // FTS: (0.8-0.7)/(0.9-0.7) * 0.4 + KNN: (0.9-0.7)/(0.9-0.7) * 0.6 = 0.2 + 0.6 = 0.8
|
||||
{ID: "a", Score: 0.7}, // FTS: (0.9-0.7)/(0.9-0.7) * 0.4 + KNN: (0.8-0.7)/(0.9-0.7) * 0.6 = 0.4 + 0.3 = 0.7
|
||||
{ID: "c", Score: 0.0}, // FTS: (0.7-0.7)/(0.9-0.7) * 0.4 + KNN: (0.7-0.7)/(0.9-0.7) * 0.6 = 0.0 + 0.0 = 0.0
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.8,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "multiple knn queries",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.8, 1: 0.6}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9, 1: 0.5}},
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.7, 1: 0.7}},
|
||||
},
|
||||
weights: []float64{0.3, 0.4, 0.3},
|
||||
windowSize: 3,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.65}, // FTS: (0.9-0.7)/(0.9-0.7)*0.3 + KNN0: (0.8-0.7)/(0.9-0.7)*0.4 + KNN1: (0.6-0.5)/(0.7-0.5)*0.3 = 1.0*0.3 + 0.5*0.4 + 0.5*0.3 = 0.65
|
||||
{ID: "b", Score: 0.55}, // FTS: (0.8-0.7)/(0.9-0.7)*0.3 + KNN0: (0.9-0.7)/(0.9-0.7)*0.4 + KNN1: (0.5-0.5)/(0.7-0.5)*0.3 = 0.5*0.3 + 1.0*0.4 + 0.0*0.3 = 0.55
|
||||
{ID: "c", Score: 0.3}, // FTS: (0.7-0.7)/(0.9-0.7)*0.3 + KNN0: (0.7-0.7)/(0.9-0.7)*0.4 + KNN1: (0.7-0.5)/(0.7-0.5)*0.3 = 0.0*0.3 + 0.0*0.4 + 1.0*0.3 = 0.3
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.65,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "all scores identical should normalize to 1.0",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
{ID: "c", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
},
|
||||
weights: []float64{0.4, 0.6},
|
||||
windowSize: 3,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 1.0}, // All scores identical: 1.0 * 0.4 + 1.0 * 0.6 = 1.0
|
||||
{ID: "b", Score: 1.0},
|
||||
{ID: "c", Score: 1.0},
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 1.0,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "window size smaller than hits",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.7}},
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{0: 0.9}},
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.8}},
|
||||
},
|
||||
weights: []float64{0.4, 0.6},
|
||||
windowSize: 2,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "b", Score: 0.6}, // Using top 2 for min/max: FTS min/max from [0.9,0.8] = [0.8,0.9], KNN min/max from [0.9,0.7] = [0.7,0.9]
|
||||
{ID: "a", Score: 0.4}, // FTS: (0.9-0.8)/(0.9-0.8) * 0.4 + KNN: (0.7-0.7)/(0.9-0.7) * 0.6 = 0.4 + 0 = 0.4
|
||||
},
|
||||
Total: 2,
|
||||
MaxScore: 0.6,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "documents with partial scores missing KNN scores",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.9, ScoreBreakdown: map[int]float64{0: 0.8}}, // has FTS and KNN query 0, missing KNN query 1
|
||||
{ID: "b", Score: 0.8, ScoreBreakdown: map[int]float64{1: 0.7}}, // has FTS and KNN query 1, missing KNN query 0
|
||||
{ID: "c", Score: 0.7, ScoreBreakdown: map[int]float64{0: 0.6, 1: 0.9}}, // has all scores
|
||||
{ID: "d", Score: 0.6, ScoreBreakdown: map[int]float64{}}, // has only FTS, missing all KNN scores
|
||||
},
|
||||
weights: []float64{0.3, 0.4, 0.3}, // FTS, KNN query 0, KNN query 1
|
||||
windowSize: 4,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.7}, // FTS: (0.9-0.6)/(0.9-0.6)*0.3 + KNN0: (0.8-0.6)/(0.8-0.6)*0.4 + KNN1: 0 = 1.0*0.3 + 1.0*0.4 + 0 = 0.7
|
||||
{ID: "c", Score: 0.4}, // FTS: (0.7-0.6)/(0.9-0.6)*0.3 + KNN0: (0.6-0.6)/(0.8-0.6)*0.4 + KNN1: (0.9-0.7)/(0.9-0.7)*0.3 = 0.33*0.3 + 0.0*0.4 + 1.0*0.3 = 0.1 + 0 + 0.3 = 0.4
|
||||
{ID: "b", Score: 0.2}, // FTS: (0.8-0.6)/(0.9-0.6)*0.3 + KNN0: 0 + KNN1: (0.7-0.7)/(0.9-0.7)*0.3 = 0.67*0.3 + 0 + 0.0*0.3 = 0.2 + 0 + 0 = 0.2
|
||||
{ID: "d", Score: 0.0}, // FTS: (0.6-0.6)/(0.9-0.6)*0.3 + KNN0: 0 + KNN1: 0 = 0.0*0.3 + 0 + 0 = 0
|
||||
},
|
||||
Total: 4,
|
||||
MaxScore: 0.7,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "documents with only KNN scores",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.0, ScoreBreakdown: map[int]float64{0: 0.9}}, // no FTS rank (Score 0.0), only KNN query 0
|
||||
{ID: "b", Score: 0.0, ScoreBreakdown: map[int]float64{1: 0.8}}, // no FTS rank (Score 0.0), only KNN query 1
|
||||
{ID: "c", Score: 0.0, ScoreBreakdown: map[int]float64{0: 0.7, 1: 0.6}}, // no FTS rank (Score 0.0), both KNN queries
|
||||
},
|
||||
weights: []float64{0.5, 0.3, 0.2}, // FTS, KNN query 0, KNN query 1
|
||||
windowSize: 3,
|
||||
numKNNQueries: 2,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.3}, // FTS: 0 + KNN0: 1.0 * 0.3 + KNN1: 0 = 0.3
|
||||
{ID: "b", Score: 0.2}, // FTS: 0 + KNN0: 0 + KNN1: 1.0 * 0.2 = 0.2
|
||||
{ID: "c", Score: 0.0}, // FTS: 0 + KNN0: 0 * 0.3 + KNN1: 0 * 0.2 = 0
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.3,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "mixed scenario with different score ranges",
|
||||
hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 1.0, ScoreBreakdown: map[int]float64{0: 0.1}}, // high FTS, low KNN
|
||||
{ID: "b", Score: 0.1, ScoreBreakdown: map[int]float64{0: 1.0}}, // low FTS, high KNN
|
||||
{ID: "c", Score: 0.5, ScoreBreakdown: map[int]float64{0: 0.5}}, // mid FTS, mid KNN
|
||||
},
|
||||
weights: []float64{0.5, 0.5}, // Equal weights
|
||||
windowSize: 3,
|
||||
numKNNQueries: 1,
|
||||
want: FusionResult{
|
||||
Hits: search.DocumentMatchCollection{
|
||||
{ID: "a", Score: 0.5}, // FTS: (1.0-0.1)/(1.0-0.1)*0.5 + KNN: (0.1-0.1)/(1.0-0.1)*0.5 = 1.0*0.5 + 0.0*0.5 = 0.5
|
||||
{ID: "b", Score: 0.5}, // FTS: (0.1-0.1)/(1.0-0.1)*0.5 + KNN: (1.0-0.1)/(1.0-0.1)*0.5 = 0.0*0.5 + 1.0*0.5 = 0.5
|
||||
{ID: "c", Score: 0.444}, // FTS: (0.5-0.1)/(1.0-0.1)*0.5 + KNN: (0.5-0.1)/(1.0-0.1)*0.5 = 0.444*0.5 + 0.444*0.5 = 0.444
|
||||
},
|
||||
Total: 3,
|
||||
MaxScore: 0.5,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
for i, hit := range tt.hits {
|
||||
hit.HitNumber = uint64(i)
|
||||
}
|
||||
|
||||
if got := RelativeScoreFusion(tt.hits, tt.weights, tt.windowSize, tt.numKNNQueries, false); !compareFusionResults(*got, tt.want) {
|
||||
t.Errorf("RelativeScoreFusion() = %v, want %v", got, tt.want)
|
||||
// Print detailed comparison for debugging
|
||||
t.Logf("Got hits:")
|
||||
for i, hit := range got.Hits {
|
||||
t.Logf(" [%d] ID: %s, Score: %.6f", i, hit.ID, hit.Score)
|
||||
}
|
||||
t.Logf("Want hits:")
|
||||
for i, hit := range tt.want.Hits {
|
||||
t.Logf(" [%d] ID: %s, Score: %.6f", i, hit.ID, hit.Score)
|
||||
}
|
||||
t.Logf("Got Total: %d, MaxScore: %.6f", got.Total, got.MaxScore)
|
||||
t.Logf("Want Total: %d, MaxScore: %.6f", tt.want.Total, tt.want.MaxScore)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
111
fusion/util.go
Normal file
111
fusion/util.go
Normal file
@@ -0,0 +1,111 @@
|
||||
// Copyright (c) 2025 Couchbase, Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package fusion
|
||||
|
||||
import (
|
||||
"sort"
|
||||
|
||||
"github.com/blevesearch/bleve/v2/search"
|
||||
)
|
||||
|
||||
// sortDocMatchesByScore orders the provided collection in-place by the primary
|
||||
// score in descending order, breaking ties with the original `HitNumber` to
|
||||
// ensure deterministic output.
|
||||
func sortDocMatchesByScore(hits search.DocumentMatchCollection) {
|
||||
if len(hits) < 2 {
|
||||
return
|
||||
}
|
||||
|
||||
sort.Slice(hits, func(a, b int) bool {
|
||||
i := hits[a]
|
||||
j := hits[b]
|
||||
if i.Score == j.Score {
|
||||
return i.HitNumber < j.HitNumber
|
||||
}
|
||||
return i.Score > j.Score
|
||||
})
|
||||
}
|
||||
|
||||
// scoreBreakdownForQuery fetches the score for a specific KNN query index from
|
||||
// the provided hit. The boolean return indicates whether the score is present.
|
||||
func scoreBreakdownForQuery(hit *search.DocumentMatch, idx int) (float64, bool) {
|
||||
if hit == nil || hit.ScoreBreakdown == nil {
|
||||
return 0, false
|
||||
}
|
||||
|
||||
score, ok := hit.ScoreBreakdown[idx]
|
||||
return score, ok
|
||||
}
|
||||
|
||||
// sortDocMatchesByBreakdown orders the hits in-place using the KNN score for
|
||||
// the supplied query index (descending), breaking ties with `HitNumber` and
|
||||
// placing hits without a score at the end.
|
||||
func sortDocMatchesByBreakdown(hits search.DocumentMatchCollection, queryIdx int) {
|
||||
if len(hits) < 2 {
|
||||
return
|
||||
}
|
||||
|
||||
sort.SliceStable(hits, func(a, b int) bool {
|
||||
left := hits[a]
|
||||
right := hits[b]
|
||||
|
||||
var leftScore float64
|
||||
leftOK := false
|
||||
if left != nil && left.ScoreBreakdown != nil {
|
||||
leftScore, leftOK = left.ScoreBreakdown[queryIdx]
|
||||
}
|
||||
|
||||
var rightScore float64
|
||||
rightOK := false
|
||||
if right != nil && right.ScoreBreakdown != nil {
|
||||
rightScore, rightOK = right.ScoreBreakdown[queryIdx]
|
||||
}
|
||||
|
||||
if leftOK && rightOK {
|
||||
if leftScore == rightScore {
|
||||
return left.HitNumber < right.HitNumber
|
||||
}
|
||||
return leftScore > rightScore
|
||||
}
|
||||
|
||||
if leftOK != rightOK {
|
||||
return leftOK
|
||||
}
|
||||
|
||||
return left.HitNumber < right.HitNumber
|
||||
})
|
||||
}
|
||||
|
||||
// getFusionExplAt copies the existing explanation child at the requested index
|
||||
// and wraps it in a new node describing how the fusion algorithm adjusted the
|
||||
// score.
|
||||
func getFusionExplAt(hit *search.DocumentMatch, i int, value float64, message string) *search.Explanation {
|
||||
return &search.Explanation{
|
||||
Value: value,
|
||||
Message: message,
|
||||
Children: []*search.Explanation{hit.Expl.Children[i]},
|
||||
}
|
||||
}
|
||||
|
||||
// finalizeFusionExpl installs the collection of fusion explanation children and
|
||||
// updates the root message so the caller sees the fused score as the sum of its
|
||||
// parts.
|
||||
func finalizeFusionExpl(hit *search.DocumentMatch, explChildren []*search.Explanation) {
|
||||
hit.Expl.Children = explChildren
|
||||
|
||||
hit.Expl.Value = hit.Score
|
||||
hit.Expl.Message = "sum of"
|
||||
}
|
||||
Reference in New Issue
Block a user