Skip to content

Adding MinScore support to Linear Retriever #124182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 21 commits into from
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
a41dac9
MinScore implementation in Linear retriever
mridula-s109 Mar 6, 2025
61dd8df
[CI] Auto commit changes from spotless
elasticsearchmachine Mar 6, 2025
6be41de
Merge remote-tracking branch 'origin/minscore-linear' into minscore-l…
mridula-s109 Mar 12, 2025
a52628c
Resolving PR comments
mridula-s109 Mar 13, 2025
a225c01
Fixed PR comments, added yaml and made changes to the markdown
mridula-s109 Mar 13, 2025
429a620
Merge branch 'main' into minscore-linear
mridula-s109 Mar 13, 2025
95710e6
Update docs/changelog/124182.yaml
mridula-s109 Mar 13, 2025
a32f947
[CI] Auto commit changes from spotless
elasticsearchmachine Mar 13, 2025
a424e77
Resolved on the PR comments
mridula-s109 Mar 20, 2025
38a9b50
[CI] Auto commit changes from spotless
elasticsearchmachine Mar 21, 2025
810d151
Added changes wrt to yaml testing from PR comments
mridula-s109 Mar 25, 2025
88703fd
Worked on kathleen comments first half
mridula-s109 Mar 26, 2025
d9b44e2
Reverted the integration test in line with the main branch
mridula-s109 Mar 26, 2025
46a1b94
Resolved comments in the PR and its in compiling state
mridula-s109 Mar 26, 2025
73a9bad
Unit tests passing
mridula-s109 Mar 26, 2025
6f93d3d
[CI] Auto commit changes from spotless
elasticsearchmachine Mar 26, 2025
1fd4a22
reverted inclusion of pit in retrierver it file
mridula-s109 Mar 26, 2025
f53cd0a
Removed transport versions
mridula-s109 Mar 26, 2025
9be3783
Modified rrfRank doc to the way main was
mridula-s109 Mar 27, 2025
7e2f732
Committing the changes done until now, will be doing a clean commit next
mridula-s109 Apr 2, 2025
3feaa3f
[CI] Auto commit changes from spotless
elasticsearchmachine Apr 3, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Resolving PR comments
  • Loading branch information
