Repository: numaproj/numaflow
Pull request #756 (2023-05-27 08:39:53)

fix: toVertexPartitions for reduce was incorrectly populated to 1
`toVertexPartition` in the edge of the `Vertex` object was incorrectly populated to 1 by the pipeline controller because we always set `DeprecatedParallelism` to 1 if it's not set in `pipeline.ListAllEdges()`, which should not be done now.


<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/pkg/apis/numaflow/v1alpha1/pipeline_types.go b/pkg/apis/numaflow/v1alpha1/pipeline_types.go
index e348735e..5591f184 100644
--- a/pkg/apis/numaflow/v1alpha1/pipeline_types.go
+++ b/pkg/apis/numaflow/v1alpha1/pipeline_types.go
@@ -28,7 +28,6 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
-	"k8s.io/utils/pointer"
 )
 
 // +kubebuilder:validation:Enum="";Running;Succeeded;Failed;Pausing;Paused;Deleting
@@ -88,19 +87,6 @@ func (p Pipeline) ListAllEdges() []Edge {
 	edges := []Edge{}
 	for _, e := range p.Spec.Edges {
 		edgeCopy := e.DeepCopy()
-		toVertex := p.GetVertex(e.To)
-		if toVertex == nil {
-			continue
-		}
-		if toVertex.UDF == nil || toVertex.UDF.GroupBy == nil {
-			// Clean up parallelism if downstream vertex is not a reduce UDF.
-			// This has been validated by the controller, harmless to do it here.
-			edgeCopy.DeprecatedParallelism = nil
-		} else if edgeCopy.DeprecatedParallelism == nil || *edgeCopy.DeprecatedParallelism < 1 || !toVertex.UDF.GroupBy.Keyed {
-			// Set parallelism = 1 if it's not set, or it's a non-keyed reduce.
-			// Already validated by the controller to make sure parallelism is not > 1 if it's not keyed, harmless to check it again.
-			edgeCopy.DeprecatedParallelism = pointer.Int32(1)
-		}
 		edges = append(edges, *edgeCopy)
 	}
 	return edges
diff --git a/pkg/reconciler/pipeline/controller.go b/pkg/reconciler/pipeline/controller.go
index aa244de6..ba41d80a 100644
--- a/pkg/reconciler/pipeline/controller.go
+++ b/pkg/reconciler/pipeline/controller.go
@@ -529,18 +529,8 @@ func mergeLimits(plLimits dfv1.PipelineLimits, vLimits *dfv1.VertexLimits) dfv1.
 }
 
 func copyEdges(pl *dfv1.Pipeline, edges []dfv1.Edge) []dfv1.CombinedEdge {
-	plLimits := pl.GetPipelineLimits()
 	result := []dfv1.CombinedEdge{}
 	for _, e := range edges {
-		if e.DeprecatedLimits == nil {
-			e.DeprecatedLimits = &dfv1.DeprecatedEdgeLimits{}
-		}
-		if e.DeprecatedLimits.BufferMaxLength == nil {
-			e.DeprecatedLimits.BufferMaxLength = plLimits.BufferMaxLength
-		}
-		if e.DeprecatedLimits.BufferUsageLimit == nil {
-			e.DeprecatedLimits.BufferUsageLimit = plLimits.BufferUsageLimit
-		}
 		vFrom := pl.GetVertex(e.From)
 		vTo := pl.GetVertex(e.To)
 		fromVertexLimits := mergeLimits(pl.GetPipelineLimits(), vFrom.Limits)
diff --git a/pkg/sources/source.go b/pkg/sources/source.go
index f54549b2..2a820c14 100644
--- a/pkg/sources/source.go
+++ b/pkg/sources/source.go
@@ -86,10 +86,17 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 				writeOpts = append(writeOpts, redisclient.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 			}
 			buffers := dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitions())
-			for _, buffer := range buffers {
-				group := buffer + "-group"
+			if e.ToVertexType == dfv1.VertexTypeReduceUDF {
+				for _, buffer := range buffers {
+					group := buffer + "-group"
+					redisClient := redisclient.NewInClusterRedisClient()
+					writer := redisisb.NewBufferWrite(ctx, redisClient, buffer, group, writeOpts...)
+					writers = append(writers, writer)
+				}
+			} else {
+				group := buffers[0] + "-group"
 				redisClient := redisclient.NewInClusterRedisClient()
-				writer := redisisb.NewBufferWrite(ctx, redisClient, buffer, group, writeOpts...)
+				writer := redisisb.NewBufferWrite(ctx, redisClient, buffers[0], group, writeOpts...)
 				writers = append(writers, writer)
 			}
 		}
@@ -122,10 +129,20 @@ func (sp *SourceProcessor) Start(ctx context.Context) error {
 				writeOpts = append(writeOpts, jetstreamisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 			}
 			buffers := dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitions())
-			for _, buffer := range buffers {
-				streamName := isbsvc.JetStreamName(buffer)
+			if e.ToVertexType == dfv1.VertexTypeReduceUDF {
+				for _, buffer := range buffers {
+					streamName := isbsvc.JetStreamName(buffer)
+					jetStreamClient := jsclient.NewInClusterJetStreamClient()
+					writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jetStreamClient, buffer, streamName, streamName, writeOpts...)
+					if err != nil {
+						return err
+					}
+					writers = append(writers, writer)
+				}
+			} else {
+				streamName := isbsvc.JetStreamName(buffers[0])
 				jetStreamClient := jsclient.NewInClusterJetStreamClient()
-				writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jetStreamClient, buffer, streamName, streamName, writeOpts...)
+				writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jetStreamClient, buffers[0], streamName, streamName, writeOpts...)
 				if err != nil {
 					return err
 				}
