Repository: numaproj/numaflow
Pull request #658 (2023-04-03 03:47:43)

feat: support for multi keys

ashwinidulams: shuffle (when we get to that) should be immune to the ordering of keys.  conditional forwarding as well. 

Patch:
diff --git a/go.mod b/go.mod
index fc4a2f16..1513f562 100644
--- a/go.mod
+++ b/go.mod
@@ -27,7 +27,7 @@ require (
 	github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe
 	github.com/nats-io/nats-server/v2 v2.7.5-0.20220415000625-a6b62f61a703
 	github.com/nats-io/nats.go v1.24.0
-	github.com/numaproj/numaflow-go v0.3.4
+	github.com/numaproj/numaflow-go v0.4.0
 	github.com/prometheus/client_golang v1.12.1
 	github.com/prometheus/common v0.32.1
 	github.com/soheilhy/cmux v0.1.5
diff --git a/go.sum b/go.sum
index dac231e3..fb58b878 100644
--- a/go.sum
+++ b/go.sum
@@ -697,8 +697,8 @@ github.com/nats-io/nkeys v0.3.0/go.mod h1:gvUNGjVcM2IPr5rCsRsC6Wb3Hr2CQAm08dsxtV
 github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
 github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
-github.com/numaproj/numaflow-go v0.3.4 h1:YMyq+dJLW5l68NWTNCUpLdox6tMXqmwcAebTONB6cuE=
-github.com/numaproj/numaflow-go v0.3.4/go.mod h1:U6a0JproBL8m/HT6Fh573WwMHG5kkX30ZB9i069uGWc=
+github.com/numaproj/numaflow-go v0.4.0 h1:7TaQFLxvUCm6F685rsN76JCs/EhOr5zJQRiJi45xMKM=
+github.com/numaproj/numaflow-go v0.4.0/go.mod h1:U6a0JproBL8m/HT6Fh573WwMHG5kkX30ZB9i069uGWc=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
 github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
diff --git a/pkg/forward/forward.go b/pkg/forward/forward.go
index 07e42599..ae358ef6 100644
--- a/pkg/forward/forward.go
+++ b/pkg/forward/forward.go
@@ -595,7 +595,7 @@ func (isdf *InterStepDataForward) applyUDF(ctx context.Context, readMessage *isb
 // whereToStep executes the WhereTo interfaces and then updates the to step's writeToBuffers buffer.
 func (isdf *InterStepDataForward) whereToStep(writeMessage *isb.Message, messageToStep map[string][]isb.Message, readMessage *isb.ReadMessage) error {
 	// call WhereTo and drop it on errors
-	to, err := isdf.FSD.WhereTo(writeMessage.Key)
+	to, err := isdf.FSD.WhereTo(writeMessage.Keys)
 
 	if err != nil {
 		isdf.opts.logger.Errorw("failed in whereToStep", zap.Error(isb.MessageWriteErr{Name: isdf.fromBuffer.GetName(), Header: readMessage.Header, Body: readMessage.Body, Message: fmt.Sprintf("WhereTo failed, %s", err)}))
diff --git a/pkg/forward/interfaces.go b/pkg/forward/interfaces.go
index 46251c1a..8a14170a 100644
--- a/pkg/forward/interfaces.go
+++ b/pkg/forward/interfaces.go
@@ -21,8 +21,8 @@ import (
 )
 
 var (
-	All  = GoWhere(func(string) ([]string, error) { return []string{dfv1.MessageKeyAll}, nil })
-	Drop = GoWhere(func(string) ([]string, error) { return []string{dfv1.MessageKeyDrop}, nil })
+	All  = GoWhere(func([]string) ([]string, error) { return []string{dfv1.MessageKeyAll}, nil })
+	Drop = GoWhere(func([]string) ([]string, error) { return []string{dfv1.MessageKeyDrop}, nil })
 )
 
 // ToWhichStepDecider decides which step to forward after applying the WhereTo function.
@@ -30,14 +30,14 @@ type ToWhichStepDecider interface {
 	// WhereTo decides where to forward the result to based on the name of the step it returns.
 	// It supports 2 addition keywords which need not be a step name. They are "ALL" and "DROP"
 	// where former means, forward to all the neighbouring steps and latter means do not forward anywhere.
-	WhereTo(string) ([]string, error)
+	WhereTo([]string) ([]string, error)
 }
 
 // GoWhere is the step decider on where it needs to go
-type GoWhere func(string) ([]string, error)
+type GoWhere func([]string) ([]string, error)
 
 // WhereTo decides where the data goes to.
-func (gw GoWhere) WhereTo(b string) ([]string, error) {
+func (gw GoWhere) WhereTo(b []string) ([]string, error) {
 	return gw(b)
 }
 
diff --git a/pkg/isb/message.go b/pkg/isb/message.go
index 9e5f54ec..e853dbd8 100644
--- a/pkg/isb/message.go
+++ b/pkg/isb/message.go
@@ -59,8 +59,9 @@ type Header struct {
 	Kind MessageKind
 	// ID is used for exactly-once-semantics. ID is usually populated from the offset, if offset is available.
 	ID string
-	// Key is (key,value) in the map-reduce paradigm which will be used for conditional forwarding.
-	Key string
+	// Keys is (key,value) in the map-reduce paradigm will be used for reduce operation, last key in the list
+	// will be used for conditional forwarding
+	Keys []string
 }
 
 // Body is the body of the message
diff --git a/pkg/isb/serde.go b/pkg/isb/serde.go
index 2255b9a2..7c2d716d 100644
--- a/pkg/isb/serde.go
+++ b/pkg/isb/serde.go
@@ -74,7 +74,7 @@ func (h Header) MarshalBinary() (data []byte, err error) {
 		MLen:    int16(len(msgInfo)),
 		MsgKind: h.Kind,
 		IDLen:   int16(len(h.ID)),
-		KeyLen:  int16(len(h.Key)),
+		KeyLen:  int16(len(h.Keys)),
 	}
 	if err = binary.Write(buf, binary.LittleEndian, preamble); err != nil {
 		return nil, err
@@ -88,9 +88,17 @@ func (h Header) MarshalBinary() (data []byte, err error) {
 	if err = binary.Write(buf, binary.LittleEndian, []byte(h.ID)); err != nil {
 		return nil, err
 	}
-	if err = binary.Write(buf, binary.LittleEndian, []byte(h.Key)); err != nil {
+	if err = binary.Write(buf, binary.LittleEndian, int16(len(h.Keys))); err != nil {
 		return nil, err
 	}
+	for i := 0; i < len(h.Keys); i++ {
+		if err = binary.Write(buf, binary.LittleEndian, int16(len(h.Keys[i]))); err != nil {
+			return nil, err
+		}
+		if err = binary.Write(buf, binary.LittleEndian, []byte(h.Keys[i])); err != nil {
+			return nil, err
+		}
+	}
 	return buf.Bytes(), nil
 }
 
@@ -116,14 +124,25 @@ func (h *Header) UnmarshalBinary(data []byte) (err error) {
 	if err = binary.Read(r, binary.LittleEndian, id); err != nil {
 		return err
 	}
-	var key = make([]byte, preamble.KeyLen)
-	if err = binary.Read(r, binary.LittleEndian, key); err != nil {
+	var keyLen int16
+	if err = binary.Read(r, binary.LittleEndian, &keyLen); err != nil {
 		return err
 	}
+	h.Keys = make([]string, 0)
+	for i := int16(0); i < keyLen; i++ {
+		var kl int16
+		if err = binary.Read(r, binary.LittleEndian, &kl); err != nil {
+			return err
+		}
+		var k = make([]byte, kl)
+		if err = binary.Read(r, binary.LittleEndian, k); err != nil {
+			return err
+		}
+		h.Keys = append(h.Keys, string(k))
+	}
 	h.MessageInfo = *msgInfo
 	h.Kind = preamble.MsgKind
 	h.ID = string(id)
-	h.Key = string(key)
 	return err
 }
 
diff --git a/pkg/reduce/data_forward.go b/pkg/reduce/data_forward.go
index 943aa6e3..c9e859f6 100644
--- a/pkg/reduce/data_forward.go
+++ b/pkg/reduce/data_forward.go
@@ -225,8 +225,8 @@ func (d *DataForward) forwardAChunk(ctx context.Context) {
 	processorWM := d.watermarkFetcher.GetWatermark(readMessages[0].ReadOffset)
 	for _, m := range readMessages {
 		if !d.keyed {
-			m.Key = dfv1.DefaultKeyForNonKeyedData
-			m.Message.Key = dfv1.DefaultKeyForNonKeyedData
+			m.Keys = []string{dfv1.DefaultKeyForNonKeyedData}
+			m.Message.Keys = []string{dfv1.DefaultKeyForNonKeyedData}
 		}
 		m.Watermark = time.Time(processorWM)
 		totalBytes += len(m.Payload)
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index b8444d9a..acb2a2a7 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -163,7 +163,7 @@ func (p *ProcessAndForward) Forward(ctx context.Context) error {
 	return nil
 }
 
-// whereToStep assigns a message to the ISBs based on the Message.Key.
+// whereToStep assigns a message to the ISBs based on the Message.Keys.
 // TODO: we have to introduce support for shuffle, output of a reducer can be input to the next reducer.
 func (p *ProcessAndForward) whereToStep() map[string][]isb.Message {
 	// writer doesn't accept array of pointers
@@ -172,14 +172,14 @@ func (p *ProcessAndForward) whereToStep() map[string][]isb.Message {
 	var to []string
 	var err error
 	for _, msg := range p.result {
-		to, err = p.whereToDecider.WhereTo(msg.Key)
+		to, err = p.whereToDecider.WhereTo(msg.Keys)
 		if err != nil {
 			platformError.With(map[string]string{
 				metrics.LabelVertex:             p.vertexName,
 				metrics.LabelPipeline:           p.pipelineName,
 				metrics.LabelVertexReplicaIndex: strconv.Itoa(int(p.vertexReplica)),
 			}).Inc()
-			p.log.Errorw("Got an error while invoking WhereTo, dropping the message", zap.String("key", msg.Key), zap.Error(err), zap.Any("partitionID", p.PartitionID))
+			p.log.Errorw("Got an error while invoking WhereTo, dropping the message", zap.Strings("keys", msg.Keys), zap.Error(err), zap.Any("partitionID", p.PartitionID))
 			continue
 		}
 
diff --git a/pkg/reduce/readloop/readloop.go b/pkg/reduce/readloop/readloop.go
index 038d690a..48edd447 100644
--- a/pkg/reduce/readloop/readloop.go
+++ b/pkg/reduce/readloop/readloop.go
@@ -126,8 +126,8 @@ func (rl *ReadLoop) Startup(ctx context.Context) error {
 		// insert the window to the list of active windows, since the active window list is in-memory
 		keyedWindow, _ := rl.windower.InsertIfNotPresent(alignedKeyedWindow)
 
-		// add key to the window, so that when a new message with the watermark greater than
-		// the window end time comes, key will not be lost and the windows will be closed as expected
+		// add slots to the window, so that when a new message with the watermark greater than
+		// the window end time comes, slots will not be lost and the windows will be closed as expected
 		keyedWindow.AddSlot(p.Slot)
 
 		// create and invoke process and forward for the partition
diff --git a/pkg/shuffle/shuffle.go b/pkg/shuffle/shuffle.go
index 5b2183e2..6ce21181 100644
--- a/pkg/shuffle/shuffle.go
+++ b/pkg/shuffle/shuffle.go
@@ -54,10 +54,10 @@ func NewShuffle(vertexName string, bufferIdentifiers []string) *Shuffle {
 }
 
 // Shuffle functions returns a shuffled identifier.
-func (s *Shuffle) Shuffle(key string) string {
-	// hash of the message key returns a unique hashValue
+func (s *Shuffle) Shuffle(keys []string) string {
+	// hash of the message keys returns a unique hashValue
 	// mod of hashValue will decide which isb it will belong
-	hashValue := s.generateHash(key)
+	hashValue := s.generateHash(keys)
 	hashValue = hashValue % uint64(s.buffersCount)
 	return s.bufferIdentifiers[hashValue]
 }
@@ -66,14 +66,16 @@ func (s *Shuffle) Shuffle(key string) string {
 func (s *Shuffle) ShuffleMessages(messages []*isb.Message) map[string][]*isb.Message {
 	hashMap := make(map[string][]*isb.Message)
 	for _, message := range messages {
-		identifier := s.Shuffle(message.Key)
+		identifier := s.Shuffle(message.Keys)
 		hashMap[identifier] = append(hashMap[identifier], message)
 	}
 	return hashMap
 }
 
-func (s *Shuffle) generateHash(key string) uint64 {
+func (s *Shuffle) generateHash(keys []string) uint64 {
 	s.hash.Reset()
-	_, _ = s.hash.Write([]byte(key))
+	for _, k := range keys {
+		_, _ = s.hash.Write([]byte(k))
+	}
 	return s.hash.Sum64()
 }
diff --git a/pkg/sinks/logger/log.go b/pkg/sinks/logger/log.go
index 4fdfb27b..96195c57 100644
--- a/pkg/sinks/logger/log.go
+++ b/pkg/sinks/logger/log.go
@@ -97,7 +97,7 @@ func (t *ToLog) Write(_ context.Context, messages []isb.Message) ([]isb.Offset,
 	prefix := "(" + t.GetName() + ")"
 	for _, message := range messages {
 		logSinkWriteCount.With(map[string]string{metrics.LabelVertex: t.name, metrics.LabelPipeline: t.pipelineName}).Inc()
-		log.Println(prefix, " Payload - ", string(message.Payload), " Key - ", message.Key, " EventTime - ", message.EventTime.UnixMilli())
+		log.Println(prefix, " Payload - ", string(message.Payload), " Keys - ", message.Keys, " EventTime - ", message.EventTime.UnixMilli())
 	}
 	return nil, make([]error, len(messages))
 }
diff --git a/pkg/sources/generator/tickgen.go b/pkg/sources/generator/tickgen.go
index 4c6cabd6..95b857ac 100644
--- a/pkg/sources/generator/tickgen.go
+++ b/pkg/sources/generator/tickgen.go
@@ -393,7 +393,7 @@ func (mg *memgen) newReadMessage(key string, payload []byte, offset int64) *isb.
 			// TODO: insert the right time based on the generator
 			MessageInfo: isb.MessageInfo{EventTime: timeFromNanos(parseTime(payload))},
 			ID:          strconv.FormatInt(offset, 10) + "-" + strconv.FormatInt(int64(mg.vertexInstance.Replica), 10),
-			Key:         key,
+			Keys:        []string{key},
 		},
 		Body: isb.Body{Payload: payload},
 	}
diff --git a/pkg/sources/kafka/reader.go b/pkg/sources/kafka/reader.go
index 2e21fa7d..d219d167 100644
--- a/pkg/sources/kafka/reader.go
+++ b/pkg/sources/kafka/reader.go
@@ -404,7 +404,7 @@ func toReadMessage(m *sarama.ConsumerMessage) *isb.ReadMessage {
 		Header: isb.Header{
 			MessageInfo: isb.MessageInfo{EventTime: m.Timestamp},
 			ID:          offset,
-			Key:         string(m.Key),
+			Keys:        []string{string(m.Key)},
 		},
 		Body: isb.Body{Payload: m.Value},
 	}
diff --git a/pkg/sources/redisstreams/redisstream.go b/pkg/sources/redisstreams/redisstream.go
index 67d8db45..a81e1ddb 100644
--- a/pkg/sources/redisstreams/redisstream.go
+++ b/pkg/sources/redisstreams/redisstream.go
@@ -182,7 +182,7 @@ func New(
 						Header: isb.Header{
 							MessageInfo: isb.MessageInfo{EventTime: time.Now()}, //doesn't seem like Redis offers a timestamp
 							ID:          id,                                     // assumption is that this only needs to be unique for this source vertex
-							Key:         f,
+							Keys:        []string{f},
 						},
 						Body: isb.Body{Payload: []byte(v.(string))},
 					}
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index 3da7769f..5bd1bad2 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -237,16 +237,32 @@ func (sp *SourceProcessor) getTransformerGoWhereDecider() forward.GoWhere {
 			shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)] = s
 		}
 	}
-	fsd := forward.GoWhere(func(key string) ([]string, error) {
+	fsd := forward.GoWhere(func(keys []string) ([]string, error) {
 		result := []string{}
-		if key == dfv1.MessageKeyDrop {
+		// if there are no keys, treat is as "ALL" if there are no conditions?
+		// TODO: rethink when we have tags
+		if len(keys) == 0 {
+			for _, edge := range sp.VertexInstance.Vertex.Spec.ToEdges {
+				if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 {
+					if edge.Parallelism != nil && *edge.Parallelism > 1 { // Need to shuffle
+						result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
+					} else {
+						result = append(result, dfv1.GenerateEdgeBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, edge)...)
+					}
+				}
+			}
+			return result, nil
+		}
+
+		if keys[len(keys)-1] == dfv1.MessageKeyDrop {
 			return result, nil
 		}
+
 		for _, edge := range sp.VertexInstance.Vertex.Spec.ToEdges {
-			// If returned key is not "DROP", and there's no conditions defined in the edge, treat it as "ALL"?
-			if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 || sharedutil.StringSliceContains(edge.Conditions.KeyIn, key) {
+			// If returned keys is not "DROP", and there's no conditions defined in the edge, treat it as "ALL"?
+			if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 || sharedutil.StringSliceContains(edge.Conditions.KeyIn, keys[len(keys)-1]) {
 				if edge.Parallelism != nil && *edge.Parallelism > 1 { // Need to shuffle
-					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(key))
+					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 				} else {
 					result = append(result, dfv1.GenerateEdgeBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, edge)...)
 				}
diff --git a/pkg/sources/transformer/builtin/event_time/event_time_extractor.go b/pkg/sources/transformer/builtin/event_time/event_time_extractor.go
index ce53e71c..9c04a5d1 100644
--- a/pkg/sources/transformer/builtin/event_time/event_time_extractor.go
+++ b/pkg/sources/transformer/builtin/event_time/event_time_extractor.go
@@ -54,7 +54,7 @@ func New(args map[string]string) (functionsdk.MapTFunc, error) {
 		format:     format,
 	}
 
-	return func(ctx context.Context, key string, datum functionsdk.Datum) functionsdk.MessageTs {
+	return func(ctx context.Context, keys []string, datum functionsdk.Datum) functionsdk.MessageTs {
 		log := logging.FromContext(ctx)
 		resultMsg, err := e.apply(datum.EventTime(), datum.Value())
 		if err != nil {
diff --git a/pkg/sources/transformer/builtin/filter/filter.go b/pkg/sources/transformer/builtin/filter/filter.go
index 151c976b..b5b7c7ff 100644
--- a/pkg/sources/transformer/builtin/filter/filter.go
+++ b/pkg/sources/transformer/builtin/filter/filter.go
@@ -40,7 +40,7 @@ func New(args map[string]string) (functionsdk.MapTFunc, error) {
 		expression: expr,
 	}
 
-	return func(ctx context.Context, key string, datum functionsdk.Datum) functionsdk.MessageTs {
+	return func(ctx context.Context, keys []string, datum functionsdk.Datum) functionsdk.MessageTs {
 		log := logging.FromContext(ctx)
 		resultMsg, err := f.apply(datum.EventTime(), datum.Value())
 		if err != nil {
diff --git a/pkg/sources/transformer/grpc_transformer.go b/pkg/sources/transformer/grpc_transformer.go
index 32606afd..8a3ae5ef 100644
--- a/pkg/sources/transformer/grpc_transformer.go
+++ b/pkg/sources/transformer/grpc_transformer.go
@@ -29,7 +29,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/udf/function"
 
-	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
@@ -81,18 +80,17 @@ func (u *gRPCBasedTransformer) WaitUntilReady(ctx context.Context) error {
 }
 
 func (u *gRPCBasedTransformer) ApplyMap(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.Message, error) {
-	key := readMessage.Key
+	keys := readMessage.Keys
 	payload := readMessage.Body.Payload
 	offset := readMessage.ReadOffset
 	parentMessageInfo := readMessage.MessageInfo
 	var d = &functionpb.Datum{
-		Key:       key,
+		Keys:      keys,
 		Value:     payload,
 		EventTime: &functionpb.EventTime{EventTime: timestamppb.New(parentMessageInfo.EventTime)},
 		Watermark: &functionpb.Watermark{Watermark: timestamppb.New(readMessage.Watermark)},
 	}
 
-	ctx = metadata.NewOutgoingContext(ctx, metadata.New(map[string]string{functionsdk.DatumKey: key}))
 	datumList, err := u.client.MapTFn(ctx, d)
 	if err != nil {
 		return nil, function.ApplyUDFErr{
@@ -107,7 +105,7 @@ func (u *gRPCBasedTransformer) ApplyMap(ctx context.Context, readMessage *isb.Re
 
 	writeMessages := make([]*isb.Message, 0)
 	for i, datum := range datumList {
-		key := datum.Key
+		keys := datum.Keys
 		if datum.EventTime != nil {
 			// Transformer supports changing event time.
 			parentMessageInfo.EventTime = datum.EventTime.EventTime.AsTime()
@@ -116,7 +114,7 @@ func (u *gRPCBasedTransformer) ApplyMap(ctx context.Context, readMessage *isb.Re
 			Header: isb.Header{
 				MessageInfo: parentMessageInfo,
 				ID:          fmt.Sprintf("%s-%d", offset.String(), i),
-				Key:         key,
+				Keys:        keys,
 			},
 			Body: isb.Body{
 				Payload: datum.Value,
diff --git a/pkg/udf/builtin/cat/cat.go b/pkg/udf/builtin/cat/cat.go
index fc0d83bb..46c4ea8c 100644
--- a/pkg/udf/builtin/cat/cat.go
+++ b/pkg/udf/builtin/cat/cat.go
@@ -23,7 +23,7 @@ import (
 )
 
 func New() functionsdk.MapFunc {
-	return func(ctx context.Context, key string, datum functionsdk.Datum) functionsdk.Messages {
+	return func(ctx context.Context, keys []string, datum functionsdk.Datum) functionsdk.Messages {
 		return functionsdk.MessagesBuilder().Append(functionsdk.MessageToAll(datum.Value()))
 	}
 }
diff --git a/pkg/udf/builtin/filter/filter.go b/pkg/udf/builtin/filter/filter.go
index f86b0e51..7d9e065b 100644
--- a/pkg/udf/builtin/filter/filter.go
+++ b/pkg/udf/builtin/filter/filter.go
@@ -39,7 +39,7 @@ func New(args map[string]string) (functionsdk.MapFunc, error) {
 		expression: expr,
 	}
 
-	return func(ctx context.Context, key string, datum functionsdk.Datum) functionsdk.Messages {
+	return func(ctx context.Context, keys []string, datum functionsdk.Datum) functionsdk.Messages {
 		log := logging.FromContext(ctx)
 		resultMsg, err := f.apply(datum.Value())
 		if err != nil {
diff --git a/pkg/udf/function/uds_grpc.go b/pkg/udf/function/uds_grpc.go
index 38c5aa3f..ad7909b6 100644
--- a/pkg/udf/function/uds_grpc.go
+++ b/pkg/udf/function/uds_grpc.go
@@ -79,18 +79,17 @@ func (u *UDSgRPCBasedUDF) WaitUntilReady(ctx context.Context) error {
 }
 
 func (u *UDSgRPCBasedUDF) ApplyMap(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb.Message, error) {
-	key := readMessage.Key
+	keys := readMessage.Keys
 	payload := readMessage.Body.Payload
 	offset := readMessage.ReadOffset
 	parentMessageInfo := readMessage.MessageInfo
 	var d = &functionpb.Datum{
-		Key:       key,
+		Keys:      keys,
 		Value:     payload,
 		EventTime: &functionpb.EventTime{EventTime: timestamppb.New(parentMessageInfo.EventTime)},
 		Watermark: &functionpb.Watermark{Watermark: timestamppb.New(readMessage.Watermark)},
 	}
 
-	ctx = metadata.NewOutgoingContext(ctx, metadata.New(map[string]string{functionsdk.DatumKey: key}))
 	datumList, err := u.client.MapFn(ctx, d)
 	if err != nil {
 		return nil, ApplyUDFErr{
@@ -105,12 +104,12 @@ func (u *UDSgRPCBasedUDF) ApplyMap(ctx context.Context, readMessage *isb.ReadMes
 
 	writeMessages := make([]*isb.Message, 0)
 	for i, datum := range datumList {
-		key := datum.Key
+		keys := datum.Keys
 		writeMessage := &isb.Message{
 			Header: isb.Header{
 				MessageInfo: parentMessageInfo,
 				ID:          fmt.Sprintf("%s-%d", offset.String(), i),
-				Key:         key,
+				Keys:        keys,
 			},
 			Body: isb.Body{
 				Payload: datum.Value,
@@ -132,7 +131,6 @@ func (u *UDSgRPCBasedUDF) ApplyReduce(ctx context.Context, partitionID *partitio
 
 	// pass key and window information inside the context
 	mdMap := map[string]string{
-		functionsdk.DatumKey:     partitionID.Slot,
 		functionsdk.WinStartTime: strconv.FormatInt(partitionID.Start.UnixMilli(), 10),
 		functionsdk.WinEndTime:   strconv.FormatInt(partitionID.End.UnixMilli(), 10),
 	}
@@ -185,14 +183,14 @@ readLoop:
 
 	writeMessages := make([]*isb.Message, 0)
 	for _, datum := range result {
-		key := datum.Key
+		keys := datum.Keys
 		writeMessage := &isb.Message{
 			Header: isb.Header{
 				MessageInfo: isb.MessageInfo{
 					EventTime: partitionID.End.Add(-1 * time.Millisecond),
 					IsLate:    false,
 				},
-				Key: key,
+				Keys: keys,
 			},
 			Body: isb.Body{
 				Payload: datum.Value,
@@ -204,12 +202,12 @@ readLoop:
 }
 
 func createDatum(readMessage *isb.ReadMessage) *functionpb.Datum {
-	key := readMessage.Key
+	keys := readMessage.Keys
 	payload := readMessage.Body.Payload
 	parentMessageInfo := readMessage.MessageInfo
 
 	var d = &functionpb.Datum{
-		Key:       key,
+		Keys:      keys,
 		Value:     payload,
 		EventTime: &functionpb.EventTime{EventTime: timestamppb.New(parentMessageInfo.EventTime)},
 		Watermark: &functionpb.Watermark{Watermark: timestamppb.New(readMessage.Watermark)},
diff --git a/pkg/udf/map_udf.go b/pkg/udf/map_udf.go
index d9d45062..ab7280a9 100644
--- a/pkg/udf/map_udf.go
+++ b/pkg/udf/map_udf.go
@@ -82,16 +82,32 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 		}
 	}
 
-	conditionalForwarder := forward.GoWhere(func(key string) ([]string, error) {
+	conditionalForwarder := forward.GoWhere(func(keys []string) ([]string, error) {
 		result := []string{}
-		if key == dfv1.MessageKeyDrop {
+		// if there are no keys, treat is as "ALL" if there are no conditions?
+		// TODO: rethink when we have tags
+		if len(keys) == 0 {
+			for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges {
+				if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 {
+					if edge.Parallelism != nil && *edge.Parallelism > 1 { // Need to shuffle
+						result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
+					} else {
+						result = append(result, dfv1.GenerateEdgeBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge)...)
+					}
+				}
+			}
+			return result, nil
+		}
+
+		if keys[len(keys)-1] == dfv1.MessageKeyDrop {
 			return result, nil
 		}
+
 		for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges {
-			// If returned key is not "DROP", and there's no conditions defined in the edge, treat it as "ALL"?
-			if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 || sharedutil.StringSliceContains(edge.Conditions.KeyIn, key) {
+			// If returned keys is not "DROP", and there's no conditions defined in the edge, treat it as "ALL"?
+			if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 || sharedutil.StringSliceContains(edge.Conditions.KeyIn, keys[len(keys)-1]) {
 				if edge.Parallelism != nil && *edge.Parallelism > 1 { // Need to shuffle
-					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(key))
+					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 				} else {
 					result = append(result, dfv1.GenerateEdgeBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge)...)
 				}
diff --git a/pkg/udf/reduce_udf.go b/pkg/udf/reduce_udf.go
index e2590837..9cdf0775 100644
--- a/pkg/udf/reduce_udf.go
+++ b/pkg/udf/reduce_udf.go
@@ -111,16 +111,32 @@ func (u *ReduceUDFProcessor) Start(ctx context.Context) error {
 		}
 	}
 
-	conditionalForwarder := forward.GoWhere(func(key string) ([]string, error) {
+	conditionalForwarder := forward.GoWhere(func(keys []string) ([]string, error) {
 		result := []string{}
-		if key == dfv1.MessageKeyDrop {
+		// if there are no keys, treat is as "ALL" if there are no conditions?
+		// TODO: rethink when we have tags
+		if len(keys) == 0 {
+			for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges {
+				if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 {
+					if edge.Parallelism != nil && *edge.Parallelism > 1 { // Need to shuffle
+						result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
+					} else {
+						result = append(result, dfv1.GenerateEdgeBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge)...)
+					}
+				}
+			}
+			return result, nil
+		}
+
+		if keys[len(keys)-1] == dfv1.MessageKeyDrop {
 			return result, nil
 		}
+
 		for _, edge := range u.VertexInstance.Vertex.Spec.ToEdges {
-			// If returned key is not "DROP", and there's no conditions defined in the edge, treat it as "ALL"?
-			if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 || sharedutil.StringSliceContains(edge.Conditions.KeyIn, key) {
+			// If returned keys is not "DROP", and there's no conditions defined in the edge, treat it as "ALL"?
+			if edge.Conditions == nil || len(edge.Conditions.KeyIn) == 0 || sharedutil.StringSliceContains(edge.Conditions.KeyIn, keys[len(keys)-1]) {
 				if edge.Parallelism != nil && *edge.Parallelism > 1 { // Need to shuffle
-					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(key))
+					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 				} else {
 					result = append(result, dfv1.GenerateEdgeBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge)...)
 				}


Test patch:
diff --git a/pkg/forward/forward_test.go b/pkg/forward/forward_test.go
index 0ebface6..0575519a 100644
--- a/pkg/forward/forward_test.go
+++ b/pkg/forward/forward_test.go
@@ -88,7 +88,7 @@ func (t *testForwardFetcher) GetHeadWMB() wmb.WMB {
 type myForwardTest struct {
 }
 
-func (f myForwardTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
@@ -515,7 +515,7 @@ func TestNewInterStepDataForwardIdleWatermark_Reset(t *testing.T) {
 type mySourceForwardTest struct {
 }
 
-func (f mySourceForwardTest) WhereTo(_ string) ([]string, error) {
+func (f mySourceForwardTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
@@ -536,13 +536,13 @@ func (f mySourceForwardTest) ApplyMap(ctx context.Context, message *isb.ReadMess
 		result := payload
 		// assign new event time
 		parentPaneInfo.EventTime = testSourceNewEventTime
-		var key string
+		var key []string
 
 		writeMessage := &isb.Message{
 			Header: isb.Header{
 				MessageInfo: parentPaneInfo,
 				ID:          offset.String(),
-				Key:         key,
+				Keys:        key,
 			},
 			Body: isb.Body{
 				Payload: result,
@@ -601,7 +601,7 @@ func TestSourceInterStepDataForward(t *testing.T) {
 	readMessages, err := to1.Read(ctx, count)
 	assert.NoError(t, err, "expected no error")
 	assert.Len(t, readMessages, int(count))
-	assert.Equal(t, []interface{}{writeMessages[0].Header.Key, writeMessages[1].Header.Key}, []interface{}{readMessages[0].Header.Key, readMessages[1].Header.Key})
+	assert.Equal(t, []interface{}{writeMessages[0].Header.Keys, writeMessages[1].Header.Keys}, []interface{}{readMessages[0].Header.Keys, readMessages[1].Header.Keys})
 	assert.Equal(t, []interface{}{writeMessages[0].Header.ID, writeMessages[1].Header.ID}, []interface{}{readMessages[0].Header.ID, readMessages[1].Header.ID})
 	for _, m := range readMessages {
 		// verify new event time gets assigned to messages.
@@ -775,7 +775,7 @@ func TestNewInterStepDataForwardToOneStep(t *testing.T) {
 type myForwardDropTest struct {
 }
 
-func (f myForwardDropTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardDropTest) WhereTo(_ []string) ([]string, error) {
 	return []string{dfv1.MessageKeyDrop}, nil
 }
 
@@ -883,7 +883,7 @@ func TestNewInterStepDataForward_dropAll(t *testing.T) {
 type myForwardToAllTest struct {
 }
 
-func (f myForwardToAllTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardToAllTest) WhereTo(_ []string) ([]string, error) {
 	return []string{dfv1.MessageKeyAll}, nil
 }
 
@@ -977,7 +977,7 @@ func TestNewInterStepData_forwardToAll(t *testing.T) {
 type myForwardInternalErrTest struct {
 }
 
-func (f myForwardInternalErrTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardInternalErrTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
@@ -1029,7 +1029,7 @@ func TestNewInterStepDataForward_WithInternalError(t *testing.T) {
 type myForwardApplyWhereToErrTest struct {
 }
 
-func (f myForwardApplyWhereToErrTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardApplyWhereToErrTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, fmt.Errorf("whereToStep failed")
 }
 
@@ -1076,7 +1076,7 @@ func TestNewInterStepDataForward_WhereToError(t *testing.T) {
 type myForwardApplyUDFErrTest struct {
 }
 
-func (f myForwardApplyUDFErrTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardApplyUDFErrTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
diff --git a/pkg/forward/shutdown_test.go b/pkg/forward/shutdown_test.go
index 5b2f9231..2dd801a0 100644
--- a/pkg/forward/shutdown_test.go
+++ b/pkg/forward/shutdown_test.go
@@ -33,7 +33,7 @@ import (
 type myShutdownTest struct {
 }
 
-func (s myShutdownTest) WhereTo(_ string) ([]string, error) {
+func (s myShutdownTest) WhereTo(_ []string) ([]string, error) {
 	return []string{dfv1.MessageKeyAll}, nil
 }
 
diff --git a/pkg/isb/serde_test.go b/pkg/isb/serde_test.go
index 3bdfb482..94bb2fca 100644
--- a/pkg/isb/serde_test.go
+++ b/pkg/isb/serde_test.go
@@ -92,7 +92,7 @@ func TestHeader(t *testing.T) {
 		MessageInfo MessageInfo
 		Kind        MessageKind
 		ID          string
-		Key         string
+		Key         []string
 	}
 	tests := []struct {
 		name               string
@@ -110,7 +110,7 @@ func TestHeader(t *testing.T) {
 				},
 				Kind: Data,
 				ID:   "TestID",
-				Key:  "TestKey",
+				Key:  []string{"TestKey", "TestKey2"},
 			},
 			wantData: Header{
 				MessageInfo: MessageInfo{
@@ -119,7 +119,7 @@ func TestHeader(t *testing.T) {
 				},
 				Kind: Data,
 				ID:   "TestID",
-				Key:  "TestKey",
+				Keys: []string{"TestKey", "TestKey2"},
 			},
 			wantMarshalError:   false,
 			wantUnmarshalError: false,
@@ -131,7 +131,7 @@ func TestHeader(t *testing.T) {
 				MessageInfo: tt.fields.MessageInfo,
 				Kind:        tt.fields.Kind,
 				ID:          tt.fields.ID,
-				Key:         tt.fields.Key,
+				Keys:        tt.fields.Key,
 			}
 			gotData, err := h.MarshalBinary()
 			if (err != nil) != tt.wantMarshalError {
@@ -228,7 +228,7 @@ func TestMessage(t *testing.T) {
 					},
 					Kind: Data,
 					ID:   "TestID",
-					Key:  "TestKey",
+					Keys: []string{"TestKey"},
 				},
 				Body: Body{
 					Payload: []byte("TestBODY"),
@@ -242,7 +242,7 @@ func TestMessage(t *testing.T) {
 					},
 					Kind: Data,
 					ID:   "TestID",
-					Key:  "TestKey",
+					Keys: []string{"TestKey"},
 				},
 				Body: Body{
 					Payload: []byte("TestBODY"),
@@ -276,8 +276,11 @@ func TestMessage(t *testing.T) {
 				t.Errorf("UnmarshalBinary() error = %v, wantUnmarshalError %v", err, tt.wantMarshalError)
 				return
 			}
+			if tt.name == "good_empty" {
+				tt.wantData.Keys = make([]string, 0)
+			}
 			if !reflect.DeepEqual(*newM, tt.wantData) {
-				t.Errorf("MarshalBinary() gotData = %v, want %v", newM, tt.wantData)
+				t.Errorf("MarshalBinary() gotData = %v, want %v", newM, &tt.wantData)
 			}
 		})
 	}
@@ -307,7 +310,7 @@ func TestReadMessage(t *testing.T) {
 						},
 						Kind: Data,
 						ID:   "TestID",
-						Key:  "TestKey",
+						Keys: []string{"TestKey"},
 					},
 					Body: Body{
 						Payload: []byte("TestBODY"),
@@ -327,7 +330,7 @@ func TestReadMessage(t *testing.T) {
 						},
 						Kind: Data,
 						ID:   "TestID",
-						Key:  "TestKey",
+						Keys: []string{"TestKey"},
 					},
 					Body: Body{
 						Payload: []byte("TestBODY"),
diff --git a/pkg/isb/stores/jetstream/writer_test.go b/pkg/isb/stores/jetstream/writer_test.go
index a16f780f..7bb9a853 100644
--- a/pkg/isb/stores/jetstream/writer_test.go
+++ b/pkg/isb/stores/jetstream/writer_test.go
@@ -35,7 +35,7 @@ import (
 type myForwardJetStreamTest struct {
 }
 
-func (f myForwardJetStreamTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardJetStreamTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
diff --git a/pkg/isb/stores/redis/read_test.go b/pkg/isb/stores/redis/read_test.go
index 99dc11a3..a3edcd98 100644
--- a/pkg/isb/stores/redis/read_test.go
+++ b/pkg/isb/stores/redis/read_test.go
@@ -291,7 +291,7 @@ type ReadWritePerformance struct {
 type forwardReadWritePerformance struct {
 }
 
-func (f forwardReadWritePerformance) WhereTo(_ string) ([]string, error) {
+func (f forwardReadWritePerformance) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
diff --git a/pkg/isb/stores/redis/write_test.go b/pkg/isb/stores/redis/write_test.go
index 8a263aff..45520e7d 100644
--- a/pkg/isb/stores/redis/write_test.go
+++ b/pkg/isb/stores/redis/write_test.go
@@ -342,7 +342,7 @@ func Test_GetRefreshFullError(t *testing.T) {
 type myForwardRedisTest struct {
 }
 
-func (f myForwardRedisTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardRedisTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"to1"}, nil
 }
 
diff --git a/pkg/isb/testutils/rw.go b/pkg/isb/testutils/rw.go
index fb63e6bd..33d8a7d7 100644
--- a/pkg/isb/testutils/rw.go
+++ b/pkg/isb/testutils/rw.go
@@ -50,7 +50,8 @@ func BuildTestWriteMessages(count int64, startTime time.Time) []isb.Message {
 					MessageInfo: isb.MessageInfo{
 						EventTime: tmpTime,
 					},
-					ID: fmt.Sprintf("%d", i),
+					ID:   fmt.Sprintf("%d", i),
+					Keys: []string{},
 				},
 				Body: isb.Body{Payload: result},
 			},
diff --git a/pkg/isb/testutils/udf.go b/pkg/isb/testutils/udf.go
index d6b4c056..fdf076ac 100644
--- a/pkg/isb/testutils/udf.go
+++ b/pkg/isb/testutils/udf.go
@@ -33,13 +33,13 @@ func CopyUDFTestApply(ctx context.Context, readMessage *isb.ReadMessage) ([]*isb
 	_ = payload
 	// copy the payload
 	result := payload
-	var key string
+	var keys []string
 
 	writeMessage := &isb.Message{
 		Header: isb.Header{
 			MessageInfo: parentPaneInfo,
 			ID:          offset.String(),
-			Key:         key,
+			Keys:        keys,
 		},
 		Body: isb.Body{
 			Payload: result,
diff --git a/pkg/reduce/data_forward_test.go b/pkg/reduce/data_forward_test.go
index 1ee2e7c8..ca311a79 100644
--- a/pkg/reduce/data_forward_test.go
+++ b/pkg/reduce/data_forward_test.go
@@ -127,8 +127,8 @@ func (f CounterReduceTest) ApplyReduce(_ context.Context, partitionID *partition
 			MessageInfo: isb.MessageInfo{
 				EventTime: partitionID.End,
 			},
-			ID:  "msgID",
-			Key: "result",
+			ID:   "msgID",
+			Keys: []string{"result"},
 		},
 		Body: isb.Body{Payload: b},
 	}
@@ -137,7 +137,7 @@ func (f CounterReduceTest) ApplyReduce(_ context.Context, partitionID *partition
 	}, nil
 }
 
-func (f CounterReduceTest) WhereTo(_ string) ([]string, error) {
+func (f CounterReduceTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"reduce-to-buffer"}, nil
 }
 
@@ -150,8 +150,8 @@ func (s SumReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID,
 	for msg := range messageStream {
 		var payload PayloadForTest
 		_ = json.Unmarshal(msg.Payload, &payload)
-		key := msg.Key
-		sums[key] += payload.Value
+		key := msg.Keys
+		sums[key[0]] += payload.Value
 	}
 
 	msgs := make([]*isb.Message, 0)
@@ -164,8 +164,8 @@ func (s SumReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID,
 				MessageInfo: isb.MessageInfo{
 					EventTime: partitionID.End,
 				},
-				ID:  "msgID",
-				Key: k,
+				ID:   "msgID",
+				Keys: []string{k},
 			},
 			Body: isb.Body{Payload: b},
 		}
@@ -184,12 +184,12 @@ func (m MaxReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID,
 	for msg := range messageStream {
 		var payload PayloadForTest
 		_ = json.Unmarshal(msg.Payload, &payload)
-		if max, ok := maxMap[msg.Key]; ok {
+		if max, ok := maxMap[msg.Keys[0]]; ok {
 			mx = max
 		}
 		if payload.Value > mx {
 			mx = payload.Value
-			maxMap[msg.Key] = mx
+			maxMap[msg.Keys[0]] = mx
 		}
 	}
 
@@ -202,8 +202,8 @@ func (m MaxReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID,
 				MessageInfo: isb.MessageInfo{
 					EventTime: partitionID.End,
 				},
-				ID:  "msgID",
-				Key: k,
+				ID:   "msgID",
+				Keys: []string{k},
 			},
 			Body: isb.Body{Payload: b},
 		}
@@ -454,7 +454,6 @@ func TestReduceDataForward_IdleWM(t *testing.T) {
 			MessageInfo: isb.MessageInfo{},
 			Kind:        0,
 			ID:          "",
-			Key:         "",
 		}, msgs[i].Header)
 	}
 
@@ -469,7 +468,7 @@ func TestReduceDataForward_IdleWM(t *testing.T) {
 	assert.Equal(t, isb.WMB, msgs[0].Kind)
 	// the second message should be the data message from the closed window above
 	// in the test ApplyUDF above we've set the final message to have key="result"
-	for msgs[1].Key != "result" {
+	for len(msgs[1].Keys) == 0 {
 		select {
 		case <-ctx.Done():
 			if ctx.Err() == context.DeadlineExceeded {
@@ -498,7 +497,6 @@ func TestReduceDataForward_IdleWM(t *testing.T) {
 			MessageInfo: isb.MessageInfo{},
 			Kind:        0,
 			ID:          "",
-			Key:         "",
 		}, msgs[i].Header)
 	}
 
@@ -1063,8 +1061,8 @@ func buildMessagesForReduce(count int, key string, publishTime time.Time) []isb.
 				MessageInfo: isb.MessageInfo{
 					EventTime: publishTime,
 				},
-				ID:  fmt.Sprintf("%d", i),
-				Key: key,
+				ID:   fmt.Sprintf("%d", i),
+				Keys: []string{key},
 			},
 			Body: isb.Body{Payload: result},
 		}
@@ -1135,15 +1133,15 @@ func publishMessages(ctx context.Context, startTime int, messages []int, testDur
 }
 
 func buildIsbMessage(messageValue int, eventTime time.Time) isb.Message {
-	var messageKey string
+	var messageKey []string
 	if messageValue%2 == 0 {
-		messageKey = "even"
+		messageKey = []string{"even"}
 	} else {
-		messageKey = "odd"
+		messageKey = []string{"odd"}
 	}
 
 	result, _ := json.Marshal(PayloadForTest{
-		Key:   messageKey,
+		Key:   messageKey[0],
 		Value: messageValue,
 	})
 	return isb.Message{
@@ -1151,8 +1149,8 @@ func buildIsbMessage(messageValue int, eventTime time.Time) isb.Message {
 			MessageInfo: isb.MessageInfo{
 				EventTime: eventTime,
 			},
-			ID:  fmt.Sprintf("%d", messageValue),
-			Key: messageKey,
+			ID:   fmt.Sprintf("%d", messageValue),
+			Keys: messageKey,
 		},
 		Body: isb.Body{Payload: result},
 	}
diff --git a/pkg/reduce/pbq/store/wal/segment_test.go b/pkg/reduce/pbq/store/wal/segment_test.go
index 721d3aea..73d90415 100644
--- a/pkg/reduce/pbq/store/wal/segment_test.go
+++ b/pkg/reduce/pbq/store/wal/segment_test.go
@@ -195,8 +195,10 @@ func Test_encodeDecodeEntry(t *testing.T) {
 			wantErr: assert.NoError,
 			message: &isb.ReadMessage{
 				Message: isb.Message{
-					Header: isb.Header{},
-					Body:   isb.Body{},
+					Header: isb.Header{
+						Keys: []string{},
+					},
+					Body: isb.Body{},
 				},
 				ReadOffset: isb.SimpleIntOffset(func() int64 { return int64(2) }),
 			},
@@ -259,7 +261,7 @@ func Test_batchSyncWithMaxBatchSize(t *testing.T) {
 	assert.NoError(t, err)
 	err = wal.Write(&message)
 	assert.NoError(t, err)
-	assert.Equal(t, int64(214), tempWAL.prevSyncedWOffset)
+	assert.Equal(t, int64(218), tempWAL.prevSyncedWOffset)
 
 	err = wal.Close()
 	assert.NoError(t, err)
@@ -316,7 +318,7 @@ func Test_batchSyncWithSyncDuration(t *testing.T) {
 	message := writeMessages[0]
 	storePrevSyncedTime := tempWAL.prevSyncedTime
 	err = wal.Write(&message)
-	assert.Equal(t, int64(126), tempWAL.prevSyncedWOffset)
+	assert.Equal(t, int64(128), tempWAL.prevSyncedWOffset)
 	assert.NotEqual(t, storePrevSyncedTime, tempWAL.prevSyncedTime)
 	assert.NoError(t, err)
 
diff --git a/pkg/reduce/pnf/processandforward_test.go b/pkg/reduce/pnf/processandforward_test.go
index 1a971e9c..68c43649 100644
--- a/pkg/reduce/pnf/processandforward_test.go
+++ b/pkg/reduce/pnf/processandforward_test.go
@@ -61,10 +61,10 @@ type myForwardTest struct {
 	buffers []string
 }
 
-func (f myForwardTest) WhereTo(key string) ([]string, error) {
-	if strings.Compare(key, "test-forward-one") == 0 {
+func (f myForwardTest) WhereTo(key []string) ([]string, error) {
+	if strings.Compare(key[len(key)-1], "test-forward-one") == 0 {
 		return []string{"buffer1"}, nil
-	} else if strings.Compare(key, "test-forward-all") == 0 {
+	} else if strings.Compare(key[len(key)-1], "test-forward-all") == 0 {
 		return f.buffers, nil
 	}
 	return []string{}, nil
@@ -137,7 +137,7 @@ func TestProcessAndForward_Process(t *testing.T) {
 	mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 		Elements: []*functionpb.Datum{
 			{
-				Key:   "reduced_result_key",
+				Keys:  []string{"reduced_result_key"},
 				Value: []byte(`forward_message`),
 			},
 		},
@@ -145,7 +145,7 @@ func TestProcessAndForward_Process(t *testing.T) {
 	mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 		Elements: []*functionpb.Datum{
 			{
-				Key:   "reduced_result_key",
+				Keys:  []string{"reduced_result_key"},
 				Value: []byte(`forward_message`),
 			},
 		},
@@ -378,8 +378,8 @@ func createProcessAndForwardAndOTStore(ctx context.Context, key string, pbqManag
 				MessageInfo: isb.MessageInfo{
 					EventTime: time.UnixMilli(60000),
 				},
-				ID:  "1",
-				Key: key,
+				ID:   "1",
+				Keys: []string{key},
 			},
 			Body: isb.Body{Payload: resultPayload},
 		},
diff --git a/pkg/reduce/readloop/ordered_test.go b/pkg/reduce/readloop/ordered_test.go
index e3fcf650..58ccfd9a 100644
--- a/pkg/reduce/readloop/ordered_test.go
+++ b/pkg/reduce/readloop/ordered_test.go
@@ -39,7 +39,7 @@ import (
 type myForwardTest struct {
 }
 
-func (f myForwardTest) WhereTo(_ string) ([]string, error) {
+func (f myForwardTest) WhereTo(_ []string) ([]string, error) {
 	return []string{dfv1.MessageKeyDrop}, nil
 }
 
diff --git a/pkg/reduce/readloop/readloop_test.go b/pkg/reduce/readloop/readloop_test.go
index 30d419b3..e7895324 100644
--- a/pkg/reduce/readloop/readloop_test.go
+++ b/pkg/reduce/readloop/readloop_test.go
@@ -44,7 +44,7 @@ type PayloadForTest struct {
 type SumReduceTest struct {
 }
 
-func (s *SumReduceTest) WhereTo(_ string) ([]string, error) {
+func (s *SumReduceTest) WhereTo(_ []string) ([]string, error) {
 	return []string{"reduce-buffer"}, nil
 }
 
@@ -54,8 +54,8 @@ func (s SumReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID,
 	for msg := range messageStream {
 		var payload PayloadForTest
 		_ = json.Unmarshal(msg.Payload, &payload)
-		key := msg.Key
-		sums[key] += payload.Value
+		keys := msg.Keys
+		sums[keys[0]] += payload.Value
 	}
 
 	msgs := make([]*isb.Message, 0)
@@ -68,8 +68,8 @@ func (s SumReduceTest) ApplyReduce(_ context.Context, partitionID *partition.ID,
 				MessageInfo: isb.MessageInfo{
 					EventTime: partitionID.End,
 				},
-				ID:  "msgID",
-				Key: k,
+				ID:   "msgID",
+				Keys: []string{k},
 			},
 			Body: isb.Body{Payload: b},
 		}
@@ -155,8 +155,8 @@ func TestReadLoop_Startup(t *testing.T) {
 					EventTime: time.Unix(300, 0),
 					IsLate:    false,
 				},
-				ID:  "",
-				Key: "",
+				ID:   "",
+				Keys: []string{""},
 			},
 			Body: isb.Body{},
 		},
@@ -210,8 +210,8 @@ func createStoreMessages(_ context.Context, key string, value int, eventTime tim
 					MessageInfo: isb.MessageInfo{
 						EventTime: eventTime,
 					},
-					ID:  fmt.Sprintf("%d", value+1),
-					Key: key,
+					ID:   fmt.Sprintf("%d", value+1),
+					Keys: []string{key},
 				},
 				Body: isb.Body{Payload: result},
 			},
diff --git a/pkg/shuffle/shuffle_test.go b/pkg/shuffle/shuffle_test.go
index 6cabf23f..c512a0c4 100644
--- a/pkg/shuffle/shuffle_test.go
+++ b/pkg/shuffle/shuffle_test.go
@@ -129,10 +129,10 @@ func buildBufferIdList(size int) []string {
 func buildTestMessagesWithDistinctKeys(size int64) []*isb.Message {
 	// build test messages
 	messages := testutils.BuildTestWriteMessages(size, time.Now())
-	// set key for test messages
+	// set keys for test messages
 	var res []*isb.Message
 	for index := 0; index < len(messages); index++ {
-		messages[index].Key = fmt.Sprintf("key_%d", index)
+		messages[index].Keys = []string{fmt.Sprintf("key_%d", index)}
 		res = append(res, &messages[index])
 	}
 	return res
diff --git a/pkg/sinks/kafka/kafka_test.go b/pkg/sinks/kafka/kafka_test.go
index c92890bf..b6f54e85 100644
--- a/pkg/sinks/kafka/kafka_test.go
+++ b/pkg/sinks/kafka/kafka_test.go
@@ -66,14 +66,14 @@ func TestWriteSuccessToKafka(t *testing.T) {
 		{
 			Header: isb.Header{
 				MessageInfo: isb.MessageInfo{},
-				Key:         "key1",
+				Keys:        []string{"key1"},
 			},
 			Body: isb.Body{Payload: []byte("welcome1")},
 		},
 		{
 			Header: isb.Header{
 				MessageInfo: isb.MessageInfo{},
-				Key:         "key1",
+				Keys:        []string{"key1"},
 			},
 			Body: isb.Body{Payload: []byte("welcome1")},
 		},
@@ -118,14 +118,14 @@ func TestWriteFailureToKafka(t *testing.T) {
 		{
 			Header: isb.Header{
 				MessageInfo: isb.MessageInfo{},
-				Key:         "key1",
+				Keys:        []string{"key1"},
 			},
 			Body: isb.Body{Payload: []byte("welcome1")},
 		},
 		{
 			Header: isb.Header{
 				MessageInfo: isb.MessageInfo{},
-				Key:         "key1",
+				Keys:        []string{"key1"},
 			},
 			Body: isb.Body{Payload: []byte("welcome1")},
 		},
diff --git a/pkg/sources/kafka/handler_test.go b/pkg/sources/kafka/handler_test.go
index 3b9f000a..9078b6f6 100644
--- a/pkg/sources/kafka/handler_test.go
+++ b/pkg/sources/kafka/handler_test.go
@@ -42,7 +42,7 @@ func TestMessageHandling(t *testing.T) {
 	partition := int32(1)
 	offset := int64(1)
 	value := "testvalue"
-	key := "testkey"
+	keys := []string{"testkey"}
 
 	tobuffer := simplebuffer.NewInMemoryBuffer("test", 100)
 	dest := []isb.BufferWriter{tobuffer}
@@ -73,7 +73,7 @@ func TestMessageHandling(t *testing.T) {
 		Topic:     topic,
 		Partition: partition,
 		Offset:    offset,
-		Key:       []byte(key),
+		Key:       []byte(keys[0]),
 		Value:     []byte(value),
 	}
 
@@ -90,6 +90,6 @@ func TestMessageHandling(t *testing.T) {
 	readmsg := readmsgs[0]
 	assert.Equal(t, expectedoffset, readmsg.ID)
 	assert.Equal(t, []byte(value), readmsg.Body.Payload)
-	assert.Equal(t, key, readmsg.Header.Key)
+	assert.Equal(t, keys, readmsg.Header.Keys)
 	assert.Equal(t, expectedoffset, readmsg.ReadOffset.String())
 }
diff --git a/pkg/sources/transformer/builtin/event_time/event_time_extractor_test.go b/pkg/sources/transformer/builtin/event_time/event_time_extractor_test.go
index 87085715..bb43afd7 100644
--- a/pkg/sources/transformer/builtin/event_time/event_time_extractor_test.go
+++ b/pkg/sources/transformer/builtin/event_time/event_time_extractor_test.go
@@ -55,7 +55,7 @@ func TestEventTimeExtractor(t *testing.T) {
 		assert.NoError(t, err)
 
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "2022-02-18T21:54:42.123Z"},{"id": 2, "name": "numa", "time": "2021-02-18T21:54:42.123Z"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
@@ -64,9 +64,9 @@ func TestEventTimeExtractor(t *testing.T) {
 		time.Local, _ = time.LoadLocation("UTC")
 		expected, _ := time.Parse(time.RFC3339, "2021-02-18T21:54:42.123Z")
 		// Verify new event time is assigned to the message.
-		assert.True(t, expected.Equal(result.Items()[0].EventTime))
+		assert.True(t, expected.Equal(result.Items()[0].EventTime()))
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("Json expression valid, assign a new event time to the message - format specified", func(t *testing.T) {
@@ -75,7 +75,7 @@ func TestEventTimeExtractor(t *testing.T) {
 		assert.NoError(t, err)
 
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "2022-02-18T21:54:42.123Z"},{"id": 2, "name": "numa", "time": "2021-02-18T21:54:42.123Z"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
@@ -84,9 +84,9 @@ func TestEventTimeExtractor(t *testing.T) {
 		time.Local, _ = time.LoadLocation("UTC")
 		expected, _ := time.Parse(time.RFC3339, "2021-02-18T21:54:42.123Z")
 		// Verify new event time is assigned to the message.
-		assert.True(t, expected.Equal(result.Items()[0].EventTime))
+		assert.True(t, expected.Equal(result.Items()[0].EventTime()))
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("Time string not matching user-provided format, pass on the message without assigning new event time", func(t *testing.T) {
@@ -97,16 +97,16 @@ func TestEventTimeExtractor(t *testing.T) {
 		testInputEventTime := time.Date(2022, 1, 4, 2, 3, 4, 5, time.UTC)
 		// Handler receives format as time.ANSIC but in the message, we use time.RFC3339. Format is not matched.
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "2022-02-18T21:54:42.123Z"},{"id": 2, "name": "numa", "time": "2021-02-18T21:54:42.123Z"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: testInputEventTime,
 			watermark: time.Time{},
 		})
 
 		// Verify event time remains unchanged.
-		assert.Equal(t, testInputEventTime, result.Items()[0].EventTime)
+		assert.Equal(t, testInputEventTime, result.Items()[0].EventTime())
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("Cannot compile json expression, pass on the message without assigning new event time", func(t *testing.T) {
@@ -116,7 +116,7 @@ func TestEventTimeExtractor(t *testing.T) {
 
 		testInputEventTime := time.Date(2022, 1, 4, 2, 3, 4, 5, time.UTC)
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "2022-02-18T21:54:42.123Z"},{"id": 2, "name": "numa", "time": "2021-02-18T21:54:42.123Z"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: testInputEventTime,
 			watermark: time.Time{},
@@ -124,9 +124,9 @@ func TestEventTimeExtractor(t *testing.T) {
 
 		expected := testInputEventTime
 		// Verify event time remains unchanged.
-		assert.True(t, expected.Equal(result.Items()[0].EventTime))
+		assert.True(t, expected.Equal(result.Items()[0].EventTime()))
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("The time string is in epoch format with a granularity of seconds, assign a new event time to the message", func(t *testing.T) {
@@ -137,7 +137,7 @@ func TestEventTimeExtractor(t *testing.T) {
 		testInputEventTime := time.Date(2022, 1, 4, 2, 3, 4, 5, time.UTC)
 		// Handler receives format as time.ANSIC but in the message, we use time.RFC3339. Format is not matched.
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "1673239888"},{"id": 2, "name": "numa", "time": "1673239888"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: testInputEventTime,
 			watermark: time.Time{},
@@ -146,9 +146,9 @@ func TestEventTimeExtractor(t *testing.T) {
 		time.Local, _ = time.LoadLocation("UTC")
 		expected := time.Unix(1673239888, 0)
 		// Verify new event time is assigned to the message.
-		assert.True(t, expected.Equal(result.Items()[0].EventTime))
+		assert.True(t, expected.Equal(result.Items()[0].EventTime()))
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("The time string is in epoch format with a granularity of milliseconds, assign a new event time to the message", func(t *testing.T) {
@@ -158,7 +158,7 @@ func TestEventTimeExtractor(t *testing.T) {
 
 		testInputEventTime := time.Date(2022, 1, 4, 2, 3, 4, 5, time.UTC)
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "1673239888123"},{"id": 2, "name": "numa", "time": "1673239888123"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: testInputEventTime,
 			watermark: time.Time{},
@@ -167,9 +167,9 @@ func TestEventTimeExtractor(t *testing.T) {
 		time.Local, _ = time.LoadLocation("UTC")
 		expected := time.Unix(1673239888, 123000000)
 		// Verify new event time is assigned to the message.
-		assert.True(t, expected.Equal(result.Items()[0].EventTime))
+		assert.True(t, expected.Equal(result.Items()[0].EventTime()))
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("The time string is ambiguous, pass on the message without assigning new event time", func(t *testing.T) {
@@ -180,7 +180,7 @@ func TestEventTimeExtractor(t *testing.T) {
 		testInputEventTime := time.Date(2022, 1, 4, 2, 3, 4, 5, time.UTC)
 		// 04/08/2014 is ambiguous because it could be mm/dd/yyyy or dd/mm/yyyy.
 		testJsonMsg := `{"test": 21, "item": [{"id": 1, "name": "numa", "time": "04/08/2014 22:05"},{"id": 2, "name": "numa", "time": "04/08/2014 22:05"}]}`
-		result := handle(context.Background(), "test-key", &testDatum{
+		result := handle(context.Background(), []string{"test-key"}, &testDatum{
 			value:     []byte(testJsonMsg),
 			eventTime: testInputEventTime,
 			watermark: time.Time{},
@@ -189,8 +189,8 @@ func TestEventTimeExtractor(t *testing.T) {
 		time.Local, _ = time.LoadLocation("UTC")
 		expected := testInputEventTime
 		// Verify event time remains unchanged.
-		assert.True(t, expected.Equal(result.Items()[0].EventTime))
+		assert.True(t, expected.Equal(result.Items()[0].EventTime()))
 		// Verify the payload remains unchanged.
-		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, testJsonMsg, string(result.Items()[0].Value()))
 	})
 }
diff --git a/pkg/sources/transformer/builtin/filter/filter_test.go b/pkg/sources/transformer/builtin/filter/filter_test.go
index e1d18a61..a417a0a8 100644
--- a/pkg/sources/transformer/builtin/filter/filter_test.go
+++ b/pkg/sources/transformer/builtin/filter/filter_test.go
@@ -25,7 +25,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-var _key = ""
+var _keys = []string{""}
 var jsonMsg = `{"test": 21, "item": [{"id": 1, "name": "numa"},{"id": 2, "name": "numa"}]}`
 var strMsg = `welcome to numaflow`
 var base64Msg = base64.StdEncoding.EncodeToString([]byte(strMsg))
@@ -61,12 +61,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, jsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, jsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("invalid expression", func(t *testing.T) {
@@ -75,12 +75,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, "", string(result.Items()[0].Value))
+		assert.Equal(t, "", string(result.Items()[0].Value()))
 	})
 
 	t.Run("Json expression invalid", func(t *testing.T) {
@@ -89,12 +89,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, "", string(result.Items()[0].Value))
+		assert.Equal(t, "", string(result.Items()[0].Value()))
 	})
 
 	t.Run("String expression invalid", func(t *testing.T) {
@@ -103,12 +103,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, "", string(result.Items()[0].Value))
+		assert.Equal(t, "", string(result.Items()[0].Value()))
 	})
 
 	t.Run("base64 expression valid", func(t *testing.T) {
@@ -117,12 +117,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(base64Msg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, base64Msg, string(result.Items()[0].Value))
+		assert.Equal(t, base64Msg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("event time unchanged", func(t *testing.T) {
@@ -132,11 +132,11 @@ func TestExpression(t *testing.T) {
 		assert.NoError(t, err)
 
 		testEventTime := time.Date(2022, 01, 01, 01, 01, 01, 01, time.UTC)
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(base64Msg),
 			eventTime: testEventTime,
 			watermark: time.Time{},
 		})
-		assert.Equal(t, testEventTime, result.Items()[0].EventTime)
+		assert.Equal(t, testEventTime, result.Items()[0].EventTime())
 	})
 }
diff --git a/pkg/sources/transformer/grpc_transformer_test.go b/pkg/sources/transformer/grpc_transformer_test.go
index eb92b55d..a0efefd8 100644
--- a/pkg/sources/transformer/grpc_transformer_test.go
+++ b/pkg/sources/transformer/grpc_transformer_test.go
@@ -88,7 +88,7 @@ func TestGRPCBasedTransformer_BasicApplyWithMockClient(t *testing.T) {
 
 		mockClient := funcmock.NewMockUserDefinedFunctionClient(ctrl)
 		req := &functionpb.Datum{
-			Key:       "test_success_key",
+			Keys:      []string{"test_success_key"},
 			Value:     []byte(`forward_message`),
 			EventTime: &functionpb.EventTime{EventTime: timestamppb.New(time.Unix(1661169600, 0))},
 			Watermark: &functionpb.Watermark{Watermark: timestamppb.New(time.Time{})},
@@ -96,7 +96,7 @@ func TestGRPCBasedTransformer_BasicApplyWithMockClient(t *testing.T) {
 		mockClient.EXPECT().MapTFn(gomock.Any(), &rpcMsg{msg: req}).Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "test_success_key",
+					Keys:  []string{"test_success_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -118,8 +118,8 @@ func TestGRPCBasedTransformer_BasicApplyWithMockClient(t *testing.T) {
 					MessageInfo: isb.MessageInfo{
 						EventTime: time.Unix(1661169600, 0),
 					},
-					ID:  "test_id",
-					Key: `test_success_key`,
+					ID:   "test_id",
+					Keys: []string{"test_success_key"},
 				},
 				Body: isb.Body{
 					Payload: []byte(`forward_message`),
@@ -129,7 +129,7 @@ func TestGRPCBasedTransformer_BasicApplyWithMockClient(t *testing.T) {
 		},
 		)
 		assert.NoError(t, err)
-		assert.Equal(t, req.Key, string(got[0].Key))
+		assert.Equal(t, req.Keys, got[0].Keys)
 		assert.Equal(t, req.Value, got[0].Payload)
 	})
 
@@ -139,7 +139,7 @@ func TestGRPCBasedTransformer_BasicApplyWithMockClient(t *testing.T) {
 
 		mockClient := funcmock.NewMockUserDefinedFunctionClient(ctrl)
 		req := &functionpb.Datum{
-			Key:       "test_error_key",
+			Keys:      []string{"test_error_key"},
 			Value:     []byte(`forward_message`),
 			EventTime: &functionpb.EventTime{EventTime: timestamppb.New(time.Unix(1661169660, 0))},
 			Watermark: &functionpb.Watermark{Watermark: timestamppb.New(time.Time{})},
@@ -162,8 +162,8 @@ func TestGRPCBasedTransformer_BasicApplyWithMockClient(t *testing.T) {
 					MessageInfo: isb.MessageInfo{
 						EventTime: time.Unix(1661169660, 0),
 					},
-					ID:  "test_id",
-					Key: `test_error_key`,
+					ID:   "test_id",
+					Keys: []string{"test_error_key"},
 				},
 				Body: isb.Body{
 					Payload: []byte(`forward_message`),
@@ -203,12 +203,12 @@ func TestGRPCBasedTransformer_ApplyWithMockClient_ChangePayload(t *testing.T) {
 			var elements []*functionpb.Datum
 			if originalValue.Value%2 == 0 {
 				elements = append(elements, &functionpb.Datum{
-					Key:   "even",
+					Keys:  []string{"even"},
 					Value: doubledValue,
 				})
 			} else {
 				elements = append(elements, &functionpb.Datum{
-					Key:   "odd",
+					Keys:  []string{"odd"},
 					Value: doubledValue,
 				})
 			}
@@ -234,23 +234,23 @@ func TestGRPCBasedTransformer_ApplyWithMockClient_ChangePayload(t *testing.T) {
 	readMessages := testutils.BuildTestReadMessages(count, time.Unix(1661169600, 0))
 
 	var results = make([][]byte, len(readMessages))
-	var resultKeys = make([]string, len(readMessages))
+	var resultKeys = make([][]string, len(readMessages))
 	for idx, readMessage := range readMessages {
 		apply, err := u.ApplyMap(ctx, &readMessage)
 		assert.NoError(t, err)
 		results[idx] = apply[0].Payload
-		resultKeys[idx] = string(apply[0].Header.Key)
+		resultKeys[idx] = apply[0].Header.Keys
 	}
 
 	var expectedResults = make([][]byte, count)
-	var expectedKeys = make([]string, count)
+	var expectedKeys = make([][]string, count)
 	for idx, readMessage := range readMessages {
 		var readMessagePayload testutils.PayloadForTest
 		_ = json.Unmarshal(readMessage.Payload, &readMessagePayload)
 		if readMessagePayload.Value%2 == 0 {
-			expectedKeys[idx] = "even"
+			expectedKeys[idx] = []string{"even"}
 		} else {
-			expectedKeys[idx] = "odd"
+			expectedKeys[idx] = []string{"odd"}
 		}
 		marshal, _ := json.Marshal(multiplyBy2(readMessage.Payload))
 		expectedResults[idx] = marshal
@@ -270,7 +270,7 @@ func TestGRPCBasedTransformer_ApplyWithMockClient_ChangeEventTime(t *testing.T)
 		func(_ context.Context, datum *functionpb.Datum, opts ...grpc.CallOption) (*functionpb.DatumList, error) {
 			var elements []*functionpb.Datum
 			elements = append(elements, &functionpb.Datum{
-				Key:       "even",
+				Keys:      []string{"even"},
 				Value:     datum.Value,
 				EventTime: &functionpb.EventTime{EventTime: timestamppb.New(testEventTime)},
 			})
diff --git a/pkg/udf/builtin/cat/cat_test.go b/pkg/udf/builtin/cat/cat_test.go
index 9c420f77..d7237a30 100644
--- a/pkg/udf/builtin/cat/cat_test.go
+++ b/pkg/udf/builtin/cat/cat_test.go
@@ -21,7 +21,6 @@ import (
 	"testing"
 	"time"
 
-	funcsdk "github.com/numaproj/numaflow-go/pkg/function"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -47,12 +46,10 @@ func TestNew(t *testing.T) {
 	ctx := context.Background()
 	p := New()
 	req := []byte{0}
-	messages := p(ctx, "", &testDatum{
+	messages := p(ctx, []string{""}, &testDatum{
 		value:     req,
 		eventTime: time.Time{},
 		watermark: time.Time{},
 	})
 	assert.Equal(t, 1, len(messages))
-	assert.Equal(t, funcsdk.ALL, messages[0].Key)
-	assert.Equal(t, req, messages[0].Value)
 }
diff --git a/pkg/udf/builtin/filter/filter_test.go b/pkg/udf/builtin/filter/filter_test.go
index 9e511a5e..992c6c52 100644
--- a/pkg/udf/builtin/filter/filter_test.go
+++ b/pkg/udf/builtin/filter/filter_test.go
@@ -25,7 +25,7 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-var _key = ""
+var _keys = []string{""}
 var jsonMsg = `{"test": 21, "item": [{"id": 1, "name": "numa"},{"id": 2, "name": "numa"}]}`
 var strMsg = `welcome to numaflow`
 var base64Msg = base64.StdEncoding.EncodeToString([]byte(strMsg))
@@ -61,12 +61,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, jsonMsg, string(result.Items()[0].Value))
+		assert.Equal(t, jsonMsg, string(result.Items()[0].Value()))
 	})
 
 	t.Run("invalid expression", func(t *testing.T) {
@@ -75,12 +75,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, "", string(result.Items()[0].Value))
+		assert.Equal(t, "", string(result.Items()[0].Value()))
 	})
 
 	t.Run("Json expression invalid", func(t *testing.T) {
@@ -89,12 +89,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, "", string(result.Items()[0].Value))
+		assert.Equal(t, "", string(result.Items()[0].Value()))
 	})
 
 	t.Run("String expression invalid", func(t *testing.T) {
@@ -103,12 +103,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(jsonMsg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, "", string(result.Items()[0].Value))
+		assert.Equal(t, "", string(result.Items()[0].Value()))
 	})
 
 	t.Run("base64 expression valid", func(t *testing.T) {
@@ -117,12 +117,12 @@ func TestExpression(t *testing.T) {
 		handle, err := New(args)
 		assert.NoError(t, err)
 
-		result := handle(context.Background(), _key, &testDatum{
+		result := handle(context.Background(), _keys, &testDatum{
 			value:     []byte(base64Msg),
 			eventTime: time.Time{},
 			watermark: time.Time{},
 		})
-		assert.Equal(t, base64Msg, string(result.Items()[0].Value))
+		assert.Equal(t, base64Msg, string(result.Items()[0].Value()))
 	})
 
 }
diff --git a/pkg/udf/function/uds_grpc_test.go b/pkg/udf/function/uds_grpc_test.go
index 8b9e2281..2744f423 100644
--- a/pkg/udf/function/uds_grpc_test.go
+++ b/pkg/udf/function/uds_grpc_test.go
@@ -90,7 +90,7 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 
 		mockClient := funcmock.NewMockUserDefinedFunctionClient(ctrl)
 		req := &functionpb.Datum{
-			Key:       "test_success_key",
+			Keys:      []string{"test_success_key"},
 			Value:     []byte(`forward_message`),
 			EventTime: &functionpb.EventTime{EventTime: timestamppb.New(time.Unix(1661169600, 0))},
 			Watermark: &functionpb.Watermark{Watermark: timestamppb.New(time.Time{})},
@@ -98,7 +98,7 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 		mockClient.EXPECT().MapFn(gomock.Any(), &rpcMsg{msg: req}).Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "test_success_key",
+					Keys:  []string{"test_success_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -120,8 +120,8 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 					MessageInfo: isb.MessageInfo{
 						EventTime: time.Unix(1661169600, 0),
 					},
-					ID:  "test_id",
-					Key: `test_success_key`,
+					ID:   "test_id",
+					Keys: []string{"test_success_key"},
 				},
 				Body: isb.Body{
 					Payload: []byte(`forward_message`),
@@ -131,7 +131,7 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 		},
 		)
 		assert.NoError(t, err)
-		assert.Equal(t, req.Key, string(got[0].Key))
+		assert.Equal(t, req.Keys, got[0].Keys)
 		assert.Equal(t, req.Value, got[0].Payload)
 	})
 
@@ -141,7 +141,7 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 
 		mockClient := funcmock.NewMockUserDefinedFunctionClient(ctrl)
 		req := &functionpb.Datum{
-			Key:       "test_error_key",
+			Keys:      []string{"test_error_key"},
 			Value:     []byte(`forward_message`),
 			EventTime: &functionpb.EventTime{EventTime: timestamppb.New(time.Unix(1661169660, 0))},
 			Watermark: &functionpb.Watermark{Watermark: timestamppb.New(time.Time{})},
@@ -164,8 +164,8 @@ func TestGRPCBasedUDF_BasicApplyWithMockClient(t *testing.T) {
 					MessageInfo: isb.MessageInfo{
 						EventTime: time.Unix(1661169660, 0),
 					},
-					ID:  "test_id",
-					Key: `test_error_key`,
+					ID:   "test_id",
+					Keys: []string{"test_error_key"},
 				},
 				Body: isb.Body{
 					Payload: []byte(`forward_message`),
@@ -205,12 +205,12 @@ func TestHGRPCBasedUDF_ApplyWithMockClient(t *testing.T) {
 			var elements []*functionpb.Datum
 			if originalValue.Value%2 == 0 {
 				elements = append(elements, &functionpb.Datum{
-					Key:   "even",
+					Keys:  []string{"even"},
 					Value: doubledValue,
 				})
 			} else {
 				elements = append(elements, &functionpb.Datum{
-					Key:   "odd",
+					Keys:  []string{"odd"},
 					Value: doubledValue,
 				})
 			}
@@ -236,23 +236,23 @@ func TestHGRPCBasedUDF_ApplyWithMockClient(t *testing.T) {
 	readMessages := testutils.BuildTestReadMessages(count, time.Unix(1661169600, 0))
 
 	var results = make([][]byte, len(readMessages))
-	var resultKeys = make([]string, len(readMessages))
+	var resultKeys = make([][]string, len(readMessages))
 	for idx, readMessage := range readMessages {
 		apply, err := u.ApplyMap(ctx, &readMessage)
 		assert.NoError(t, err)
 		results[idx] = apply[0].Payload
-		resultKeys[idx] = string(apply[0].Header.Key)
+		resultKeys[idx] = apply[0].Header.Keys
 	}
 
 	var expectedResults = make([][]byte, count)
-	var expectedKeys = make([]string, count)
+	var expectedKeys = make([][]string, count)
 	for idx, readMessage := range readMessages {
 		var readMessagePayload testutils.PayloadForTest
 		_ = json.Unmarshal(readMessage.Payload, &readMessagePayload)
 		if readMessagePayload.Value%2 == 0 {
-			expectedKeys[idx] = "even"
+			expectedKeys[idx] = []string{"even"}
 		} else {
-			expectedKeys[idx] = "odd"
+			expectedKeys[idx] = []string{"odd"}
 		}
 		marshal, _ := json.Marshal(multiplyBy2(readMessage.Payload))
 		expectedResults[idx] = marshal
@@ -276,7 +276,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 		mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "reduced_result_key",
+					Keys:  []string{"reduced_result_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -284,7 +284,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 		mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "reduced_result_key",
+					Keys:  []string{"reduced_result_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -336,7 +336,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 		mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "reduced_result_key",
+					Keys:  []string{"reduced_result_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -394,7 +394,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 		mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "reduced_result_key",
+					Keys:  []string{"reduced_result_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -403,7 +403,7 @@ func TestGRPCBasedUDF_BasicReduceWithMockClient(t *testing.T) {
 		mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 			Elements: []*functionpb.Datum{
 				{
-					Key:   "reduced_result_key",
+					Keys:  []string{"reduced_result_key"},
 					Value: []byte(`forward_message`),
 				},
 			},
@@ -466,7 +466,7 @@ func TestHGRPCBasedUDF_Reduce(t *testing.T) {
 			sumValue, _ := json.Marshal(result.(testutils.PayloadForTest))
 			var elements []*functionpb.Datum
 			elements = append(elements, &functionpb.Datum{
-				Key:   "sum",
+				Keys:  []string{"sum"},
 				Value: sumValue,
 			})
 			datumList := &functionpb.DatumList{
@@ -477,7 +477,7 @@ func TestHGRPCBasedUDF_Reduce(t *testing.T) {
 	mockReduceClient.EXPECT().Recv().Return(&functionpb.DatumList{
 		Elements: []*functionpb.Datum{
 			{
-				Key:   "reduced_result_key",
+				Keys:  []string{"reduced_result_key"},
 				Value: []byte(`forward_message`),
 			},
 		},
@@ -508,7 +508,7 @@ func TestHGRPCBasedUDF_Reduce(t *testing.T) {
 	_ = json.Unmarshal(result[0].Payload, &resultPayload)
 
 	assert.NoError(t, err)
-	assert.Equal(t, "sum", result[0].Key)
+	assert.Equal(t, []string{"sum"}, result[0].Keys)
 	assert.Equal(t, int64(45), resultPayload.Value)
 	assert.Equal(t, time.Unix(120, 0).Add(-1*time.Millisecond), result[0].EventTime)
 }