Add read and write optional enum sets methods to stream input and output (#17556)

* Add read and write optional enum sets to stream

Signed-off-by: Andy Qin <qinandy@amazon.com>

* Only write set if it is non-empty

Signed-off-by: Andy Qin <qinandy@amazon.com>

* Update javadoc

Signed-off-by: Andy Qin <qinandy@amazon.com>

* Read false presence flag as empty set

Co-authored-by: Daniel Widdis <widdis@gmail.com>
Signed-off-by: Andy <qinandy@amazon.com>

* Update javadoc

Signed-off-by: Andy Qin <qinandy@amazon.com>

* Add unit tests

Signed-off-by: Andy Qin <qinandy@amazon.com>

* Fix typos

Signed-off-by: Andy Qin <qinandy@amazon.com>

* Update changelog

Signed-off-by: Andy Qin <qinandy@amazon.com>

---------

Signed-off-by: Andy Qin <qinandy@amazon.com>
Signed-off-by: Andy <qinandy@amazon.com>
Co-authored-by: Daniel Widdis <widdis@gmail.com>

Author: q-andy

CHANGELOG-3.0.md

@@ -51,6 +51,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Use Lucene `BM25Similarity` as default since the `LegacyBM25Similarity` is marked as deprecated ([#17306](https://github.com/opensearch-project/OpenSearch/pull/17306))
 - Wildcard field index only 3gram of the input data [#17349](https://github.com/opensearch-project/OpenSearch/pull/17349)
 - Use BC libraries to parse PEM files, increase key length, allow general use of known cryptographic binary extensions, remove unused BC dependencies ([#3420](https://github.com/opensearch-project/OpenSearch/pull/14912))
+- Add optional enum set read / write functionality to stream input / output ([#17556](https://github.com/opensearch-project/OpenSearch/pull/17556))
 
 ### Deprecated
 

libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamInput.java

@@ -1305,6 +1305,21 @@ public <E extends Enum<E>> EnumSet<E> readEnumSet(Class<E> enumClass) throws IOE
         return res;
     }
 
+    /**
+     * Reads an optional enum set with type E that was serialized based on the value of each enum's ordinal
+     * The set is expected to have been written using {@link StreamOutput#writeOptionalEnumSet(EnumSet)}
+     *
+     * @return the enum set of strings
+     * @throws IOException if an I/O exception occurs reading the set
+     */
+    public <E extends Enum<E>> EnumSet<E> readOptionalEnumSet(Class<E> enumClass) throws IOException {
+        if (readBoolean()) {
+            return readEnumSet(enumClass);
+        } else {
+            return EnumSet.noneOf(enumClass);
+        }
+    }
+
     public static StreamInput wrap(byte[] bytes) {
         return wrap(bytes, 0, bytes.length);
     }

libs/core/src/main/java/org/opensearch/core/common/io/stream/StreamOutput.java

@@ -1255,6 +1255,19 @@ public <E extends Enum<E>> void writeEnumSet(EnumSet<E> enumSet) throws IOExcept
         }
     }
 
+    /**
+     * Writes an optional EnumSet with type E that serializes each enum based on its ordinal value
+     * For null or empty enum set, writes false;
+     */
+    public <E extends Enum<E>> void writeOptionalEnumSet(@Nullable EnumSet<E> enumSet) throws IOException {
+        if (enumSet != null && enumSet.size() > 0) {
+            writeBoolean(true);
+            writeEnumSet(enumSet);
+        } else {
+            writeBoolean(false);
+        }
+    }
+
     /**
      * Write a {@link TimeValue} to the stream
      */

libs/core/src/test/java/org/opensearch/core/common/io/stream/BaseStreamTests.java

@@ -53,6 +53,7 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
@@ -354,6 +355,32 @@ private <T> void runWriteReadCollectionTest(
         }
     }
 
+    public void testOptionalEnumSet() throws IOException {
+        EnumSet<TestEnum> enumSet = EnumSet.allOf(TestEnum.class);
+        BytesStreamOutput out = new BytesStreamOutput();
+        out.writeOptionalEnumSet(enumSet);
+        EnumSet<TestEnum> targetSet = getStreamInput(out.bytes()).readOptionalEnumSet(TestEnum.class);
+        assertEquals(enumSet, targetSet);
+
+        enumSet = EnumSet.of(TestEnum.A, TestEnum.C, TestEnum.E);
+        out = new BytesStreamOutput();
+        out.writeOptionalEnumSet(enumSet);
+        targetSet = getStreamInput(out.bytes()).readOptionalEnumSet(TestEnum.class);
+        assertEquals(enumSet, targetSet);
+
+        enumSet = EnumSet.noneOf(TestEnum.class);
+        out = new BytesStreamOutput();
+        out.writeOptionalEnumSet(enumSet);
+        targetSet = getStreamInput(out.bytes()).readOptionalEnumSet(TestEnum.class);
+        assertEquals(enumSet, targetSet);
+
+        enumSet = null;
+        out = new BytesStreamOutput();
+        out.writeOptionalEnumSet(enumSet);
+        targetSet = getStreamInput(out.bytes()).readOptionalEnumSet(TestEnum.class);
+        assertEquals(EnumSet.noneOf(TestEnum.class), targetSet);
+    }
+
     public void testSetOfLongs() throws IOException {
         final int size = randomIntBetween(0, 6);
         final Set<Long> sourceSet = new HashSet<>(size);
@@ -540,4 +567,11 @@ private void assertGenericRoundtrip(Object original) throws IOException {
         });
     }
 
+    private enum TestEnum {
+        A,
+        B,
+        C,
+        D,
+        E;
+    }
 }