@@ -260,16 +277,16 @@ func (sp *SourceProcessor) getTransformerGoWhereDecider() forward.GoWhere {
 				if edge.ToVertexType == dfv1.VertexTypeReduceUDF && edge.GetToVertexPartitions() > 1 { // Need to shuffle
 					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 				} else {
-					// TODO: need to shuffle for partitioned map vertex
-					result = append(result, dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, edge.To, edge.GetToVertexPartitions())...)
+					// TODO: need to shuffle for partitioned map vertex, for now write only to the first partition
+					result = append(result, dfv1.GenerateBufferName(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, edge.To, 0))
 				}
 			} else {
 				if sharedutil.CompareSlice(edge.Conditions.Tags.GetOperator(), tags, edge.Conditions.Tags.Values) {
 					if edge.ToVertexType == dfv1.VertexTypeReduceUDF && edge.GetToVertexPartitions() > 1 { // Need to shuffle
 						result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 					} else {
-						// TODO: need to shuffle for partitioned map vertex
-						result = append(result, dfv1.GenerateBufferNames(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, edge.To, edge.GetToVertexPartitions())...)
+						// TODO: need to shuffle for partitioned map vertex, for now write only to the first partition
+						result = append(result, dfv1.GenerateBufferName(sp.VertexInstance.Vertex.Namespace, sp.VertexInstance.Vertex.Spec.PipelineName, edge.To, 0))
 					}
 				}
 			}
diff --git a/pkg/udf/common.go b/pkg/udf/common.go
index f3b4f32a..49e0ed85 100644
--- a/pkg/udf/common.go
+++ b/pkg/udf/common.go
@@ -56,10 +56,16 @@ func buildRedisBufferIO(ctx context.Context, fromBufferName string, vertexInstan
 			// TODO: remove this after deprecation period
 			writeOpts = append(writeOpts, redisclient.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 		}
+		// TODO: support multiple partitions
 		buffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitions())
-		for _, buffer := range buffers {
-			writer := redisisb.NewBufferWrite(ctx, redisClient, buffer, buffer+"-group", writeOpts...)
-			writers[buffer] = writer
+		if e.ToVertexType == dfv1.VertexTypeReduceUDF {
+			for _, buffer := range buffers {
+				writer := redisisb.NewBufferWrite(ctx, redisClient, buffer, buffer+"-group", writeOpts...)
+				writers[buffer] = writer
+			}
+		} else {
+			writer := redisisb.NewBufferWrite(ctx, redisClient, buffers[0], buffers[0]+"-group", writeOpts...)
+			writers[buffers[0]] = writer
 		}
 	}
 
@@ -96,14 +102,24 @@ func buildJetStreamBufferIO(ctx context.Context, fromBufferName string, vertexIn
 			// TODO: remove this after deprecation period
 			writeOpts = append(writeOpts, jetstreamisb.WithBufferUsageLimit(float64(*x.BufferUsageLimit)/100))
 		}
