Add getEstimatedSize API to OutputCodec (opensearch-project#5593)

* Add getEstimatedSize API to OutputCodec

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

* Fixed build failure

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

* Addressed review comments

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

---------

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

Author: kkondaka

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/OutputCodec.java

@@ -15,6 +15,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.io.ByteArrayOutputStream;
 import java.util.Map;
 
 public interface OutputCodec {
@@ -28,7 +29,7 @@ public interface OutputCodec {
      * @param outputStream  outputStream param for wrapping
      * @param event         Event to auto-generate schema
      * @param context       Extra Context used in Codec.
-     * @throws IOException throws IOException when invalid input is received or not able to create wrapping
+     * @throws IOException  throws IOException when invalid input is received or not able to create wrapping
      */
     void start(OutputStream outputStream, Event event, OutputCodecContext context) throws IOException;
 
@@ -51,6 +52,21 @@ public interface OutputCodec {
      */
     void complete(OutputStream outputStream) throws IOException;
 
+    /**
+     * this method get called from {@link Sink} to estimate size of event in {@link OutputStream}
+     *
+     * @param event        event Record event
+     * @return long        size of the serialized event
+     * @throws IOException throws IOException when invalid input is received or not able to create wrapping
+     */
+    default long getEstimatedSize(Event event, OutputCodecContext codecContext) throws IOException {
+        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+        start(outputStream, event, codecContext);
+        writeEvent(event, outputStream);
+        complete(outputStream);
+        return outputStream.toByteArray().length;
+    }
+
     /**
      * used to get extension of file
      *

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/OutputCodecTest.java

@@ -45,7 +45,7 @@ void validateAgainstCodecContext_does_not_throw_or_interact_with_outputCodecCont
     }
 
     @Test
-    public void testWriteMetrics() throws JsonProcessingException {
+    public void testWriteMetrics() throws JsonProcessingException, IOException {
         OutputCodec outputCodec = new OutputCodec() {
             @Override
             public void start(OutputStream outputStream, Event event, final OutputCodecContext codecContext) throws IOException {
@@ -73,6 +73,7 @@ public String getExtension() {
         final JacksonEvent event = JacksonLog.builder().withData(json).withEventMetadata(defaultEventMetadata).build();
         Event tagsToEvent = outputCodec.addTagsToEvent(event, "Tag");
         assertNotEquals(event.toJsonString(), tagsToEvent.toJsonString());
+        assertThat(outputCodec.getEstimatedSize(event, new OutputCodecContext()), equalTo(0L));
     }
 
     private static Map<String, Object> generateJson() {

data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/codec/json/JsonOutputCodecTest.java

@@ -24,6 +24,7 @@
 import java.util.Set;
 
 import static org.hamcrest.CoreMatchers.equalTo;
+import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
 
@@ -147,6 +148,19 @@ void writeEvent_with_exclude_keys(final int numberOfRecords) throws IOException
         assertThat(index, equalTo(numberOfRecords));
     }
 
+    @Test
+    void testGetEstimatedSize() throws Exception {
+        int numberOfRecords = 1;
+        JsonOutputCodec jsonOutputCodec = createObjectUnderTest();
+        outputStream = new ByteArrayOutputStream();
+        OutputCodecContext codecContext = new OutputCodecContext();
+        final List<Map<String, Object>> expectedData = generateRecords(numberOfRecords);
+        final Event event = convertToEvent(expectedData.get(0));
+        jsonOutputCodec.start(outputStream, null, codecContext);
+        String expectedEventString = "{\"events\":[{\"name\":\"Person0\",\"age\":0}]";
+        assertThat(jsonOutputCodec.getEstimatedSize(event, codecContext), greaterThanOrEqualTo((long)(expectedEventString.length())));
+    }
+
     @Test
     void testGetExtension() {
         JsonOutputCodec jsonOutputCodec = createObjectUnderTest();
@@ -185,4 +199,4 @@ private Object getValue(JsonNode jsonNode) {
 
         throw new RuntimeException("Test not setup correctly.");
     }
-}
+}

data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/codec/json/NewlineDelimitedOutputCodecTest.java

@@ -89,4 +89,17 @@ void testGetExtension() {
 
         assertThat("ndjson", equalTo(ndjsonOutputCodec.getExtension()));
     }
-}
+
+    @Test
+    void testGetEstimatedSize() throws Exception {
+        NewlineDelimitedOutputCodecTest.numberOfRecords = 1;
+        NdjsonOutputCodec ndjsonOutputCodec = createObjectUnderTest();
+        outputStream = new ByteArrayOutputStream();
+        OutputCodecContext codecContext = new OutputCodecContext();
+        ndjsonOutputCodec.start(outputStream, null, codecContext);
+
+        Record<Event> record = getRecord(0);
+        String expectedEventString = "{\"name\":\"Person0\",\"age\":0}\n";
+        assertThat(ndjsonOutputCodec.getEstimatedSize(record.getData(), codecContext), equalTo((long)expectedEventString.length()));
+    }
+}