Repository: numaproj/numaflow
Pull request #438 (2023-01-03 15:28:28)

chore: move pbq under reduce
Signed-off-by: Vigith Maurice <vigith@gmail.com>

Move `pbq` under `reduce`. `pbq` is only used by the reduce


Patch:
diff --git a/pkg/pbq/interfaces.go b/pkg/reduce/pbq/interfaces.go
similarity index 100%
rename from pkg/pbq/interfaces.go
rename to pkg/reduce/pbq/interfaces.go
diff --git a/pkg/pbq/metrics.go b/pkg/reduce/pbq/metrics.go
similarity index 100%
rename from pkg/pbq/metrics.go
rename to pkg/reduce/pbq/metrics.go
diff --git a/pkg/pbq/options.go b/pkg/reduce/pbq/options.go
similarity index 100%
rename from pkg/pbq/options.go
rename to pkg/reduce/pbq/options.go
diff --git a/pkg/pbq/options_test.go b/pkg/reduce/pbq/options_test.go
similarity index 100%
rename from pkg/pbq/options_test.go
rename to pkg/reduce/pbq/options_test.go
diff --git a/pkg/pbq/partition/partition.go b/pkg/reduce/pbq/partition/partition.go
similarity index 100%
rename from pkg/pbq/partition/partition.go
rename to pkg/reduce/pbq/partition/partition.go
diff --git a/pkg/pbq/pbq.go b/pkg/reduce/pbq/pbq.go
similarity index 97%
rename from pkg/pbq/pbq.go
rename to pkg/reduce/pbq/pbq.go
index d5d6fb38..7c0d6867 100644
--- a/pkg/pbq/pbq.go
+++ b/pkg/reduce/pbq/pbq.go
@@ -21,11 +21,12 @@ import (
 	"errors"
 
 	"github.com/numaproj/numaflow/pkg/metrics"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
+
 	"go.uber.org/zap"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
 )
 
 var ErrCOB = errors.New("error while writing to pbq, pbq is closed")
diff --git a/pkg/pbq/pbqmanager.go b/pkg/reduce/pbq/pbqmanager.go
similarity index 98%
rename from pkg/pbq/pbqmanager.go
rename to pkg/reduce/pbq/pbqmanager.go
index ad8cf9e0..d71493cd 100644
--- a/pkg/pbq/pbqmanager.go
+++ b/pkg/reduce/pbq/pbqmanager.go
@@ -24,12 +24,13 @@ import (
 	"time"
 
 	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
+
 	"go.uber.org/zap"
 	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 )
 
