[ML] Change result format to suit rank feature mapping (#93460)

elastic · Feb 7, 2023 · d52b6f9 · d52b6f9
1 parent e6bbfa2
commit d52b6f9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/...gin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/results/SlimResults.java b/...gin/core/src/main/java/org/elasticsearch/xpack/core/ml/inference/results/SlimResults.java
@@ -26,9 +26,6 @@ public class SlimResults extends NlpInferenceResults {
 
     public record WeightedToken(int token, float weight) implements Writeable, ToXContentObject {
 
-        public static final String TOKEN = "token";
-        public static final String WEIGHT = "weight";
-
         public WeightedToken(StreamInput in) throws IOException {
             this(in.readVInt(), in.readFloat());
         }
@@ -42,14 +39,13 @@ public void writeTo(StreamOutput out) throws IOException {
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
-            builder.field(TOKEN, token);
-            builder.field(WEIGHT, weight);
+            builder.field(Integer.toString(token), weight);
             builder.endObject();
             return builder;
         }
 
         public Map<String, Object> asMap() {
-            return Map.of(TOKEN, token, WEIGHT, weight);
+            return Map.of(Integer.toString(token), weight);
         }
 
         @Override

diff --git a/...ore/src/test/java/org/elasticsearch/xpack/core/ml/inference/results/SlimResultsTests.java b/...ore/src/test/java/org/elasticsearch/xpack/core/ml/inference/results/SlimResultsTests.java
@@ -32,7 +32,7 @@ protected SlimResults createTestInstance() {
 
     @Override
     protected SlimResults mutateInstance(SlimResults instance) {
-        return null;// TODO implement https://github.com/elastic/elasticsearch/issues/25929
+        return new SlimResults(instance.getResultsField() + "-FOO", instance.getWeightedTokens(), instance.isTruncated() == false);
     }
 
     @Override
@@ -45,8 +45,11 @@ void assertFieldValues(SlimResults createdInstance, IngestDocument document, Str
         var originalTokens = createdInstance.getWeightedTokens();
         assertEquals(originalTokens.size(), ingestedTokens.size());
         for (int i = 0; i < createdInstance.getWeightedTokens().size(); i++) {
-            assertEquals(originalTokens.get(i).token(), (int) ingestedTokens.get(i).get("token"));
-            assertEquals(originalTokens.get(i).weight(), (float) ingestedTokens.get(i).get("weight"), 0.0001);
+            assertEquals(
+                originalTokens.get(i).weight(),
+                (float) ingestedTokens.get(i).get(Integer.toString(originalTokens.get(i).token())),
+                0.0001
+            );
         }
     }
 }