+		// TODO: support multiple partitions
 		buffers := dfv1.GenerateBufferNames(vertexInstance.Vertex.Namespace, vertexInstance.Vertex.Spec.PipelineName, e.To, e.GetToVertexPartitions())
-		for _, buffer := range buffers {
-			streamName := isbsvc.JetStreamName(buffer)
-			writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jsclient.NewInClusterJetStreamClient(), buffer, streamName, streamName, writeOpts...)
+		if e.ToVertexType == dfv1.VertexTypeReduceUDF {
+			for _, buffer := range buffers {
+				streamName := isbsvc.JetStreamName(buffer)
+				writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jsclient.NewInClusterJetStreamClient(), buffer, streamName, streamName, writeOpts...)
+				if err != nil {
+					return nil, nil, err
+				}
+				writers[buffer] = writer
+			}
+		} else {
+			streamName := isbsvc.JetStreamName(buffers[0])
+			writer, err := jetstreamisb.NewJetStreamBufferWriter(ctx, jsclient.NewInClusterJetStreamClient(), buffers[0], streamName, streamName, writeOpts...)
 			if err != nil {
 				return nil, nil, err
 			}
-			writers[buffer] = writer
+			writers[buffers[0]] = writer
 		}
 	}
 	return reader, writers, nil
diff --git a/pkg/udf/function/uds_grpc.go b/pkg/udf/function/uds_grpc.go
index 51c23ae2..9111613b 100644
--- a/pkg/udf/function/uds_grpc.go
+++ b/pkg/udf/function/uds_grpc.go
@@ -19,13 +19,13 @@ package function
 import (
 	"context"
 	"fmt"
-	"golang.org/x/sync/errgroup"
 	"strconv"
 	"sync"
 	"time"
 
 	functionpb "github.com/numaproj/numaflow-go/pkg/apis/proto/function/v1"
 	functionsdk "github.com/numaproj/numaflow-go/pkg/function"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/timestamppb"
diff --git a/pkg/udf/map_udf.go b/pkg/udf/map_udf.go
index b3a48708..2aeac0db 100644
--- a/pkg/udf/map_udf.go
+++ b/pkg/udf/map_udf.go
@@ -92,16 +92,16 @@ func (u *MapUDFProcessor) Start(ctx context.Context) error {
 				if edge.ToVertexType == dfv1.VertexTypeReduceUDF && edge.GetToVertexPartitions() > 1 { // Need to shuffle
 					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 				} else {
-					// TODO: need to shuffle for partitioned map vertex
-					result = append(result, dfv1.GenerateBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, edge.GetToVertexPartitions())...)
+					// TODO: need to shuffle for partitioned map vertex, for now write only to the first partition
+					result = append(result, dfv1.GenerateBufferName(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, 0))
 				}
 			} else {
 				if sharedutil.CompareSlice(edge.Conditions.Tags.GetOperator(), tags, edge.Conditions.Tags.Values) {
 					if edge.ToVertexType == dfv1.VertexTypeReduceUDF && edge.GetToVertexPartitions() > 1 { // Need to shuffle
 						result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 					} else {
-						// TODO: need to shuffle for partitioned map vertex
-						result = append(result, dfv1.GenerateBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, edge.GetToVertexPartitions())...)
+						// TODO: need to shuffle for partitioned map vertex, for now write only to the first partition
+						result = append(result, dfv1.GenerateBufferName(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, 0))
 					}
 				}
 			}
