Simplify LearnToRankRescorerBuilder serialization code.

afoucret · Nov 21, 2023 · 5560737 · 5560737
1 parent 275c934
commit 5560737
Showing 1 changed file with 3 additions and 12 deletions.
diff --git a/...ck/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/ltr/LearnToRankRescorerBuilder.java b/...ck/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/ltr/LearnToRankRescorerBuilder.java
@@ -21,6 +21,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.ml.action.GetTrainedModelsAction;
+import org.elasticsearch.xpack.core.ml.inference.trainedmodel.InferenceConfig;
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.LearnToRankConfig;
 import org.elasticsearch.xpack.ml.inference.loadingservice.LocalModel;
 
@@ -89,11 +90,7 @@ public LearnToRankRescorerBuilder(StreamInput input, LearnToRankService learnToR
         super(input);
         this.modelId = input.readString();
         this.params = input.readMap();
-        if (input.readBoolean()) {
-            this.learnToRankConfig = new LearnToRankConfig(input);
-        } else {
-            this.learnToRankConfig = null;
-        }
+        this.learnToRankConfig = (LearnToRankConfig) input.readOptionalNamedWriteable(InferenceConfig.class);
         this.learnToRankService = learnToRankService;
 
         this.localModel = null;
@@ -258,13 +255,7 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         assert localModel == null || rescoreOccurred : "Unnecessarily populated local model object";
         out.writeString(modelId);
         out.writeGenericMap(params);
-
-        if (learnToRankConfig != null) {
-            out.writeBoolean(true);
-            learnToRankConfig.writeTo(out);
-        } else {
-            out.writeBoolean(false);
-        }
+        out.writeOptionalNamedWriteable(learnToRankConfig);
     }
 
     @Override