mridula-s109 committed Mar 13, 2025
commit a52628c609e876a97aca2d11a89a12f0bc7c3031
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,7 @@ static TransportVersion def(int id) {
public static final TransportVersion ESQL_SERIALIZE_BLOCK_TYPE_CODE = def(9_026_0_00);
public static final TransportVersion ESQL_THREAD_NAME_IN_DRIVER_PROFILE = def(9_027_0_00);
public static final TransportVersion INFERENCE_CONTEXT = def(9_028_0_00);
public static final TransportVersion RANK_DOCS_MIN_SCORE = def(9_029_0_00);

/*
* STOP! READ THIS FIRST! No, really,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -28,12 +28,17 @@
public class RankDocsQueryBuilder extends AbstractQueryBuilder<RankDocsQueryBuilder> {

public static final String NAME = "rank_docs_query";
public static final float DEFAULT_MIN_SCORE = Float.MIN_VALUE;

private final RankDoc[] rankDocs;
private final QueryBuilder[] queryBuilders;
private final boolean onlyRankDocs;
private final float minScore;

public RankDocsQueryBuilder(RankDoc[] rankDocs, QueryBuilder[] queryBuilders, boolean onlyRankDocs) {
this(rankDocs, queryBuilders, onlyRankDocs, DEFAULT_MIN_SCORE);
}

public RankDocsQueryBuilder(RankDoc[] rankDocs, QueryBuilder[] queryBuilders, boolean onlyRankDocs, float minScore) {
this.rankDocs = rankDocs;
this.queryBuilders = queryBuilders;
Expand All @@ -43,16 +48,23 @@ public RankDocsQueryBuilder(RankDoc[] rankDocs, QueryBuilder[] queryBuilders, bo

public RankDocsQueryBuilder(StreamInput in) throws IOException {
super(in);
this.rankDocs = in.readArray(c -> c.readNamedWriteable(RankDoc.class), RankDoc[]::new);
RankDoc[] rankDocs = in.readArray(c -> c.readNamedWriteable(RankDoc.class), RankDoc[]::new);
QueryBuilder[] queryBuilders = null;
boolean onlyRankDocs = false;

if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_16_0)) {
this.queryBuilders = in.readOptionalArray(c -> c.readNamedWriteable(QueryBuilder.class), QueryBuilder[]::new);
this.onlyRankDocs = in.readBoolean();
this.minScore = in.readFloat();
} else {
this.queryBuilders = null;
this.onlyRankDocs = false;
this.minScore = Float.MIN_VALUE;
queryBuilders = in.readOptionalArray(c -> c.readNamedWriteable(QueryBuilder.class), QueryBuilder[]::new);
onlyRankDocs = in.readBoolean();
}

float minScore = in.getTransportVersion().onOrAfter(TransportVersions.RANK_DOCS_MIN_SCORE)
? in.readFloat()
: DEFAULT_MIN_SCORE;

this.rankDocs = rankDocs;
this.queryBuilders = queryBuilders;
this.onlyRankDocs = onlyRankDocs;
this.minScore = minScore;
}

@Override
Expand Down Expand Up @@ -92,7 +104,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_16_0)) {
out.writeOptionalArray(StreamOutput::writeNamedWriteable, queryBuilders);
out.writeBoolean(onlyRankDocs);
out.writeFloat(minScore);
if (out.getTransportVersion().onOrAfter(TransportVersions.RANK_DOCS_MIN_SCORE)) {
out.writeFloat(minScore);
}
}
}

Expand Down Expand Up @@ -151,6 +165,6 @@ protected int doHashCode() {

@Override
public TransportVersion getMinimalSupportedVersion() {
return TransportVersions.V_8_16_0;
return TransportVersions.RANK_DOCS_MIN_SCORE;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -49,14 +49,16 @@ public static class TopQuery extends Query {
private final String[] queryNames;
private final int[] segmentStarts;
private final Object contextIdentity;
private final float minScore;

TopQuery(RankDoc[] docs, Query[] sources, String[] queryNames, int[] segmentStarts, Object contextIdentity) {
TopQuery(RankDoc[] docs, Query[] sources, String[] queryNames, int[] segmentStarts, Object contextIdentity, float minScore) {
assert sources.length == queryNames.length;
this.docs = docs;
this.sources = sources;
this.queryNames = queryNames;
this.segmentStarts = segmentStarts;
this.contextIdentity = contextIdentity;
this.minScore = minScore;
for (RankDoc doc : docs) {
if (false == doc.score >= 0) {
throw new IllegalArgumentException("RankDoc scores must be positive values. Missing a normalization step?");
Expand All @@ -76,7 +78,7 @@ public Query rewrite(IndexSearcher searcher) throws IOException {
changed |= newSources[i] != sources[i];
}
if (changed) {
return new TopQuery(docs, newSources, queryNames, segmentStarts, contextIdentity);
return new TopQuery(docs, newSources, queryNames, segmentStarts, contextIdentity, minScore);
}
return this;
}
Expand Down Expand Up @@ -165,7 +167,7 @@ public float getMaxScore(int docId) {

@Override
public float score() throws IOException {
return Math.max(docs[upTo].score, Float.MIN_VALUE);
return Math.max(docs[upTo].score, minScore);
}

@Override
Expand Down Expand Up @@ -254,7 +256,7 @@ public RankDocsQuery(
this.docs = rankDocs.clone();
// sort rank docs by doc id
Arrays.sort(docs, Comparator.comparingInt(a -> a.doc));
this.topQuery = new TopQuery(docs, sources, queryNames, findSegmentStarts(reader, docs), reader.getContext().id());
this.topQuery = new TopQuery(docs, sources, queryNames, findSegmentStarts(reader, docs), reader.getContext().id(), minScore);
if (sources.length > 0 && false == onlyRankDocs) {
var bq = new BooleanQuery.Builder();
for (var source : sources) {
Expand All @@ -268,12 +270,12 @@ public RankDocsQuery(
this.minScore = minScore;
}

private RankDocsQuery(RankDoc[] docs, Query topQuery, Query tailQuery, boolean onlyRankDocs) {
private RankDocsQuery(RankDoc[] docs, Query topQuery, Query tailQuery, boolean onlyRankDocs, float minScore) {
this.docs = docs;
this.topQuery = topQuery;
this.tailQuery = tailQuery;
this.onlyRankDocs = onlyRankDocs;
this.minScore = Float.MIN_VALUE;
this.minScore = minScore;
}

private static int binarySearch(RankDoc[] docs, int fromIndex, int toIndex, int key) {
Expand Down Expand Up @@ -317,7 +319,7 @@ public Query rewrite(IndexSearcher searcher) throws IOException {
if (tailRewrite != tailQuery) {
hasChanged = true;
}
return hasChanged ? new RankDocsQuery(docs, topRewrite, tailRewrite, onlyRankDocs) : this;
return hasChanged ? new RankDocsQuery(docs, topRewrite, tailRewrite, onlyRankDocs, minScore) : this;
}

@Override
Expand Down