feat: Unify MapStream and Unary Map Operations Using a Shared gRPC Protocol

src/main/java/io/numaproj/numaflow/batchmapper/Service.java

@@ -149,7 +149,8 @@ private void buildAndStreamResponse(
         // Send an EOT message to indicate the end of the transmission for the batch.
         MapOuterClass.MapResponse eotResponse = MapOuterClass.MapResponse
                 .newBuilder()
-                .setStatus(MapOuterClass.Status.newBuilder().setEot(true).build()).build();
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true).build())
+                .build();
         responseObserver.onNext(eotResponse);
         responseObserver.onCompleted();
     }

src/main/java/io/numaproj/numaflow/mapstreamer/OutputObserverImpl.java

@@ -2,7 +2,7 @@
 
 import com.google.protobuf.ByteString;
 import io.grpc.stub.StreamObserver;
-import io.numaproj.numaflow.mapstream.v1.Mapstream;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
 import lombok.AllArgsConstructor;
 
 import java.util.ArrayList;
@@ -14,17 +14,17 @@
  */
 @AllArgsConstructor
 class OutputObserverImpl implements OutputObserver {
-    StreamObserver<Mapstream.MapStreamResponse> responseObserver;
+    StreamObserver<MapOuterClass.MapResponse> responseObserver;
 
     @Override
     public void send(Message message) {
-        Mapstream.MapStreamResponse response = buildResponse(message);
+        MapOuterClass.MapResponse response = buildResponse(message);
         responseObserver.onNext(response);
     }
 
-    private Mapstream.MapStreamResponse buildResponse(Message message) {
-        return Mapstream.MapStreamResponse.newBuilder()
-                .setResult(Mapstream.MapStreamResponse.Result.newBuilder()
+    private MapOuterClass.MapResponse buildResponse(Message message) {
+        return MapOuterClass.MapResponse.newBuilder()
+                .addResults(MapOuterClass.MapResponse.Result.newBuilder()
                         .setValue(
                                 message.getValue() == null ? ByteString.EMPTY : ByteString.copyFrom(
                                         message.getValue()))
@@ -33,6 +33,5 @@ private Mapstream.MapStreamResponse buildResponse(Message message) {
                         .addAllTags(message.getTags()
                                 == null ? new ArrayList<>() : List.of(message.getTags()))
                         .build()).build();
-
     }
 }

src/main/java/io/numaproj/numaflow/mapstreamer/Service.java

@@ -1,64 +1,105 @@
 package io.numaproj.numaflow.mapstreamer;
 
-import com.google.protobuf.Empty;
+import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
-import io.numaproj.numaflow.mapstream.v1.MapStreamGrpc;
-import io.numaproj.numaflow.mapstream.v1.Mapstream;
+import io.numaproj.numaflow.map.v1.MapGrpc;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
 import lombok.AllArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 import java.time.Instant;
 
-import static io.numaproj.numaflow.mapstream.v1.MapStreamGrpc.getMapStreamFnMethod;
-
+import static io.numaproj.numaflow.map.v1.MapGrpc.getMapFnMethod;
 
 @Slf4j
 @AllArgsConstructor
-class Service extends MapStreamGrpc.MapStreamImplBase {
+class Service extends MapGrpc.MapImplBase {
 
     private final MapStreamer mapStreamer;
 
-    /**
-     * Applies a map stream function to each request.
-     */
     @Override
-    public void mapStreamFn(
-            Mapstream.MapStreamRequest request,
-            StreamObserver<Mapstream.MapStreamResponse> responseObserver) {
+    public StreamObserver<MapOuterClass.MapRequest> mapFn(StreamObserver<MapOuterClass.MapResponse> responseObserver) {
 
         if (this.mapStreamer == null) {
-            io.grpc.stub.ServerCalls.asyncUnimplementedUnaryCall(
-                    getMapStreamFnMethod(),
+            return io.grpc.stub.ServerCalls.asyncUnimplementedStreamingCall(
+                    getMapFnMethod(),
                     responseObserver);
-            return;
         }
 
-        HandlerDatum handlerDatum = new HandlerDatum(
-                request.getValue().toByteArray(),
-                Instant.ofEpochSecond(
-                        request.getWatermark().getSeconds(),
-                        request.getWatermark().getNanos()),
-                Instant.ofEpochSecond(
-                        request.getEventTime().getSeconds(),
-                        request.getEventTime().getNanos()),
-                request.getHeadersMap()
-        );
+        return new StreamObserver<>() {
+            private boolean handshakeDone = false;
 
-        // process Datum
-        this.mapStreamer.processMessage(request
-                .getKeysList()
-                .toArray(new String[0]), handlerDatum, new OutputObserverImpl(responseObserver));
+            @Override
+            public void onNext(MapOuterClass.MapRequest request) {
+                // make sure the handshake is done before processing the messages
+                if (!handshakeDone) {
+                    if (!request.hasHandshake() || !request.getHandshake().getSot()) {
+                        responseObserver.onError(Status.INVALID_ARGUMENT
+                                .withDescription("Handshake request not received")
+                                .asException());
+                        return;
+                    }
+                    responseObserver.onNext(MapOuterClass.MapResponse.newBuilder()
+                            .setHandshake(request.getHandshake())
+                            .build());
+                    handshakeDone = true;
+                    return;
+                }
 
-        responseObserver.onCompleted();
-    }
+                try {
+                    // process the message
+                    mapStreamer.processMessage(
+                            request
+                                    .getRequest()
+                                    .getKeysList()
+                                    .toArray(new String[0]),
+                            constructHandlerDatum(request),
+                            new OutputObserverImpl(responseObserver));
+                } catch (Exception e) {
+                    log.error("Error processing message", e);
+                    responseObserver.onError(Status.UNKNOWN
+                            .withDescription(e.getMessage())
+                            .asException());
+                    return;
+                }
 
-    /**
-     * IsReady is the heartbeat endpoint for gRPC.
-     */
-    @Override
-    public void isReady(Empty request, StreamObserver<Mapstream.ReadyResponse> responseObserver) {
-        responseObserver.onNext(Mapstream.ReadyResponse.newBuilder().setReady(true).build());
-        responseObserver.onCompleted();
+                // Send an EOT message to indicate the end of the transmission for the batch.
+                MapOuterClass.MapResponse eotResponse = MapOuterClass.MapResponse
+                        .newBuilder()
+                        .setStatus(MapOuterClass.TransmissionStatus
+                                .newBuilder()
+                                .setEot(true)
+                                .build()).build();
+                responseObserver.onNext(eotResponse);
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+                log.error("Error Encountered in mapStream Stream", throwable);
+                var status = Status.UNKNOWN
+                        .withDescription(throwable.getMessage())
+                        .withCause(throwable);
+                responseObserver.onError(status.asException());
+            }
+
+            @Override
+            public void onCompleted() {
+                responseObserver.onCompleted();
+            }
+        };
     }
 
+    // Construct a HandlerDatum from a MapRequest
+    private HandlerDatum constructHandlerDatum(MapOuterClass.MapRequest d) {
+        return new HandlerDatum(
+                d.getRequest().getValue().toByteArray(),
+                Instant.ofEpochSecond(
+                        d.getRequest().getWatermark().getSeconds(),
+                        d.getRequest().getWatermark().getNanos()),
+                Instant.ofEpochSecond(
+                        d.getRequest().getEventTime().getSeconds(),
+                        d.getRequest().getEventTime().getNanos()),
+                d.getRequest().getHeadersMap()
+        );
+    }
 }

src/main/java/io/numaproj/numaflow/sinker/Service.java

@@ -78,6 +78,16 @@ public void onNext(SinkOuterClass.SinkRequest request) {
                             responseObserver.onNext(sinkResponse);
                         });
 
+                        // send eot response to indicate end of transmission for the batch
+                        SinkOuterClass.SinkResponse eotResponse = SinkOuterClass.SinkResponse
+                                .newBuilder()
+                                .setStatus(SinkOuterClass.TransmissionStatus
+                                        .newBuilder()
+                                        .setEot(true)
+                                        .build())
+                                .build();
+                        responseObserver.onNext(eotResponse);
+
                         // reset the startOfStream flag, since the stream has ended
                         // so that the next request will be treated as the start of the stream
                         startOfStream = true;

src/main/java/io/numaproj/numaflow/sinker/SinkerTestKit.java

@@ -181,7 +181,7 @@
             }
             // send end of transmission message
             requestObserver.onNext(SinkOuterClass.SinkRequest.newBuilder().setStatus(
-                    SinkOuterClass.SinkRequest.Status.newBuilder().setEot(true)).build());
+                    SinkOuterClass.TransmissionStatus.newBuilder().setEot(true)).build());
 
             requestObserver.onCompleted();
 
@@ -197,6 +197,9 @@
                 if (result.getHandshake().getSot()) {
                     continue;
                 }
+                if (result.hasStatus() && result.getStatus().getEot()) {
+                    continue;
+                }
                 if (result.getResult().getStatus() == SinkOuterClass.Status.SUCCESS) {
                     responseListBuilder.addResponse(Response.responseOK(result
                             .getResult()

src/main/proto/map/v1/map.proto

@@ -30,7 +30,7 @@ message MapRequest {
   // This ID is used to uniquely identify a map request
   string id = 2;
   optional Handshake handshake = 3;
-  optional Status status = 4;
+  optional TransmissionStatus status = 4;
 }
 
 /*
@@ -44,7 +44,7 @@ message Handshake {
 /*
  * Status message to indicate the status of the message.
  */
-message Status {
+message TransmissionStatus {
   bool eot = 1;
 }
 
@@ -61,7 +61,7 @@ message MapResponse {
   // This ID is used to refer the responses to the request it corresponds to.
   string id = 2;
   optional Handshake handshake = 3;
-  optional Status status = 4;
+  optional TransmissionStatus status = 4;
 }
 
 /**

src/main/proto/sink/v1/sink.proto

@@ -27,14 +27,11 @@ message SinkRequest {
     string id = 5;
     map<string, string> headers = 6;
   }
-  message Status {
-    bool eot = 1;
-  }
   // Required field indicating the request.
   Request request = 1;
   // Required field indicating the status of the request.
   // If eot is set to true, it indicates the end of transmission.
-  Status status = 2;
+  TransmissionStatus status = 2;
   // optional field indicating the handshake message.
   optional Handshake handshake = 3;
 }
@@ -54,6 +51,13 @@ message ReadyResponse {
   bool ready = 1;
 }
 
+/**
+ * TransmissionStatus is the status of the transmission.
+ */
+message TransmissionStatus {
+  bool eot = 1;
+}
+
 /*
   * Status is the status of the response.
  */
@@ -77,4 +81,5 @@ message SinkResponse {
   }
   Result result = 1;
   optional Handshake handshake = 2;
+  optional TransmissionStatus status = 3;
 }

src/test/java/io/numaproj/numaflow/batchmapper/ServerErrTest.java

@@ -127,7 +127,7 @@ public void testErrorFromUDF() {
         inputStreamObserver.onNext(request);
         inputStreamObserver.onNext(MapOuterClass.MapRequest
                 .newBuilder()
-                .setStatus(MapOuterClass.Status.newBuilder().setEot(true))
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true))
                 .build());
         inputStreamObserver.onCompleted();
         try {

src/test/java/io/numaproj/numaflow/batchmapper/ServerTest.java

@@ -90,7 +90,7 @@ public void testBatchMapHappyPath() {
 
         inputStreamObserver.onNext(MapOuterClass.MapRequest
                 .newBuilder()
-                .setStatus(MapOuterClass.Status.newBuilder().setEot(true))
+                .setStatus(MapOuterClass.TransmissionStatus.newBuilder().setEot(true))
                 .build());
         inputStreamObserver.onCompleted();
 

src/test/java/io/numaproj/numaflow/mapstreamer/MapStreamOutputStreamObserver.java

@@ -0,0 +1,41 @@
+package io.numaproj.numaflow.mapstreamer;
+
+import io.grpc.stub.StreamObserver;
+import io.numaproj.numaflow.map.v1.MapOuterClass;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+public class MapStreamOutputStreamObserver implements StreamObserver<MapOuterClass.MapResponse> {
+    List<MapOuterClass.MapResponse> mapResponses = new ArrayList<>();
+    CompletableFuture<Void> done = new CompletableFuture<>();
+    Integer responseCount;
+
+    public MapStreamOutputStreamObserver(Integer responseCount) {
+        this.responseCount = responseCount;
+    }
+
+    @Override
+    public void onNext(MapOuterClass.MapResponse mapResponse) {
+        System.out.println("Received response: " + mapResponse);
+        mapResponses.add(mapResponse);
+        if (mapResponses.size() == responseCount) {
+            done.complete(null);
+        }
+    }
+
+    @Override
+    public void onError(Throwable throwable) {
+        done.completeExceptionally(throwable);
+    }
+
+    @Override
+    public void onCompleted() {
+        done.complete(null);
+    }
+
+    public List<MapOuterClass.MapResponse> getMapResponses() {
+        return mapResponses;
+    }
+}