From 2847ebdb75adcb183f0e3090adfa21f91015837d Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Thu, 27 Feb 2025 17:01:14 -0500 Subject: [PATCH 01/21] Add rrf plugin as ent-search plugin dependency --- x-pack/plugin/ent-search/build.gradle | 1 + x-pack/plugin/ent-search/src/main/java/module-info.java | 1 + 2 files changed, 2 insertions(+) diff --git a/x-pack/plugin/ent-search/build.gradle b/x-pack/plugin/ent-search/build.gradle index 52634ad788d97..69df705cd98f9 100644 --- a/x-pack/plugin/ent-search/build.gradle +++ b/x-pack/plugin/ent-search/build.gradle @@ -15,6 +15,7 @@ base { dependencies { compileOnly project(path: xpackModule('core')) + compileOnly project(xpackModule('rank-rrf')) implementation project(xpackModule('search-business-rules')) api project(':modules:lang-mustache') diff --git a/x-pack/plugin/ent-search/src/main/java/module-info.java b/x-pack/plugin/ent-search/src/main/java/module-info.java index 2acf0654dcdc3..1ee9cf63c9aa8 100644 --- a/x-pack/plugin/ent-search/src/main/java/module-info.java +++ b/x-pack/plugin/ent-search/src/main/java/module-info.java @@ -21,6 +21,7 @@ requires org.elasticsearch.xcore; requires org.elasticsearch.searchbusinessrules; requires org.apache.lucene.suggest; + requires org.elasticsearch.rank.rrf; exports org.elasticsearch.xpack.application; exports org.elasticsearch.xpack.application.analytics; From 46db23d78e414ee50181c6a3d4cc3add37f36abb Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 10:51:13 -0500 Subject: [PATCH 02/21] Added HybridRetrieverBuilder class --- .../rank/hybrid/HybridRetrieverBuilder.java | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java new file mode 100644 index 0000000000000..bd65bdbdee0b2 --- /dev/null +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -0,0 +1,50 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.rank.hybrid; + +import org.apache.lucene.search.ScoreDoc; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.search.rank.RankDoc; +import org.elasticsearch.search.retriever.CompoundRetrieverBuilder; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.List; + +public class HybridRetrieverBuilder extends CompoundRetrieverBuilder { + public static final String NAME = "hybrid"; + public static final ParseField FIELDS_FIELD = new ParseField("fields"); + + private final List fields; + + public HybridRetrieverBuilder(List fields, int rankWindowSize) { + super(List.of(), rankWindowSize); + this.fields = List.copyOf(fields); + } + + @Override + protected HybridRetrieverBuilder clone(List newChildRetrievers, List newPreFilterQueryBuilders) { + return null; + } + + @Override + protected RankDoc[] combineInnerRetrieverResults(List rankResults, boolean explain) { + return new RankDoc[0]; + } + + @Override + public String getName() { + return NAME; + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException { + + } +} From e74cbe8f1a525799760d4785f5a1ad0002500766 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 11:18:40 -0500 Subject: [PATCH 03/21] Added WrappedStandardRetrieverBuilder class --- .../WrappedStandardRetrieverBuilder.java | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java new file mode 100644 index 0000000000000..0ac2395823cfc --- /dev/null +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java @@ -0,0 +1,62 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.rank.hybrid; + +import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.search.builder.SearchSourceBuilder; +import org.elasticsearch.search.retriever.RetrieverBuilder; +import org.elasticsearch.search.retriever.StandardRetrieverBuilder; +import org.elasticsearch.xcontent.XContentBuilder; + +import java.io.IOException; +import java.util.Objects; + +public class WrappedStandardRetrieverBuilder extends RetrieverBuilder { + private final StandardRetrieverBuilder wrappedRetriever; + private final QueryBuilder queryBuilder; + private final String field; + + public WrappedStandardRetrieverBuilder(String field) { + Objects.requireNonNull(field); + this.field = field; + this.queryBuilder = new BoolQueryBuilder(); + this.wrappedRetriever = new StandardRetrieverBuilder(this.queryBuilder); + } + + @Override + public QueryBuilder topDocsQuery() { + return wrappedRetriever.topDocsQuery(); + } + + @Override + public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder, boolean compoundUsed) { + wrappedRetriever.extractToSearchSourceBuilder(searchSourceBuilder, compoundUsed); + } + + @Override + public String getName() { + return wrappedRetriever.getName(); + } + + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException { + wrappedRetriever.doToXContent(builder, params); + } + + @Override + protected boolean doEquals(Object o) { + WrappedStandardRetrieverBuilder that = (WrappedStandardRetrieverBuilder) o; + return field.equals(that.field) && wrappedRetriever.equals(that.wrappedRetriever); + } + + @Override + protected int doHashCode() { + return Objects.hash(field, wrappedRetriever); + } +} From b58a022180e8b2ccb71a1afda2e3d2a16b737e3a Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 12:39:05 -0500 Subject: [PATCH 04/21] Refactored wrapper class --- .../StandardRetrieverBuilderWrapper.java | 47 ++++++++++++++ .../WrappedStandardRetrieverBuilder.java | 62 ------------------- 2 files changed, 47 insertions(+), 62 deletions(-) create mode 100644 x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java delete mode 100644 x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java new file mode 100644 index 0000000000000..1bc1d34a92ca5 --- /dev/null +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java @@ -0,0 +1,47 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.rank.hybrid; + +import org.elasticsearch.index.query.BoolQueryBuilder; +import org.elasticsearch.search.retriever.RetrieverBuilder; +import org.elasticsearch.search.retriever.RetrieverBuilderWrapper; +import org.elasticsearch.search.retriever.StandardRetrieverBuilder; + +import java.util.Objects; + +public class StandardRetrieverBuilderWrapper extends RetrieverBuilderWrapper { + private final String field; + + public StandardRetrieverBuilderWrapper(String field) { + super(new StandardRetrieverBuilder(new BoolQueryBuilder())); + + Objects.requireNonNull(field); + this.field = field; + } + + private StandardRetrieverBuilderWrapper(String field, RetrieverBuilder retrieverBuilder) { + super(retrieverBuilder); + this.field = field; + } + + @Override + protected StandardRetrieverBuilderWrapper clone(RetrieverBuilder sub) { + return new StandardRetrieverBuilderWrapper(field, sub); + } + + @Override + protected boolean doEquals(Object o) { + StandardRetrieverBuilderWrapper that = (StandardRetrieverBuilderWrapper) o; + return field.equals(that.field) && super.doEquals(o); + } + + @Override + protected int doHashCode() { + return Objects.hash(field, super.doHashCode()); + } +} diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java deleted file mode 100644 index 0ac2395823cfc..0000000000000 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/WrappedStandardRetrieverBuilder.java +++ /dev/null @@ -1,62 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.xpack.rank.hybrid; - -import org.elasticsearch.index.query.BoolQueryBuilder; -import org.elasticsearch.index.query.QueryBuilder; -import org.elasticsearch.search.builder.SearchSourceBuilder; -import org.elasticsearch.search.retriever.RetrieverBuilder; -import org.elasticsearch.search.retriever.StandardRetrieverBuilder; -import org.elasticsearch.xcontent.XContentBuilder; - -import java.io.IOException; -import java.util.Objects; - -public class WrappedStandardRetrieverBuilder extends RetrieverBuilder { - private final StandardRetrieverBuilder wrappedRetriever; - private final QueryBuilder queryBuilder; - private final String field; - - public WrappedStandardRetrieverBuilder(String field) { - Objects.requireNonNull(field); - this.field = field; - this.queryBuilder = new BoolQueryBuilder(); - this.wrappedRetriever = new StandardRetrieverBuilder(this.queryBuilder); - } - - @Override - public QueryBuilder topDocsQuery() { - return wrappedRetriever.topDocsQuery(); - } - - @Override - public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder, boolean compoundUsed) { - wrappedRetriever.extractToSearchSourceBuilder(searchSourceBuilder, compoundUsed); - } - - @Override - public String getName() { - return wrappedRetriever.getName(); - } - - @Override - protected void doToXContent(XContentBuilder builder, Params params) throws IOException { - wrappedRetriever.doToXContent(builder, params); - } - - @Override - protected boolean doEquals(Object o) { - WrappedStandardRetrieverBuilder that = (WrappedStandardRetrieverBuilder) o; - return field.equals(that.field) && wrappedRetriever.equals(that.wrappedRetriever); - } - - @Override - protected int doHashCode() { - return Objects.hash(field, wrappedRetriever); - } -} From 83d0d9a3aea6f74f87c2db8f6f55c3550bab8013 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 12:39:29 -0500 Subject: [PATCH 05/21] Fix equality/hashcode bug --- .../elasticsearch/search/retriever/RetrieverBuilder.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java index ce852a44c28ec..c639dd1b7f2cf 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java @@ -282,8 +282,8 @@ public final boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; RetrieverBuilder that = (RetrieverBuilder) o; - return Objects.equals(preFilterQueryBuilders, that.preFilterQueryBuilders) - && Objects.equals(minScore, that.minScore) + return Objects.equals(getPreFilterQueryBuilders(), that.getPreFilterQueryBuilders()) + && Objects.equals(minScore(), that.minScore()) && doEquals(o); } @@ -291,7 +291,7 @@ public final boolean equals(Object o) { @Override public final int hashCode() { - return Objects.hash(getClass(), preFilterQueryBuilders, minScore, doHashCode()); + return Objects.hash(getClass(), getPreFilterQueryBuilders(), minScore(), doHashCode()); } protected abstract int doHashCode(); From c41af7de0726d2e875d8ee0ecf4bd4028df43251 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 12:39:49 -0500 Subject: [PATCH 06/21] Add query builder getter --- .../search/retriever/StandardRetrieverBuilder.java | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java index 3ca74dc133d47..c7ebc94216a76 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java @@ -104,6 +104,10 @@ private StandardRetrieverBuilder(StandardRetrieverBuilder clone) { this.terminateAfter = clone.terminateAfter; } + public QueryBuilder getQueryBuilder() { + return queryBuilder; + } + @Override public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { boolean changed = false; From 2d9c0cd51c0bee3cf7b9f2a332865728e2551e38 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 14:12:40 -0500 Subject: [PATCH 07/21] Updated HybridRetrieverBuilder to be a wrapper --- .../rank/hybrid/HybridRetrieverBuilder.java | 88 +++++++++++++++---- 1 file changed, 73 insertions(+), 15 deletions(-) diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index bd65bdbdee0b2..b783e4a1f0532 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -7,35 +7,52 @@ package org.elasticsearch.xpack.rank.hybrid; -import org.apache.lucene.search.ScoreDoc; -import org.elasticsearch.index.query.QueryBuilder; -import org.elasticsearch.search.rank.RankDoc; +import org.elasticsearch.index.query.MatchQueryBuilder; import org.elasticsearch.search.retriever.CompoundRetrieverBuilder; +import org.elasticsearch.search.retriever.RetrieverBuilder; +import org.elasticsearch.search.retriever.RetrieverBuilderWrapper; +import org.elasticsearch.search.retriever.StandardRetrieverBuilder; import org.elasticsearch.xcontent.ParseField; -import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xpack.rank.linear.LinearRetrieverBuilder; +import org.elasticsearch.xpack.rank.linear.MinMaxScoreNormalizer; +import org.elasticsearch.xpack.rank.linear.ScoreNormalizer; -import java.io.IOException; +import java.util.ArrayList; +import java.util.Arrays; import java.util.List; +import java.util.Objects; -public class HybridRetrieverBuilder extends CompoundRetrieverBuilder { +public class HybridRetrieverBuilder extends RetrieverBuilderWrapper { public static final String NAME = "hybrid"; public static final ParseField FIELDS_FIELD = new ParseField("fields"); + public static final ParseField QUERY_FIELD = new ParseField("query"); private final List fields; + private final String query; - public HybridRetrieverBuilder(List fields, int rankWindowSize) { - super(List.of(), rankWindowSize); - this.fields = List.copyOf(fields); + public HybridRetrieverBuilder(List fields, String query, int rankWindowSize) { + super( + new LinearRetrieverBuilder( + generateInnerRetrievers(fields, query), + rankWindowSize, + generateWeights(fields), + generateScoreNormalizers(fields) + ) + ); + this.fields = fields == null ? List.of() : List.copyOf(fields); + this.query = query; } - @Override - protected HybridRetrieverBuilder clone(List newChildRetrievers, List newPreFilterQueryBuilders) { - return null; + private HybridRetrieverBuilder(List fields, String query, RetrieverBuilder retrieverBuilder) { + super(retrieverBuilder); + this.fields = fields; + this.query = query; } @Override - protected RankDoc[] combineInnerRetrieverResults(List rankResults, boolean explain) { - return new RankDoc[0]; + protected HybridRetrieverBuilder clone(RetrieverBuilder sub) { + assert sub instanceof LinearRetrieverBuilder; + return new HybridRetrieverBuilder(fields, query, sub); } @Override @@ -44,7 +61,48 @@ public String getName() { } @Override - protected void doToXContent(XContentBuilder builder, Params params) throws IOException { + protected boolean doEquals(Object o) { + HybridRetrieverBuilder that = (HybridRetrieverBuilder) o; + return Objects.equals(fields, that.fields) && Objects.equals(query, that.query) && super.doEquals(o); + } + + @Override + protected int doHashCode() { + return Objects.hash(fields, query, super.doHashCode()); + } + + private static List generateInnerRetrievers(List fields, String query) { + if (fields == null) { + return List.of(); + } + + List innerRetrievers = new ArrayList<>(fields.size()); + for (String field : fields) { + MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(field, query); + innerRetrievers.add(new CompoundRetrieverBuilder.RetrieverSource(new StandardRetrieverBuilder(matchQueryBuilder), null)); + } + + return innerRetrievers; + } + + private static float[] generateWeights(List fields) { + if (fields == null) { + return new float[0]; + } + + // TODO: Parse field strings for weights + float[] weights = new float[fields.size()]; + Arrays.fill(weights, 1.0f); + return weights; + } + + private static ScoreNormalizer[] generateScoreNormalizers(List fields) { + if (fields == null) { + return new ScoreNormalizer[0]; + } + ScoreNormalizer[] scoreNormalizers = new ScoreNormalizer[fields.size()]; + Arrays.fill(scoreNormalizers, MinMaxScoreNormalizer.INSTANCE); + return scoreNormalizers; } } From 1e083fd2fac272f905645b23c6497394e9ba8a26 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 14:44:30 -0500 Subject: [PATCH 08/21] Implemented doToXContent --- .../rank/hybrid/HybridRetrieverBuilder.java | 26 ++++++++++++++----- 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index b783e4a1f0532..f3e28d287efbe 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -13,15 +13,19 @@ import org.elasticsearch.search.retriever.RetrieverBuilderWrapper; import org.elasticsearch.search.retriever.StandardRetrieverBuilder; import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xpack.rank.linear.LinearRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.MinMaxScoreNormalizer; import org.elasticsearch.xpack.rank.linear.ScoreNormalizer; +import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Objects; +import static org.elasticsearch.search.retriever.CompoundRetrieverBuilder.RANK_WINDOW_SIZE_FIELD; + public class HybridRetrieverBuilder extends RetrieverBuilderWrapper { public static final String NAME = "hybrid"; public static final ParseField FIELDS_FIELD = new ParseField("fields"); @@ -29,9 +33,13 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper fields; private final String query; + private final int rankWindowSize; public HybridRetrieverBuilder(List fields, String query, int rankWindowSize) { - super( + this( + fields == null ? List.of() : List.copyOf(fields), + query, + rankWindowSize, new LinearRetrieverBuilder( generateInnerRetrievers(fields, query), rankWindowSize, @@ -39,20 +47,18 @@ public HybridRetrieverBuilder(List fields, String query, int rankWindowS generateScoreNormalizers(fields) ) ); - this.fields = fields == null ? List.of() : List.copyOf(fields); - this.query = query; } - private HybridRetrieverBuilder(List fields, String query, RetrieverBuilder retrieverBuilder) { + private HybridRetrieverBuilder(List fields, String query, int rankWindowSize, RetrieverBuilder retrieverBuilder) { super(retrieverBuilder); this.fields = fields; this.query = query; + this.rankWindowSize = rankWindowSize; } @Override protected HybridRetrieverBuilder clone(RetrieverBuilder sub) { - assert sub instanceof LinearRetrieverBuilder; - return new HybridRetrieverBuilder(fields, query, sub); + return new HybridRetrieverBuilder(fields, query, rankWindowSize, sub); } @Override @@ -60,8 +66,16 @@ public String getName() { return NAME; } + @Override + protected void doToXContent(XContentBuilder builder, Params params) throws IOException { + builder.field(FIELDS_FIELD.getPreferredName(), fields); + builder.field(QUERY_FIELD.getPreferredName(), query); + builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize); + } + @Override protected boolean doEquals(Object o) { + // TODO: Check rankWindowSize? It should be checked by the wrapped retriever. HybridRetrieverBuilder that = (HybridRetrieverBuilder) o; return Objects.equals(fields, that.fields) && Objects.equals(query, that.query) && super.doEquals(o); } From 848614a7da23fb88d11feca05fee5b38dc5c8625 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 14:44:55 -0500 Subject: [PATCH 09/21] Added TODO --- .../elasticsearch/xpack/rank/linear/LinearRetrieverBuilder.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverBuilder.java index 66bbbf95bc9d6..e83d688e806da 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverBuilder.java @@ -205,4 +205,6 @@ public void doToXContent(XContentBuilder builder, Params params) throws IOExcept } builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize); } + + // TODO: Need doEquals & doHashCode to check weights and normalizers } From f5fc29ece607f648843949d0a396590960f7c54a Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 15:02:51 -0500 Subject: [PATCH 10/21] Implemented fromXContent --- .../rank/hybrid/HybridRetrieverBuilder.java | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index f3e28d287efbe..fe6615f0805be 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -8,12 +8,16 @@ package org.elasticsearch.xpack.rank.hybrid; import org.elasticsearch.index.query.MatchQueryBuilder; +import org.elasticsearch.search.rank.RankBuilder; import org.elasticsearch.search.retriever.CompoundRetrieverBuilder; import org.elasticsearch.search.retriever.RetrieverBuilder; import org.elasticsearch.search.retriever.RetrieverBuilderWrapper; +import org.elasticsearch.search.retriever.RetrieverParserContext; import org.elasticsearch.search.retriever.StandardRetrieverBuilder; +import org.elasticsearch.xcontent.ConstructingObjectParser; import org.elasticsearch.xcontent.ParseField; import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParser; import org.elasticsearch.xpack.rank.linear.LinearRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.MinMaxScoreNormalizer; import org.elasticsearch.xpack.rank.linear.ScoreNormalizer; @@ -25,6 +29,11 @@ import java.util.Objects; import static org.elasticsearch.search.retriever.CompoundRetrieverBuilder.RANK_WINDOW_SIZE_FIELD; +import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; + +// TODO: +// - Retriever name support public class HybridRetrieverBuilder extends RetrieverBuilderWrapper { public static final String NAME = "hybrid"; @@ -35,6 +44,25 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper PARSER = new ConstructingObjectParser<>( + NAME, + false, + (args, context) -> { + List fields = (List) args[0]; + String query = (String) args[1]; + int rankWindowSize = args[2] == null ? RankBuilder.DEFAULT_RANK_WINDOW_SIZE : (int) args[2]; + return new HybridRetrieverBuilder(fields, query, rankWindowSize); + } + ); + + static { + PARSER.declareStringArray(constructorArg(), FIELDS_FIELD); + PARSER.declareString(constructorArg(), QUERY_FIELD); + PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD); + RetrieverBuilder.declareBaseParserFields(NAME, PARSER); + } + public HybridRetrieverBuilder(List fields, String query, int rankWindowSize) { this( fields == null ? List.of() : List.copyOf(fields), @@ -85,6 +113,10 @@ protected int doHashCode() { return Objects.hash(fields, query, super.doHashCode()); } + public static HybridRetrieverBuilder fromXContent(XContentParser parser, RetrieverParserContext context) throws IOException { + return PARSER.apply(parser, context); + } + private static List generateInnerRetrievers(List fields, String query) { if (fields == null) { return List.of(); From de1b185935baf8d715ce3666563c4f2c476589a2 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 15:06:34 -0500 Subject: [PATCH 11/21] Add hybrid retriever to enterprise search plugin --- .../elasticsearch/xpack/application/EnterpriseSearch.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java index db6ee3a621d84..674d79879938c 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java @@ -202,6 +202,7 @@ import org.elasticsearch.xpack.core.XPackSettings; import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction; import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction; +import org.elasticsearch.xpack.rank.hybrid.HybridRetrieverBuilder; import java.util.ArrayList; import java.util.Arrays; @@ -521,6 +522,9 @@ public List> getQueries() { @Override public List> getRetrievers() { - return List.of(new RetrieverSpec<>(new ParseField(QueryRuleRetrieverBuilder.NAME), QueryRuleRetrieverBuilder::fromXContent)); + return List.of( + new RetrieverSpec<>(new ParseField(QueryRuleRetrieverBuilder.NAME), QueryRuleRetrieverBuilder::fromXContent), + new RetrieverSpec<>(new ParseField(HybridRetrieverBuilder.NAME), HybridRetrieverBuilder::fromXContent) + ); } } From 1f41e8d016d4860d9786e42ebbf24f4ffb095f33 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Fri, 28 Feb 2025 16:13:05 -0500 Subject: [PATCH 12/21] Make rank-rrf plugin extensible --- x-pack/plugin/ent-search/build.gradle | 2 +- .../java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/ent-search/build.gradle b/x-pack/plugin/ent-search/build.gradle index 69df705cd98f9..d11b64475394a 100644 --- a/x-pack/plugin/ent-search/build.gradle +++ b/x-pack/plugin/ent-search/build.gradle @@ -6,7 +6,7 @@ esplugin { name = 'x-pack-ent-search' description = 'Elasticsearch Expanded Pack Plugin - Enterprise Search' classname = 'org.elasticsearch.xpack.application.EnterpriseSearch' - extendedPlugins = ['x-pack-core'] + extendedPlugins = ['x-pack-core', 'rank-rrf'] } base { diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java index 251015b21ff50..58ab23d0fd313 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java @@ -10,6 +10,7 @@ import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.license.License; import org.elasticsearch.license.LicensedFeature; +import org.elasticsearch.plugins.ExtensiblePlugin; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.SearchPlugin; import org.elasticsearch.search.rank.RankBuilder; @@ -22,7 +23,7 @@ import java.util.List; -public class RRFRankPlugin extends Plugin implements SearchPlugin { +public class RRFRankPlugin extends Plugin implements SearchPlugin, ExtensiblePlugin { public static final LicensedFeature.Momentary RANK_RRF_FEATURE = LicensedFeature.momentary( null, From 578028c9e70b83d0ac8b0867f29a8a32d79262e0 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Thu, 6 Mar 2025 17:26:01 -0500 Subject: [PATCH 13/21] Updated MinMaxScoreNormalizer to make the initial minimum configurable --- .../xpack/rank/hybrid/HybridRetrieverBuilder.java | 2 +- .../xpack/rank/linear/MinMaxScoreNormalizer.java | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index fe6615f0805be..f64ebfa8b9164 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -148,7 +148,7 @@ private static ScoreNormalizer[] generateScoreNormalizers(List fields) { } ScoreNormalizer[] scoreNormalizers = new ScoreNormalizer[fields.size()]; - Arrays.fill(scoreNormalizers, MinMaxScoreNormalizer.INSTANCE); + Arrays.fill(scoreNormalizers, new MinMaxScoreNormalizer(0)); return scoreNormalizers; } } diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/MinMaxScoreNormalizer.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/MinMaxScoreNormalizer.java index 56b42b48a5d47..7ca27eef28717 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/MinMaxScoreNormalizer.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/MinMaxScoreNormalizer.java @@ -17,7 +17,15 @@ public class MinMaxScoreNormalizer extends ScoreNormalizer { private static final float EPSILON = 1e-6f; - public MinMaxScoreNormalizer() {} + private final float initialMin; + + public MinMaxScoreNormalizer() { + this(Float.MAX_VALUE); + } + + public MinMaxScoreNormalizer(float initialMin) { + this.initialMin = initialMin; + } @Override public String getName() { @@ -31,7 +39,7 @@ public ScoreDoc[] normalizeScores(ScoreDoc[] docs) { } // create a new array to avoid changing ScoreDocs in place ScoreDoc[] scoreDocs = new ScoreDoc[docs.length]; - float min = Float.MAX_VALUE; + float min = initialMin; float max = Float.MIN_VALUE; boolean atLeastOneValidScore = false; for (ScoreDoc rd : docs) { From 4f06a7f11ce148c6b1602a3d97cca8fbd2aa5e96 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Mon, 10 Mar 2025 10:19:47 -0400 Subject: [PATCH 14/21] Removed StandardRetrieverBuilderWrapper --- .../StandardRetrieverBuilderWrapper.java | 47 ------------------- 1 file changed, 47 deletions(-) delete mode 100644 x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java deleted file mode 100644 index 1bc1d34a92ca5..0000000000000 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/StandardRetrieverBuilderWrapper.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.xpack.rank.hybrid; - -import org.elasticsearch.index.query.BoolQueryBuilder; -import org.elasticsearch.search.retriever.RetrieverBuilder; -import org.elasticsearch.search.retriever.RetrieverBuilderWrapper; -import org.elasticsearch.search.retriever.StandardRetrieverBuilder; - -import java.util.Objects; - -public class StandardRetrieverBuilderWrapper extends RetrieverBuilderWrapper { - private final String field; - - public StandardRetrieverBuilderWrapper(String field) { - super(new StandardRetrieverBuilder(new BoolQueryBuilder())); - - Objects.requireNonNull(field); - this.field = field; - } - - private StandardRetrieverBuilderWrapper(String field, RetrieverBuilder retrieverBuilder) { - super(retrieverBuilder); - this.field = field; - } - - @Override - protected StandardRetrieverBuilderWrapper clone(RetrieverBuilder sub) { - return new StandardRetrieverBuilderWrapper(field, sub); - } - - @Override - protected boolean doEquals(Object o) { - StandardRetrieverBuilderWrapper that = (StandardRetrieverBuilderWrapper) o; - return field.equals(that.field) && super.doEquals(o); - } - - @Override - protected int doHashCode() { - return Objects.hash(field, super.doHashCode()); - } -} From c2b7ffeeb9978706e5aa5f0971dbaebdcf157001 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Mon, 10 Mar 2025 12:20:05 -0400 Subject: [PATCH 15/21] Move hybrid retriever to rank-rrf plugin --- x-pack/plugin/ent-search/build.gradle | 3 +-- x-pack/plugin/ent-search/src/main/java/module-info.java | 1 - .../elasticsearch/xpack/application/EnterpriseSearch.java | 6 +----- .../xpack/rank/hybrid/HybridRetrieverBuilder.java | 0 .../org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java | 7 ++++--- 5 files changed, 6 insertions(+), 11 deletions(-) rename x-pack/plugin/{ent-search => rank-rrf}/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java (100%) diff --git a/x-pack/plugin/ent-search/build.gradle b/x-pack/plugin/ent-search/build.gradle index d11b64475394a..52634ad788d97 100644 --- a/x-pack/plugin/ent-search/build.gradle +++ b/x-pack/plugin/ent-search/build.gradle @@ -6,7 +6,7 @@ esplugin { name = 'x-pack-ent-search' description = 'Elasticsearch Expanded Pack Plugin - Enterprise Search' classname = 'org.elasticsearch.xpack.application.EnterpriseSearch' - extendedPlugins = ['x-pack-core', 'rank-rrf'] + extendedPlugins = ['x-pack-core'] } base { @@ -15,7 +15,6 @@ base { dependencies { compileOnly project(path: xpackModule('core')) - compileOnly project(xpackModule('rank-rrf')) implementation project(xpackModule('search-business-rules')) api project(':modules:lang-mustache') diff --git a/x-pack/plugin/ent-search/src/main/java/module-info.java b/x-pack/plugin/ent-search/src/main/java/module-info.java index 1ee9cf63c9aa8..2acf0654dcdc3 100644 --- a/x-pack/plugin/ent-search/src/main/java/module-info.java +++ b/x-pack/plugin/ent-search/src/main/java/module-info.java @@ -21,7 +21,6 @@ requires org.elasticsearch.xcore; requires org.elasticsearch.searchbusinessrules; requires org.apache.lucene.suggest; - requires org.elasticsearch.rank.rrf; exports org.elasticsearch.xpack.application; exports org.elasticsearch.xpack.application.analytics; diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java index 674d79879938c..db6ee3a621d84 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/EnterpriseSearch.java @@ -202,7 +202,6 @@ import org.elasticsearch.xpack.core.XPackSettings; import org.elasticsearch.xpack.core.action.XPackInfoFeatureAction; import org.elasticsearch.xpack.core.action.XPackUsageFeatureAction; -import org.elasticsearch.xpack.rank.hybrid.HybridRetrieverBuilder; import java.util.ArrayList; import java.util.Arrays; @@ -522,9 +521,6 @@ public List> getQueries() { @Override public List> getRetrievers() { - return List.of( - new RetrieverSpec<>(new ParseField(QueryRuleRetrieverBuilder.NAME), QueryRuleRetrieverBuilder::fromXContent), - new RetrieverSpec<>(new ParseField(HybridRetrieverBuilder.NAME), HybridRetrieverBuilder::fromXContent) - ); + return List.of(new RetrieverSpec<>(new ParseField(QueryRuleRetrieverBuilder.NAME), QueryRuleRetrieverBuilder::fromXContent)); } } diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java similarity index 100% rename from x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java rename to x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java index 58ab23d0fd313..5cd9dcaad5339 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankPlugin.java @@ -10,7 +10,6 @@ import org.elasticsearch.common.io.stream.NamedWriteableRegistry; import org.elasticsearch.license.License; import org.elasticsearch.license.LicensedFeature; -import org.elasticsearch.plugins.ExtensiblePlugin; import org.elasticsearch.plugins.Plugin; import org.elasticsearch.plugins.SearchPlugin; import org.elasticsearch.search.rank.RankBuilder; @@ -18,12 +17,13 @@ import org.elasticsearch.search.rank.RankShardResult; import org.elasticsearch.xcontent.NamedXContentRegistry; import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xpack.rank.hybrid.HybridRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.LinearRankDoc; import org.elasticsearch.xpack.rank.linear.LinearRetrieverBuilder; import java.util.List; -public class RRFRankPlugin extends Plugin implements SearchPlugin, ExtensiblePlugin { +public class RRFRankPlugin extends Plugin implements SearchPlugin { public static final LicensedFeature.Momentary RANK_RRF_FEATURE = LicensedFeature.momentary( null, @@ -58,7 +58,8 @@ public List getNamedXContent() { public List> getRetrievers() { return List.of( new RetrieverSpec<>(new ParseField(NAME), RRFRetrieverBuilder::fromXContent), - new RetrieverSpec<>(new ParseField(LinearRetrieverBuilder.NAME), LinearRetrieverBuilder::fromXContent) + new RetrieverSpec<>(new ParseField(LinearRetrieverBuilder.NAME), LinearRetrieverBuilder::fromXContent), + new RetrieverSpec<>(new ParseField(HybridRetrieverBuilder.NAME), HybridRetrieverBuilder::fromXContent) ); } } From a7b9e86230ca47102f8c518f2fe08da380987606 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Mon, 10 Mar 2025 12:50:57 -0400 Subject: [PATCH 16/21] Made inference plugin a dependency of the rank-rrf plugin --- x-pack/plugin/inference/src/main/java/module-info.java | 1 + x-pack/plugin/rank-rrf/build.gradle | 3 ++- x-pack/plugin/rank-rrf/src/main/java/module-info.java | 1 + .../xpack/rank/hybrid/HybridRetrieverBuilder.java | 2 +- 4 files changed, 5 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/module-info.java b/x-pack/plugin/inference/src/main/java/module-info.java index 78f30e7da0670..a9a3f8bc9ec39 100644 --- a/x-pack/plugin/inference/src/main/java/module-info.java +++ b/x-pack/plugin/inference/src/main/java/module-info.java @@ -43,6 +43,7 @@ exports org.elasticsearch.xpack.inference; exports org.elasticsearch.xpack.inference.action.task; exports org.elasticsearch.xpack.inference.telemetry; + exports org.elasticsearch.xpack.inference.rank.textsimilarity; provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.inference.InferenceFeatures; } diff --git a/x-pack/plugin/rank-rrf/build.gradle b/x-pack/plugin/rank-rrf/build.gradle index 216e85f48f56f..0588140b47ce8 100644 --- a/x-pack/plugin/rank-rrf/build.gradle +++ b/x-pack/plugin/rank-rrf/build.gradle @@ -13,11 +13,12 @@ esplugin { name = 'rank-rrf' description = 'Reciprocal rank fusion in search.' classname ='org.elasticsearch.xpack.rank.rrf.RRFRankPlugin' - extendedPlugins = ['x-pack-core'] + extendedPlugins = ['x-pack-core', 'x-pack-inference'] } dependencies { compileOnly project(path: xpackModule('core')) + compileOnly project(path: xpackModule('inference')) testImplementation(testArtifact(project(xpackModule('core')))) testImplementation(testArtifact(project(':server'))) diff --git a/x-pack/plugin/rank-rrf/src/main/java/module-info.java b/x-pack/plugin/rank-rrf/src/main/java/module-info.java index fbe467fdf3eae..42bfcdcdc19e4 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/module-info.java +++ b/x-pack/plugin/rank-rrf/src/main/java/module-info.java @@ -13,6 +13,7 @@ requires org.elasticsearch.xcontent; requires org.elasticsearch.server; requires org.elasticsearch.xcore; + requires org.elasticsearch.inference; exports org.elasticsearch.xpack.rank; exports org.elasticsearch.xpack.rank.rrf; diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index f64ebfa8b9164..b729bcba4712a 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -60,7 +60,7 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper fields, String query, int rankWindowSize) { From 4e0f5c5ba682c1a4bbc222e0e76cbc2f159ad33c Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Mon, 10 Mar 2025 13:03:10 -0400 Subject: [PATCH 17/21] Added rerank and rerank_inference_id fields --- .../rank/hybrid/HybridRetrieverBuilder.java | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index b729bcba4712a..5d01228939dd0 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -39,9 +39,13 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper fields; private final String query; + private final boolean rerank; + private final String rerankInferenceId; private final int rankWindowSize; @SuppressWarnings("unchecked") @@ -51,22 +55,28 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper { List fields = (List) args[0]; String query = (String) args[1]; - int rankWindowSize = args[2] == null ? RankBuilder.DEFAULT_RANK_WINDOW_SIZE : (int) args[2]; - return new HybridRetrieverBuilder(fields, query, rankWindowSize); + boolean rerank = args[2] != null && (boolean) args[2]; + String rerankInferenceId = (String) args[3]; + int rankWindowSize = args[4] == null ? RankBuilder.DEFAULT_RANK_WINDOW_SIZE : (int) args[4]; + return new HybridRetrieverBuilder(fields, query, rerank, rerankInferenceId, rankWindowSize); } ); static { PARSER.declareStringArray(constructorArg(), FIELDS_FIELD); PARSER.declareString(constructorArg(), QUERY_FIELD); + PARSER.declareBoolean(optionalConstructorArg(), RERANK_FIELD); + PARSER.declareString(optionalConstructorArg(), RERANK_INFERENCE_ID_FIELD); PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD); RetrieverBuilder.declareBaseParserFields(PARSER); } - public HybridRetrieverBuilder(List fields, String query, int rankWindowSize) { + public HybridRetrieverBuilder(List fields, String query, boolean rerank, String rerankInferenceId, int rankWindowSize) { this( fields == null ? List.of() : List.copyOf(fields), query, + rerank, + rerankInferenceId, rankWindowSize, new LinearRetrieverBuilder( generateInnerRetrievers(fields, query), @@ -77,16 +87,25 @@ public HybridRetrieverBuilder(List fields, String query, int rankWindowS ); } - private HybridRetrieverBuilder(List fields, String query, int rankWindowSize, RetrieverBuilder retrieverBuilder) { + private HybridRetrieverBuilder( + List fields, + String query, + boolean rerank, + String rerankInferenceId, + int rankWindowSize, + RetrieverBuilder retrieverBuilder + ) { super(retrieverBuilder); this.fields = fields; this.query = query; + this.rerank = rerank; + this.rerankInferenceId = rerankInferenceId; this.rankWindowSize = rankWindowSize; } @Override protected HybridRetrieverBuilder clone(RetrieverBuilder sub) { - return new HybridRetrieverBuilder(fields, query, rankWindowSize, sub); + return new HybridRetrieverBuilder(fields, query, rerank, rerankInferenceId, rankWindowSize, sub); } @Override From 27e349738fe6a3cb0509258509e245b311487be7 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Mon, 10 Mar 2025 14:39:35 -0400 Subject: [PATCH 18/21] Integrated text similarity reranker into hybrid retriever --- .../TextSimilarityRankRetrieverBuilder.java | 25 ++++-- .../rank/hybrid/HybridRetrieverBuilder.java | 90 +++++++++++++++---- 2 files changed, 91 insertions(+), 24 deletions(-) diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java index d6883d3743a1d..51fd95b263be9 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java @@ -23,6 +23,7 @@ import org.elasticsearch.xcontent.XContentParser; import java.io.IOException; +import java.util.ArrayList; import java.util.List; import java.util.Objects; @@ -49,7 +50,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder public static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>(TextSimilarityRankBuilder.NAME, args -> { RetrieverBuilder retrieverBuilder = (RetrieverBuilder) args[0]; - String inferenceId = args[1] == null ? DEFAULT_RERANK_ID : (String) args[1]; + String inferenceId = (String) args[1]; String inferenceText = (String) args[2]; String field = (String) args[3]; int rankWindowSize = args[4] == null ? DEFAULT_RANK_WINDOW_SIZE : (int) args[4]; @@ -104,11 +105,17 @@ public TextSimilarityRankRetrieverBuilder( int rankWindowSize, boolean failuresAllowed ) { - super(List.of(new RetrieverSource(retrieverBuilder, null)), rankWindowSize); - this.inferenceId = inferenceId; - this.inferenceText = inferenceText; - this.field = field; - this.failuresAllowed = failuresAllowed; + this( + List.of(new RetrieverSource(retrieverBuilder, null)), + inferenceId, + inferenceText, + field, + rankWindowSize, + null, + failuresAllowed, + null, + new ArrayList<>() + ); } public TextSimilarityRankRetrieverBuilder( @@ -124,9 +131,11 @@ public TextSimilarityRankRetrieverBuilder( ) { super(retrieverSource, rankWindowSize); if (retrieverSource.size() != 1) { - throw new IllegalArgumentException("[" + getName() + "] retriever should have exactly one inner retriever"); + throw new IllegalArgumentException( + "[" + TextSimilarityRankBuilder.NAME + "] retriever should have exactly one inner retriever" + ); } - this.inferenceId = inferenceId; + this.inferenceId = inferenceId == null ? DEFAULT_RERANK_ID : inferenceId; this.inferenceText = inferenceText; this.field = field; this.minScore = minScore; diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index 5d01228939dd0..963187166c51b 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -18,6 +18,7 @@ import org.elasticsearch.xcontent.ParseField; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentParser; +import org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.LinearRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.MinMaxScoreNormalizer; import org.elasticsearch.xpack.rank.linear.ScoreNormalizer; @@ -40,11 +41,13 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper fields; private final String query; - private final boolean rerank; + private final Boolean rerank; + private final String rerankField; private final String rerankInferenceId; private final int rankWindowSize; @@ -55,10 +58,11 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper { List fields = (List) args[0]; String query = (String) args[1]; - boolean rerank = args[2] != null && (boolean) args[2]; - String rerankInferenceId = (String) args[3]; - int rankWindowSize = args[4] == null ? RankBuilder.DEFAULT_RANK_WINDOW_SIZE : (int) args[4]; - return new HybridRetrieverBuilder(fields, query, rerank, rerankInferenceId, rankWindowSize); + Boolean rerank = (Boolean) args[2]; + String rerankField = (String) args[3]; + String rerankInferenceId = (String) args[4]; + int rankWindowSize = args[5] == null ? RankBuilder.DEFAULT_RANK_WINDOW_SIZE : (int) args[5]; + return new HybridRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, rankWindowSize); } ); @@ -66,31 +70,36 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper fields, String query, boolean rerank, String rerankInferenceId, int rankWindowSize) { + public HybridRetrieverBuilder( + List fields, + String query, + Boolean rerank, + String rerankField, + String rerankInferenceId, + int rankWindowSize + ) { this( fields == null ? List.of() : List.copyOf(fields), query, rerank, + rerankField, rerankInferenceId, rankWindowSize, - new LinearRetrieverBuilder( - generateInnerRetrievers(fields, query), - rankWindowSize, - generateWeights(fields), - generateScoreNormalizers(fields) - ) + generateRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, rankWindowSize) ); } private HybridRetrieverBuilder( List fields, String query, - boolean rerank, + Boolean rerank, + String rerankField, String rerankInferenceId, int rankWindowSize, RetrieverBuilder retrieverBuilder @@ -99,13 +108,14 @@ private HybridRetrieverBuilder( this.fields = fields; this.query = query; this.rerank = rerank; + this.rerankField = rerankField; this.rerankInferenceId = rerankInferenceId; this.rankWindowSize = rankWindowSize; } @Override protected HybridRetrieverBuilder clone(RetrieverBuilder sub) { - return new HybridRetrieverBuilder(fields, query, rerank, rerankInferenceId, rankWindowSize, sub); + return new HybridRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, rankWindowSize, sub); } @Override @@ -117,6 +127,15 @@ public String getName() { protected void doToXContent(XContentBuilder builder, Params params) throws IOException { builder.field(FIELDS_FIELD.getPreferredName(), fields); builder.field(QUERY_FIELD.getPreferredName(), query); + if (rerank != null) { + builder.field(RERANK_FIELD.getPreferredName(), rerank); + } + if (rerankField != null) { + builder.field(RERANK_FIELD_FIELD.getPreferredName(), rerankField); + } + if (rerankInferenceId != null) { + builder.field(RERANK_INFERENCE_ID_FIELD.getPreferredName(), rerankInferenceId); + } builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize); } @@ -124,18 +143,57 @@ protected void doToXContent(XContentBuilder builder, Params params) throws IOExc protected boolean doEquals(Object o) { // TODO: Check rankWindowSize? It should be checked by the wrapped retriever. HybridRetrieverBuilder that = (HybridRetrieverBuilder) o; - return Objects.equals(fields, that.fields) && Objects.equals(query, that.query) && super.doEquals(o); + return Objects.equals(fields, that.fields) + && Objects.equals(query, that.query) + && Objects.equals(rerank, that.rerank) + && Objects.equals(rerankField, that.rerankField) + && Objects.equals(rerankInferenceId, that.rerankInferenceId) + && super.doEquals(o); } @Override protected int doHashCode() { - return Objects.hash(fields, query, super.doHashCode()); + return Objects.hash(fields, query, rerank, rerankField, rerankInferenceId, super.doHashCode()); } public static HybridRetrieverBuilder fromXContent(XContentParser parser, RetrieverParserContext context) throws IOException { return PARSER.apply(parser, context); } + private static RetrieverBuilder generateRetrieverBuilder( + List fields, + String query, + Boolean rerank, + String rerankField, + String rerankInferenceId, + int rankWindowSize + ) { + LinearRetrieverBuilder linearRetrieverBuilder = new LinearRetrieverBuilder( + generateInnerRetrievers(fields, query), + rankWindowSize, + generateWeights(fields), + generateScoreNormalizers(fields) + ); + + RetrieverBuilder rootRetriever = linearRetrieverBuilder; + if (rerank != null && rerank) { + if (rerankField == null) { + throw new IllegalArgumentException("[" + RERANK_FIELD_FIELD.getPreferredName() + "] is required when reranking is enabled"); + } + + rootRetriever = new TextSimilarityRankRetrieverBuilder( + linearRetrieverBuilder, + rerankInferenceId, + query, + rerankField, + rankWindowSize, + false + ); + } + + return rootRetriever; + } + private static List generateInnerRetrievers(List fields, String query) { if (fields == null) { return List.of(); From 94713c833c0045eaf02ccb4c7b0d8426dde53a6c Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Wed, 12 Mar 2025 15:18:07 -0400 Subject: [PATCH 19/21] Parse field strings for weights --- .../rank/hybrid/HybridRetrieverBuilder.java | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index 963187166c51b..999e12a11d616 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -26,6 +26,7 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.Objects; @@ -168,10 +169,12 @@ private static RetrieverBuilder generateRetrieverBuilder( String rerankInferenceId, int rankWindowSize ) { + FieldsAndsWeights fieldsAndsWeights = generateFieldsAndWeights(fields); + LinearRetrieverBuilder linearRetrieverBuilder = new LinearRetrieverBuilder( - generateInnerRetrievers(fields, query), + generateInnerRetrievers(fieldsAndsWeights.fields(), query), rankWindowSize, - generateWeights(fields), + fieldsAndsWeights.weights(), generateScoreNormalizers(fields) ); @@ -208,15 +211,29 @@ private static List generateInnerRetri return innerRetrievers; } - private static float[] generateWeights(List fields) { + private static FieldsAndsWeights generateFieldsAndWeights(List fields) { if (fields == null) { - return new float[0]; + return new FieldsAndsWeights(List.of(), new float[0]); + } + + int fieldCount = fields.size(); + List parsedFields = new ArrayList<>(fieldCount); + float[] parsedWeights = new float[fieldCount]; + for (int i = 0; i < fieldCount; i++) { + String[] fieldSplit = fields.get(i).split("\\^"); + + float weight = 1.0f; + if (fieldSplit.length > 2) { + throw new IllegalArgumentException("Invalid field name [" + fields.get(i) + "]"); + } else if (fieldSplit.length == 2) { + weight = Float.parseFloat(fieldSplit[1]); + } + + parsedFields.add(fieldSplit[0]); + parsedWeights[i] = weight; } - // TODO: Parse field strings for weights - float[] weights = new float[fields.size()]; - Arrays.fill(weights, 1.0f); - return weights; + return new FieldsAndsWeights(Collections.unmodifiableList(parsedFields), parsedWeights); } private static ScoreNormalizer[] generateScoreNormalizers(List fields) { @@ -228,4 +245,6 @@ private static ScoreNormalizer[] generateScoreNormalizers(List fields) { Arrays.fill(scoreNormalizers, new MinMaxScoreNormalizer(0)); return scoreNormalizers; } + + private record FieldsAndsWeights(List fields, float[] weights) {} } From 70c01bf486b973f6af50c80db5ef3e22a4ce5fc2 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Thu, 13 Mar 2025 15:32:23 -0400 Subject: [PATCH 20/21] Added per-field query settings --- .../rank/hybrid/HybridRetrieverBuilder.java | 134 ++++++++++++++++-- .../xpack/rank/hybrid/MatchQuerySettings.java | 84 +++++++++++ .../xpack/rank/hybrid/QuerySettings.java | 25 ++++ .../xpack/rank/hybrid/QueryType.java | 23 +++ 4 files changed, 258 insertions(+), 8 deletions(-) create mode 100644 x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/MatchQuerySettings.java create mode 100644 x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QuerySettings.java create mode 100644 x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QueryType.java diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index 999e12a11d616..0b0898ec58c5c 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -7,7 +7,11 @@ package org.elasticsearch.xpack.rank.hybrid; +import org.elasticsearch.common.collect.ImmutableOpenMap; +import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; +import org.elasticsearch.common.xcontent.support.XContentMapValues; import org.elasticsearch.index.query.MatchQueryBuilder; +import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.search.rank.RankBuilder; import org.elasticsearch.search.retriever.CompoundRetrieverBuilder; import org.elasticsearch.search.retriever.RetrieverBuilder; @@ -15,9 +19,12 @@ import org.elasticsearch.search.retriever.RetrieverParserContext; import org.elasticsearch.search.retriever.StandardRetrieverBuilder; import org.elasticsearch.xcontent.ConstructingObjectParser; +import org.elasticsearch.xcontent.NamedXContentRegistry; import org.elasticsearch.xcontent.ParseField; import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParseException; import org.elasticsearch.xcontent.XContentParser; +import org.elasticsearch.xcontent.support.MapXContentParser; import org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.LinearRetrieverBuilder; import org.elasticsearch.xpack.rank.linear.MinMaxScoreNormalizer; @@ -27,12 +34,15 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Objects; import static org.elasticsearch.search.retriever.CompoundRetrieverBuilder.RANK_WINDOW_SIZE_FIELD; import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; +import static org.elasticsearch.xpack.rank.hybrid.QuerySettings.TYPE_FIELD; // TODO: // - Retriever name support @@ -44,12 +54,14 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper fields; private final String query; private final Boolean rerank; private final String rerankField; private final String rerankInferenceId; + private final Map> querySettingsMap; private final int rankWindowSize; @SuppressWarnings("unchecked") @@ -63,17 +75,60 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper> querySettingsMap = (Map>) args[6]; + + return new HybridRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, querySettingsMap, rankWindowSize); } ); + private static final NamedXContentRegistry NAMED_X_CONTENT_REGISTRY; + static { + List xContentRegistryEntries = List.of( + new NamedXContentRegistry.Entry( + QuerySettings.class, + new ParseField(MatchQuerySettings.QUERY_TYPE.getQueryName()), + MatchQuerySettings::fromXContent + ) + ); + + NAMED_X_CONTENT_REGISTRY = new NamedXContentRegistry(xContentRegistryEntries); + PARSER.declareStringArray(constructorArg(), FIELDS_FIELD); PARSER.declareString(constructorArg(), QUERY_FIELD); PARSER.declareBoolean(optionalConstructorArg(), RERANK_FIELD); PARSER.declareString(optionalConstructorArg(), RERANK_FIELD_FIELD); PARSER.declareString(optionalConstructorArg(), RERANK_INFERENCE_ID_FIELD); PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD); + PARSER.declareObject(optionalConstructorArg(), (p, c) -> { + Map> querySettingsMap = new HashMap<>(); + + Map unparsedMap = p.map(); + for (var entry : unparsedMap.entrySet()) { + String field = entry.getKey(); + Object value = entry.getValue(); + + if (value instanceof List list) { + List querySettings = querySettingsMap.computeIfAbsent(field, f -> new ArrayList<>(list.size())); + for (Object listValue : list) { + if (listValue instanceof Map map) { + querySettings.add(parseQuerySettings(map)); + } else { + throw new IllegalArgumentException( + "Query settings for field [" + field + "] must be an object or list of objects" + ); + } + } + } else if (value instanceof Map map) { + List querySettings = querySettingsMap.computeIfAbsent(field, f -> new ArrayList<>()); + querySettings.add(parseQuerySettings(map)); + } else { + throw new IllegalArgumentException("Query settings for field [" + field + "] must be an object or list of objects"); + } + } + + return querySettingsMap; + }, QUERY_SETTINGS_FIELD); RetrieverBuilder.declareBaseParserFields(PARSER); } @@ -83,6 +138,7 @@ public HybridRetrieverBuilder( Boolean rerank, String rerankField, String rerankInferenceId, + Map> querySettingsMap, int rankWindowSize ) { this( @@ -91,8 +147,9 @@ public HybridRetrieverBuilder( rerank, rerankField, rerankInferenceId, + copyQuerySettingsMap(querySettingsMap), rankWindowSize, - generateRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, rankWindowSize) + generateRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, querySettingsMap, rankWindowSize) ); } @@ -102,6 +159,7 @@ private HybridRetrieverBuilder( Boolean rerank, String rerankField, String rerankInferenceId, + Map> querySettingsMap, int rankWindowSize, RetrieverBuilder retrieverBuilder ) { @@ -111,12 +169,13 @@ private HybridRetrieverBuilder( this.rerank = rerank; this.rerankField = rerankField; this.rerankInferenceId = rerankInferenceId; + this.querySettingsMap = querySettingsMap; this.rankWindowSize = rankWindowSize; } @Override protected HybridRetrieverBuilder clone(RetrieverBuilder sub) { - return new HybridRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, rankWindowSize, sub); + return new HybridRetrieverBuilder(fields, query, rerank, rerankField, rerankInferenceId, querySettingsMap, rankWindowSize, sub); } @Override @@ -161,18 +220,35 @@ public static HybridRetrieverBuilder fromXContent(XContentParser parser, Retriev return PARSER.apply(parser, context); } + private static Map> copyQuerySettingsMap(Map> querySettingsMap) { + if (querySettingsMap == null) { + return Map.of(); + } + + ImmutableOpenMap.Builder> copyBuilder = new ImmutableOpenMap.Builder<>(querySettingsMap.size()); + for (var entry : querySettingsMap.entrySet()) { + String field = entry.getKey(); + List querySettings = entry.getValue(); + + copyBuilder.put(field, querySettings != null ? List.copyOf(querySettings) : List.of()); + } + + return copyBuilder.build(); + } + private static RetrieverBuilder generateRetrieverBuilder( List fields, String query, Boolean rerank, String rerankField, String rerankInferenceId, + Map> querySettingsMap, int rankWindowSize ) { FieldsAndsWeights fieldsAndsWeights = generateFieldsAndWeights(fields); LinearRetrieverBuilder linearRetrieverBuilder = new LinearRetrieverBuilder( - generateInnerRetrievers(fieldsAndsWeights.fields(), query), + generateInnerRetrievers(fieldsAndsWeights.fields(), query, querySettingsMap), rankWindowSize, fieldsAndsWeights.weights(), generateScoreNormalizers(fields) @@ -197,15 +273,31 @@ private static RetrieverBuilder generateRetrieverBuilder( return rootRetriever; } - private static List generateInnerRetrievers(List fields, String query) { + private static List generateInnerRetrievers( + List fields, + String query, + Map> querySettingsMap + ) { if (fields == null) { return List.of(); } - List innerRetrievers = new ArrayList<>(fields.size()); + List innerRetrievers = new ArrayList<>(); for (String field : fields) { - MatchQueryBuilder matchQueryBuilder = new MatchQueryBuilder(field, query); - innerRetrievers.add(new CompoundRetrieverBuilder.RetrieverSource(new StandardRetrieverBuilder(matchQueryBuilder), null)); + List fieldQueryBuilders = new ArrayList<>(); + List fieldQuerySettings = querySettingsMap != null ? querySettingsMap.get(field) : null; + if (fieldQuerySettings == null || fieldQuerySettings.isEmpty()) { + // Default to match query + fieldQueryBuilders.add(new MatchQueryBuilder(field, query)); + } else { + for (QuerySettings querySettings : fieldQuerySettings) { + fieldQueryBuilders.add(querySettings.constructQueryBuilder(field, query)); + } + } + + for (QueryBuilder queryBuilder : fieldQueryBuilders) { + innerRetrievers.add(new CompoundRetrieverBuilder.RetrieverSource(new StandardRetrieverBuilder(queryBuilder), null)); + } } return innerRetrievers; @@ -247,4 +339,30 @@ private static ScoreNormalizer[] generateScoreNormalizers(List fields) { } private record FieldsAndsWeights(List fields, float[] weights) {} + + // TODO: Probably a better way to do this, but this is quick & dirty for POC purposes + private static QuerySettings parseQuerySettings(Map map) { + Map querySettingsMap = XContentMapValues.nodeMapValue(map, "query settings"); + + Object typeObject = querySettingsMap.get(TYPE_FIELD.getPreferredName()); + if (typeObject == null) { + throw new IllegalArgumentException("[" + TYPE_FIELD.getPreferredName() + "] must be provided in query settings"); + } else if (typeObject instanceof String == false) { + throw new IllegalArgumentException("[" + TYPE_FIELD.getPreferredName() + "] must have a string value"); + } + + String typeString = (String) typeObject; + MapXContentParser mapXContentParser = new MapXContentParser( + NAMED_X_CONTENT_REGISTRY, + LoggingDeprecationHandler.INSTANCE, + querySettingsMap, + null + ); + + try (mapXContentParser) { + return mapXContentParser.namedObject(QuerySettings.class, typeString, null); + } catch (IOException e) { + throw new XContentParseException(mapXContentParser.getTokenLocation(), "Failed to parse query settings"); + } + } } diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/MatchQuerySettings.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/MatchQuerySettings.java new file mode 100644 index 0000000000000..e9d5beaf236b4 --- /dev/null +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/MatchQuerySettings.java @@ -0,0 +1,84 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.rank.hybrid; + +import org.elasticsearch.index.query.MatchQueryBuilder; +import org.elasticsearch.index.query.Operator; +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.xcontent.ConstructingObjectParser; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; + +import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; + +public class MatchQuerySettings implements QuerySettings { + public static final QueryType QUERY_TYPE = QueryType.MATCH; + + public static final ParseField OPERATOR_FIELD = new ParseField("operator"); + + static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + "match_query_settings", + false, + args -> { + String typeString = (String) args[0]; + String operatorString = (String) args[1]; + + QueryType queryType = QueryType.fromString(typeString); + if (queryType != QueryType.MATCH) { + throw new IllegalStateException("Query type must be " + QueryType.MATCH); + } + + return new MatchQuerySettings( + operatorString != null ? Operator.fromString(operatorString) : MatchQueryBuilder.DEFAULT_OPERATOR + ); + } + ); + + static { + PARSER.declareString(constructorArg(), TYPE_FIELD); + PARSER.declareString(optionalConstructorArg(), OPERATOR_FIELD); + } + + private final Operator operator; + + public MatchQuerySettings(Operator operator) { + this.operator = operator; + } + + public Operator getOperator() { + return operator; + } + + @Override + public QueryType getQueryType() { + return QUERY_TYPE; + } + + @Override + public QueryBuilder constructQueryBuilder(String field, String query) { + return new MatchQueryBuilder(field, query).operator(getOperator()); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field(TYPE_FIELD.getPreferredName(), getName()); + builder.field(OPERATOR_FIELD.getPreferredName(), operator); + builder.endObject(); + + return builder; + } + + public static MatchQuerySettings fromXContent(XContentParser parser) { + return PARSER.apply(parser, null); + } +} diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QuerySettings.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QuerySettings.java new file mode 100644 index 0000000000000..9061af538737f --- /dev/null +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QuerySettings.java @@ -0,0 +1,25 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.rank.hybrid; + +import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xpack.core.ml.utils.NamedXContentObject; + +public interface QuerySettings extends NamedXContentObject { + ParseField TYPE_FIELD = new ParseField("type"); + + QueryType getQueryType(); + + QueryBuilder constructQueryBuilder(String field, String query); + + @Override + default String getName() { + return getQueryType().getQueryName(); + } +} diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QueryType.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QueryType.java new file mode 100644 index 0000000000000..af719efb5af01 --- /dev/null +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/QueryType.java @@ -0,0 +1,23 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.rank.hybrid; + +import java.util.Locale; + +public enum QueryType { + MATCH, + MATCH_PHRASE; + + public String getQueryName() { + return name().toLowerCase(Locale.ROOT); + } + + public static QueryType fromString(String queryName) { + return valueOf(queryName.toUpperCase(Locale.ROOT)); + } +} From ea83a3ac21c3e84deed92df21808c8645c0b57c8 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Thu, 13 Mar 2025 16:26:08 -0400 Subject: [PATCH 21/21] Added match phrase query settings --- .../rank/hybrid/HybridRetrieverBuilder.java | 5 ++ .../rank/hybrid/MatchPhraseQuerySettings.java | 88 +++++++++++++++++++ 2 files changed, 93 insertions(+) create mode 100644 x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/MatchPhraseQuerySettings.java diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java index 0b0898ec58c5c..341f578d52480 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/hybrid/HybridRetrieverBuilder.java @@ -89,6 +89,11 @@ public class HybridRetrieverBuilder extends RetrieverBuilderWrapper PARSER = new ConstructingObjectParser<>( + "match_phrase_query_settings", + false, + args -> { + String typeString = (String) args[0]; + Integer slop = (Integer) args[1]; + + QueryType queryType = QueryType.fromString(typeString); + if (queryType != QUERY_TYPE) { + throw new IllegalStateException("Query type must be " + QUERY_TYPE); + } + + return new MatchPhraseQuerySettings(slop); + } + ); + + static { + PARSER.declareString(constructorArg(), TYPE_FIELD); + PARSER.declareInt(optionalConstructorArg(), SLOP_FIELD); + } + + private final Integer slop; + + public MatchPhraseQuerySettings(Integer slop) { + this.slop = slop; + } + + public Integer getSlop() { + return slop; + } + + @Override + public QueryType getQueryType() { + return QUERY_TYPE; + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + builder.field(TYPE_FIELD.getPreferredName(), getName()); + if (slop != null) { + builder.field(SLOP_FIELD.getPreferredName(), slop); + } + builder.endObject(); + + return builder; + } + + @Override + public QueryBuilder constructQueryBuilder(String field, String query) { + MatchPhraseQueryBuilder queryBuilder = new MatchPhraseQueryBuilder(field, query); + if (slop != null) { + queryBuilder.slop(getSlop()); + } + + return queryBuilder; + } + + public static MatchPhraseQuerySettings fromXContent(XContentParser parser) { + return PARSER.apply(parser, null); + } +}