propagate

Author: thecoop

server/src/internalClusterTest/java/org/elasticsearch/search/rank/FieldBasedRerankerIT.java

@@ -100,7 +100,7 @@ public static FieldBasedRankBuilder fromXContent(XContentParser parser) throws I
         }
 
         public FieldBasedRankBuilder(final int rankWindowSize, final String field) {
-            super(rankWindowSize);
+            super(rankWindowSize, false);
             this.field = field;
         }
 
@@ -205,7 +205,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
 
         @Override
         public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorContext(int size, int from, Client client) {
-            return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize()) {
+            return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize(), isLenient()) {
                 @Override
                 protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                     float[] scores = new float[featureDocs.length];
@@ -346,7 +346,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
         @Override
         public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorContext(int size, int from, Client client) {
             if (this.throwingRankBuilderType == ThrowingRankBuilderType.THROWING_RANK_FEATURE_PHASE_COORDINATOR_CONTEXT)
-                return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize()) {
+                return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize(), isLenient()) {
                     @Override
                     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                         throw new UnsupportedOperationException("rfc - simulated failure");

server/src/internalClusterTest/java/org/elasticsearch/search/rank/MockedRequestActionBasedRerankerIT.java

@@ -249,7 +249,7 @@ public static class TestRerankingRankFeaturePhaseRankCoordinatorContext extends
             String inferenceText,
             float minScore
         ) {
-            super(size, from, windowSize);
+            super(size, from, windowSize, false);
             this.client = client;
             this.inferenceId = inferenceId;
             this.inferenceText = inferenceText;
@@ -337,7 +337,7 @@ public MockRequestActionBasedRankBuilder(
             final String inferenceText,
             final float minScore
         ) {
-            super(rankWindowSize);
+            super(rankWindowSize, false);
             this.field = field;
             this.inferenceId = inferenceId;
             this.inferenceText = inferenceText;

server/src/main/java/org/elasticsearch/search/rank/RankBuilder.java

@@ -152,16 +152,15 @@ public final boolean equals(Object obj) {
         if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
-        @SuppressWarnings("unchecked")
         RankBuilder other = (RankBuilder) obj;
-        return Objects.equals(rankWindowSize, other.rankWindowSize()) && doEquals(other);
+        return rankWindowSize == other.rankWindowSize && lenient == other.lenient && doEquals(other);
     }
 
     protected abstract boolean doEquals(RankBuilder other);
 
     @Override
     public final int hashCode() {
-        return Objects.hash(getClass(), rankWindowSize, doHashCode());
+        return Objects.hash(getClass(), rankWindowSize, lenient, doHashCode());
     }
 
     protected abstract int doHashCode();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java

@@ -44,6 +44,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
     public static final ParseField INFERENCE_ID_FIELD = new ParseField("inference_id");
     public static final ParseField INFERENCE_TEXT_FIELD = new ParseField("inference_text");
     public static final ParseField FIELD_FIELD = new ParseField("field");
+    public static final ParseField LENIENT_FIELD = new ParseField("lenient");
 
     public static final ConstructingObjectParser<TextSimilarityRankRetrieverBuilder, RetrieverParserContext> PARSER =
         new ConstructingObjectParser<>(TextSimilarityRankBuilder.NAME, args -> {
@@ -54,7 +55,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
             int rankWindowSize = args[4] == null ? DEFAULT_RANK_WINDOW_SIZE : (int) args[4];
             boolean lenient = args[5] != null && (Boolean) args[5];
 
-            return new TextSimilarityRankRetrieverBuilder(retrieverBuilder, inferenceId, inferenceText, field, rankWindowSize);
+            return new TextSimilarityRankRetrieverBuilder(retrieverBuilder, inferenceId, inferenceText, field, rankWindowSize, lenient);
         });
 
     static {
@@ -67,6 +68,7 @@ public class TextSimilarityRankRetrieverBuilder extends CompoundRetrieverBuilder
         PARSER.declareString(constructorArg(), INFERENCE_TEXT_FIELD);
         PARSER.declareString(constructorArg(), FIELD_FIELD);
         PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
+        PARSER.declareBoolean(optionalConstructorArg(), LENIENT_FIELD);
 
         RetrieverBuilder.declareBaseParserFields(TextSimilarityRankBuilder.NAME, PARSER);
     }
@@ -85,18 +87,21 @@ public static TextSimilarityRankRetrieverBuilder fromXContent(
     private final String inferenceId;
     private final String inferenceText;
     private final String field;
+    private final boolean lenient;
 
     public TextSimilarityRankRetrieverBuilder(
         RetrieverBuilder retrieverBuilder,
         String inferenceId,
         String inferenceText,
         String field,
-        int rankWindowSize
+        int rankWindowSize,
+        boolean lenient
     ) {
         super(List.of(new RetrieverSource(retrieverBuilder, null)), rankWindowSize);
         this.inferenceId = inferenceId;
         this.inferenceText = inferenceText;
         this.field = field;
+        this.lenient = lenient;
     }
 
     public TextSimilarityRankRetrieverBuilder(
@@ -106,6 +111,7 @@ public TextSimilarityRankRetrieverBuilder(
         String field,
         int rankWindowSize,
         Float minScore,
+        boolean lenient,
         String retrieverName,
         List<QueryBuilder> preFilterQueryBuilders
     ) {
@@ -117,6 +123,7 @@ public TextSimilarityRankRetrieverBuilder(
         this.inferenceText = inferenceText;
         this.field = field;
         this.minScore = minScore;
+        this.lenient = lenient;
         this.retrieverName = retrieverName;
         this.preFilterQueryBuilders = preFilterQueryBuilders;
     }
@@ -133,6 +140,7 @@ protected TextSimilarityRankRetrieverBuilder clone(
             field,
             rankWindowSize,
             minScore,
+            lenient,
             retrieverName,
             newPreFilterQueryBuilders
         );
@@ -163,9 +171,7 @@ protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, b
 
     @Override
     protected SearchSourceBuilder finalizeSourceBuilder(SearchSourceBuilder sourceBuilder) {
-        sourceBuilder.rankBuilder(
-            new TextSimilarityRankBuilder(this.field, this.inferenceId, this.inferenceText, this.rankWindowSize, this.minScore, false)
-        );
+        sourceBuilder.rankBuilder(new TextSimilarityRankBuilder(field, inferenceId, inferenceText, rankWindowSize, minScore, lenient));
         return sourceBuilder;
     }
 
@@ -189,6 +195,9 @@ protected void doToXContent(XContentBuilder builder, Params params) throws IOExc
         builder.field(INFERENCE_TEXT_FIELD.getPreferredName(), inferenceText);
         builder.field(FIELD_FIELD.getPreferredName(), field);
         builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize);
+        if (lenient) {
+            builder.field(LENIENT_FIELD.getPreferredName(), lenient);
+        }
     }
 
     @Override
@@ -198,12 +207,13 @@ public boolean doEquals(Object other) {
             && Objects.equals(inferenceId, that.inferenceId)
             && Objects.equals(inferenceText, that.inferenceText)
             && Objects.equals(field, that.field)
-            && Objects.equals(rankWindowSize, that.rankWindowSize)
-            && Objects.equals(minScore, that.minScore);
+            && rankWindowSize == that.rankWindowSize
+            && Objects.equals(minScore, that.minScore)
+            && lenient == that.lenient;
     }
 
     @Override
     public int doHashCode() {
-        return Objects.hash(inferenceId, inferenceText, field, rankWindowSize, minScore);
+        return Objects.hash(inferenceId, inferenceText, field, rankWindowSize, minScore, lenient);
     }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilderTests.java

@@ -57,7 +57,8 @@ public static TextSimilarityRankRetrieverBuilder createRandomTextSimilarityRankR
             randomAlphaOfLength(10),
             randomAlphaOfLength(20),
             randomAlphaOfLength(50),
-            randomIntBetween(100, 10000)
+            randomIntBetween(100, 10000),
+            randomBoolean()
         );
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverTelemetryTests.java

@@ -139,7 +139,8 @@ public void testTelemetryForRRFRetriever() throws IOException {
                         "some_inference_id",
                         "some_inference_text",
                         "some_field",
-                        10
+                        10,
+                        false
                     )
                 )
             );