update raw payload docID to include shardID

Signed-off-by: Varun Bharadwaj <[email protected]>

Author: varunbharadwaj

server/src/main/java/org/opensearch/indices/pollingingest/DefaultStreamPoller.java

@@ -114,7 +114,7 @@ private DefaultStreamPoller(
             pollTimeout,
             pointerBasedLagUpdateIntervalMs,
             ingestionEngine.config().getIndexSettings(),
-            IngestionMessageMapper.create(mapperType.getName())
+            IngestionMessageMapper.create(mapperType.getName(), shardId)
         );
     }
 

server/src/main/java/org/opensearch/indices/pollingingest/mappers/IngestionMessageMapper.java

@@ -68,15 +68,16 @@ public static MapperType fromString(String value) {
      * Factory method to create a mapper instance based on type string.
      *
      * @param mapperTypeString the type of mapper to create as a string
+     * @param shardId the shard ID
      * @return the mapper instance
      */
-    static IngestionMessageMapper create(String mapperTypeString) {
+    static IngestionMessageMapper create(String mapperTypeString, int shardId) {
         MapperType mapperType = MapperType.fromString(mapperTypeString);
         switch (mapperType) {
             case DEFAULT:
                 return new DefaultIngestionMessageMapper();
             case RAW_PAYLOAD:
-                return new RawPayloadIngestionMessageMapper();
+                return new RawPayloadIngestionMessageMapper(shardId);
             default:
                 throw new IllegalArgumentException("Unknown mapper type: " + mapperType);
         }

server/src/main/java/org/opensearch/indices/pollingingest/mappers/RawPayloadIngestionMessageMapper.java

@@ -22,22 +22,28 @@
 
 /**
  * Mapper implementation that uses the entire message payload as the document source.
- * The message pointer (Kafka offset, Kinesis sequence number, etc) is used as the document ID, and the operation type is "index".
+ * The message pointer (Kafka offset, Kinesis sequence number, etc) along with shardID is used as the document ID, and the operation type is "index".
+ * For the document ID, the shard ID is prefixed with the shard pointer to ensure uniqueness across all shards.
  *
  * <p> Note that raw payload will not support document versioning. Eventually consistent view of the documents can be expected
  * if there is a shard recovery resulting in message replay.
  */
 public class RawPayloadIngestionMessageMapper implements IngestionMessageMapper {
 
     private static final String OP_TYPE_INDEX = "index";
+    private final int shardId;
+
+    public RawPayloadIngestionMessageMapper(int shardId) {
+        this.shardId = shardId;
+    }
 
     @Override
     public ShardUpdateMessage mapAndProcess(IngestionShardPointer pointer, Message message) throws IllegalArgumentException {
         // Parse the raw payload - this will be the _source content
         Map<String, Object> sourceMap = IngestionUtils.getParsedPayloadMap((byte[]) message.getPayload());
 
-        // Use pointer as the document ID
-        String id = pointer.asString();
+        // Use shard ID prefix + pointer as the document ID to ensure uniqueness across shards
+        String id = shardId + "-" + pointer.asString();
 
         // No auto-generated ID timestamp since we're using the pointer as ID
         long autoGeneratedIdTimestamp = UNSET_AUTO_GENERATED_TIMESTAMP;

server/src/test/java/org/opensearch/indices/pollingingest/mappers/IngestionMessageMapperTests.java

@@ -36,15 +36,15 @@ public void testMapperTypeFromStringInvalid() {
     }
 
     public void testMapperCreation() {
-        IngestionMessageMapper defaultMapper = IngestionMessageMapper.create("default");
+        IngestionMessageMapper defaultMapper = IngestionMessageMapper.create("default", 0);
         assertNotNull(defaultMapper);
         assertTrue(defaultMapper instanceof DefaultIngestionMessageMapper);
 
-        IngestionMessageMapper rawPayloadMapper = IngestionMessageMapper.create("raw_payload");
+        IngestionMessageMapper rawPayloadMapper = IngestionMessageMapper.create("raw_payload", 0);
         assertNotNull(rawPayloadMapper);
         assertTrue(rawPayloadMapper instanceof RawPayloadIngestionMessageMapper);
 
-        expectThrows(IllegalArgumentException.class, () -> IngestionMessageMapper.create("unknown"));
+        expectThrows(IllegalArgumentException.class, () -> IngestionMessageMapper.create("unknown", 0));
     }
 
     public void testDefaultMapperWithIdPresent() {
@@ -113,7 +113,7 @@ public void testDefaultMapperWithVersion() {
     }
 
     public void testRawPayloadMapper() {
-        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper();
+        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(0);
         String payload = "{\"name\":\"alice\",\"age\":30,\"city\":\"Seattle\"}";
         byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);
 
@@ -129,8 +129,8 @@ public void testRawPayloadMapper() {
 
         Map<String, Object> parsedMap = result.parsedPayloadMap();
 
-        // Verify _id is set to pointer value
-        assertEquals("100", parsedMap.get("_id"));
+        // Verify _id is set to shard ID + pointer value
+        assertEquals("0-100", parsedMap.get("_id"));
 
         // Verify _op_type is set to "index"
         assertEquals("index", parsedMap.get("_op_type"));
@@ -151,7 +151,7 @@ public void testRawPayloadMapper() {
     }
 
     public void testRawPayloadMapperWithComplexObject() {
-        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper();
+        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(1);
         String payload = "{\"user\":{\"name\":\"bob\",\"email\":\"[email protected]\"},\"tags\":[\"tag1\",\"tag2\"],\"count\":42}";
         byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);
 
@@ -163,7 +163,7 @@ public void testRawPayloadMapperWithComplexObject() {
         assertNotNull(result);
         Map<String, Object> parsedMap = result.parsedPayloadMap();
 
-        assertEquals("200", parsedMap.get("_id"));
+        assertEquals("1-200", parsedMap.get("_id"));
         assertEquals("index", parsedMap.get("_op_type"));
 
         @SuppressWarnings("unchecked")
@@ -175,7 +175,7 @@ public void testRawPayloadMapperWithComplexObject() {
     }
 
     public void testRawPayloadMapperWithEmptyObject() {
-        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper();
+        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(2);
         String payload = "{}";
         byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);
 
@@ -187,7 +187,7 @@ public void testRawPayloadMapperWithEmptyObject() {
         assertNotNull(result);
         Map<String, Object> parsedMap = result.parsedPayloadMap();
 
-        assertEquals("300", parsedMap.get("_id"));
+        assertEquals("2-300", parsedMap.get("_id"));
         assertEquals("index", parsedMap.get("_op_type"));
 
         @SuppressWarnings("unchecked")
@@ -207,7 +207,7 @@ public void testDefaultMapperWithInvalidJson() {
     }
 
     public void testRawPayloadMapperWithInvalidJson() {
-        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper();
+        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(3);
         String payload = "not a json";
         byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);