Resolved merge conflict

Author: mridula-s109

x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/search/retriever/AbstractRetrieverBuilderTests.java

@@ -0,0 +1,241 @@
+/*
+ * 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.search.retriever;
+
+import org.elasticsearch.action.MockResolvedIndices;
+import org.elasticsearch.action.OriginalIndices;
+import org.elasticsearch.action.ResolvedIndices;
+import org.elasticsearch.action.support.IndicesOptions;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.metadata.InferenceFieldMetadata;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.Tuple;
+import org.elasticsearch.index.Index;
+import org.elasticsearch.index.IndexVersion;
+import org.elasticsearch.index.query.BoolQueryBuilder;
+import org.elasticsearch.index.query.MatchQueryBuilder;
+import org.elasticsearch.index.query.MultiMatchQueryBuilder;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryRewriteContext;
+import org.elasticsearch.index.query.TermsQueryBuilder;
+import org.elasticsearch.search.retriever.CompoundRetrieverBuilder.RetrieverSource;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.rank.linear.ScoreNormalizer;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public abstract class AbstractRetrieverBuilderTests<T extends CompoundRetrieverBuilder<T>> extends ESTestCase {
+
+    protected abstract float[] getWeights(T builder);
+
+    protected abstract ScoreNormalizer[] getScoreNormalizers(T builder);
+
+    protected abstract void assertCompoundRetriever(T originalRetriever, RetrieverBuilder rewrittenRetriever);
+
+    protected static ResolvedIndices createMockResolvedIndices(
+        Map<String, List<String>> localIndexInferenceFields,
+        Map<String, String> remoteIndexNames,
+        Map<String, String> commonInferenceIds
+    ) {
+        Map<Index, IndexMetadata> indexMetadata = new HashMap<>();
+
+        for (var indexEntry : localIndexInferenceFields.entrySet()) {
+            String indexName = indexEntry.getKey();
+            List<String> inferenceFields = indexEntry.getValue();
+
+            Index index = new Index(indexName, randomAlphaOfLength(10));
+
+            IndexMetadata.Builder indexMetadataBuilder = IndexMetadata.builder(index.getName())
+                .settings(
+                    Settings.builder()
+                        .put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current())
+                        .put(IndexMetadata.SETTING_INDEX_UUID, index.getUUID())
+                )
+                .numberOfShards(1)
+                .numberOfReplicas(0);
+
+            for (String inferenceField : inferenceFields) {
+                String inferenceId = commonInferenceIds.containsKey(inferenceField)
+                    ? commonInferenceIds.get(inferenceField)
+                    : randomAlphaOfLengthBetween(3, 5);
+
+                indexMetadataBuilder.putInferenceField(
+                    new InferenceFieldMetadata(inferenceField, inferenceId, new String[] { inferenceField }, null)
+                );
+            }
+
+            indexMetadata.put(index, indexMetadataBuilder.build());
+        }
+
+        Map<String, OriginalIndices> remoteIndices = new HashMap<>();
+        if (remoteIndexNames != null) {
+            for (Map.Entry<String, String> entry : remoteIndexNames.entrySet()) {
+                remoteIndices.put(entry.getKey(), new OriginalIndices(new String[] { entry.getValue() }, IndicesOptions.DEFAULT));
+            }
+        }
+
+        return new MockResolvedIndices(
+            remoteIndices,
+            new OriginalIndices(localIndexInferenceFields.keySet().toArray(new String[0]), IndicesOptions.DEFAULT),
+            indexMetadata
+        );
+    }
+
+    protected void assertMultiFieldsParamsRewrite(
+        T retriever,
+        QueryRewriteContext ctx,
+        Map<String, Float> expectedNonInferenceFields,
+        Map<String, Float> expectedInferenceFields,
+        String expectedQuery,
+        ScoreNormalizer expectedNormalizer
+    ) {
+        Map<Tuple<String, List<String>>, Float> inferenceFields = new HashMap<>();
+        expectedInferenceFields.forEach((key, value) -> inferenceFields.put(new Tuple<>(key, List.of()), value));
+
+        assertMultiIndexMultiFieldsParamsRewrite(
+            retriever,
+            ctx,
+            Map.of(expectedNonInferenceFields, List.of()),
+            inferenceFields,
+            expectedQuery,
+            expectedNormalizer
+        );
+    }
+
+    @SuppressWarnings("unchecked")
+    protected void assertMultiIndexMultiFieldsParamsRewrite(
+        T retriever,
+        QueryRewriteContext ctx,
+        Map<Map<String, Float>, List<String>> expectedNonInferenceFields,
+        Map<Tuple<String, List<String>>, Float> expectedInferenceFields,
+        String expectedQuery,
+        ScoreNormalizer expectedNormalizer
+    ) {
+        Set<QueryBuilder> expectedLexicalQueryBuilders = expectedNonInferenceFields.entrySet().stream().map(entry -> {
+            Map<String, Float> fields = entry.getKey();
+            List<String> indices = entry.getValue();
+
+            QueryBuilder queryBuilder = new MultiMatchQueryBuilder(expectedQuery).type(MultiMatchQueryBuilder.Type.MOST_FIELDS)
+                .fields(fields);
+
+            if (indices.isEmpty() == false) {
+                queryBuilder = new BoolQueryBuilder().must(queryBuilder).filter(new TermsQueryBuilder("_index", indices));
+            }
+            return queryBuilder;
+        }).collect(Collectors.toSet());
+
+        Set<InnerRetriever> expectedInnerSemanticRetrievers = expectedInferenceFields.entrySet().stream().map(entry -> {
+            var groupedInferenceField = entry.getKey();
+            var fieldName = groupedInferenceField.v1();
+            var indices = groupedInferenceField.v2();
+            var weight = entry.getValue();
+            QueryBuilder queryBuilder = new MatchQueryBuilder(fieldName, expectedQuery);
+            if (indices.isEmpty() == false) {
+                queryBuilder = new BoolQueryBuilder().must(queryBuilder).filter(new TermsQueryBuilder("_index", indices));
+            }
+            return new InnerRetriever(new StandardRetrieverBuilder(queryBuilder), weight, expectedNormalizer);
+        }).collect(Collectors.toSet());
+
+        RetrieverBuilder rewritten = retriever.doRewrite(ctx);
+        assertNotSame(retriever, rewritten);
+        assertCompoundRetriever(retriever, rewritten);
+
+        boolean assertedLexical = false;
+        boolean assertedSemantic = false;
+
+        for (InnerRetriever topInnerRetriever : getInnerRetrieversAsSet(retriever, (T) rewritten)) {
+            assertEquals(expectedNormalizer, topInnerRetriever.normalizer);
+            assertEquals(1.0f, topInnerRetriever.weight, 0.0f);
+
+            if (topInnerRetriever.retriever instanceof StandardRetrieverBuilder standardRetrieverBuilder) {
+                assertFalse("the lexical retriever is only asserted once", assertedLexical);
+                assertFalse(expectedNonInferenceFields.isEmpty());
+
+                QueryBuilder topDocsQueryBuilder = standardRetrieverBuilder.topDocsQuery();
+                if (expectedLexicalQueryBuilders.size() == 1) {
+                    assertEquals(topDocsQueryBuilder, expectedLexicalQueryBuilders.iterator().next());
+                } else {
+                    assertTrue(topDocsQueryBuilder instanceof BoolQueryBuilder);
+                    BoolQueryBuilder boolQueryBuilder = (BoolQueryBuilder) topDocsQueryBuilder;
+                    assertEquals(new HashSet<>(expectedLexicalQueryBuilders), new HashSet<>(boolQueryBuilder.should()));
+                }
+                assertedLexical = true;
+            } else {
+                assertFalse("the semantic retriever is only asserted once", assertedSemantic);
+                assertFalse(expectedInferenceFields.isEmpty());
+                assertEquals(expectedInnerSemanticRetrievers, topInnerRetriever.retriever);
+                assertedSemantic = true;
+            }
+        }
+    }
+
+    @SuppressWarnings("unchecked")
+    private Set<InnerRetriever> getInnerRetrieversAsSet(T originalRetriever, T rewrittenRetriever) {
+        float[] weights = getWeights(rewrittenRetriever);
+        ScoreNormalizer[] normalizers = getScoreNormalizers(rewrittenRetriever);
+
+        int i = 0;
+        Set<InnerRetriever> innerRetrieversSet = new HashSet<>();
+        for (RetrieverSource innerRetriever : rewrittenRetriever.innerRetrievers()) {
+            float weight = weights[i];
+            ScoreNormalizer normalizer = normalizers != null ? normalizers[i] : null;
+
+            if (innerRetriever.retriever() instanceof CompoundRetrieverBuilder<?> compoundRetriever) {
+                assertCompoundRetriever(originalRetriever, compoundRetriever);
+                innerRetrieversSet.add(
+                    new InnerRetriever(getInnerRetrieversAsSet(originalRetriever, (T) compoundRetriever), weight, normalizer)
+                );
+            } else {
+                innerRetrieversSet.add(new InnerRetriever(innerRetriever.retriever(), weight, normalizer));
+            }
+
+            i++;
+        }
+
+        return innerRetrieversSet;
+    }
+
+    private static class InnerRetriever {
+        private final Object retriever;
+        private final float weight;
+        private final ScoreNormalizer normalizer;
+
+        InnerRetriever(RetrieverBuilder retriever, float weight, ScoreNormalizer normalizer) {
+            this.retriever = retriever;
+            this.weight = weight;
+            this.normalizer = normalizer;
+        }
+
+        InnerRetriever(Set<InnerRetriever> innerRetrievers, float weight, ScoreNormalizer normalizer) {
+            this.retriever = innerRetrievers;
+            this.weight = weight;
+            this.normalizer = normalizer;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            InnerRetriever that = (InnerRetriever) o;
+            return Float.compare(weight, that.weight) == 0
+                && Objects.equals(retriever, that.retriever)
+                && Objects.equals(normalizer, that.normalizer);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(retriever, weight, normalizer);
+        }
+    }
+}