add S3 Enrich processor to merge ml batch job output with source inputs by Zhangxunmt · Pull Request #5992 · opensearch-project/data-prepper

Zhangxunmt · 2025-08-13T18:38:17Z

Description

This PR implements the proposed S3 Enricher processor to merge the data from S3 into the events in the data prepper pipeline.

This processor can be used to merge the ML batch job results in the S3Key "<source_file_basename-xxxxx>.jsonl.out" with certain fields from source data in the S3Key "source_file_basename.jsonl".

This is the config of this processor:

  - s3_enrich:
      bucket:
        name: offlinebatch                              # Name of the S3 bucket containing enrichment source files
        filter:
          include_prefix: bedrockbatch/originsource/    # Only scan objects under this key prefix within the bucket
      codec:
        ndjson:                                         # Codec used to parse enrichment source files (newline-delimited JSON)
      default_bucket_owner: 802041417063                # AWS account ID that owns the bucket; used to validate bucket ownership on access
      aws:
        region: us-east-1
      s3_object_size_limit: 100mb                      # Maximum file size to load into memory; objects exceeding this limit are skipped (min: 100mb, max: 300mb)
      s3_key_path: "s3/key"                             # Event field path whose value contains the S3 output object key (used to derive the enrichment source object)

      # Regex applied to the output S3 key to extract the base name of the enrichment source file.
      # Capture group 1 becomes the source file stem; the codec extension is appended to form the final lookup key.
      # Example: "test_batch_50k-2025-11-06T21-19-15Z-xxx.jsonl.out" → capture group 1 = "test_batch_50k"
      #          → looks up "bedrockbatch/originsource/test_batch_50k.jsonl" in the bucket
      s3_object_name_pattern: ^(.*?)-\d{4}-\d{2}-\d{2}T\d{2}-\d{2}-\d{2}Z-.*\.jsonl\.out$

      correlation_keys: 
         - "recordId"                       # Field present in both the pipeline event and the enrichment source record; used to match and join the two
      keys_to_merge:                                    # Fields to copy from the matched enrichment source record into the pipeline event
        - "field_A"
        - "field_B"
        - "field_C"
      enrich_when: /s3/key != null                      # Data Prepper expression; processor is skipped for events where this evaluates to false``` 

 
### Issues Resolved
Resolves #[Issue number to be closed when this PR is merged]
 
### Check List
- [ ] New functionality includes testing.
- [ ] New functionality has a documentation issue. Please link to it in this PR.
  - [ ] New functionality has javadoc added
