Pinned retriever

docs/changelog/126401.yaml

@@ -0,0 +1,5 @@
+pr: 126401
+summary: Add pinned retriever
+area: Relevance
+type: enhancement
+issues: []

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,

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

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')
 }

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 }
+
+

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"))
+}
+

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;
 }

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);

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));
     }
 
 }

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);
+    }
+}

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])
     );

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;
+    }
+}