diff --git a/plugin/src/main/java/org/opensearch/ml/stats/MLModelStats.java b/plugin/src/main/java/org/opensearch/ml/stats/MLModelStats.java index 592011938e..df3ceccd6a 100644 --- a/plugin/src/main/java/org/opensearch/ml/stats/MLModelStats.java +++ b/plugin/src/main/java/org/opensearch/ml/stats/MLModelStats.java @@ -27,7 +27,9 @@ public class MLModelStats implements ToXContentFragment, Writeable { private Map modelStats; public MLModelStats(StreamInput in) throws IOException { - this.modelStats = in.readMap(stream -> stream.readEnum(ActionName.class), MLActionStats::new); + if (in.readBoolean()) { + this.modelStats = in.readMap(stream -> stream.readEnum(ActionName.class), MLActionStats::new); + } } public MLModelStats(Map modelStats) { @@ -37,7 +39,10 @@ public MLModelStats(Map modelStats) { @Override public void writeTo(StreamOutput out) throws IOException { if (modelStats != null && modelStats.size() > 0) { + out.writeBoolean(true); out.writeMap(modelStats, (stream, v) -> stream.writeEnum(v), (stream, stats) -> stats.writeTo(stream)); + } else { + out.writeBoolean(false); } } diff --git a/plugin/src/test/java/org/opensearch/ml/stats/MLModelStatsTests.java b/plugin/src/test/java/org/opensearch/ml/stats/MLModelStatsTests.java index a38ccfde51..8be4b56f05 100644 --- a/plugin/src/test/java/org/opensearch/ml/stats/MLModelStatsTests.java +++ b/plugin/src/test/java/org/opensearch/ml/stats/MLModelStatsTests.java @@ -57,7 +57,10 @@ public void testEmptySerializationDeserialization() throws IOException { MLModelStats mlModelEmptyStats = new MLModelStats(modelStats); BytesStreamOutput output = new BytesStreamOutput(); mlModelEmptyStats.writeTo(output); - assertEquals(0, output.bytes().length()); + MLModelStats parsedMLModelStats = new MLModelStats(output.bytes().streamInput()); + MLActionStats parsedMLActionStats = parsedMLModelStats.getActionStats(ActionName.PREDICT); + assertNull(parsedMLActionStats); + // assertEquals(0, output.bytes().length()); } public void testToXContent() throws IOException {