Pinned retriever (elastic#126401)

mridula-s109 · elasticsearchmachine · kderusso · web-flow · commit b742b026cc09 · 2025-05-01T13:50:33.000+01:00
* Created PinnedQueryRetriever Builder * Added retriever build test file * Changed the test to accept all the retrievers as its not licensed * [CI] Auto commit changes from spotless * Added integration test and fixed the code style issues * Fixed merge issues * Added cluster test feature to the integration tests * Update docs/changelog/126401.yaml * Registered retriever plugin * Enhanced changelog description * Added validation to the constructor creation * Introduced validations in the code and incorporated in tests - compiling version * Included PinnedRankDoc to enhance the explain query, validations for inputs anfd sorting. Fixed the tests after these changes * Working on improving integration test and introducing cluster features * Removed cluster test temporarily * Got pinnedretriever yaml to working state without cluster * [CI] Auto commit changes from spotless * Resolved no source yaml error * Cluster loaded successfully * [CI] Auto commit changes from spotless * Fixed integration test * Made validate sort less strict * Included shareddoc sorting in the validate sort * All yaml issues resolved * Explanation yaml files added * Trying to add clustering to the test * [CI] Auto commit changes from spotless * Removed explicit feature specification * Deleted the empty file * Put the node feature in the proper place * Added additional validation in ids and docs * Cleaned the create pinned query validation * made unit tests robust * Remove query.es from version control * Remove result.json from version control * [CI] Auto commit changes from spotless * Everything except explanation is fixed * added duplicate doc test * Applied checkstyle fix, spotless and a failing yaml * Improvements based on PR comments * [CI] Auto commit changes from spotless * Modified the unit test to accomodate the change in createPinnedQuery * Split the yaml to test for basic and trial, cleanedup and acted on all the comments * Remove result.json and query.es from version control * Removed unnecessary comments * Removed redundant file * Fixing CI build error * [CI] Auto commit changes from spotless * Removed pinnedBy as it wasnt necessary * Removed unnecessary ToXContent Information * Fixed transport version charges and cleaned up null checks: " * [CI] Auto commit changes from spotless * Retriever status changed to 9.1 version * cleaned up 2 yamltestsuite for different licenses * Trying to get the clustering works * Cleaned up the yaml clustering and also reorganised the yaml tests * Unnecessary file introduction removed * Reverted the plugins to the previous state as the changes werent necessary * Removed unnecessary transport versioning from pinnedrankdoc * reverted * [CI] Auto commit changes from spotless * Edited the SearchBusinessRules to remove the class from getFeatures * [CI] Auto commit changes from spotless * Cleaned up Pinned retriever to allow only id or docs * Added more test to the pinned retriever * did spotless * Introduced new transport versioning * Cleaning it up * Resolved validate module error * BWT issues * [CI] Auto commit changes from spotless * fix NPE occuring in ci * trying to fix duplicate feature issue * [CI] Auto commit changes from spotless * modified the tests * Playing with pinned retriever CI (elastic#127530) * Fix compilation error * Remove SearchPlugin from META-INF * Remove duplicate FeatureSpecification * Add service as test resource * Move to test * More file moving * Remove from module-info, remove file * Make search business rules plugin extensible * fix ent search plugin * Delete dup files * [CI] Auto commit changes from spotless * Add module info back in --------- Co-authored-by: elasticsearchmachine <infra-root+elasticsearchmachine@elastic.co> * top document no pinned * Removing this to see if it works without this transport version check, as i had made changes related to the cluster error * Fixed the retriever builder comments and error message * removed the unnencessary explain * Fixed all the tests in the yaml file * added extra test case * [CI] Auto commit changes from spotless --------- Co-authored-by: elasticsearchmachine <infra-root+elasticsearchmachine@elastic.co> Co-authored-by: Kathleen DeRusso <kathleen.derusso@elastic.co>
diff --git a/docs/changelog/126401.yaml b/docs/changelog/126401.yaml
@@ -0,0 +1,5 @@
+pr: 126401
+summary: Add pinned retriever
+area: Relevance
+type: enhancement
+issues: []
diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java
@@ -231,6 +231,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INTRODUCE_FAILURES_LIFECYCLE = def(9_065_0_00);
     public static final TransportVersion PROJECT_METADATA_SETTINGS = def(9_066_00_0);
     public static final TransportVersion AGGREGATE_METRIC_DOUBLE_BLOCK = def(9_067_00_0);
+    public static final TransportVersion PINNED_RETRIEVER = def(9_068_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,
diff --git a/x-pack/plugin/ent-search/build.gradle b/x-pack/plugin/ent-search/build.gradle
@@ -8,7 +8,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', 'search-business-rules']
 }
 
 base {
@@ -17,7 +17,7 @@ base {
 
 dependencies {
   compileOnly project(path: xpackModule('core'))
-  implementation project(xpackModule('search-business-rules'))
+  compileOnly project(xpackModule('search-business-rules'))
   api project(':modules:lang-mustache')
 
   // JSON Schema dependencies
diff --git a/x-pack/plugin/rank-rrf/build.gradle b/x-pack/plugin/rank-rrf/build.gradle
@@ -26,6 +26,7 @@ dependencies {
   clusterModules project(xpackModule('rank-rrf'))
   clusterModules project(xpackModule('inference'))
   clusterModules project(':modules:lang-painless')
+  clusterModules project(xpackModule('search-business-rules'))
 
   clusterPlugins project(':x-pack:plugin:inference:qa:test-service-plugin')
 }
diff --git a/x-pack/plugin/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/950_pinned_interaction.yml b/x-pack/plugin/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/950_pinned_interaction.yml
@@ -0,0 +1,98 @@
+setup:
+  - requires:
+      cluster_features: 'pinned_retriever_supported'
+      reason: 'test requires pinned retriever implementation'
+  - do:
+      indices.create:
+        index: test-index1
+
+  - do:
+      bulk:
+        refresh: true
+        index: test-index1
+        body:
+          - index:
+              _id: doc1
+          - { "text": "document one" }
+          - index:
+              _id: doc2
+          - { "text": "document two" }
+          - index:
+              _id: doc3
+          - { "text": "document three" }
+          - index:
+              _id: doc4
+          - { "text": "document four" }
+          - index:
+              _id: doc5
+          - { "text": "document five" }
+
+---
+"rrf combined with pinned retriever":
+  - skip: { features: headers }
+  - do:
+      headers:
+        Content-Type: application/json
+      search:
+        index: test-index1
+        body:
+          retriever:
+            pinned:
+              ids: ["doc1"]
+              retriever:
+                rrf:
+                  retrievers: [
+                    {
+                      standard: {
+                        query: {
+                          term: { text: "document" }
+                        }
+                      }
+                    },
+                    {
+                      standard: {
+                        query: {
+                          term: { text: "three" }
+                        }
+                      }
+                    }
+                  ]
+              rank_window_size: 10
+
+  - match: { hits.total.value: 5 }
+  - match: { hits.hits.0._id: doc1 }
+  - match: { hits.hits.0._score: 1.7014122E38 }
+  - match: { hits.hits.1._id: doc3 }
+  - match: { hits.hits.1._score < 100.0 }
+  - match: { hits.hits.2._id: doc2 }
+
+---
+"rrf with pinned retriever as a sub-retriever":
+  - skip: { features: headers }
+  - do:
+      headers: { Content-Type: application/json }
+      search:
+        index: test-index1
+        body:
+          retriever:
+            rrf:
+              retrievers:
+                -
+                  standard:
+                    query:
+                      match: { text: "document" }
+                -
+                  pinned:
+                    ids: ["doc4", "doc5"]
+                    retriever:
+                      standard:
+                        query:
+                          match: { text: "document" }
+
+  - match: { hits.total.value: 5 }
+  - match: { hits.hits.0._id: doc1 }
+  - lt: { hits.hits.0._score: 100.0 }
+  - match: { hits.hits.1._id: doc4 }
+  - match: { hits.hits.2._id: doc5 }
+
+
diff --git a/x-pack/plugin/search-business-rules/build.gradle b/x-pack/plugin/search-business-rules/build.gradle
@@ -1,5 +1,6 @@
 apply plugin: 'elasticsearch.internal-es-plugin'
 apply plugin: 'elasticsearch.internal-cluster-test'
+apply plugin: 'elasticsearch.internal-yaml-rest-test'
 
 esplugin {
   name = 'search-business-rules'
@@ -13,6 +14,19 @@ base {
 
 dependencies {
   compileOnly project(path: xpackModule('core'))
+  compileOnly project(':server')
   testImplementation(testArtifact(project(xpackModule('core'))))
-  testImplementation project(":test:framework")
+  testImplementation(testArtifact(project(':server')))
+  clusterModules project(xpackModule('search-business-rules'))
+  clusterModules project(':modules:mapper-extras')
+  clusterModules project(':modules:lang-painless')
+  clusterModules project(xpackModule('inference'))
 }
+
+tasks.named("yamlRestTest") {
+  usesDefaultDistribution("uses search business rules plugin")
+}
+artifacts {
+  restXpackTests(new File(projectDir, "src/yamlRestTest/resources/rest-api-spec/test"))
+}
+
diff --git a/x-pack/plugin/search-business-rules/src/main/java/module-info.java b/x-pack/plugin/search-business-rules/src/main/java/module-info.java
@@ -16,4 +16,7 @@
     requires org.elasticsearch.xcore;
 
     exports org.elasticsearch.xpack.searchbusinessrules;
+
+    provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.searchbusinessrules.SearchBusinessRulesFeatures;
+    provides org.elasticsearch.plugins.SearchPlugin with org.elasticsearch.xpack.searchbusinessrules.SearchBusinessRules;
 }
diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/PinnedQueryBuilder.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/PinnedQueryBuilder.java
@@ -59,7 +59,7 @@ public class PinnedQueryBuilder extends AbstractQueryBuilder<PinnedQueryBuilder>
 
     // Organic queries will have their scores capped to this number range,
     // We reserve the highest float exponent for scores of pinned queries
-    private static final float MAX_ORGANIC_SCORE = Float.intBitsToFloat((0xfe << 23)) - 1;
+    public static final float MAX_ORGANIC_SCORE = Float.intBitsToFloat((0xfe << 23)) - 1;
 
     public PinnedQueryBuilder(QueryBuilder organicQuery, String... ids) {
         this(organicQuery, Arrays.asList(ids), null);
diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SearchBusinessRules.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SearchBusinessRules.java
@@ -7,18 +7,28 @@
 
 package org.elasticsearch.xpack.searchbusinessrules;
 
+import org.elasticsearch.plugins.ExtensiblePlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.SearchPlugin;
+import org.elasticsearch.plugins.SearchPlugin.QuerySpec;
+import org.elasticsearch.plugins.SearchPlugin.RetrieverSpec;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xpack.searchbusinessrules.retriever.PinnedRetrieverBuilder;
 
 import java.util.List;
 
 import static java.util.Collections.singletonList;
 
-public class SearchBusinessRules extends Plugin implements SearchPlugin {
+public class SearchBusinessRules extends Plugin implements SearchPlugin, ExtensiblePlugin {
 
     @Override
     public List<QuerySpec<?>> getQueries() {
-        return singletonList(new QuerySpec<>(PinnedQueryBuilder.NAME, PinnedQueryBuilder::new, PinnedQueryBuilder::fromXContent));
+        return List.of(new QuerySpec<>(PinnedQueryBuilder.NAME, PinnedQueryBuilder::new, PinnedQueryBuilder::fromXContent));
+    }
+
+    @Override
+    public List<RetrieverSpec<?>> getRetrievers() {
+        return singletonList(new RetrieverSpec<>(new ParseField(PinnedRetrieverBuilder.NAME), PinnedRetrieverBuilder::fromXContent));
     }
 
 }
diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SearchBusinessRulesFeatures.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SearchBusinessRulesFeatures.java
@@ -0,0 +1,28 @@
+/*
+ * 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.searchbusinessrules;
+
+import org.elasticsearch.features.FeatureSpecification;
+import org.elasticsearch.features.NodeFeature;
+
+import java.util.Set;
+
+import static org.elasticsearch.xpack.searchbusinessrules.retriever.PinnedRetrieverBuilder.PINNED_RETRIEVER_FEATURE;
+
+public class SearchBusinessRulesFeatures implements FeatureSpecification {
+
+    @Override
+    public Set<NodeFeature> getFeatures() {
+        return Set.of();
+    }
+
+    @Override
+    public Set<NodeFeature> getTestFeatures() {
+        return Set.of(PINNED_RETRIEVER_FEATURE);
+    }
+}
diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SpecifiedDocument.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/SpecifiedDocument.java
@@ -94,7 +94,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         return builder.endObject();
     }
 
-    static final ConstructingObjectParser<SpecifiedDocument, Void> PARSER = new ConstructingObjectParser<>(
+    public static final ConstructingObjectParser<SpecifiedDocument, Void> PARSER = new ConstructingObjectParser<>(
         NAME,
         a -> new SpecifiedDocument((String) a[0], (String) a[1])
     );
diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRankDoc.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRankDoc.java
@@ -0,0 +1,81 @@
+/*
+ * 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.searchbusinessrules.retriever;
+
+import org.apache.lucene.search.Explanation;
+import org.elasticsearch.TransportVersion;
+import org.elasticsearch.TransportVersions;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.search.rank.RankDoc;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class PinnedRankDoc extends RankDoc {
+    public static final String NAME = "pinned_rank_doc";
+
+    private final boolean isPinned;
+
+    public PinnedRankDoc(int docId, float score, int shardIndex, boolean isPinned) {
+        super(docId, score, shardIndex);
+        this.isPinned = isPinned;
+    }
+
+    public PinnedRankDoc(StreamInput in) throws IOException {
+        super(in);
+        this.isPinned = in.readBoolean();
+    }
+
+    public boolean isPinned() {
+        return isPinned;
+    }
+
+    @Override
+    public Explanation explain(Explanation[] sources, String[] queryNames) {
+        if (isPinned) {
+            return Explanation.match(score, "Pinned document, original explanation:", sources);
+        } else {
+            return super.explain(sources, queryNames);
+        }
+    }
+
+    @Override
+    public String toString() {
+        return super.toString() + ", isPinned=" + isPinned;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    protected void doWriteTo(StreamOutput out) throws IOException {
+        out.writeBoolean(isPinned);
+    }
+
+    @Override
+    protected boolean doEquals(RankDoc rd) {
+        if (rd instanceof PinnedRankDoc other) {
+            return this.isPinned == other.isPinned;
+        } else {
+            return false;
+        }
+    }
+
+    @Override
+    protected int doHashCode() {
+        return Objects.hash(super.doHashCode(), isPinned);
+    }
+
+    @Override
+    public TransportVersion getMinimalSupportedVersion() {
+        return TransportVersions.PINNED_RETRIEVER;
+    }
+}
diff --git a/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java b/x-pack/plugin/search-business-rules/src/main/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilder.java
diff --git a/x-pack/plugin/search-business-rules/src/main/resources/META-INF/services/org.elasticsearch.features.FeatureSpecification b/x-pack/plugin/search-business-rules/src/main/resources/META-INF/services/org.elasticsearch.features.FeatureSpecification
diff --git a/x-pack/plugin/search-business-rules/src/test/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilderTests.java b/x-pack/plugin/search-business-rules/src/test/java/org/elasticsearch/xpack/searchbusinessrules/retriever/PinnedRetrieverBuilderTests.java
diff --git a/x-pack/plugin/search-business-rules/src/yamlRestTest/java/org/elasticsearch/xpack/searchbusinessrules/SearchBusinessRulesClientYamlTestSuiteIT.java b/x-pack/plugin/search-business-rules/src/yamlRestTest/java/org/elasticsearch/xpack/searchbusinessrules/SearchBusinessRulesClientYamlTestSuiteIT.java
diff --git a/x-pack/plugin/search-business-rules/src/yamlRestTest/resources/rest-api-spec/test/search-business-rules/10_pinned_retriever.yml b/x-pack/plugin/search-business-rules/src/yamlRestTest/resources/rest-api-spec/test/search-business-rules/10_pinned_retriever.yml

Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@ dependencies {`
`26`	`26`	`clusterModules project(xpackModule('rank-rrf'))`
`27`	`27`	`clusterModules project(xpackModule('inference'))`
`28`	`28`	`clusterModules project(':modules:lang-painless')`
	`29`	`+ clusterModules project(xpackModule('search-business-rules'))`
`29`	`30`
`30`	`31`	`clusterPlugins project(':x-pack:plugin:inference:qa:test-service-plugin')`
`31`	`32`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,4 +16,7 @@`
`16`	`16`	`requires org.elasticsearch.xcore;`
`17`	`17`
`18`	`18`	`exports org.elasticsearch.xpack.searchbusinessrules;`
	`19`	`+`
	`20`	`+ provides org.elasticsearch.features.FeatureSpecification with org.elasticsearch.xpack.searchbusinessrules.SearchBusinessRulesFeatures;`
	`21`	`+ provides org.elasticsearch.plugins.SearchPlugin with org.elasticsearch.xpack.searchbusinessrules.SearchBusinessRules;`
`19`	`22`	`}`