diff --git a/pkg/pbq/store/errors.go b/pkg/reduce/pbq/store/errors.go
similarity index 100%
rename from pkg/pbq/store/errors.go
rename to pkg/reduce/pbq/store/errors.go
diff --git a/pkg/pbq/store/interfaces.go b/pkg/reduce/pbq/store/interfaces.go
similarity index 96%
rename from pkg/pbq/store/interfaces.go
rename to pkg/reduce/pbq/store/interfaces.go
index 912782c7..76b3dcb3 100644
--- a/pkg/pbq/store/interfaces.go
+++ b/pkg/reduce/pbq/store/interfaces.go
@@ -20,7 +20,7 @@ import (
 	"context"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 // Store provides methods to read, write and delete data from the store.
diff --git a/pkg/pbq/store/memory/options.go b/pkg/reduce/pbq/store/memory/options.go
similarity index 94%
rename from pkg/pbq/store/memory/options.go
rename to pkg/reduce/pbq/store/memory/options.go
index 2696d469..32454aaa 100644
--- a/pkg/pbq/store/memory/options.go
+++ b/pkg/reduce/pbq/store/memory/options.go
@@ -19,7 +19,7 @@ package memory
 import (
 	"context"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 type Option func(stores *memoryStores)
diff --git a/pkg/pbq/store/memory/store.go b/pkg/reduce/pbq/store/memory/store.go
similarity index 95%
rename from pkg/pbq/store/memory/store.go
rename to pkg/reduce/pbq/store/memory/store.go
index 5128d254..877ecc76 100644
--- a/pkg/pbq/store/memory/store.go
+++ b/pkg/reduce/pbq/store/memory/store.go
@@ -20,8 +20,9 @@ import (
 	"math"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
+
 	"go.uber.org/zap"
 )
 
diff --git a/pkg/pbq/store/memory/stores.go b/pkg/reduce/pbq/store/memory/stores.go
similarity index 95%
rename from pkg/pbq/store/memory/stores.go
rename to pkg/reduce/pbq/store/memory/stores.go
index 264e2bb9..b6a63bf9 100644
--- a/pkg/pbq/store/memory/stores.go
+++ b/pkg/reduce/pbq/store/memory/stores.go
@@ -22,8 +22,8 @@ import (
 	"sync"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 )
 
diff --git a/pkg/pbq/store/noop/store.go b/pkg/reduce/pbq/store/noop/store.go
similarity index 95%
rename from pkg/pbq/store/noop/store.go
rename to pkg/reduce/pbq/store/noop/store.go
index 3f97bdf9..6b353163 100644
--- a/pkg/pbq/store/noop/store.go
+++ b/pkg/reduce/pbq/store/noop/store.go
@@ -18,7 +18,7 @@ package noop
 
 import (
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
 )
 
 // PBQNoOpStore is a no-op pbq store which does not do any operation but can be safely invoked.
diff --git a/pkg/pbq/store/noop/stores.go b/pkg/reduce/pbq/store/noop/stores.go
similarity index 90%
rename from pkg/pbq/store/noop/stores.go
rename to pkg/reduce/pbq/store/noop/stores.go
index 9e392718..c50c4a18 100644
--- a/pkg/pbq/store/noop/stores.go
+++ b/pkg/reduce/pbq/store/noop/stores.go
@@ -19,8 +19,8 @@ package noop
 import (
 	"context"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
 )
 
 type noopStores struct {
diff --git a/pkg/pbq/store/wal/bootstrap.go b/pkg/reduce/pbq/store/wal/bootstrap.go
similarity index 98%
rename from pkg/pbq/store/wal/bootstrap.go
rename to pkg/reduce/pbq/store/wal/bootstrap.go
index b07cb702..29bb3afb 100644
--- a/pkg/pbq/store/wal/bootstrap.go
+++ b/pkg/reduce/pbq/store/wal/bootstrap.go
@@ -26,7 +26,7 @@ import (
 	"time"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 var location *time.Location
diff --git a/pkg/pbq/store/wal/doc.go b/pkg/reduce/pbq/store/wal/doc.go
similarity index 100%
rename from pkg/pbq/store/wal/doc.go
rename to pkg/reduce/pbq/store/wal/doc.go
diff --git a/pkg/pbq/store/wal/metrics.go b/pkg/reduce/pbq/store/wal/metrics.go
similarity index 99%
rename from pkg/pbq/store/wal/metrics.go
rename to pkg/reduce/pbq/store/wal/metrics.go
index 0f0d8945..31d2256e 100644
--- a/pkg/pbq/store/wal/metrics.go
+++ b/pkg/reduce/pbq/store/wal/metrics.go
@@ -17,9 +17,10 @@ limitations under the License.
 package wal
 
 import (
-	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+
+	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
 )
 
 const (
diff --git a/pkg/pbq/store/wal/options.go b/pkg/reduce/pbq/store/wal/options.go
similarity index 100%
rename from pkg/pbq/store/wal/options.go
rename to pkg/reduce/pbq/store/wal/options.go
diff --git a/pkg/pbq/store/wal/segment.go b/pkg/reduce/pbq/store/wal/segment.go
similarity index 99%
rename from pkg/pbq/store/wal/segment.go
rename to pkg/reduce/pbq/store/wal/segment.go
index b715d655..c68d83a3 100644
--- a/pkg/pbq/store/wal/segment.go
+++ b/pkg/reduce/pbq/store/wal/segment.go
@@ -29,7 +29,7 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 const (
diff --git a/pkg/pbq/store/wal/stores.go b/pkg/reduce/pbq/store/wal/stores.go
similarity index 98%
rename from pkg/pbq/store/wal/stores.go
rename to pkg/reduce/pbq/store/wal/stores.go
index 62489a31..a1ba0957 100644
--- a/pkg/pbq/store/wal/stores.go
+++ b/pkg/reduce/pbq/store/wal/stores.go
@@ -27,8 +27,8 @@ import (
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
 )
 
 type walStores struct {
diff --git a/pkg/reduce/pnf/processandforward.go b/pkg/reduce/pnf/processandforward.go
index f36763d8..8e196ee7 100644
--- a/pkg/reduce/pnf/processandforward.go
+++ b/pkg/reduce/pnf/processandforward.go
@@ -28,6 +28,9 @@ import (
 	"time"
 
 	metricspkg "github.com/numaproj/numaflow/pkg/metrics"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+
 	"go.uber.org/zap"
 	"k8s.io/apimachinery/pkg/util/wait"
 
@@ -39,8 +42,6 @@ import (
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 )
 
diff --git a/pkg/reduce/readloop/ordered.go b/pkg/reduce/readloop/ordered.go
index afbf480e..3f2a818d 100644
--- a/pkg/reduce/readloop/ordered.go
+++ b/pkg/reduce/readloop/ordered.go
@@ -25,13 +25,13 @@ import (
 	"go.uber.org/zap"
 
 	"github.com/numaproj/numaflow/pkg/metrics"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
 
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/reduce/pnf"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
diff --git a/pkg/reduce/readloop/readloop.go b/pkg/reduce/readloop/readloop.go
index 49fb617f..01a6e537 100644
--- a/pkg/reduce/readloop/readloop.go
+++ b/pkg/reduce/readloop/readloop.go
@@ -37,11 +37,11 @@ import (
 	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/numaproj/numaflow/pkg/metrics"
+	pbq2 "github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
@@ -55,7 +55,7 @@ type ReadLoop struct {
 	vertexName            string
 	pipelineName          string
 	UDF                   applier.ReduceApplier
-	pbqManager            *pbq.Manager
+	pbqManager            *pbq2.Manager
 	windower              window.Windower
 	op                    *orderedForwarder
 	log                   *zap.SugaredLogger
@@ -70,7 +70,7 @@ func NewReadLoop(ctx context.Context,
 	vertexName string,
 	pipelineName string,
 	udf applier.ReduceApplier,
-	pbqManager *pbq.Manager,
+	pbqManager *pbq2.Manager,
 	windowingStrategy window.Windower,
 	toBuffers map[string]isb.BufferWriter,
 	whereToDecider forward.ToWhichStepDecider,
@@ -296,7 +296,7 @@ func (rl *ReadLoop) ackMessages(ctx context.Context, messages []*isb.ReadMessage
 
 // associatePBQAndPnF associates a PBQ with the partition if a PBQ exists, else creates a new one and then associates
 // it to the partition.
-func (rl *ReadLoop) associatePBQAndPnF(ctx context.Context, partitionID partition.ID) pbq.ReadWriteCloser {
+func (rl *ReadLoop) associatePBQAndPnF(ctx context.Context, partitionID partition.ID) pbq2.ReadWriteCloser {
 	// look for existing pbq
 	q := rl.pbqManager.GetPBQ(partitionID)
 
diff --git a/pkg/reduce/reduce.go b/pkg/reduce/reduce.go
index 6aa75a52..c49228c7 100644
--- a/pkg/reduce/reduce.go
+++ b/pkg/reduce/reduce.go
@@ -24,12 +24,13 @@ import (
 	"time"
 
 	"github.com/numaproj/numaflow/pkg/metrics"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq"
+
 	"go.uber.org/zap"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
-	"github.com/numaproj/numaflow/pkg/pbq"
 	"github.com/numaproj/numaflow/pkg/reduce/readloop"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
diff --git a/pkg/udf/applier/reducer_applier.go b/pkg/udf/applier/reducer_applier.go
index 5e212486..03f48983 100644
--- a/pkg/udf/applier/reducer_applier.go
+++ b/pkg/udf/applier/reducer_applier.go
@@ -20,7 +20,7 @@ import (
 	"context"
 
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 // ReduceApplier applies the HTTPBasedUDF on the read message and gives back a new message. Any UserError will be retried here, while
diff --git a/pkg/udf/function/uds_grpc.go b/pkg/udf/function/uds_grpc.go
index 034bb9aa..e9152160 100644
--- a/pkg/udf/function/uds_grpc.go
+++ b/pkg/udf/function/uds_grpc.go
@@ -23,7 +23,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 
 	functionpb "github.com/numaproj/numaflow-go/pkg/apis/proto/function/v1"
diff --git a/pkg/udf/reduce_udf.go b/pkg/udf/reduce_udf.go
index c0434733..c7953c88 100644
--- a/pkg/udf/reduce_udf.go
+++ b/pkg/udf/reduce_udf.go
@@ -23,18 +23,19 @@ import (
 	"sync"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
+	"go.uber.org/zap"
+
+	"github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
+	wal2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/wal"
 	"github.com/numaproj/numaflow/pkg/window/strategy/fixed"
 	"github.com/numaproj/numaflow/pkg/window/strategy/sliding"
-	"go.uber.org/zap"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/forward"
 	"github.com/numaproj/numaflow/pkg/metrics"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
-	"github.com/numaproj/numaflow/pkg/pbq/store/wal"
 	"github.com/numaproj/numaflow/pkg/reduce"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	sharedutil "github.com/numaproj/numaflow/pkg/shared/util"
@@ -148,7 +149,7 @@ func (u *ReduceUDFProcessor) Start(ctx context.Context) error {
 
 	var storeProvider store.StoreProvider
 	if storage := u.VertexInstance.Vertex.Spec.UDF.GroupBy.Storage; storage != nil && storage.PersistentVolumeClaim != nil {
-		storeProvider = wal.NewWALStores(u.VertexInstance, wal.WithStorePath(dfv1.DefaultStorePath), wal.WithMaxBufferSize(dfv1.DefaultStoreMaxBufferSize), wal.WithSyncDuration(dfv1.DefaultStoreSyncDuration))
+		storeProvider = wal2.NewWALStores(u.VertexInstance, wal2.WithStorePath(dfv1.DefaultStorePath), wal2.WithMaxBufferSize(dfv1.DefaultStoreMaxBufferSize), wal2.WithSyncDuration(dfv1.DefaultStoreSyncDuration))
 	} else {
 		// use in memory type by default
 		storeProvider = memory.NewMemoryStores()
diff --git a/pkg/window/keyed/keyed.go b/pkg/window/keyed/keyed.go
index f352fb9d..a8ed992d 100644
--- a/pkg/window/keyed/keyed.go
+++ b/pkg/window/keyed/keyed.go
@@ -22,7 +22,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 	"github.com/numaproj/numaflow/pkg/window"
 )
 
diff --git a/pkg/window/windower.go b/pkg/window/windower.go
index b2a57801..58275ed9 100644
--- a/pkg/window/windower.go
+++ b/pkg/window/windower.go
@@ -19,7 +19,7 @@ package window
 import (
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 // AlignedKeyedWindower represents a bounded window (i.e., it will have a definite start and end time), and the


Test patch:
diff --git a/pkg/pbq/pbq_test.go b/pkg/reduce/pbq/pbq_test.go
similarity index 88%
rename from pkg/pbq/pbq_test.go
rename to pkg/reduce/pbq/pbq_test.go
index 832defda..28b83015 100644
--- a/pkg/pbq/pbq_test.go
+++ b/pkg/reduce/pbq/pbq_test.go
@@ -24,9 +24,10 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store"
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store"
+	memory2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
+
 	"github.com/stretchr/testify/assert"
 )
 
@@ -40,7 +41,7 @@ func TestPBQ_ReadWrite(t *testing.T) {
 
 	ctx := context.Background()
 
-	qManager, _ := NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(storeSize)), WithChannelBufferSize(int64(buffSize)), WithReadTimeout(1*time.Second))
+	qManager, _ := NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(storeSize)), WithChannelBufferSize(int64(buffSize)), WithReadTimeout(1*time.Second))
 
 	// write 10 isb messages to persisted store
 	msgCount := 10
@@ -101,7 +102,7 @@ func Test_PBQReadWithCanceledContext(t *testing.T) {
 
 	ctx := context.Background()
 
-	qManager, err = NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(storeSize)), WithChannelBufferSize(int64(bufferSize)), WithReadTimeout(1*time.Second))
+	qManager, err = NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(storeSize)), WithChannelBufferSize(int64(bufferSize)), WithReadTimeout(1*time.Second))
 
 	assert.NoError(t, err)
 
@@ -167,7 +168,7 @@ func TestPBQ_WriteWithStoreFull(t *testing.T) {
 	var err error
 	ctx := context.Background()
 
-	qManager, err = NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(storeSize)), WithChannelBufferSize(int64(buffSize)), WithReadTimeout(1*time.Second))
+	qManager, err = NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(storeSize)), WithChannelBufferSize(int64(buffSize)), WithReadTimeout(1*time.Second))
 	assert.NoError(t, err)
 
 	// write 101 isb messages to pbq, but the store size is 100, we should get store is full error
diff --git a/pkg/pbq/pbqmanager_test.go b/pkg/reduce/pbq/pbqmanager_test.go
similarity index 93%
rename from pkg/pbq/pbqmanager_test.go
rename to pkg/reduce/pbq/pbqmanager_test.go
index 2c20c7dc..1779dce5 100644
--- a/pkg/pbq/pbqmanager_test.go
+++ b/pkg/reduce/pbq/pbqmanager_test.go
@@ -26,9 +26,9 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
-	"github.com/numaproj/numaflow/pkg/pbq/store/noop"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	memory2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/store/noop"
 )
 
 // tests for pbqManager (store type - in-memory)
@@ -37,7 +37,7 @@ func TestManager_ListPartitions(t *testing.T) {
 	size := int64(100)
 
 	ctx := context.Background()
-	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(size)),
+	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(size)),
 		WithReadTimeout(1*time.Second), WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
@@ -75,7 +75,7 @@ func TestManager_GetPBQ(t *testing.T) {
 	size := int64(100)
 	var pb1, pb2 ReadWriteCloser
 	ctx := context.Background()
-	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(size)),
+	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(size)),
 		WithReadTimeout(1*time.Second), WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
@@ -99,7 +99,7 @@ func TestPBQFlow(t *testing.T) {
 	size := int64(100)
 
 	ctx := context.Background()
-	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(size)),
+	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(size)),
 		WithReadTimeout(1*time.Second), WithChannelBufferSize(10))
 	assert.NoError(t, err)
 	testPartition := partition.ID{
@@ -217,7 +217,7 @@ func TestManager_Replay(t *testing.T) {
 	size := int64(100)
 
 	ctx := context.Background()
-	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(size)),
+	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(size)),
 		WithReadTimeout(1*time.Second), WithChannelBufferSize(10), WithReadBatchSize(10))
 	assert.NoError(t, err)
 	testPartition := partition.ID{
@@ -289,7 +289,7 @@ func TestManager_StartUp(t *testing.T) {
 		return []partition.ID{pID1,
 			pID2}, nil
 	}
-	stores := memory.NewMemoryStores(memory.WithStoreSize(size), memory.WithDiscoverer(dp))
+	stores := memory2.NewMemoryStores(memory2.WithStoreSize(size), memory2.WithDiscoverer(dp))
 	ctx := context.Background()
 	pbqManager, err := NewManager(ctx, "reduce", "test-pipeline", stores,
 		WithReadTimeout(1*time.Second), WithChannelBufferSize(10), WithReadBatchSize(10))
diff --git a/pkg/pbq/store/memory/store_test.go b/pkg/reduce/pbq/store/memory/store_test.go
similarity index 98%
rename from pkg/pbq/store/memory/store_test.go
rename to pkg/reduce/pbq/store/memory/store_test.go
index 8c537ff6..6be711d2 100644
--- a/pkg/pbq/store/memory/store_test.go
+++ b/pkg/reduce/pbq/store/memory/store_test.go
@@ -23,7 +23,8 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+
 	"github.com/stretchr/testify/assert"
 )
 
diff --git a/pkg/pbq/store/memory/stores_test.go b/pkg/reduce/pbq/store/memory/stores_test.go
similarity index 95%
rename from pkg/pbq/store/memory/stores_test.go
rename to pkg/reduce/pbq/store/memory/stores_test.go
index 6d8ae399..528d8741 100644
--- a/pkg/pbq/store/memory/stores_test.go
+++ b/pkg/reduce/pbq/store/memory/stores_test.go
@@ -5,8 +5,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
 	"github.com/stretchr/testify/assert"
+
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 func TestMemoryStores(t *testing.T) {
diff --git a/pkg/pbq/store/wal/segment_test.go b/pkg/reduce/pbq/store/wal/segment_test.go
similarity index 99%
rename from pkg/pbq/store/wal/segment_test.go
rename to pkg/reduce/pbq/store/wal/segment_test.go
index 517579a4..9e6419f3 100644
--- a/pkg/pbq/store/wal/segment_test.go
+++ b/pkg/reduce/pbq/store/wal/segment_test.go
@@ -28,7 +28,7 @@ import (
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 var vi = &dfv1.VertexInstance{
diff --git a/pkg/pbq/store/wal/stores_test.go b/pkg/reduce/pbq/store/wal/stores_test.go
similarity index 96%
rename from pkg/pbq/store/wal/stores_test.go
rename to pkg/reduce/pbq/store/wal/stores_test.go
index c94b3e16..e6afe162 100644
--- a/pkg/pbq/store/wal/stores_test.go
+++ b/pkg/reduce/pbq/store/wal/stores_test.go
@@ -2,11 +2,12 @@ package wal
 
 import (
 	"context"
-	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"testing"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+
 	"github.com/stretchr/testify/assert"
 )
 
diff --git a/pkg/reduce/pnf/processandforward_test.go b/pkg/reduce/pnf/processandforward_test.go
index 38edd6fb..754794c5 100644
--- a/pkg/reduce/pnf/processandforward_test.go
+++ b/pkg/reduce/pnf/processandforward_test.go
@@ -25,6 +25,9 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
+	pbq2 "github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	memory2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/shared/logging"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
@@ -39,9 +42,6 @@ import (
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
 	udfcall "github.com/numaproj/numaflow/pkg/udf/function"
 	wmstore "github.com/numaproj/numaflow/pkg/watermark/store"
 )
@@ -88,14 +88,14 @@ func TestProcessAndForward_Process(t *testing.T) {
 		Key:   "partition-1",
 	}
 	var err error
-	var pbqManager *pbq.Manager
+	var pbqManager *pbq2.Manager
 
-	pbqManager, err = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(100)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	pbqManager, err = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(100)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create a pbq for a partition
-	var simplePbq pbq.ReadWriteCloser
+	var simplePbq pbq2.ReadWriteCloser
 	simplePbq, err = pbqManager.CreateNewPBQ(ctx, testPartition)
 	assert.NoError(t, err)
 
@@ -145,9 +145,9 @@ func TestProcessAndForward_Process(t *testing.T) {
 func TestProcessAndForward_Forward(t *testing.T) {
 	ctx := context.Background()
 
-	var pbqManager *pbq.Manager
+	var pbqManager *pbq2.Manager
 
-	pbqManager, _ = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores())
+	pbqManager, _ = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores())
 
 	test1Buffer1 := simplebuffer.NewInMemoryBuffer("buffer1", 10)
 	test1Buffer2 := simplebuffer.NewInMemoryBuffer("buffer2", 10)
@@ -245,7 +245,7 @@ func TestProcessAndForward_Forward(t *testing.T) {
 	}
 }
 
-func createProcessAndForward(ctx context.Context, key string, pbqManager *pbq.Manager, toBuffers map[string]isb.BufferWriter) ProcessAndForward {
+func createProcessAndForward(ctx context.Context, key string, pbqManager *pbq2.Manager, toBuffers map[string]isb.BufferWriter) ProcessAndForward {
 
 	testPartition := partition.ID{
 		Start: time.UnixMilli(60000),
@@ -255,7 +255,7 @@ func createProcessAndForward(ctx context.Context, key string, pbqManager *pbq.Ma
 
 	// create a pbq for a partition
 	pw := buildPublisherMap(toBuffers)
-	var simplePbq pbq.Reader
+	var simplePbq pbq2.Reader
 	simplePbq, _ = pbqManager.CreateNewPBQ(ctx, testPartition)
 
 	resultPayload, _ := json.Marshal(PayloadForTest{
diff --git a/pkg/reduce/readloop/ordered_test.go b/pkg/reduce/readloop/ordered_test.go
index 188ae9d4..1a2b77eb 100644
--- a/pkg/reduce/readloop/ordered_test.go
+++ b/pkg/reduce/readloop/ordered_test.go
@@ -24,15 +24,15 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
+	pbq2 "github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	memory2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 
 	"github.com/stretchr/testify/assert"
 
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/udf/applier"
 )
 
@@ -107,8 +107,8 @@ func TestOrderedProcessing(t *testing.T) {
 			op := newOrderedForwarder(ctx, "reduce", "test-pipeline")
 			// although this could be declared outside, since we are using common naming scheme for partitions,
 			// things will go haywire.
-			pbqManager, _ := pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(100)),
-				pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+			pbqManager, _ := pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(100)),
+				pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 			cCtx, cancelFn := context.WithCancel(ctx)
 			defer cancelFn()
 			for _, _partition := range tt.partitions {
diff --git a/pkg/reduce/readloop/readloop_test.go b/pkg/reduce/readloop/readloop_test.go
index ddcb4306..ea6f1ee8 100644
--- a/pkg/reduce/readloop/readloop_test.go
+++ b/pkg/reduce/readloop/readloop_test.go
@@ -11,9 +11,9 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
+	pbq2 "github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	memory2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/watermark/generic"
 	"github.com/numaproj/numaflow/pkg/window/strategy/fixed"
 )
@@ -87,7 +87,7 @@ func TestReadLoop_Startup(t *testing.T) {
 		},
 	}
 
-	memStoreProvider := memory.NewMemoryStores(memory.WithStoreSize(100))
+	memStoreProvider := memory2.NewMemoryStores(memory2.WithStoreSize(100))
 
 	for _, id := range partitionIds {
 		memStore, err := memStoreProvider.CreateStore(ctx, id)
@@ -108,7 +108,7 @@ func TestReadLoop_Startup(t *testing.T) {
 		}
 	}
 
-	pManager, _ := pbq.NewManager(ctx, "reduce", "test-pipeline", memStoreProvider, pbq.WithChannelBufferSize(10))
+	pManager, _ := pbq2.NewManager(ctx, "reduce", "test-pipeline", memStoreProvider, pbq2.WithChannelBufferSize(10))
 
 	to1 := simplebuffer.NewInMemoryBuffer("reduce-buffer", 4)
 	toSteps := map[string]isb.BufferWriter{
diff --git a/pkg/reduce/reduce_test.go b/pkg/reduce/reduce_test.go
index c0844ec9..eadd4657 100644
--- a/pkg/reduce/reduce_test.go
+++ b/pkg/reduce/reduce_test.go
@@ -31,9 +31,9 @@ import (
 	dfv1 "github.com/numaproj/numaflow/pkg/apis/numaflow/v1alpha1"
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/stores/simplebuffer"
-	"github.com/numaproj/numaflow/pkg/pbq"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-	"github.com/numaproj/numaflow/pkg/pbq/store/memory"
+	pbq2 "github.com/numaproj/numaflow/pkg/reduce/pbq"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
+	memory2 "github.com/numaproj/numaflow/pkg/reduce/pbq/store/memory"
 	"github.com/numaproj/numaflow/pkg/watermark/fetch"
 	"github.com/numaproj/numaflow/pkg/watermark/processor"
 	"github.com/numaproj/numaflow/pkg/watermark/publish"
@@ -225,11 +225,11 @@ func TestDataForward_StartWithNoOpWM(t *testing.T) {
 	}
 
 	var err error
-	var pbqManager *pbq.Manager
+	var pbqManager *pbq2.Manager
 
 	// create pbqManager
-	pbqManager, err = pbq.NewManager(child, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(100)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	pbqManager, err = pbq2.NewManager(child, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(100)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	publisher := map[string]publish.Publisher{
@@ -297,9 +297,9 @@ func TestReduceDataForward_Count(t *testing.T) {
 	}
 
 	// create pbq manager
-	var pbqManager *pbq.Manager
-	pbqManager, err = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(1000)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	var pbqManager *pbq2.Manager
+	pbqManager, err = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(1000)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create in memory watermark publisher and fetcher
@@ -370,9 +370,9 @@ func TestReduceDataForward_Sum(t *testing.T) {
 	}
 
 	// create pbq manager
-	var pbqManager *pbq.Manager
-	pbqManager, err = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(1000)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	var pbqManager *pbq2.Manager
+	pbqManager, err = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(1000)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create in memory watermark publisher and fetcher
@@ -443,9 +443,9 @@ func TestReduceDataForward_Max(t *testing.T) {
 	}
 
 	// create pbq manager
-	var pbqManager *pbq.Manager
-	pbqManager, err = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(1000)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	var pbqManager *pbq2.Manager
+	pbqManager, err = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(1000)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create in memory watermark publisher and fetcher
@@ -516,9 +516,9 @@ func TestReduceDataForward_SumWithDifferentKeys(t *testing.T) {
 	}
 
 	// create pbq manager
-	var pbqManager *pbq.Manager
-	pbqManager, err = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(1000)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	var pbqManager *pbq2.Manager
+	pbqManager, err = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(1000)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create in memory watermark publisher and fetcher
@@ -597,9 +597,9 @@ func TestReduceDataForward_NonKeyed(t *testing.T) {
 	}
 
 	// create pbq manager
-	var pbqManager *pbq.Manager
-	pbqManager, err = pbq.NewManager(ctx, "reduce", "test-pipeline", memory.NewMemoryStores(memory.WithStoreSize(1000)),
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	var pbqManager *pbq2.Manager
+	pbqManager, err = pbq2.NewManager(ctx, "reduce", "test-pipeline", memory2.NewMemoryStores(memory2.WithStoreSize(1000)),
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create in memory watermark publisher and fetcher
@@ -674,12 +674,12 @@ func TestDataForward_WithContextClose(t *testing.T) {
 	}
 
 	// create a store provider
-	storeProvider := memory.NewMemoryStores(memory.WithStoreSize(1000))
+	storeProvider := memory2.NewMemoryStores(memory2.WithStoreSize(1000))
 
 	// create pbq manager
-	var pbqManager *pbq.Manager
-	pbqManager, err = pbq.NewManager(cctx, "reduce", "test-pipeline", storeProvider,
-		pbq.WithReadTimeout(1*time.Second), pbq.WithChannelBufferSize(10))
+	var pbqManager *pbq2.Manager
+	pbqManager, err = pbq2.NewManager(cctx, "reduce", "test-pipeline", storeProvider,
+		pbq2.WithReadTimeout(1*time.Second), pbq2.WithChannelBufferSize(10))
 	assert.NoError(t, err)
 
 	// create in memory watermark publisher and fetcher
diff --git a/pkg/udf/function/uds_grpc_test.go b/pkg/udf/function/uds_grpc_test.go
index bb41dc74..974a110b 100644
--- a/pkg/udf/function/uds_grpc_test.go
+++ b/pkg/udf/function/uds_grpc_test.go
@@ -35,7 +35,7 @@ import (
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/isb/testutils"
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 func NewMockUDSGRPCBasedUDF(mockClient *funcmock.MockUserDefinedFunctionClient) *udsGRPCBasedUDF {
diff --git a/pkg/window/keyed/keyed_test.go b/pkg/window/keyed/keyed_test.go
index 4a29a2a5..ba71fef4 100644
--- a/pkg/window/keyed/keyed_test.go
+++ b/pkg/window/keyed/keyed_test.go
@@ -21,9 +21,9 @@ import (
 	"testing"
 	"time"
 
-	"github.com/numaproj/numaflow/pkg/pbq/partition"
-
 	"github.com/stretchr/testify/assert"
+
+	"github.com/numaproj/numaflow/pkg/reduce/pbq/partition"
 )
 
 func TestKeyedWindow_AddKey(t *testing.T) {