Write results to output field

davidkyle · Oct 3, 2023 · 13e2264 · 13e2264
1 parent fc99bf5
commit 13e2264
Show file tree

Hide file tree

Showing 17 changed files with 468 additions and 36 deletions.
diff --git a/server/src/main/java/org/elasticsearch/inference/InferenceResults.java b/server/src/main/java/org/elasticsearch/inference/InferenceResults.java
@@ -9,6 +9,7 @@
 package org.elasticsearch.inference;
 
 import org.elasticsearch.common.io.stream.NamedWriteable;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.ingest.IngestDocument;
 import org.elasticsearch.xcontent.ToXContentFragment;
 
@@ -24,16 +25,58 @@ static void writeResult(InferenceResults results, IngestDocument ingestDocument,
         Objects.requireNonNull(resultField, "resultField");
         Map<String, Object> resultMap = results.asMap();
         resultMap.put(MODEL_ID_RESULTS_FIELD, modelId);
-        if (ingestDocument.hasField(resultField)) {
-            ingestDocument.appendFieldValue(resultField, resultMap);
+        setOrAppendValue(resultField, resultMap, ingestDocument);
+    }
+
+    static void writeResultToField(
+        InferenceResults results,
+        IngestDocument ingestDocument,
+        @Nullable String basePath,
+        String outputField,
+        String modelId,
+        boolean includeModelId
+    ) {
+        Objects.requireNonNull(results, "results");
+        Objects.requireNonNull(ingestDocument, "ingestDocument");
+        Objects.requireNonNull(outputField, "outputField");
+        Map<String, Object> resultMap = results.asMap(outputField);
+        if (includeModelId) {
+            resultMap.put(MODEL_ID_RESULTS_FIELD, modelId);
+        }
+        if (basePath == null) {
+            // insert the results into the root of the document
+            for (var entry : resultMap.entrySet()) {
+                setOrAppendValue(entry.getKey(), entry.getValue(), ingestDocument);
+            }
         } else {
-            ingestDocument.setFieldValue(resultField, resultMap);
+            for (var entry : resultMap.entrySet()) {
+                setOrAppendValue(basePath + "." + entry.getKey(), entry.getValue(), ingestDocument);
+            }
+        }
+    }
+
+    private static void setOrAppendValue(String path, Object value, IngestDocument ingestDocument) {
+        if (ingestDocument.hasField(path)) {
+            ingestDocument.appendFieldValue(path, value);
+        } else {
+            ingestDocument.setFieldValue(path, value);
         }
     }
 
     String getResultsField();
 
+    /**
+     * Convert to a map
+     * @return Map representation of the InferenceResult
+     */
     Map<String, Object> asMap();
 
+    /**
+     * Convert to a map placing the inference result in {@code outputField}
+     * @param outputField Write the inference result to this field
+     * @return Map representation of the InferenceResult
+     */
+    Map<String, Object> asMap(String outputField);
+
     Object predictedValue();
 }
diff --git a/...ava/org/elasticsearch/xpack/core/ml/inference/results/ClassificationInferenceResults.java b/...ava/org/elasticsearch/xpack/core/ml/inference/results/ClassificationInferenceResults.java
@@ -220,6 +220,19 @@ public String getResultsField() {
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
         map.put(resultsField, predictionFieldType.transformPredictedValue(value(), valueAsString()));
+        addSupportingFieldsToMap(map);
+        return map;
+    }
+
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        Map<String, Object> map = new LinkedHashMap<>();
+        map.put(outputField, predictionFieldType.transformPredictedValue(value(), valueAsString()));
+        addSupportingFieldsToMap(map);
+        return map;
+    }
+
+    private void addSupportingFieldsToMap(Map<String, Object> map) {
         if (topClasses.isEmpty() == false) {
             map.put(topNumClassesField, topClasses.stream().map(TopClassEntry::asValueMap).collect(Collectors.toList()));
         }
@@ -235,7 +248,6 @@ public Map<String, Object> asMap() {
                 featureImportance.stream().map(ClassificationFeatureImportance::toMap).collect(Collectors.toList())
             );
         }
-        return map;
     }
 
     @Override

diff --git a/...rc/main/java/org/elasticsearch/xpack/core/ml/inference/results/ErrorInferenceResults.java b/...rc/main/java/org/elasticsearch/xpack/core/ml/inference/results/ErrorInferenceResults.java
@@ -74,6 +74,12 @@ public Map<String, Object> asMap() {
         return asMap;
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        // errors do not have a result
+        return asMap();
+    }
+
     @Override
     public String toString() {
         return Strings.toString(this);

diff --git a/...core/src/main/java/org/elasticsearch/xpack/core/ml/inference/results/FillMaskResults.java b/...core/src/main/java/org/elasticsearch/xpack/core/ml/inference/results/FillMaskResults.java
@@ -55,6 +55,13 @@ void addMapFields(Map<String, Object> map) {
         map.put(resultsField + "_sequence", predictedSequence);
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        var map = super.asMap(outputField);
+        map.put(outputField + "_sequence", predictedSequence);
+        return map;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;

diff --git a/.../src/main/java/org/elasticsearch/xpack/core/ml/inference/results/NlpInferenceResults.java b/.../src/main/java/org/elasticsearch/xpack/core/ml/inference/results/NlpInferenceResults.java
@@ -59,10 +59,21 @@ public final XContentBuilder toXContent(XContentBuilder builder, Params params)
     public final Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
         addMapFields(map);
+        addSupportingFieldsToMap(map);
+        return map;
+    }
+
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        Map<String, Object> map = new LinkedHashMap<>();
+        addSupportingFieldsToMap(map);
+        return map;
+    }
+
+    private void addSupportingFieldsToMap(Map<String, Object> map) {
         if (isTruncated) {
             map.put("is_truncated", isTruncated);
         }
-        return map;
     }
 
     @Override

