Use FieldInferenceMetadata structure in dependencies

carlosdelest · Mar 13, 2024 · 588e67f · 588e67f
1 parent 547ea9a
commit 588e67f
Show file tree

Hide file tree

Showing 11 changed files with 48 additions and 34 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java b/server/src/main/java/org/elasticsearch/index/query/QueryRewriteContext.java
@@ -9,10 +9,12 @@
 
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.cluster.metadata.FieldInferenceMetadata;
 import org.elasticsearch.common.collect.Iterators;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.util.concurrent.CountDown;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexSettings;
 import org.elasticsearch.index.analysis.IndexAnalyzers;
@@ -64,7 +66,7 @@ public class QueryRewriteContext {
     protected Predicate<String> allowedFields;
     private final ModelRegistry modelRegistry;
     private final InferenceServiceRegistry inferenceServiceRegistry;
-    private final Map<String, Set<String>> modelsForFields;
+    private final Map<String, FieldInferenceMetadata> fieldInferenceMetadataForIndices;
 
     public QueryRewriteContext(
         final XContentParserConfiguration parserConfiguration,
@@ -83,8 +85,8 @@ public QueryRewriteContext(
         final ScriptCompiler scriptService,
         final ModelRegistry modelRegistry,
         final InferenceServiceRegistry inferenceServiceRegistry,
-        final Map<String, Set<String>> modelsForFields
-    ) {
+        @Nullable final Map<String, FieldInferenceMetadata> fieldInferenceMetadataForIndices
+        ) {
 
         this.parserConfiguration = parserConfiguration;
         this.client = client;
@@ -103,9 +105,7 @@ public QueryRewriteContext(
         this.scriptService = scriptService;
         this.modelRegistry = modelRegistry;
         this.inferenceServiceRegistry = inferenceServiceRegistry;
-        this.modelsForFields = modelsForFields != null ?
-            modelsForFields.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, e -> Set.copyOf(e.getValue()))) :
-            Collections.emptyMap();
+        this.fieldInferenceMetadataForIndices = Objects.requireNonNullElse(fieldInferenceMetadataForIndices, Map.of());
     }
 
     public QueryRewriteContext(final XContentParserConfiguration parserConfiguration, final Client client, final LongSupplier nowInMillis) {
@@ -136,7 +136,7 @@ public QueryRewriteContext(
         final LongSupplier nowInMillis,
         final ModelRegistry modelRegistry,
         final InferenceServiceRegistry inferenceServiceRegistry,
-        final Map<String, Set<String>> modelsForFields
+        final Map<String, FieldInferenceMetadata> fieldInferenceMetadataMap
     ) {
         this(
             parserConfiguration,
@@ -155,7 +155,7 @@ public QueryRewriteContext(
             null,
             modelRegistry,
             inferenceServiceRegistry,
-            modelsForFields
+            fieldInferenceMetadataMap
         );
     }
 
@@ -400,8 +400,7 @@ public InferenceServiceRegistry getInferenceServiceRegistry() {
         return inferenceServiceRegistry;
     }
 
-    public Set<String> getModelsForField(String fieldName) {
-        Set<String> models = modelsForFields.get(fieldName);
-        return models != null ? models : Collections.emptySet();
+    public Set<String> getInferenceIdsForField(String fieldName) {
+        return fieldInferenceMetadataForIndices.values().stream().map(v -> v.getInferenceIdForField(fieldName)).collect(Collectors.toSet());
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/indices/IndicesService.java b/server/src/main/java/org/elasticsearch/indices/IndicesService.java
@@ -32,6 +32,7 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.metadata.DataStream;
+import org.elasticsearch.cluster.metadata.FieldInferenceMetadata;
 import org.elasticsearch.cluster.metadata.IndexAbstraction;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
@@ -1702,18 +1703,21 @@ public AliasFilter buildAliasFilter(ClusterState state, String index, Set<String
     public QueryRewriteContext getRewriteContext(LongSupplier nowInMillis, IndicesRequest indicesRequest, ModelRegistry modelRegistry,
                                                  InferenceServiceRegistry inferenceServiceRegistry) {
         Index[] indices = indexNameExpressionResolver.concreteIndices(clusterService.state(), indicesRequest);
-        Map<String, Set<String>> modelsForFields = new HashMap<>();
+        Map<String, FieldInferenceMetadata> fieldInferenceMetadataMap = Arrays.stream(indices)
+            .collect(Collectors.toMap(Index::getName, index -> indexServiceSafe(index).getMetadata().getFieldInferenceMetadata()));
+
+        Map<String, Set<String>> inferenceIdsForFields = new HashMap<>();
+        // Collect all index inference ids for each field
         for (Index index : indices) {
-            Map<String, Set<String>> fieldsForModels = indexService(index).getMetadata().getFieldsForModels();
-            for (Map.Entry<String, Set<String>> entry : fieldsForModels.entrySet()) {
-                for (String fieldName : entry.getValue()) {
-                    Set<String> models = modelsForFields.computeIfAbsent(fieldName, v -> new HashSet<>());
-                    models.add(entry.getKey());
-                }
-            }
+            Map<String, String> inferenceIdForFieldsInIndex = indexService(index).getMetadata()
+                .getFieldInferenceMetadata()
+                .getInferenceIdForFields();
+            inferenceIdForFieldsInIndex.entrySet().forEach(entry -> {
+                inferenceIdsForFields.computeIfAbsent(entry.getKey(), k -> new HashSet<>()).add(entry.getValue());
+            });
         }
 
-        return new QueryRewriteContext(parserConfig, client, nowInMillis, modelRegistry, inferenceServiceRegistry, modelsForFields);
+        return new QueryRewriteContext(parserConfig, client, nowInMillis, modelRegistry, inferenceServiceRegistry, fieldInferenceMetadataMap);
     }
 
     public DataRewriteContext getDataRewriteContext(LongSupplier nowInMillis) {

diff --git a/server/src/test/java/org/elasticsearch/action/bulk/BulkOperationTests.java b/server/src/test/java/org/elasticsearch/action/bulk/BulkOperationTests.java
@@ -472,7 +472,7 @@ private static BulkShardRequest runBulkOperation(
     ) {
         Settings settings = Settings.builder().put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current()).build();
         IndexMetadata indexMetadata = IndexMetadata.builder(INDEX_NAME)
-            .fieldsForModels(fieldsForModels)
+            .fieldInferenceMetadata(fieldsForModels)
             .settings(settings)
             .numberOfShards(1)
             .numberOfReplicas(0)

diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexMetadataTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexMetadataTests.java
@@ -108,7 +108,7 @@ public void testIndexMetadataSerialization() throws IOException {
             .stats(indexStats)
             .indexWriteLoadForecast(indexWriteLoadForecast)
             .shardSizeInBytesForecast(shardSizeInBytesForecast)
-            .fieldsForModels(fieldsForModels)
+            .fieldInferenceMetadata(fieldsForModels)
             .build();
         assertEquals(system, metadata.isSystem());
 
@@ -556,7 +556,7 @@ public void testFieldsForModels() {
         assertThat(idxMeta1.getFieldsForModels(), equalTo(Map.of()));
 
         Map<String, Set<String>> fieldsForModels = randomFieldsForModels(false);
-        IndexMetadata idxMeta2 = IndexMetadata.builder(idxMeta1).fieldsForModels(fieldsForModels).build();
+        IndexMetadata idxMeta2 = IndexMetadata.builder(idxMeta1).fieldInferenceMetadata(fieldsForModels).build();
         assertThat(idxMeta2.getFieldsForModels(), equalTo(fieldsForModels));
     }
 

diff --git a/server/src/test/java/org/elasticsearch/index/mapper/FieldTypeLookupTests.java b/server/src/test/java/org/elasticsearch/index/mapper/FieldTypeLookupTests.java
@@ -37,7 +37,7 @@ public void testEmpty() {
         assertNotNull(names);
         assertThat(names, hasSize(0));
 
-        Map<String, Set<String>> fieldsForModels = lookup.getFieldsForModels();
+        Map<String, Set<String>> fieldsForModels = lookup.getInferenceIdsForFields();
         assertNotNull(fieldsForModels);
         assertTrue(fieldsForModels.isEmpty());
     }
@@ -48,7 +48,7 @@ public void testAddNewField() {
         assertNull(lookup.get("bar"));
         assertEquals(f.fieldType(), lookup.get("foo"));
 
-        Map<String, Set<String>> fieldsForModels = lookup.getFieldsForModels();
+        Map<String, Set<String>> fieldsForModels = lookup.getInferenceIdsForFields();
         assertNotNull(fieldsForModels);
         assertTrue(fieldsForModels.isEmpty());
     }
@@ -440,7 +440,7 @@ public void testInferenceModelFieldType() {
         assertEquals(f2.fieldType(), lookup.get("foo2"));
         assertEquals(f3.fieldType(), lookup.get("foo3"));
 
-        Map<String, Set<String>> fieldsForModels = lookup.getFieldsForModels();
+        Map<String, Set<String>> fieldsForModels = lookup.getInferenceIdsForFields();
         assertNotNull(fieldsForModels);
         assertEquals(2, fieldsForModels.size());
         assertEquals(Set.of("foo1", "foo2"), fieldsForModels.get("bar1"));

diff --git a/server/src/test/java/org/elasticsearch/index/mapper/MappingLookupTests.java b/server/src/test/java/org/elasticsearch/index/mapper/MappingLookupTests.java
@@ -122,8 +122,8 @@ public void testEmptyMappingLookup() {
         assertEquals(0, mappingLookup.getMapping().getMetadataMappersMap().size());
         assertFalse(mappingLookup.fieldMappers().iterator().hasNext());
         assertEquals(0, mappingLookup.getMatchingFieldNames("*").size());
-        assertNotNull(mappingLookup.getFieldsForModels());
-        assertTrue(mappingLookup.getFieldsForModels().isEmpty());
+        assertNotNull(mappingLookup.getInferenceIdsForFields());
+        assertTrue(mappingLookup.getInferenceIdsForFields().isEmpty());
     }
 
     public void testValidateDoesNotShadow() {
@@ -201,7 +201,7 @@ public void testFieldsForModels() {
         assertEquals(1, size(mappingLookup.fieldMappers()));
         assertEquals(fieldType, mappingLookup.getFieldType("test_field_name"));
 
-        Map<String, Set<String>> fieldsForModels = mappingLookup.getFieldsForModels();
+        Map<String, Set<String>> fieldsForModels = mappingLookup.getInferenceIdsForFields();
         assertNotNull(fieldsForModels);
         assertEquals(1, fieldsForModels.size());
         assertEquals(Collections.singleton("test_field_name"), fieldsForModels.get("test_model_id"));

diff --git a/test/framework/src/main/java/org/elasticsearch/index/mapper/MockInferenceModelFieldType.java b/test/framework/src/main/java/org/elasticsearch/index/mapper/MockInferenceModelFieldType.java
@@ -39,7 +39,7 @@ public ValueFetcher valueFetcher(SearchExecutionContext context, String format)
     }
 
     @Override
-    public String getInferenceModel() {
+    public String getInferenceId() {
         return modelId;
     }
 }
diff --git a/test/framework/src/main/java/org/elasticsearch/node/MockNode.java b/test/framework/src/main/java/org/elasticsearch/node/MockNode.java
@@ -28,6 +28,8 @@
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
 import org.elasticsearch.indices.recovery.RecoverySettings;
+import org.elasticsearch.inference.InferenceServiceRegistry;
+import org.elasticsearch.inference.ModelRegistry;
 import org.elasticsearch.plugins.MockPluginsService;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.PluginsService;
@@ -101,7 +103,9 @@ SearchService newSearchService(
             ResponseCollectorService responseCollectorService,
             CircuitBreakerService circuitBreakerService,
             ExecutorSelector executorSelector,
-            Tracer tracer
+            Tracer tracer,
+            ModelRegistry modelRegistry,
+            InferenceServiceRegistry inferenceServiceRegistry
         ) {
             if (pluginsService.filterPlugins(MockSearchService.TestPlugin.class).findAny().isEmpty()) {
                 return super.newSearchService(
@@ -115,7 +119,9 @@ SearchService newSearchService(
                     responseCollectorService,
                     circuitBreakerService,
                     executorSelector,
-                    tracer
+                    tracer,
+                    modelRegistry,
+                    inferenceServiceRegistry
                 );
             }
             return new MockSearchService(

diff --git a/test/framework/src/main/java/org/elasticsearch/search/MockSearchService.java b/test/framework/src/main/java/org/elasticsearch/search/MockSearchService.java
@@ -15,6 +15,8 @@
 import org.elasticsearch.indices.ExecutorSelector;
 import org.elasticsearch.indices.IndicesService;
 import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.inference.InferenceServiceRegistry;
+import org.elasticsearch.inference.ModelRegistry;
 import org.elasticsearch.node.MockNode;
 import org.elasticsearch.node.ResponseCollectorService;
 import org.elasticsearch.plugins.Plugin;
@@ -97,7 +99,9 @@ public MockSearchService(
             responseCollectorService,
             circuitBreakerService,
             executorSelector,
-            tracer
+            tracer,
+            null,
+            null
         );
     }
 

diff --git a/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/AbstractBuilderTestCase.java
@@ -603,6 +603,7 @@ QueryRewriteContext createQueryRewriteContext() {
                 () -> true,
                 scriptService,
                 null,
+                null,
                 null
             );
         }

diff --git a/...ference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java b/...ference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java
@@ -109,7 +109,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
             return this;
         }
 
-        Set<String> modelsForField = queryRewriteContext.getModelsForField(fieldName);
+        Set<String> modelsForField = queryRewriteContext.getInferenceIdsForField(fieldName);
         if (modelsForField.isEmpty()) {
             throw new IllegalArgumentException("Field [" + fieldName + "] is not a semantic_text field type");
         }
-Original file line number
+Diff line change
@@ Expand Up / @@ -603,6 +603,7 @@ QueryRewriteContext createQueryRewriteContext() { @@
                     () -> true,
                     scriptService,
                     null,
+                    null,
                     null
                 );
             }
@@ Expand Down @@