From 4abc1d1eb144d6a5eafbaa1533311e7c39db1e7e Mon Sep 17 00:00:00 2001 From: Dhrubo Saha Date: Wed, 4 Oct 2023 12:58:08 -0700 Subject: [PATCH] adding boolean check if model stats exists Signed-off-by: Dhrubo Saha --- .../main/java/org/opensearch/ml/stats/MLModelStats.java | 7 ++++++- .../java/org/opensearch/ml/stats/MLModelStatsTests.java | 5 ++++- 2 files changed, 10 insertions(+), 2 deletions(-) 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 {