Repository: numaproj/numaflow (go)

Issue:
Refactor appliers to use option when there are multiple UDFs
              I am not fully in favor of this way of sending appliers to forward. This is because even if we need just `MapUDF`, we still have to set the other `UDFs`. So, I'm thinking, should we use an options approach for all appliers? This way, we just pass the option that we need.

_Originally posted by @yhl25 in https://github.com/numaproj/numaflow/pull/1778#discussion_r1665118156_
            

Pull request:
chore: update map appliers to use options
fixes #1809 

Patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index ea28d9dd..58621af6 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -48,15 +48,6 @@ jobs:
         run: |
           echo /home/runner/go/bin >> $GITHUB_PATH
           echo /usr/local/bin >> $GITHUB_PATH
-      - name: Install pandoc
-        run: |
-          set -eux -o pipefail
-          PANDOC_VERSION=3.2.1
-          PANDOC_ZIP=pandoc-$PANDOC_VERSION-linux-amd64.tar.gz
-          curl -OL https://github.com/jgm/pandoc/releases/download/$PANDOC_VERSION/$PANDOC_ZIP
-          sudo tar xvzf $PANDOC_ZIP --strip-components 1 -C /usr/local
-          rm -f $PANDOC_ZIP
-          echo /usr/local/pandoc-$PANDOC_VERSION/bin >> $GITHUB_PATH
       - name: Get dependencies
         run: go mod download
       - name: Make codegen
diff --git a/hack/library.sh b/hack/library.sh
index 45928775..4b499dde 100644
--- a/hack/library.sh
+++ b/hack/library.sh
@@ -38,28 +38,15 @@ ensure_vendor() {
   go mod vendor
 }	
 
