Adds Internal mechanisms and retriever for MMR based result diversification

server/src/main/java/org/elasticsearch/search/SearchModule.java

@@ -208,6 +208,7 @@
 import org.elasticsearch.search.aggregations.pipeline.StatsBucketPipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.pipeline.SumBucketPipelineAggregationBuilder;
 import org.elasticsearch.search.aggregations.support.ValuesSourceRegistry;
+import org.elasticsearch.search.diversification.ResultDiversificationRetrieverBuilder;
 import org.elasticsearch.search.fetch.FetchPhase;
 import org.elasticsearch.search.fetch.FetchSubPhase;
 import org.elasticsearch.search.fetch.subphase.ExplainPhase;
@@ -1087,6 +1088,9 @@ private void registerRetrieverParsers(List<SearchPlugin> plugins) {
         registerRetriever(new RetrieverSpec<>(StandardRetrieverBuilder.NAME, StandardRetrieverBuilder::fromXContent));
         registerRetriever(new RetrieverSpec<>(KnnRetrieverBuilder.NAME, KnnRetrieverBuilder::fromXContent));
         registerRetriever(new RetrieverSpec<>(RescorerRetrieverBuilder.NAME, RescorerRetrieverBuilder::fromXContent));
+        registerRetriever(
+            new RetrieverSpec<>(ResultDiversificationRetrieverBuilder.NAME, ResultDiversificationRetrieverBuilder::fromXContent)
+        );
 
         registerFromPlugin(plugins, SearchPlugin::getRetrievers, this::registerRetriever);
     }

server/src/main/java/org/elasticsearch/search/diversification/ResultDiversification.java

@@ -0,0 +1,35 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.diversification;
+
+import org.apache.lucene.index.VectorSimilarityFunction;
+import org.elasticsearch.search.rank.RankDoc;
+import org.elasticsearch.search.vectors.VectorData;
+
+import java.io.IOException;
+
+/**
+ * Base interface for result diversification.
+ */
+public abstract class ResultDiversification {
+
+    public abstract RankDoc[] diversify(RankDoc[] docs, ResultDiversificationContext diversificationContext) throws IOException;
+
+    protected float getVectorComparisonScore(
+        VectorSimilarityFunction similarityFunction,
+        boolean useFloat,
+        VectorData thisDocVector,
+        VectorData comparisonVector
+    ) {
+        return useFloat
+            ? similarityFunction.compare(thisDocVector.floatVector(), comparisonVector.floatVector())
+            : similarityFunction.compare(thisDocVector.byteVector(), comparisonVector.byteVector());
+    }
+}

server/src/main/java/org/elasticsearch/search/diversification/ResultDiversificationContext.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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.diversification;
+
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
+import org.elasticsearch.search.vectors.VectorData;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+
+public abstract class ResultDiversificationContext {
+    private final String field;
+    private final int numCandidates;
+    private final DenseVectorFieldMapper fieldMapper;
+    private final IndexVersion indexVersion;
+    private final VectorData queryVector;
+    private Map<Integer, VectorData> fieldVectors;
+
+    // Field _must_ be a dense_vector type
+    protected ResultDiversificationContext(
+        String field,
+        int numCandidates,
+        DenseVectorFieldMapper fieldMapper,
+        IndexVersion indexVersion,
+        @Nullable VectorData queryVector,
+        @Nullable Map<Integer, VectorData> fieldVectors
+    ) {
+        this.field = field;
+        this.numCandidates = numCandidates;
+        this.fieldMapper = fieldMapper;
+        this.indexVersion = indexVersion;
+        this.queryVector = queryVector;
+        this.fieldVectors = fieldVectors == null ? new HashMap<>() : fieldVectors;
+    }
+
+    public String getField() {
+        return field;
+    }
+
+    public int getNumCandidates() {
+        return numCandidates;
+    }
+
+    public DenseVectorFieldMapper getFieldMapper() {
+        return fieldMapper;
+    }
+
+    public DenseVectorFieldMapper.ElementType getElementType() {
+        return fieldMapper.fieldType().getElementType();
+    }
+
+    public IndexVersion getIndexVersion() {
+        return indexVersion;
+    }
+
+    public void setFieldVectors(Map<Integer, VectorData> fieldVectors) {
+        this.fieldVectors = fieldVectors;
+    }
+
+    public VectorData getQueryVector() {
+        return queryVector;
+    }
+
+    public VectorData getFieldVector(int docId) {
+        return fieldVectors.getOrDefault(docId, null);
+    }
+
+    public Set<Map.Entry<Integer, VectorData>> getFieldVectorsEntrySet() {
+        return fieldVectors.entrySet();
+    }
+}