- [ ] Commits are signed with a real name per the DCO

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check [here](https://github.com/opensearch-project/data-prepper/blob/main/CONTRIBUTING.md).

dlvenable

Thank you @Zhangxunmt for this change!

I left some comments. One thing that might help reduce the amount of work in this PR is to create a different PR that moves all the common code from S3 source into s3-common as its own PR.

dlvenable

Thanks @Zhangxunmt , I have a few more comments.

dlvenable · 2026-03-20T20:49:47Z

+     */
+    public Cache<String, Event> getOrLoadRecordCache(String s3Url, Supplier<Cache<String, Event>> loader) {
+        return s3Cache.get(s3Url, key -> {
+            LOG.info("Loading S3 data for: {}", key);


This should probably be a debug log. It will happen too often for normal logging.

dlvenable · 2026-03-20T20:52:08Z

+            if (matcher.matches() && matcher.groupCount() >= 1) {
+                final String baseName = matcher.group(1);
+                if (baseName != null && !baseName.isBlank()) {
+                    return baseName + ".jsonl";


Should this jsonl extension be configurable?

let's reuse this config in the extension to make it parametric. This is used in the S3 sink to define the extension of the generated file.

codec: ndjson: extension: jsonl

dlvenable · 2026-03-20T20:56:52Z

+     * @param targetEvent the event to enrich (output record)
+     * @param sourceEvent the event containing source data (from cache)
+     */
+    private void mergeData(Event targetEvent, Event sourceEvent) {


The Event class has a merge method. Can it work for you? If not, is there a way it could be configured for your needs?

data-prepper/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java

Lines 432 to 447 in 45b0e86

@Override

public void merge(final Event other) {

if(!(other instanceof JacksonEvent))

throw new IllegalArgumentException("Unable to merge the Event. The input Event must be a JacksonEvent.");

final JacksonEvent otherJacksonEvent = (JacksonEvent) other;

if(!(otherJacksonEvent.jsonNode instanceof ObjectNode)) {

throw new IllegalArgumentException("Unable to merge the Event. The input Event must be a JacksonEvent with object data.");

}

final ObjectNode otherObjectNode = (ObjectNode) otherJacksonEvent.jsonNode;

if(!(jsonNode instanceof ObjectNode)) {

throw new UnsupportedOperationException("Unable to merge the Event. The current Event must have object data.");

}

((ObjectNode) jsonNode).setAll(otherObjectNode);

}

dlvenable · 2026-03-20T20:57:59Z

+@ExtendWith(MockitoExtension.class)
+class S3ObjectReferenceResolverTest {
+
+    private static final String BUCKET_NAME = "test-bucket";


Make this and KEY_PATH into a non-static fields. Initialize them as random values in @BeforeEach.

dlvenable · 2026-03-20T20:59:06Z

+
+    @Test
+    void resolve_throws_when_s3_key_is_blank() {
+        when(event.get(KEY_PATH, String.class)).thenReturn("   ");


I think you also need this:

when(bucketOption.getName()).thenReturn(bucketName);

Otherwise you may not testing what you mean to.

dlvenable · 2026-03-20T20:59:11Z

+
+    @Test
+    void resolve_throws_when_s3_key_is_null() {
+        when(event.get(KEY_PATH, String.class)).thenReturn(null);


I think you also need this:

when(bucketOption.getName()).thenReturn(bucketName);

Otherwise you may not testing what you mean to.

dlvenable · 2026-03-20T21:05:43Z

+    private final Counter numberOfRecordsSuccessCounter;
+    private final Counter numberOfRecordsFailedCounter;
+
+    @DataPrepperPluginConstructor


Maybe we should also make this @Experimental while we work on it in case we want to make some changes.

Yes agree. I will add the Experimental annotation to it.

dlvenable · 2026-03-23T21:29:37Z

+     * @throws UnsupportedOperationException if the current Event does not support merging.
+     * @since 2.11
+     */
+    void merge(Event other, List<String> keys);


Make this Collection<String>

dlvenable · 2026-03-23T21:30:54Z

    private final List<ScanObjectWorker> workers;

-    public S3ScanService(final S3SourceConfig s3SourceConfig,
+    public  S3ScanService(final S3SourceConfig s3SourceConfig,


Please revert this change.

dlvenable · 2026-03-23T21:54:18Z

+     * @param keys  the list of keys to selectively merge
+     * @throws IllegalArgumentException if the input event is not compatible to merge.
+     * @throws UnsupportedOperationException if the current Event does not support merging.
+     * @since 2.11


Suggested change

* @since 2.11

* @since 2.15

The next version is 2.15.

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Zhangxunmt · 2026-03-21T00:54:42Z

+            if (matcher.matches() && matcher.groupCount() >= 1) {
+                final String baseName = matcher.group(1);
+                if (baseName != null && !baseName.isBlank()) {
+                    return baseName + ".jsonl";


let's reuse this config in the extension to make it parametric. This is used in the S3 sink to define the extension of the generated file.

codec: ndjson: extension: jsonl

Zhangxunmt · 2026-03-23T21:38:04Z

+    private final Counter numberOfRecordsSuccessCounter;
+    private final Counter numberOfRecordsFailedCounter;
+
+    @DataPrepperPluginConstructor


Yes agree. I will add the Experimental annotation to it.

Zhangxunmt · 2026-03-24T19:36:47Z

+
+import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;
+
+@Experimental


added @experimental here for the new processor class.

dlvenable

Thanks @Zhangxunmt for this great new processor!

graytaylor0 · 2026-03-24T21:55:10Z

+        for (final String key : keys) {
+            final Object value = other.get(key, Object.class);
+            if (value != null) {
+                put(key, value);


If the data has something like this { "key": null }, do you thing users will want it merged?

Good question. For the S3 enrich use case(adding data from a lookup), skipping nulls makes sense — you don't want a missing/null enrichment field to overwrite a valid value in the target event. But in general, a user might explicitly want to merge a null to clear a field.

Given that the only current consumer is S3EnrichProcessor and the semantics are enrichment, I think we can keep skipping nulls. But I'll add a comment in the code in the next PR to make the decision explicit so in the future this can enhanced for other cases.

Zhangxunmt requested review from KarstenSchnitter, chenqi0805, dinujoh, dlvenable, engechas, graytaylor0, kkondaka, oeyh, san81, sb2k16 and srikanthjg as code owners August 13, 2025 18:38

Zhangxunmt force-pushed the mainline branch from eb3fda3 to fc37ac1 Compare December 16, 2025 07:01

Zhangxunmt changed the title ~~add ml merge processor to merge S3 inputs~~ add S3 Enricher processor to merge ml batch job output with source inputs Dec 16, 2025

Zhangxunmt force-pushed the mainline branch 2 times, most recently from 05e390c to a6cddc2 Compare December 16, 2025 07:52

dlvenable requested changes Dec 18, 2025

View reviewed changes

Zhangxunmt mentioned this pull request Jan 13, 2026

create s3 common module #6404

Merged

4 tasks

Zhangxunmt force-pushed the mainline branch 2 times, most recently from a7e549a to 32673a9 Compare March 10, 2026 04:08

Zhangxunmt changed the title ~~add S3 Enricher processor to merge ml batch job output with source inputs~~ add S3 Enrich processor to merge ml batch job output with source inputs Mar 19, 2026

dlvenable requested changes Mar 20, 2026

View reviewed changes

Zhangxunmt force-pushed the mainline branch from 32673a9 to 56c4ff7 Compare March 23, 2026 07:32

dlvenable reviewed Mar 23, 2026

View reviewed changes

dlvenable requested changes Mar 23, 2026

View reviewed changes

Zhangxunmt added 4 commits March 24, 2026 12:21

add s3 enricher processor

a70d6c3

Signed-off-by: Xun Zhang <xunzh@amazon.com>

rename from S3Enricher to S3Enrich

43f6507

Signed-off-by: Xun Zhang <xunzh@amazon.com>

add unit tests for s3 enrich processor

478e15c

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Add selective merge to Event API and refactor S3EnrichProcessor

b92c00c

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Zhangxunmt force-pushed the mainline branch from 56c4ff7 to 3b31dbf Compare March 24, 2026 19:22

add experimental to the S3 enrich processor

6e9681a

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Zhangxunmt force-pushed the mainline branch from 3b31dbf to 6e9681a Compare March 24, 2026 19:28

Zhangxunmt commented Mar 24, 2026

View reviewed changes

dlvenable approved these changes Mar 24, 2026

View reviewed changes

dlvenable added this to the v2.15 milestone Mar 24, 2026

graytaylor0 approved these changes Mar 24, 2026

View reviewed changes

graytaylor0 merged commit 627a6b7 into opensearch-project:main Mar 24, 2026
101 of 108 checks passed

dlvenable mentioned this pull request Apr 29, 2026

Adds Xun Zhang (Zhangxunmt) to the maintainers. #6793

Merged

4 tasks

	@Override
	public void merge(final Event other) {
	if(!(other instanceof JacksonEvent))
	throw new IllegalArgumentException("Unable to merge the Event. The input Event must be a JacksonEvent.");
	final JacksonEvent otherJacksonEvent = (JacksonEvent) other;
	if(!(otherJacksonEvent.jsonNode instanceof ObjectNode)) {
	throw new IllegalArgumentException("Unable to merge the Event. The input Event must be a JacksonEvent with object data.");
	}
	final ObjectNode otherObjectNode = (ObjectNode) otherJacksonEvent.jsonNode;

	if(!(jsonNode instanceof ObjectNode)) {
	throw new UnsupportedOperationException("Unable to merge the Event. The current Event must have object data.");
	}

	((ObjectNode) jsonNode).setAll(otherObjectNode);
	}


		import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;

		@Experimental

Conversation

Zhangxunmt commented Aug 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

dlvenable left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

dlvenable left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dlvenable left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Zhangxunmt commented Aug 13, 2025 •

edited

Loading