From cda89ffd407be450c367750b20f7f28146dff6cb Mon Sep 17 00:00:00 2001
From: Yaliang Wu <ylwu@amazon.com>
Date: Thu, 28 Sep 2023 10:49:20 -0700
Subject: [PATCH] fix null exception in text docs data set (#1403)

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
---
 .../common/dataset/TextDocsInputDataSet.java  | 12 ++++-
 .../dataset/TextDocsInputDataSetTest.java     | 53 +++++++++++++++++++
 2 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 common/src/test/java/org/opensearch/ml/common/dataset/TextDocsInputDataSetTest.java

diff --git a/common/src/main/java/org/opensearch/ml/common/dataset/TextDocsInputDataSet.java b/common/src/main/java/org/opensearch/ml/common/dataset/TextDocsInputDataSet.java
index 93a26f649c..98672841d7 100644
--- a/common/src/main/java/org/opensearch/ml/common/dataset/TextDocsInputDataSet.java
+++ b/common/src/main/java/org/opensearch/ml/common/dataset/TextDocsInputDataSet.java
@@ -15,6 +15,7 @@
 import org.opensearch.ml.common.output.model.ModelResultFilter;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 
@@ -40,7 +41,11 @@ public TextDocsInputDataSet(List<String> docs, ModelResultFilter resultFilter) {
 
     public TextDocsInputDataSet(StreamInput streamInput) throws IOException {
         super(MLInputDataType.TEXT_DOCS);
-        docs = streamInput.readStringList();
+        docs = new ArrayList<>();
+        int size = streamInput.readInt();
+        for (int i=0; i<size; i++) {
+            docs.add(streamInput.readOptionalString());
+        }
         if (streamInput.readBoolean()) {
             resultFilter = new ModelResultFilter(streamInput);
         } else {
@@ -51,7 +56,10 @@ public TextDocsInputDataSet(StreamInput streamInput) throws IOException {
     @Override
     public void writeTo(StreamOutput streamOutput) throws IOException {
         super.writeTo(streamOutput);
-        streamOutput.writeStringCollection(docs);
+        streamOutput.writeInt(docs.size());
+        for (String doc : docs) {
+            streamOutput.writeOptionalString(doc);
+        }
         if (resultFilter != null) {
             streamOutput.writeBoolean(true);
             resultFilter.writeTo(streamOutput);
diff --git a/common/src/test/java/org/opensearch/ml/common/dataset/TextDocsInputDataSetTest.java b/common/src/test/java/org/opensearch/ml/common/dataset/TextDocsInputDataSetTest.java
new file mode 100644
index 0000000000..89f629e7c2
--- /dev/null
+++ b/common/src/test/java/org/opensearch/ml/common/dataset/TextDocsInputDataSetTest.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.dataset;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.io.stream.StreamInput;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+import static org.junit.Assert.assertEquals;
+
+public class TextDocsInputDataSetTest {
+
+    @Rule
+    public ExpectedException exceptionRule = ExpectedException.none();
+
+    @Test
+    public void writeTo_Success() throws IOException {
+        TextDocsInputDataSet inputDataSet = TextDocsInputDataSet.builder().docs(Arrays.asList("doc1", "doc2")).build();
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        inputDataSet.writeTo(bytesStreamOutput);
+
+        StreamInput streamInput = bytesStreamOutput.bytes().streamInput();
+        MLInputDataType inputDataType = streamInput.readEnum(MLInputDataType.class);
+        assertEquals(MLInputDataType.TEXT_DOCS, inputDataType);
+        TextDocsInputDataSet inputDataSet1 = new TextDocsInputDataSet((streamInput));
+        assertEquals(2, inputDataSet1.getDocs().size());
+        assertEquals("doc1", inputDataSet1.getDocs().get(0));
+        assertEquals("doc2", inputDataSet1.getDocs().get(1));
+    }
+
+    @Test
+    public void writeTo_Success_NullValue() throws IOException {
+        TextDocsInputDataSet inputDataSet = TextDocsInputDataSet.builder().docs(Arrays.asList("doc1", null)).build();
+        BytesStreamOutput bytesStreamOutput = new BytesStreamOutput();
+        inputDataSet.writeTo(bytesStreamOutput);
+
+        StreamInput streamInput = bytesStreamOutput.bytes().streamInput();
+        MLInputDataType inputDataType = streamInput.readEnum(MLInputDataType.class);
+        assertEquals(MLInputDataType.TEXT_DOCS, inputDataType);
+        TextDocsInputDataSet inputDataSet1 = new TextDocsInputDataSet((streamInput));
+        assertEquals(2, inputDataSet1.getDocs().size());
+        assertEquals("doc1", inputDataSet1.getDocs().get(0));
+        assertEquals(null, inputDataSet1.getDocs().get(1));
+    }
+}