Aggregator processor should evaluate aggregate_when condition before forwarding events to remote peer

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/peerforwarder/RequiresPeerForwarding.java

@@ -4,6 +4,8 @@
  */
 
 package org.opensearch.dataprepper.model.peerforwarder;
+import org.opensearch.dataprepper.model.event.Event;
+import org.opensearch.dataprepper.model.record.Record;
 
 import java.util.Collection;
 
@@ -18,4 +20,15 @@ public interface RequiresPeerForwarding {
      * @return A set of keys
      */
     Collection<String> getIdentificationKeys();
+
+    /**
+     * Returns events that are applicable for peer forwarding.
+     *
+     * @param records collection of input records
+     *
+     * @return a collection of output records
+     */
+    default Collection<Record<Event>> applicableEventsForPeerForwarding(Collection<Record<Event>> records) {
+        return records;
+    }
 }

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/peerforwarder/RequiresPeerForwardingTest.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.dataprepper.model.peerforwarder;
+
+import org.junit.jupiter.api.Test;
+import org.opensearch.dataprepper.model.record.Record;
+import org.opensearch.dataprepper.model.event.Event;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
+import static org.mockito.Mockito.mock;
+
+import java.util.Collection;
+
+class RequiresPeerForwardingTest {
+
+    public class SimpleRequiresPeerForwarding implements RequiresPeerForwarding {
+        @Override
+        public Collection<String> getIdentificationKeys() {
+            return null;
+        }
+    }
+
+    @Test
+    void testRequiresPeerForwardingTest() {
+        Collection<Record<Event>> records = mock(Collection.class);
+        RequiresPeerForwarding requiresPeerForwarding = new SimpleRequiresPeerForwarding();
+        assertThat(requiresPeerForwarding.applicableEventsForPeerForwarding(records), equalTo(records));
+    }
+}
+
+

data-prepper-core/src/main/java/org/opensearch/dataprepper/peerforwarder/PeerForwardingProcessorDecorator.java

@@ -15,6 +15,7 @@
 
 import java.util.Collection;
 import java.util.Collections;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
@@ -79,7 +80,9 @@ private PeerForwardingProcessorDecorator(final PeerForwarder peerForwarder, fina
 
     @Override
     public Collection<Record<Event>> execute(final Collection<Record<Event>> records) {
-        final Collection<Record<Event>> recordsToProcessOnLocalPeer = peerForwarder.forwardRecords(records);
+        final Collection<Record<Event>> recordsToProcess = ((RequiresPeerForwarding)innerProcessor).applicableEventsForPeerForwarding(records);
+        final Collection<Record<Event>> recordsToProcessOnLocalPeer = peerForwarder.forwardRecords(recordsToProcess);
+
         final Collection<Record<Event>> receivedRecordsFromBuffer = peerForwarder.receiveRecords();
 
         final Collection<Record<Event>> recordsToProcessLocally = CollectionUtils.union(

data-prepper-core/src/main/java/org/opensearch/dataprepper/pipeline/ProcessWorker.java

@@ -93,8 +93,8 @@ public void run() {
         }
     }
 
-    private void processAcknowledgements(List<Event> inputEvents, Collection outputRecords) {
-        Set<Event> outputEventsSet = ((ArrayList<Record<Event>>)outputRecords).stream().map(Record::getData).collect(Collectors.toSet());
+    private void processAcknowledgements(List<Event> inputEvents, Collection<Record<Event>> outputRecords) {
+        Set<Event> outputEventsSet = outputRecords.stream().map(Record::getData).collect(Collectors.toSet());
         // For each event in the input events list that is not present in the output events, send positive acknowledgement, if acknowledgements are enabled for it
         inputEvents.forEach(event -> {
             EventHandle eventHandle = event.getEventHandle();

data-prepper-plugins/aggregate-processor/src/main/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateProcessor.java

@@ -20,6 +20,7 @@
 import io.micrometer.core.instrument.Counter;
 import org.opensearch.dataprepper.plugins.hasher.IdentificationKeysHasher;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedList;
 import java.util.List;
@@ -148,6 +149,21 @@ public void shutdown() {
 
     }
 
+    @Override
+    public Collection<Record<Event>> applicableEventsForPeerForwarding(Collection<Record<Event>> records) {
+        if (whenCondition == null) {
+            return records;
+        }
+        final Collection<Record<Event>> recordsOut = new ArrayList<>();
+        for (Record<Event> record: records) {
+            Event event = record.getData();
+            if (expressionEvaluator.evaluateConditional(whenCondition, event)) {
+                recordsOut.add(record);
+            }
+        }
+        return recordsOut;
+    }
+
     @Override
     public Collection<String> getIdentificationKeys() {
         return aggregateProcessorConfig.getIdentificationKeys();

data-prepper-plugins/aggregate-processor/src/test/java/org/opensearch/dataprepper/plugins/processor/aggregate/AggregateProcessorTest.java

@@ -228,6 +228,8 @@ void handleEvent_returning_with_condition_eliminates_one_record() {
             recordsIn.add(new Record<Event>(secondEvent));
             recordsIn.add(new Record<Event>(event));
             Collection<Record<Event>> c = recordsIn;
+            Collection<Record<Event>> applicableRecords = objectUnderTest.applicableEventsForPeerForwarding(recordsIn);
+            assertThat(applicableRecords.size(), equalTo(2));
             final List<Record<Event>> recordsOut = (List<Record<Event>>) objectUnderTest.doExecute(c);
 
             assertThat(recordsOut.size(), equalTo(2));