diff --git a/...ain/java/org/elasticsearch/xpack/core/ml/inference/results/PyTorchPassThroughResults.java b/...ain/java/org/elasticsearch/xpack/core/ml/inference/results/PyTorchPassThroughResults.java
@@ -65,6 +65,13 @@ void addMapFields(Map<String, Object> map) {
         map.put(resultsField, inference);
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        var map = super.asMap(outputField);
+        map.put(outputField, inference);
+        return map;
+    }
+
     @Override
     public Object predictedValue() {
         throw new UnsupportedOperationException("[" + NAME + "] does not support a single predicted value");

diff --git a/.../org/elasticsearch/xpack/core/ml/inference/results/QuestionAnsweringInferenceResults.java b/.../org/elasticsearch/xpack/core/ml/inference/results/QuestionAnsweringInferenceResults.java
@@ -121,6 +121,18 @@ public String predictedValue() {
     @Override
     void addMapFields(Map<String, Object> map) {
         map.put(resultsField, answer);
+        addSupportingFieldsToMap(map);
+    }
+
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        var map = super.asMap(outputField);
+        map.put(outputField, answer);
+        addSupportingFieldsToMap(map);
+        return map;
+    }
+
+    private void addSupportingFieldsToMap(Map<String, Object> map) {
         map.put(START_OFFSET.getPreferredName(), startOffset);
         map.put(END_OFFSET.getPreferredName(), endOffset);
         if (topClasses.isEmpty() == false) {

diff --git a/.../src/main/java/org/elasticsearch/xpack/core/ml/inference/results/RawInferenceResults.java b/.../src/main/java/org/elasticsearch/xpack/core/ml/inference/results/RawInferenceResults.java
@@ -69,6 +69,11 @@ public Map<String, Object> asMap() {
         throw new UnsupportedOperationException("[raw] does not support map conversion");
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        throw new UnsupportedOperationException("[raw] does not support map conversion");
+    }
+
     @Override
     public Object predictedValue() {
         return null;

diff --git a/...in/java/org/elasticsearch/xpack/core/ml/inference/results/RegressionInferenceResults.java b/...in/java/org/elasticsearch/xpack/core/ml/inference/results/RegressionInferenceResults.java
@@ -121,11 +121,23 @@ public String getResultsField() {
     @Override
     public Map<String, Object> asMap() {
         Map<String, Object> map = new LinkedHashMap<>();
+        addSupportingFieldsToMap(map);
         map.put(resultsField, value());
+        return map;
+    }
+
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        Map<String, Object> map = new LinkedHashMap<>();
+        addSupportingFieldsToMap(map);
+        map.put(outputField, value());
+        return map;
+    }
+
+    private void addSupportingFieldsToMap(Map<String, Object> map) {
         if (featureImportance.isEmpty() == false) {
             map.put(FEATURE_IMPORTANCE, featureImportance.stream().map(RegressionFeatureImportance::toMap).collect(Collectors.toList()));
         }
-        return map;
     }
 
     @Override

diff --git a/...src/main/java/org/elasticsearch/xpack/core/ml/inference/results/TextEmbeddingResults.java b/...src/main/java/org/elasticsearch/xpack/core/ml/inference/results/TextEmbeddingResults.java
@@ -72,6 +72,13 @@ void addMapFields(Map<String, Object> map) {
         map.put(resultsField, inference);
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        var map = super.asMap(outputField);
+        map.put(outputField, inference);
+        return map;
+    }
+
     @Override
     public Object predictedValue() {
         throw new UnsupportedOperationException("[" + NAME + "] does not support a single predicted value");

diff --git a/...src/main/java/org/elasticsearch/xpack/core/ml/inference/results/TextExpansionResults.java b/...src/main/java/org/elasticsearch/xpack/core/ml/inference/results/TextExpansionResults.java
@@ -119,4 +119,11 @@ void doWriteTo(StreamOutput out) throws IOException {
     void addMapFields(Map<String, Object> map) {
         map.put(resultsField, weightedTokens.stream().collect(Collectors.toMap(WeightedToken::token, WeightedToken::weight)));
     }
+
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        var map = super.asMap(outputField);
+        map.put(outputField, weightedTokens.stream().collect(Collectors.toMap(WeightedToken::token, WeightedToken::weight)));
+        return map;
+    }
 }
diff --git a/...ava/org/elasticsearch/xpack/core/ml/inference/results/TextSimilarityInferenceResults.java b/...ava/org/elasticsearch/xpack/core/ml/inference/results/TextSimilarityInferenceResults.java
@@ -68,6 +68,13 @@ void addMapFields(Map<String, Object> map) {
         map.put(resultsField, score);
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        var map = super.asMap(outputField);
+        map.put(outputField, score);
+        return map;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;

diff --git a/.../main/java/org/elasticsearch/xpack/core/ml/inference/results/WarningInferenceResults.java b/.../main/java/org/elasticsearch/xpack/core/ml/inference/results/WarningInferenceResults.java
@@ -75,6 +75,12 @@ public Map<String, Object> asMap() {
         return asMap;
     }
 
+    @Override
+    public Map<String, Object> asMap(String outputField) {
+        // warnings do not have a result
+        return asMap();
+    }
+
     @Override
     public Object predictedValue() {
         return null;