diff --git a/pkg/udf/reduce_udf.go b/pkg/udf/reduce_udf.go
index 333def7d..23c7322a 100644
--- a/pkg/udf/reduce_udf.go
+++ b/pkg/udf/reduce_udf.go
@@ -123,16 +123,16 @@ func (u *ReduceUDFProcessor) Start(ctx context.Context) error {
 				if edge.ToVertexType == dfv1.VertexTypeReduceUDF && edge.GetToVertexPartitions() > 1 { // Need to shuffle
 					result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 				} else {
-					// TODO: need to shuffle for partitioned map vertex
-					result = append(result, dfv1.GenerateBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, edge.GetToVertexPartitions())...)
+					// TODO: need to shuffle for partitioned map vertex, for now write only to the first partition
+					result = append(result, dfv1.GenerateBufferName(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, 0))
 				}
 			} else {
 				if sharedutil.CompareSlice(edge.Conditions.Tags.GetOperator(), tags, edge.Conditions.Tags.Values) {
 					if edge.ToVertexType == dfv1.VertexTypeReduceUDF && edge.GetToVertexPartitions() > 1 { // Need to shuffle
 						result = append(result, shuffleFuncMap[fmt.Sprintf("%s:%s", edge.From, edge.To)].Shuffle(keys))
 					} else {
-						// TODO: need to shuffle for partitioned map vertex
-						result = append(result, dfv1.GenerateBufferNames(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, edge.GetToVertexPartitions())...)
+						// TODO: need to shuffle for partitioned map vertex, for now write only to the first partition
+						result = append(result, dfv1.GenerateBufferName(u.VertexInstance.Vertex.Namespace, u.VertexInstance.Vertex.Spec.PipelineName, edge.To, 0))
 					}
 				}
 			}