-# TODO: Install it in Makefile
-ensure_pandoc() {
-  if [ "`command -v pandoc`" = "" ]; then
-    warning "Please install pandoc with - brew install pandoc"
-    exit 1
-  fi
-
-  if [ "`pandoc -v | head -1 | awk '{print $2}'`" != "3.2.1" ]; then
-    warning "Please upgrade pandoc to version 3.2.1 with - brew upgrade pandoc"
-    exit 1
-  fi
-}
-
 ensure_node(){
-  if [ "`command -v node`" = "" ]; then
+  if [[ "`command -v node`" = "" ]]; then
       warning "Please install node with - brew install node"
       exit 1
   fi
 }
 
 ensure_yarn(){
-  if [ "`command -v yarn`" = "" ]; then
+  if [[ "`command -v yarn`" = "" ]]; then
       warning "Please install yarn with - brew install yarn"
       exit 1
   fi
diff --git a/hack/update-api-docs.sh b/hack/update-api-docs.sh
index 2db980d9..6c7d7d08 100755
--- a/hack/update-api-docs.sh
+++ b/hack/update-api-docs.sh
@@ -7,7 +7,6 @@ set -o pipefail
 source $(dirname $0)/library.sh
 header "updating api docs"
 
-ensure_pandoc
 ensure_vendor
 make_fake_paths
 
@@ -24,9 +23,58 @@ go run ${FAKE_REPOPATH}/vendor/github.com/ahmetb/gen-crd-api-reference-docs/main
  -out-file "${FAKE_REPOPATH}/docs/APIs.html" \
  -template-dir "${FAKE_REPOPATH}/hack/api-docs-template"
 
-# Setup - https://pandoc.org/installing.html
+install-pandoc() {
+  # pandoc version
+  PANDOC_VERSION=3.2.1
 
-pandoc --from markdown --to gfm ${FAKE_REPOPATH}/docs/APIs.html > ${FAKE_REPOPATH}/docs/APIs.md
+  if [[ "`command -v pandoc`" != "" ]]; then
+    if [[ "`pandoc -v | head -1 | awk '{print $2}'`" != "${PANDOC_VERSION}" ]]; then
+      warning "Existing pandoc version does not match the requirement (${PANDOC_VERSION}), will download a new one..."
+    else
+      PANDOC_BINARY="`command -v pandoc`"
+      return 
+    fi
+  fi
+
+  PD_REL="https://github.com/jgm/pandoc/releases"
+  OS=$(uname_os)
+  ARCH=$(uname_arch)
+
+  echo "OS: $OS  ARCH: $ARCH"
+
+  BINARY_NAME="pandoc-${PANDOC_VERSION}-${OS}-${ARCH}.zip"
+  if [[ "$OS" = "darwin" ]]; then
+    if [[ "$ARCH" = "arm64" ]]; then
+      BINARY_NAME="pandoc-${PANDOC_VERSION}-arm64-macOS.zip"
+    elif [[ "$ARCH" = "amd64" ]]; then
+      BINARY_NAME="pandoc-${PANDOC_VERSION}-x86_64-macOS.zip"
+    fi
+  elif [[ "$OS" = "linux" ]]; then
+    if [[ "$ARCH" = "arm64" ]]; then
+      BINARY_NAME="pandoc-${PANDOC_VERSION}-linux-arm64.tar.gz"
+    elif [[ "$ARCH" = "amd64" ]]; then
+      BINARY_NAME="pandoc-${PANDOC_VERSION}-linux-amd64.tar.gz"
+    fi
+  fi
+  BINARY_URL=$PD_REL/download/${PANDOC_VERSION}/${BINARY_NAME}
+  echo "Downloading $BINARY_URL"
+
+  tmp=$(mktemp -d)
+  trap 'rm -rf ${tmp}' EXIT
+
+  curl -sL -o ${tmp}/${BINARY_NAME} $BINARY_URL
+  if [[ "$BINARY_NAME" =~ .zip$ ]]; then
+    unzip ${tmp}/${BINARY_NAME} -d ${tmp}
+    for a in `ls -d -1 ${tmp}/* | grep pandoc | grep -v zip`; do mv $a/* ${tmp}; rmdir $a; done
+  elif [[ "$BINARY_NAME" =~ .tar.gz$ ]]; then
+    tar xvzf ${tmp}/${BINARY_NAME} --strip-components 1 -C ${tmp}/
+  fi
+  PANDOC_BINARY="${tmp}/bin/pandoc"
+}
+
+install-pandoc
+
+${PANDOC_BINARY} --from markdown --to gfm ${FAKE_REPOPATH}/docs/APIs.html > ${FAKE_REPOPATH}/docs/APIs.md
 
 rm ${FAKE_REPOPATH}/docs/APIs.html
 
diff --git a/pkg/udf/forward/forward.go b/pkg/udf/forward/forward.go
index dc09bb41..5e667e7f 100644
--- a/pkg/udf/forward/forward.go
+++ b/pkg/udf/forward/forward.go
@@ -38,19 +38,11 @@ import (
 	"github.com/numaproj/numaflow/pkg/metrics"
 	"github.com/numaproj/numaflow/pkg/shared/idlehandler"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
-	"github.com/numaproj/numaflow/pkg/udf/forward/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 	"github.com/numaproj/numaflow/pkg/watermark/wmb"
 )
 
-// MapAppliers consolidates the appliers for different map modes together
-type MapAppliers struct {
-	MapUDF       applier.MapApplier
-	MapStreamUDF applier.MapStreamApplier
-	BatchMapUDF  applier.BatchMapApplier
-}
-
 // InterStepDataForward forwards the data from previous step to the current step via inter-step buffer.
 type InterStepDataForward struct {
 	// I have my reasons for overriding the default principle https://github.com/golang/go/issues/22602
@@ -60,10 +52,9 @@ type InterStepDataForward struct {
 	cancelFn            context.CancelFunc
 	fromBufferPartition isb.BufferReader
 	// toBuffers is a map of toVertex name to the toVertex's owned buffers.
-	toBuffers   map[string][]isb.BufferWriter
-	FSD         forwarder.ToWhichStepDecider
-	mapAppliers MapAppliers
-	wmFetcher   fetch.Fetcher
+	toBuffers map[string][]isb.BufferWriter
+	FSD       forwarder.ToWhichStepDecider
+	wmFetcher fetch.Fetcher
 	// wmPublishers stores the vertex to publisher mapping
 	wmPublishers  map[string]publish.Publisher
 	opts          options
@@ -78,7 +69,7 @@ type InterStepDataForward struct {
 }
 
 // NewInterStepDataForward creates an inter-step forwarder.
-func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.BufferReader, toSteps map[string][]isb.BufferWriter, fsd forwarder.ToWhichStepDecider, mapApplier MapAppliers, fetchWatermark fetch.Fetcher, publishWatermark map[string]publish.Publisher, idleManager wmb.IdleManager, opts ...Option) (*InterStepDataForward, error) {
+func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.BufferReader, toSteps map[string][]isb.BufferWriter, fsd forwarder.ToWhichStepDecider, fetchWatermark fetch.Fetcher, publishWatermark map[string]publish.Publisher, idleManager wmb.IdleManager, opts ...Option) (*InterStepDataForward, error) {
 
 	options := DefaultOptions()
 	for _, o := range opts {
@@ -86,6 +77,12 @@ func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.B
 			return nil, err
 		}
 	}
+
+	// we can have all modes empty if no option was enabled, this is an invalid case
+	if !isValidMapMode(options) {
+		return nil, fmt.Errorf("no valid map mode selected")
+	}
+
 	// creating a context here which is managed by the forwarder's lifecycle
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -95,7 +92,6 @@ func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.B
 		fromBufferPartition: fromStep,
 		toBuffers:           toSteps,
 		FSD:                 fsd,
-		mapAppliers:         mapApplier,
 		wmFetcher:           fetchWatermark,
 		wmPublishers:        publishWatermark,
 		// should we do a check here for the values not being null?
@@ -113,7 +109,7 @@ func NewInterStepDataForward(vertexInstance *dfv1.VertexInstance, fromStep isb.B
 	// Add logger from parent ctx to child context.
 	isdf.ctx = logging.WithLogger(ctx, options.logger)
 
-	if isdf.opts.enableMapUdfStream && isdf.opts.readBatchSize != 1 {
+	if (isdf.opts.streamMapUdfApplier != nil) && isdf.opts.readBatchSize != 1 {
 		return nil, fmt.Errorf("batch size is not 1 with map UDF streaming")
 	}
 
@@ -251,8 +247,8 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 
 	var udfResults []isb.ReadWriteMessagePair
 	var writeOffsets map[string][][]isb.Offset
-	// Check if map streaming mode is enabled
-	if isdf.opts.enableMapUdfStream {
+	// Check if map streaming mode is enabled, if the applier is not nil that means we have enabled the required mode
+	if isdf.opts.streamMapUdfApplier != nil {
 		writeOffsets, err = isdf.streamMessage(ctx, dataMessages)
 		if err != nil {
 			isdf.opts.logger.Errorw("failed to streamMessage", zap.Error(err))
@@ -270,7 +266,7 @@ func (isdf *InterStepDataForward) forwardAChunk(ctx context.Context) {
 		// Trigger the UDF processing based on the mode enabled for map
 		// ie Batch Map or unary map
 		// This will be a blocking call until the all the UDF results for the batch are received.
-		if isdf.opts.enableBatchMapUdf {
+		if isdf.opts.batchMapUdfApplier != nil {
 			udfResults = isdf.processBatchMessages(ctx, dataMessages)
 		} else {
 			udfResults = isdf.processConcurrentMap(ctx, dataMessages)
@@ -415,7 +411,7 @@ func (isdf *InterStepDataForward) processConcurrentMap(ctx context.Context, data
 // response for all of them.
 func (isdf *InterStepDataForward) processBatchMessages(ctx context.Context, dataMessages []*isb.ReadMessage) []isb.ReadWriteMessagePair {
 	concurrentUDFProcessingStart := time.Now()
-	udfResults, err := isdf.mapAppliers.BatchMapUDF.ApplyBatchMap(ctx, dataMessages)
+	udfResults, err := isdf.opts.batchMapUdfApplier.ApplyBatchMap(ctx, dataMessages)
 	if err != nil {
 		// In case of an error received while processing the UDF call, we would do not handle partial failures,
 		// and hence would want to replay the whole batch currently.
@@ -468,7 +464,7 @@ func (isdf *InterStepDataForward) streamMessage(ctx context.Context, dataMessage
 		writeMessageCh := make(chan isb.WriteMessage)
 		errs, ctx := errgroup.WithContext(ctx)
 		errs.Go(func() error {
-			return isdf.mapAppliers.MapStreamUDF.ApplyMapStream(ctx, dataMessages[0], writeMessageCh)
+			return isdf.opts.streamMapUdfApplier.ApplyMapStream(ctx, dataMessages[0], writeMessageCh)
 		})
 
 		// Stream the message to the next vertex. First figure out which vertex
@@ -700,7 +696,7 @@ func (isdf *InterStepDataForward) concurrentApplyUDF(ctx context.Context, readMe
 // The UserError retry will be done on the ApplyUDF.
 func (isdf *InterStepDataForward) applyUDF(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.WriteMessage, error) {
 	for {
-		writeMessages, err := isdf.mapAppliers.MapUDF.ApplyMap(ctx, readMessage)
+		writeMessages, err := isdf.opts.unaryMapUdfApplier.ApplyMap(ctx, readMessage)
 		if err != nil {
 			isdf.opts.logger.Errorw("mapUDF.Apply error", zap.Error(err))
 			// TODO: implement retry with backoff etc.
@@ -753,3 +749,10 @@ func errorArrayToMap(errs []error) map[string]int64 {
 	}
 	return result
 }
+
+// check if the options provided are for a valid map mode
+// exactly one of the appliers should not be nil as only one mode can be active at a time, not more not less only 1
+func isValidMapMode(opts *options) bool {
+	// if all the appliers are empty, then it is an invalid scenario
+	return !((opts.batchMapUdfApplier == nil) && (opts.unaryMapUdfApplier == nil) && (opts.streamMapUdfApplier == nil))
+}
diff --git a/pkg/udf/forward/options.go b/pkg/udf/forward/options.go
index acd4d4c8..31374d64 100644
--- a/pkg/udf/forward/options.go
+++ b/pkg/udf/forward/options.go
@@ -17,6 +17,7 @@ limitations under the License.
 package forward
 
 import (
+	"fmt"
 	"time"
 
 	"go.uber.org/zap"
@@ -24,6 +25,7 @@ import (
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/shared/callback"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
+	"github.com/numaproj/numaflow/pkg/udf/forward/applier"
 )
 
 // options for forwarding the message
@@ -36,23 +38,27 @@ type options struct {
 	retryInterval time.Duration
 	// logger is used to pass the logger variable
 	logger *zap.SugaredLogger
-	// enableMapUdfStream indicates whether the message streaming is enabled or not for map UDF processing
-	enableMapUdfStream bool
 	// cbPublisher is the callback publisher for the vertex.
 	cbPublisher *callback.Uploader
-	// enableBatchMapUdf indicates whether the batch map mode is enabled for the map UDF
-	enableBatchMapUdf bool
+	// unaryMapUdfApplier is the UDF applier for a unary map mode
+	unaryMapUdfApplier applier.MapApplier
+	// streamMapUdfApplier is the UDF applier for a server streaming map mode
+	streamMapUdfApplier applier.MapStreamApplier
+	// batchMapUdfApplier is the UDF applier for a batch map mode
+	batchMapUdfApplier applier.BatchMapApplier
 }
 
 type Option func(*options) error
 
 func DefaultOptions() *options {
 	return &options{
-		readBatchSize:      dfv1.DefaultReadBatchSize,
-		udfConcurrency:     dfv1.DefaultReadBatchSize,
-		retryInterval:      time.Millisecond,
-		logger:             logging.NewLogger(),
-		enableMapUdfStream: false,
+		readBatchSize:       dfv1.DefaultReadBatchSize,
+		udfConcurrency:      dfv1.DefaultReadBatchSize,
+		retryInterval:       time.Millisecond,
+		logger:              logging.NewLogger(),
+		unaryMapUdfApplier:  nil,
+		batchMapUdfApplier:  nil,
+		streamMapUdfApplier: nil,
 	}
 }
 
@@ -88,26 +94,58 @@ func WithLogger(l *zap.SugaredLogger) Option {
 	}
 }
 
-// WithUDFStreaming sets streaming for map UDF processing
-func WithUDFStreaming(f bool) Option {
+// WithCallbackUploader sets the callback uploader for the vertex
+func WithCallbackUploader(cp *callback.Uploader) Option {
 	return func(o *options) error {
-		o.enableMapUdfStream = f
+		o.cbPublisher = cp
 		return nil
 	}
 }
 
-// WithCallbackUploader sets the callback uploader for the vertex
-func WithCallbackUploader(cp *callback.Uploader) Option {
+// Options to set the map mode to be used, as all of them are mutually exclusive, at one point of time
+// one of them can be enabled, and others are set to nil
+// NOTE: if multiple of these are passed with valid appliers, we consider that as an invalid scenario
+// There 3 are mutually exclusive by virtue of the fields that they control, and thus only one of them should be used.
+// Currently, cannot think of a valid case where multiple of them might be required to be passed together, if later
+// on we feel the need, we can allow that and check for mutual exclusivity later on.
+//
+// Overwrites are allowed for the same applier, and hence a single With call can be added multiple times with different
+// values of the same applier.
+// exactly one of the appliers should not be nil as only one mode can be active at a time, not more not less only 1
+// There is a case where none of them is set which cannot be allowed
+
+// WithUDFBatchMap enables the batch map for UDF if provided with a non-nil applier
+func WithUDFBatchMap(f applier.BatchMapApplier) Option {
 	return func(o *options) error {
-		o.cbPublisher = cp
-		return nil
+		// only overwrite for the same option is allowed, other two cannot be set
+		if f != nil && o.unaryMapUdfApplier == nil && o.streamMapUdfApplier == nil {
+			o.batchMapUdfApplier = f
+			return nil
+		}
+		return fmt.Errorf("invalid option")
 	}
 }
 
-// WithUDFBatchMap enables the batch map for UDF
-func WithUDFBatchMap(f bool) Option {
+// WithUDFUnaryMap enables the unary map for UDF if provided with a non-nil applier
+func WithUDFUnaryMap(f applier.MapApplier) Option {
 	return func(o *options) error {
-		o.enableBatchMapUdf = f
-		return nil
+		// only overwrite for the same option is allowed, other two cannot be set
+		if f != nil && o.batchMapUdfApplier == nil && o.streamMapUdfApplier == nil {
+			o.unaryMapUdfApplier = f
+			return nil
+		}
+		return fmt.Errorf("invalid option")
+	}
+}
+
+// WithUDFStreamingMap sets streaming for map UDF processing if provided with a non-nil applier
+func WithUDFStreamingMap(f applier.MapStreamApplier) Option {
+	return func(o *options) error {
+		// only overwrite for the same option is allowed, other two cannot be set
+		if f != nil && o.unaryMapUdfApplier == nil && o.batchMapUdfApplier == nil {
+			o.streamMapUdfApplier = f
+			return nil
+		}
+		return fmt.Errorf("invalid option")
 	}
 }
diff --git a/pkg/udf/map_udf.go b/pkg/udf/map_udf.go
index 8b690360..94fcc4aa 100644
--- a/pkg/udf/map_udf.go
+++ b/pkg/udf/map_udf.go
@@ -131,6 +131,7 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 		return fmt.Errorf("unrecognized isbsvc type %q", u.ISBSvcType)
 	}
 
+	opts := []forward.Option{forward.WithLogger(log)}
 	maxMessageSize := sharedutil.LookupEnvIntOr(dfv1.EnvGRPCMaxMessageSize, sdkclient.DefaultGRPCMaxMessageSize)
 	enableMapUdfStream := sharedutil.LookupEnvBoolOr(dfv1.EnvMapStreaming, false)
 	enableBatchMapUdf := sharedutil.LookupEnvBoolOr(dfv1.EnvBatchMap, false)
@@ -164,6 +165,7 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 				log.Warnw("Failed to close gRPC client conn", zap.Error(err))
 			}
 		}()
+		opts = append(opts, forward.WithUDFStreamingMap(mapStreamHandler))
 
 	} else if enableBatchMapUdf {
 		// if Batch Map mode is enabled create the client and handler for that accordingly
@@ -190,7 +192,10 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 				log.Warnw("Failed to close gRPC client conn", zap.Error(err))
 			}
 		}()
+		opts = append(opts, forward.WithUDFBatchMap(batchMapHandler))
 	} else {
+		// Default is to enable unary map mode
+
 		// Wait for server info to be ready
 		serverInfo, err := sdkserverinfo.SDKServerInfo(sdkserverinfo.WithServerInfoFilePath(sdkclient.MapServerInfoFile))
 		if err != nil {
@@ -214,7 +219,7 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 				log.Warnw("Failed to close gRPC client conn", zap.Error(err))
 			}
 		}()
-
+		opts = append(opts, forward.WithUDFUnaryMap(mapHandler))
 	}
 
 	for index, bufferPartition := range fromBuffer {
@@ -264,11 +269,6 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 
 			return result, nil
 		})
-
-		// add the options to check which mode for the map UDF is enabled
-		opts := []forward.Option{forward.WithLogger(log),
-			forward.WithUDFStreaming(enableMapUdfStream),
-			forward.WithUDFBatchMap(enableBatchMapUdf)}
 		if x := u.VertexInstance.Vertex.Spec.Limits; x != nil {
 			if x.ReadBatchSize != nil {
 				opts = append(opts, forward.WithReadBatchSize(int64(*x.ReadBatchSize)))
@@ -288,14 +288,8 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 			opts = append(opts, forward.WithCallbackUploader(cbPublisher))
 		}
 
-		mapAppliers := forward.MapAppliers{
-			MapUDF:       mapHandler,
-			MapStreamUDF: mapStreamHandler,
-			BatchMapUDF:  batchMapHandler,
-		}
-
 		// create a forwarder for each partition
-		df, err := forward.NewInterStepDataForward(u.VertexInstance, readers[index], writers, conditionalForwarder, mapAppliers, fetchWatermark, publishWatermark, idleManager, opts...)
+		df, err := forward.NewInterStepDataForward(u.VertexInstance, readers[index], writers, conditionalForwarder, fetchWatermark, publishWatermark, idleManager, opts...)
 		if err != nil {
 			return err
 		}


Test patch:
diff --git a/pkg/isb/stores/jetstream/writer_test.go b/pkg/isb/stores/jetstream/writer_test.go
index 8408f5b2..4928e211 100644
--- a/pkg/isb/stores/jetstream/writer_test.go
+++ b/pkg/isb/stores/jetstream/writer_test.go
@@ -63,16 +63,29 @@ func TestForwarderJetStreamBuffer(t *testing.T) {
 		name          string
 		batchSize     int64
 		streamEnabled bool
+		unaryEnabled  bool
+		batchEnabled  bool
 	}{
 		{
-			name:          "batch",
+			name:          "unary",
 			batchSize:     10,
 			streamEnabled: false,
+			unaryEnabled:  true,
+			batchEnabled:  false,
 		},
 		{
 			name:          "stream",
 			batchSize:     1,
 			streamEnabled: true,
+			unaryEnabled:  false,
+			batchEnabled:  false,
+		},
+		{
+			name:          "batch_map",
+			batchSize:     10,
+			streamEnabled: false,
+			unaryEnabled:  false,
+			batchEnabled:  true,
 		},
 	}
 	for _, tt := range tests {
@@ -154,12 +167,19 @@ func TestForwarderJetStreamBuffer(t *testing.T) {
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 			idleManager, err := wmb.NewIdleManager(1, len(toSteps))
 			assert.NoError(t, err)
-			appliers := forward.MapAppliers{
-				MapUDF:       myForwardJetStreamTest{},
-				MapStreamUDF: myForwardJetStreamTest{},
-				BatchMapUDF:  myForwardJetStreamTest{},
+
+			opts := []forward.Option{forward.WithReadBatchSize(tt.batchSize)}
+			if tt.batchEnabled {
+				opts = append(opts, forward.WithUDFBatchMap(myForwardJetStreamTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, forward.WithUDFStreamingMap(myForwardJetStreamTest{}))
 			}
-			f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardJetStreamTest{}, appliers, fetchWatermark, publishWatermark, idleManager, forward.WithReadBatchSize(tt.batchSize), forward.WithUDFStreaming(tt.streamEnabled))
+			if tt.unaryEnabled {
+				opts = append(opts, forward.WithUDFUnaryMap(myForwardJetStreamTest{}))
+			}
+
+			f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardJetStreamTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 			assert.NoError(t, err)
 
 			stopped := f.Start()
diff --git a/pkg/isb/stores/redis/read_test.go b/pkg/isb/stores/redis/read_test.go
index e32c3688..562fa827 100644
--- a/pkg/isb/stores/redis/read_test.go
+++ b/pkg/isb/stores/redis/read_test.go
@@ -45,11 +45,6 @@ var (
 		// Addrs: []string{":7000", ":7001", ":7002", ":7003", ":7004", ":7005"}
 		Addrs: []string{":6379"},
 	}
-	readTestAppliers = forward.MapAppliers{
-		MapUDF:       forwardReadWritePerformance{},
-		MapStreamUDF: forwardReadWritePerformance{},
-		BatchMapUDF:  forwardReadWritePerformance{},
-	}
 
 	testStartTime = time.Unix(1636470000, 0).UTC()
 )
@@ -145,7 +140,7 @@ func TestRedisCheckBacklog(t *testing.T) {
 	}
 
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-	f, err := forward.NewInterStepDataForward(vertexInstance, rqr, toSteps, forwardReadWritePerformance{}, readTestAppliers, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager(), forward.WithReadBatchSize(10))
+	f, err := forward.NewInterStepDataForward(vertexInstance, rqr, toSteps, forwardReadWritePerformance{}, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager(), forward.WithReadBatchSize(10), forward.WithUDFUnaryMap(forwardReadWritePerformance{}))
 
 	stopped := f.Start()
 	// validate the length of the toStep stream.
@@ -353,7 +348,7 @@ func (suite *ReadWritePerformance) SetupSuite() {
 	}
 
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-	isdf, _ := forward.NewInterStepDataForward(vertexInstance, rqr, toSteps, forwardReadWritePerformance{}, readTestAppliers, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager())
+	isdf, _ := forward.NewInterStepDataForward(vertexInstance, rqr, toSteps, forwardReadWritePerformance{}, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager(), forward.WithUDFUnaryMap(forwardReadWritePerformance{}))
 
 	suite.ctx = ctx
 	suite.rclient = client
@@ -447,7 +442,7 @@ func (suite *ReadWritePerformance) TestReadWriteLatencyPipelining() {
 	}
 
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-	suite.isdf, _ = forward.NewInterStepDataForward(vertexInstance, suite.rqr, toSteps, forwardReadWritePerformance{}, readTestAppliers, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager())
+	suite.isdf, _ = forward.NewInterStepDataForward(vertexInstance, suite.rqr, toSteps, forwardReadWritePerformance{}, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager(), forward.WithUDFUnaryMap(forwardReadWritePerformance{}))
 
 	suite.False(suite.rqw.IsFull())
 	var writeMessages = make([]isb.Message, 0, suite.count)
diff --git a/pkg/isb/stores/redis/write_test.go b/pkg/isb/stores/redis/write_test.go
index db03d4f9..e3fa3e62 100644
--- a/pkg/isb/stores/redis/write_test.go
+++ b/pkg/isb/stores/redis/write_test.go
@@ -37,14 +37,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/wmb"
 )
 
-var (
-	appliers = forward.MapAppliers{
-		MapUDF:       myForwardRedisTest{},
-		MapStreamUDF: myForwardRedisTest{},
-		BatchMapUDF:  myForwardRedisTest{},
-	}
-)
-
 func TestRedisQWrite_Write(t *testing.T) {
 	client := redisclient.NewRedisClient(redisOptions)
 	ctx, cancel := context.WithTimeout(context.Background(), time.Minute*10)
@@ -422,7 +414,7 @@ func TestNewInterStepDataForwardRedis(t *testing.T) {
 	}
 
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-	f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardRedisTest{}, appliers, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager())
+	f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardRedisTest{}, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager(), forward.WithUDFUnaryMap(myForwardRedisTest{}))
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 
@@ -471,7 +463,7 @@ func TestReadTimeout(t *testing.T) {
 	}
 
 	fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
-	f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardRedisTest{}, appliers, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager())
+	f, err := forward.NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardRedisTest{}, fetchWatermark, publishWatermark, wmb.NewNoOpIdleManager(), forward.WithUDFUnaryMap(myForwardRedisTest{}))
 	assert.NoError(t, err)
 	stopped := f.Start()
 	// Call stop to end the test as we have a blocking read. The forwarder is up and running with no messages written
diff --git a/pkg/udf/forward/forward_test.go b/pkg/udf/forward/forward_test.go
index d615d730..e4ca74e3 100644
--- a/pkg/udf/forward/forward_test.go
+++ b/pkg/udf/forward/forward_test.go
@@ -101,30 +101,136 @@ func (f myForwardTest) ApplyMapStream(ctx context.Context, message *isb.ReadMess
 	return testutils.CopyUDFTestApplyStream(ctx, "", writeMessageCh, message)
 }
 
-func TestNewInterStepDataForward(t *testing.T) {
+type testingOpts struct {
+	name            string
+	batchSize       int64
+	streamEnabled   bool
+	batchMapEnabled bool
+	unaryMapEnabled bool
+}
+
+// Check valid initialization for map modes in forwarder
+func TestValidMapModeInit(t *testing.T) {
 	tests := []struct {
 		name            string
-		batchSize       int64
+		valid           bool
 		streamEnabled   bool
 		batchMapEnabled bool
+		unaryMapEnabled bool
 	}{
+		{
+			name:            "valid_stream",
+			streamEnabled:   true,
+			batchMapEnabled: false,
+			unaryMapEnabled: false,
+			valid:           true,
+		},
+		{
+			name:            "valid_batch",
+			streamEnabled:   false,
+			batchMapEnabled: false,
+			unaryMapEnabled: true,
+			valid:           true,
+		},
+		{
+			name:            "valid_unary",
+			streamEnabled:   false,
+			batchMapEnabled: true,
+			unaryMapEnabled: false,
+			valid:           true,
+		},
+		{
+			name:            "all_disabled",
+			streamEnabled:   false,
+			batchMapEnabled: false,
+			unaryMapEnabled: false,
+			valid:           false,
+		},
+		{
+			name:            "all_enabled",
+			streamEnabled:   true,
+			batchMapEnabled: true,
+			unaryMapEnabled: true,
+			valid:           false,
+		},
+		{
+			name:            "two_enabled",
+			streamEnabled:   true,
+			batchMapEnabled: true,
+			unaryMapEnabled: false,
+			valid:           false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name+"_map_mode", func(t *testing.T) {
+			batchSize := int64(1)
+			fromStep := simplebuffer.NewInMemoryBuffer("from", 5*batchSize, 0)
+			to11 := simplebuffer.NewInMemoryBuffer("to1-1", 2*batchSize, 0)
+			to12 := simplebuffer.NewInMemoryBuffer("to1-2", 2*batchSize, 1)
+			toSteps := map[string][]isb.BufferWriter{
+				"to1": {to11, to12},
+			}
+
+			vertex := &dfv1.Vertex{Spec: dfv1.VertexSpec{
+				PipelineName: "testPipeline",
+				AbstractVertex: dfv1.AbstractVertex{
+					Name: "testVertex",
+				},
+			}}
+			vertexInstance := &dfv1.VertexInstance{
+				Vertex:  vertex,
+				Replica: 0,
+			}
+
+			_, cancel := context.WithTimeout(context.Background(), time.Second*10)
+			defer cancel()
+
+			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
+
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardTest{}))
+			}
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			_, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, fetchWatermark, publishWatermark, idleManager, opts...)
+			if tt.valid {
+				assert.NoError(t, err, "expected no error")
+			} else {
+				assert.Error(t, err, "expected error")
+			}
+		})
+	}
+}
+
+func TestNewInterStepDataForward(t *testing.T) {
+	tests := []testingOpts{
 		{
 			name:            "stream_forward",
 			batchSize:       1,
 			streamEnabled:   true,
 			batchMapEnabled: false,
+			unaryMapEnabled: false,
 		},
 		{
 			name:            "batch_forward",
 			batchSize:       10,
 			streamEnabled:   false,
 			batchMapEnabled: false,
+			unaryMapEnabled: true,
 		},
 		{
 			name:            "batch_map_forward",
 			batchSize:       10,
 			streamEnabled:   false,
 			batchMapEnabled: true,
+			unaryMapEnabled: false,
 		},
 	}
 	for _, tt := range tests {
@@ -156,13 +262,19 @@ func TestNewInterStepDataForward(t *testing.T) {
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 
-			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
-			appliers := MapAppliers{
-				MapUDF:       myForwardTest{},
-				MapStreamUDF: myForwardTest{},
-				BatchMapUDF:  myForwardTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardTest{}))
 			}
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+
+			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -265,12 +377,17 @@ func TestNewInterStepDataForward(t *testing.T) {
 			}()
 
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
-			appliers := MapAppliers{
-				MapUDF:       &myForwardToAllTest{},
-				MapStreamUDF: &myForwardToAllTest{},
-				BatchMapUDF:  &myForwardToAllTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(&myForwardToAllTest{}))
 			}
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(&myForwardToAllTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(&myForwardToAllTest{}))
+			}
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &myForwardToAllTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -437,13 +554,18 @@ func TestNewInterStepDataForward(t *testing.T) {
 
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myForwardDropTest{},
-				MapStreamUDF: myForwardDropTest{},
-				BatchMapUDF:  myForwardDropTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardDropTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardDropTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardDropTest{}))
 			}
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardDropTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -622,13 +744,18 @@ func TestNewInterStepDataForward(t *testing.T) {
 
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myForwardTest{},
-				MapStreamUDF: myForwardTest{},
-				BatchMapUDF:  myForwardTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardTest{}))
 			}
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.False(t, to11.IsFull())
@@ -764,13 +891,18 @@ func TestNewInterStepDataForward(t *testing.T) {
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myForwardApplyUDFErrTest{},
-				MapStreamUDF: myForwardApplyUDFErrTest{},
-				BatchMapUDF:  myForwardApplyUDFErrTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardApplyUDFErrTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardApplyUDFErrTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardApplyUDFErrTest{}))
 			}
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyUDFErrTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyUDFErrTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.False(t, to1.IsFull())
@@ -820,13 +952,18 @@ func TestNewInterStepDataForward(t *testing.T) {
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myForwardApplyWhereToErrTest{},
-				MapStreamUDF: myForwardApplyWhereToErrTest{},
-				BatchMapUDF:  myForwardApplyWhereToErrTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardApplyWhereToErrTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardApplyWhereToErrTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardApplyWhereToErrTest{}))
 			}
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardApplyWhereToErrTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.True(t, to1.IsEmpty())
@@ -869,13 +1006,18 @@ func TestNewInterStepDataForward(t *testing.T) {
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myForwardInternalErrTest{},
-				MapStreamUDF: myForwardInternalErrTest{},
-				BatchMapUDF:  myForwardInternalErrTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardInternalErrTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardInternalErrTest{}))
+			}
+			if tt.unaryMapEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardInternalErrTest{}))
 			}
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled), WithUDFBatchMap(tt.batchMapEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardInternalErrTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 
 			assert.NoError(t, err)
 			assert.False(t, to1.IsFull())
@@ -1019,13 +1161,8 @@ func TestNewInterStepDataForwardIdleWatermark(t *testing.T) {
 
 	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-	appliers := MapAppliers{
-		MapUDF:       myForwardTest{},
-		MapStreamUDF: myForwardTest{},
-		BatchMapUDF:  myForwardTest{},
-	}
-
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(2))
+	opts := []Option{WithReadBatchSize(2), WithUDFUnaryMap(myForwardTest{})}
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -1204,14 +1341,9 @@ func TestNewInterStepDataForwardIdleWatermark_Reset(t *testing.T) {
 	}()
 
 	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
+	opts := []Option{WithReadBatchSize(2), WithUDFUnaryMap(myForwardTest{})}
 
-	appliers := MapAppliers{
-		MapUDF:       myForwardTest{},
-		MapStreamUDF: myForwardTest{},
-		BatchMapUDF:  myForwardTest{},
-	}
-
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(2))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -1450,13 +1582,9 @@ func TestInterStepDataForwardSinglePartition(t *testing.T) {
 	// create a forwarder
 	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-	appliers := MapAppliers{
-		MapUDF:       mySourceForwardTest{},
-		MapStreamUDF: mySourceForwardTest{},
-		BatchMapUDF:  mySourceForwardTest{},
-	}
+	opts := []Option{WithReadBatchSize(5), WithUDFUnaryMap(mySourceForwardTest{})}
 
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(5))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, mySourceForwardTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 	assert.NoError(t, err)
 	assert.False(t, to1.IsFull())
 	assert.True(t, to1.IsEmpty())
@@ -1522,13 +1650,9 @@ func TestInterStepDataForwardMultiplePartition(t *testing.T) {
 
 	idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-	appliers := MapAppliers{
-		MapUDF:       mySourceForwardTest{},
-		MapStreamUDF: mySourceForwardTest{},
-		BatchMapUDF:  mySourceForwardTest{},
-	}
+	opts := []Option{WithReadBatchSize(5), WithUDFUnaryMap(mySourceForwardTest{})}
 
-	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(5))
+	f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, &mySourceForwardTestRoundRobin{}, fetchWatermark, publishWatermark, idleManager, opts...)
 	assert.NoError(t, err)
 	assert.False(t, to11.IsFull())
 	assert.False(t, to12.IsFull())
@@ -1596,6 +1720,8 @@ func TestWriteToBuffer(t *testing.T) {
 		batchSize     int64
 		strategy      dfv1.BufferFullWritingStrategy
 		streamEnabled bool
+		unaryEnabled  bool
+		batchEnabled  bool
 		throwError    bool
 	}{
 		{
@@ -1603,6 +1729,8 @@ func TestWriteToBuffer(t *testing.T) {
 			batchSize:     10,
 			strategy:      dfv1.DiscardLatest,
 			streamEnabled: false,
+			batchEnabled:  false,
+			unaryEnabled:  true,
 			// should not throw any error as we drop messages and finish writing before context is cancelled
 			throwError: false,
 		},
@@ -1611,6 +1739,8 @@ func TestWriteToBuffer(t *testing.T) {
 			batchSize:     10,
 			strategy:      dfv1.RetryUntilSuccess,
 			streamEnabled: false,
+			batchEnabled:  false,
+			unaryEnabled:  true,
 			// should throw context closed error as we keep retrying writing until context is cancelled
 			throwError: true,
 		},
@@ -1619,6 +1749,8 @@ func TestWriteToBuffer(t *testing.T) {
 			batchSize:     1,
 			strategy:      dfv1.DiscardLatest,
 			streamEnabled: true,
+			batchEnabled:  false,
+			unaryEnabled:  false,
 			// should not throw any error as we drop messages and finish writing before context is cancelled
 			throwError: false,
 		},
@@ -1627,6 +1759,28 @@ func TestWriteToBuffer(t *testing.T) {
 			batchSize:     1,
 			strategy:      dfv1.RetryUntilSuccess,
 			streamEnabled: true,
+			batchEnabled:  false,
+			unaryEnabled:  false,
+			// should throw context closed error as we keep retrying writing until context is cancelled
+			throwError: true,
+		},
+		{
+			name:          "test-discard-latest",
+			batchSize:     10,
+			strategy:      dfv1.DiscardLatest,
+			streamEnabled: false,
+			batchEnabled:  true,
+			unaryEnabled:  false,
+			// should not throw any error as we drop messages and finish writing before context is cancelled
+			throwError: false,
+		},
+		{
+			name:          "test-retry-until-success",
+			batchSize:     10,
+			strategy:      dfv1.RetryUntilSuccess,
+			streamEnabled: false,
+			batchEnabled:  true,
+			unaryEnabled:  false,
 			// should throw context closed error as we keep retrying writing until context is cancelled
 			throwError: true,
 		},
@@ -1655,12 +1809,18 @@ func TestWriteToBuffer(t *testing.T) {
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myForwardTest{},
-				MapStreamUDF: myForwardTest{},
-				BatchMapUDF:  myForwardTest{},
+			opts := []Option{WithReadBatchSize(value.batchSize)}
+			if value.batchEnabled {
+				opts = append(opts, WithUDFBatchMap(myForwardApplyWhereToErrTest{}))
 			}
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(value.batchSize), WithUDFStreaming(value.streamEnabled))
+			if value.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myForwardApplyWhereToErrTest{}))
+			}
+			if value.unaryEnabled {
+				opts = append(opts, WithUDFUnaryMap(myForwardApplyWhereToErrTest{}))
+			}
+
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myForwardTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 			assert.NoError(t, err)
 			assert.False(t, buffer.IsFull())
 			assert.True(t, buffer.IsEmpty())
diff --git a/pkg/udf/forward/shutdown_test.go b/pkg/udf/forward/shutdown_test.go
index 33b7d9f8..42f547fb 100644
--- a/pkg/udf/forward/shutdown_test.go
+++ b/pkg/udf/forward/shutdown_test.go
@@ -5,7 +5,7 @@ Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at
 
-    http://www.apache.org/licenses/LICENSE-2.0
+	http://www.apache.org/licenses/LICENSE-2.0
 
 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
@@ -13,7 +13,6 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
-
 package forward
 
 import (
@@ -53,19 +52,32 @@ func (s myShutdownTest) ApplyMapStream(ctx context.Context, message *isb.ReadMes
 
 func TestInterStepDataForward(t *testing.T) {
 	tests := []struct {
-		name          string
-		batchSize     int64
-		streamEnabled bool
+		name            string
+		batchSize       int64
+		streamEnabled   bool
+		batchMapEnabled bool
+		unaryEnabled    bool
 	}{
 		{
-			name:          "stream_forward",
-			batchSize:     1,
-			streamEnabled: true,
+			name:            "stream_forward",
+			batchSize:       1,
+			streamEnabled:   true,
+			batchMapEnabled: false,
+			unaryEnabled:    false,
+		},
+		{
+			name:            "batch_forward",
+			batchSize:       5,
+			streamEnabled:   false,
+			batchMapEnabled: false,
+			unaryEnabled:    true,
 		},
 		{
-			name:          "batch_forward",
-			batchSize:     5,
-			streamEnabled: false,
+			name:            "batch_map_forward",
+			batchSize:       5,
+			streamEnabled:   false,
+			batchMapEnabled: true,
+			unaryEnabled:    false,
 		},
 	}
 	for _, tt := range tests {
@@ -96,14 +108,19 @@ func TestInterStepDataForward(t *testing.T) {
 
 			fetchWatermark, publishWatermark := generic.BuildNoOpWatermarkProgressorsFromBufferMap(toSteps)
 
-			appliers := MapAppliers{
-				MapUDF:       myShutdownTest{},
-				MapStreamUDF: myShutdownTest{},
-				BatchMapUDF:  myShutdownTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myShutdownTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myShutdownTest{}))
+			}
+			if tt.unaryEnabled {
+				opts = append(opts, WithUDFUnaryMap(myShutdownTest{}))
 			}
 
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data but buffer is not full even though we are not reading
@@ -144,13 +161,18 @@ func TestInterStepDataForward(t *testing.T) {
 
 			idleManager, _ := wmb.NewIdleManager(1, len(toSteps))
 
-			appliers := MapAppliers{
-				MapUDF:       myShutdownTest{},
-				MapStreamUDF: myShutdownTest{},
-				BatchMapUDF:  myShutdownTest{},
+			opts := []Option{WithReadBatchSize(batchSize)}
+			if tt.batchMapEnabled {
+				opts = append(opts, WithUDFBatchMap(myShutdownTest{}))
+			}
+			if tt.streamEnabled {
+				opts = append(opts, WithUDFStreamingMap(myShutdownTest{}))
+			}
+			if tt.unaryEnabled {
+				opts = append(opts, WithUDFUnaryMap(myShutdownTest{}))
 			}
 
-			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, appliers, fetchWatermark, publishWatermark, idleManager, WithReadBatchSize(batchSize), WithUDFStreaming(tt.streamEnabled))
+			f, err := NewInterStepDataForward(vertexInstance, fromStep, toSteps, myShutdownTest{}, fetchWatermark, publishWatermark, idleManager, opts...)
 			assert.NoError(t, err)
 			stopped := f.Start()
 			// write some data such that the fromBufferPartition can be empty, that is toBuffer gets full