diff --git a/pkg/watermark/fetch/edge_fetcher.go b/pkg/watermark/fetch/edge_fetcher.go
index e188eef5..2af99cdc 100644
--- a/pkg/watermark/fetch/edge_fetcher.go
+++ b/pkg/watermark/fetch/edge_fetcher.go
@@ -41,19 +41,19 @@ import (
 // edgeFetcher is a fetcher between two vertices.
 type edgeFetcher struct {
 	ctx              context.Context
-	bufferName       string
+	bucketName       string
 	storeWatcher     store.WatermarkStoreWatcher
 	processorManager *processor.ProcessorManager
 	log              *zap.SugaredLogger
 }
 
 // NewEdgeFetcher returns a new edge fetcher.
-func NewEdgeFetcher(ctx context.Context, bufferName string, storeWatcher store.WatermarkStoreWatcher, manager *processor.ProcessorManager) Fetcher {
-	log := logging.FromContext(ctx).With("bufferName", bufferName)
+func NewEdgeFetcher(ctx context.Context, bucketName string, storeWatcher store.WatermarkStoreWatcher, manager *processor.ProcessorManager) Fetcher {
+	log := logging.FromContext(ctx).With("bucketName", bucketName)
 	log.Info("Creating a new edge watermark fetcher")
 	return &edgeFetcher{
 		ctx:              ctx,
-		bufferName:       bufferName,
+		bucketName:       bucketName,
 		storeWatcher:     storeWatcher,
 		processorManager: manager,
 		log:              log,
@@ -87,7 +87,7 @@ func (e *edgeFetcher) GetWatermark(inputOffset isb.Offset) wmb.Watermark {
 	if epoch == math.MaxInt64 {
 		epoch = -1
 	}
-	e.log.Debugf("%s[%s] get watermark for offset %d: %+v", debugString.String(), e.bufferName, offset, epoch)
+	e.log.Debugf("%s[%s] get watermark for offset %d: %+v", debugString.String(), e.bucketName, offset, epoch)
 
 	return wmb.Watermark(time.UnixMilli(epoch))
 }
@@ -144,7 +144,7 @@ func (e *edgeFetcher) GetHeadWMB() wmb.WMB {
 		// we only consider the latest wmb in the offset timeline
 		var curHeadWMB = p.GetOffsetTimeline().GetHeadWMB()
 		if !curHeadWMB.Idle {
-			e.log.Debugf("[%s] GetHeadWMB finds an active head wmb for offset, return early", e.bufferName)
+			e.log.Debugf("[%s] GetHeadWMB finds an active head wmb for offset, return early", e.bucketName)
 			return wmb.WMB{}
 		}
 		if curHeadWMB.Watermark != -1 {
@@ -160,7 +160,7 @@ func (e *edgeFetcher) GetHeadWMB() wmb.WMB {
 		// there is no valid watermark yet
 		return wmb.WMB{}
 	}
-	e.log.Debugf("GetHeadWMB: %s[%s] get idle head wmb for offset", debugString.String(), e.bufferName)
+	e.log.Debugf("GetHeadWMB: %s[%s] get idle head wmb for offset", debugString.String(), e.bucketName)
 	return headWMB
 }
 


Test patch:
diff --git a/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go b/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go
index 2030da06..111227f8 100644
--- a/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go
+++ b/pkg/apis/numaflow/v1alpha1/pipeline_types_test.go
@@ -49,23 +49,13 @@ var (
 func Test_ListAllEdges(t *testing.T) {
 	es := testPipeline.ListAllEdges()
 	assert.Equal(t, 2, len(es))
-	assert.Nil(t, es[0].DeprecatedParallelism)
-	assert.Nil(t, es[1].DeprecatedParallelism)
 	pl := testPipeline.DeepCopy()
 	pl.Spec.Vertices[1].UDF.GroupBy = &GroupBy{}
 	es = pl.ListAllEdges()
 	assert.Equal(t, 2, len(es))
-	assert.NotNil(t, es[0].DeprecatedParallelism)
-	assert.Equal(t, int32(1), *es[0].DeprecatedParallelism)
-	assert.Nil(t, es[1].DeprecatedParallelism)
 	pl.Spec.Edges[0].DeprecatedParallelism = pointer.Int32(3)
 	es = pl.ListAllEdges()
 	assert.Equal(t, 2, len(es))
-	assert.NotNil(t, es[0].DeprecatedParallelism)
-	assert.Equal(t, int32(1), *es[0].DeprecatedParallelism)
-	pl.Spec.Vertices[1].UDF.GroupBy.Keyed = true
-	es = pl.ListAllEdges()
-	assert.Equal(t, int32(3), *es[0].DeprecatedParallelism)
 }
 
 func Test_GetToEdges(t *testing.T) {
diff --git a/pkg/reconciler/pipeline/controller_test.go b/pkg/reconciler/pipeline/controller_test.go
index d4b63506..d77816f2 100644
--- a/pkg/reconciler/pipeline/controller_test.go
+++ b/pkg/reconciler/pipeline/controller_test.go
@@ -194,34 +194,58 @@ func Test_copyVertexLimits(t *testing.T) {
 }
 
 func Test_copyEdges(t *testing.T) {
-	pl := testPipeline.DeepCopy()
-	edges := []dfv1.Edge{{From: "input", To: "p1"}}
-	result := copyEdges(pl, edges)
-	for _, e := range result {
-		assert.NotNil(t, e.ToVertexLimits)
-		assert.Equal(t, int64(dfv1.DefaultBufferLength), int64(*e.ToVertexLimits.BufferMaxLength))
-	}
-	onethouand := uint64(1000)
-	eighty := uint32(80)
-	pl.Spec.Limits = &dfv1.PipelineLimits{BufferMaxLength: &onethouand, BufferUsageLimit: &eighty}
-	result = copyEdges(pl, edges)
-	for _, e := range result {
-		assert.NotNil(t, e.ToVertexLimits)
-		assert.NotNil(t, e.ToVertexLimits.BufferMaxLength)
-		assert.NotNil(t, e.ToVertexLimits.BufferUsageLimit)
-	}
+	t.Run("test copy map", func(t *testing.T) {
+		pl := testPipeline.DeepCopy()
+		edges := []dfv1.Edge{{From: "input", To: "p1"}}
+		result := copyEdges(pl, edges)
+		for _, e := range result {
+			assert.NotNil(t, e.ToVertexLimits)
+			assert.Equal(t, int64(dfv1.DefaultBufferLength), int64(*e.ToVertexLimits.BufferMaxLength))
+		}
+		onethouand := uint64(1000)
+		eighty := uint32(80)
+		pl.Spec.Limits = &dfv1.PipelineLimits{BufferMaxLength: &onethouand, BufferUsageLimit: &eighty}
+		result = copyEdges(pl, edges)
+		for _, e := range result {
+			assert.NotNil(t, e.ToVertexLimits)
+			assert.NotNil(t, e.ToVertexLimits.BufferMaxLength)
+			assert.NotNil(t, e.ToVertexLimits.BufferUsageLimit)
+		}
+
+		twothouand := uint64(2000)
+		pl.Spec.Vertices[2].Limits = &dfv1.VertexLimits{BufferMaxLength: &twothouand}
+		edges = []dfv1.Edge{{From: "p1", To: "output"}}
+		result = copyEdges(pl, edges)
+		for _, e := range result {
+			assert.NotNil(t, e.ToVertexLimits)
+			assert.NotNil(t, e.ToVertexLimits.BufferMaxLength)
+			assert.Equal(t, twothouand, *e.ToVertexLimits.BufferMaxLength)
+			assert.NotNil(t, e.ToVertexLimits.BufferUsageLimit)
+			assert.Equal(t, eighty, *e.ToVertexLimits.BufferUsageLimit)
+		}
+	})
+
+	t.Run("test copy reduce", func(t *testing.T) {
+		pl := testReducePipeline.DeepCopy()
+		edges := []dfv1.Edge{{From: "p1", To: "p2"}}
+		result := copyEdges(pl, edges)
+		assert.Equal(t, 1, len(result))
+		assert.Equal(t, "p1", result[0].From)
+		assert.Equal(t, "p2", result[0].To)
+		assert.NotNil(t, result[0].ToVertexLimits)
+		assert.Equal(t, int64(dfv1.DefaultBufferLength), int64(*result[0].ToVertexLimits.BufferMaxLength))
+		assert.Equal(t, int32(2), *result[0].ToVertexPartitions)
+		assert.Equal(t, int32(1), *result[0].FromVertexPartitions)
+
+		edges = []dfv1.Edge{{From: "p2", To: "p3"}}
+		result = copyEdges(pl, edges)
+		assert.Equal(t, 1, len(result))
+		assert.Equal(t, "p2", result[0].From)
+		assert.Equal(t, "p3", result[0].To)
+		assert.Equal(t, int32(1), *result[0].ToVertexPartitions)
+		assert.Equal(t, int32(2), *result[0].FromVertexPartitions)
+	})
 
-	twothouand := uint64(2000)
-	pl.Spec.Vertices[2].Limits = &dfv1.VertexLimits{BufferMaxLength: &twothouand}
-	edges = []dfv1.Edge{{From: "p1", To: "output"}}
-	result = copyEdges(pl, edges)
-	for _, e := range result {
-		assert.NotNil(t, e.ToVertexLimits)
-		assert.NotNil(t, e.ToVertexLimits.BufferMaxLength)
-		assert.Equal(t, twothouand, *e.ToVertexLimits.BufferMaxLength)
-		assert.NotNil(t, e.ToVertexLimits.BufferUsageLimit)
-		assert.Equal(t, eighty, *e.ToVertexLimits.BufferUsageLimit)
-	}
 }
 
 func Test_buildISBBatchJob(t *testing.T) {
diff --git a/pkg/reconciler/pipeline/validate_test.go b/pkg/reconciler/pipeline/validate_test.go
index 27bc0da3..2bbc1d6c 100644
--- a/pkg/reconciler/pipeline/validate_test.go
+++ b/pkg/reconciler/pipeline/validate_test.go
@@ -102,7 +102,8 @@ var (
 					},
 				},
 				{
-					Name: "p2",
+					Name:       "p2",
+					Partitions: pointer.Int32(2),
 					UDF: &dfv1.UDF{
 						Container: &dfv1.Container{
 							Image: "my-image",
@@ -115,6 +116,7 @@ var (
 									},
 								},
 							},
+							Keyed: true,
 							Storage: &dfv1.PBQStorage{
 								EmptyDir: &corev1.EmptyDirVolumeSource{},
 							},
diff --git a/pkg/watermark/fetch/edge_fetcher_test.go b/pkg/watermark/fetch/edge_fetcher_test.go
index 89182b36..9b56a539 100644
--- a/pkg/watermark/fetch/edge_fetcher_test.go
+++ b/pkg/watermark/fetch/edge_fetcher_test.go
@@ -144,7 +144,7 @@ func TestBuffer_GetWatermark(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			b := &edgeFetcher{
 				ctx:              ctx,
-				bufferName:       "testBuffer",
+				bucketName:       "testBucket",
 				processorManager: tt.processorManager,
 				log:              zaptest.NewLogger(t).Sugar(),
 			}
@@ -161,7 +161,7 @@ func TestBuffer_GetWatermark(t *testing.T) {
 func Test_edgeFetcher_GetHeadWatermark(t *testing.T) {
 	var (
 		ctx               = context.Background()
-		bufferName        = "testBuffer"
+		bucketName        = "testBucket"
 		hbWatcher         = noop.NewKVOpWatch()
 		otWatcher         = noop.NewKVOpWatch()
 		storeWatcher      = store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher)
@@ -192,7 +192,7 @@ func Test_edgeFetcher_GetHeadWatermark(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			e := &edgeFetcher{
 				ctx:              ctx,
-				bufferName:       bufferName,
+				bucketName:       bucketName,
 				storeWatcher:     storeWatcher,
 				processorManager: tt.processorManager,
 				log:              zaptest.NewLogger(t).Sugar(),
@@ -289,7 +289,7 @@ func getHeadWMTest2(ctx context.Context, processorManager1 *processor.ProcessorM
 func Test_edgeFetcher_GetHeadWMB(t *testing.T) {
 	var (
 		ctx               = context.Background()
-		bufferName        = "testBuffer"
+		bucketName        = "testBucket"
 		hbWatcher         = noop.NewKVOpWatch()
 		otWatcher         = noop.NewKVOpWatch()
 		storeWatcher      = store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher)
@@ -338,7 +338,7 @@ func Test_edgeFetcher_GetHeadWMB(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			e := &edgeFetcher{
 				ctx:              ctx,
-				bufferName:       bufferName,
+				bucketName:       bucketName,
 				storeWatcher:     storeWatcher,
 				processorManager: tt.processorManager,
 				log:              zaptest.NewLogger(t).Sugar(),
@@ -534,7 +534,7 @@ func TestFetcherWithSameOTBucket_InMem(t *testing.T) {
 	assert.NoError(t, err)
 	storeWatcher := store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher)
 	var processorManager = processor.NewProcessorManager(ctx, storeWatcher)
-	var fetcher = NewEdgeFetcher(ctx, "testBuffer", storeWatcher, processorManager)
+	var fetcher = NewEdgeFetcher(ctx, "testBucket", storeWatcher, processorManager)
 	// start p1 heartbeat for 3 loops
 	wg.Add(1)
 	go func() {
@@ -791,7 +791,7 @@ func TestFetcherWithSameOTBucket(t *testing.T) {
 	assert.NoError(t, err)
 	storeWatcher := store.BuildWatermarkStoreWatcher(hbWatcher, otWatcher)
 	processorManager := processor.NewProcessorManager(ctx, storeWatcher)
-	fetcher := NewEdgeFetcher(ctx, "testBuffer", storeWatcher, processorManager)
+	fetcher := NewEdgeFetcher(ctx, "testBucket", storeWatcher, processorManager)
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
diff --git a/pkg/watermark/processor/processor_manager_test.go b/pkg/watermark/processor/processor_manager_test.go
index e12d1c3a..ed59d36d 100644
--- a/pkg/watermark/processor/processor_manager_test.go
+++ b/pkg/watermark/processor/processor_manager_test.go
@@ -213,7 +213,7 @@ func TestProcessorManagerWatchForMap(t *testing.T) {
 		err = otStore.PutKV(ctx, "p1", otValueByte)
 		assert.NoError(t, err)
 	}
-	for processorManager.GetProcessor("p1").GetOffsetTimeline().GetHeadOffset() == -1 || processorManager.GetProcessor("p2").GetOffsetTimeline().GetHeadOffset() == -1 {
+	for processorManager.GetProcessor("p1").GetOffsetTimeline().GetHeadOffset() != 115 || processorManager.GetProcessor("p2").GetOffsetTimeline().GetHeadOffset() != 115 {
 		select {
 		case <-ctx.Done():
 			if ctx.Err() == context.DeadlineExceeded {
@@ -333,7 +333,7 @@ func TestProcessorManagerWatchForReduce(t *testing.T) {
 		err = otStore.PutKV(ctx, "p1", otValueByte)
 		assert.NoError(t, err)
 	}
-	for processorManager.GetProcessor("p1").GetOffsetTimeline().GetHeadOffset() == -1 || processorManager.GetProcessor("p2").GetOffsetTimeline().GetHeadOffset() == -1 {
+	for processorManager.GetProcessor("p1").GetOffsetTimeline().GetHeadOffset() != 115 || processorManager.GetProcessor("p2").GetOffsetTimeline().GetHeadOffset() != 115 {
 		select {
 		case <-ctx.Done():
 			if ctx.Err() == context.DeadlineExceeded {