From 47b5b59ed9499ef92916e433f1324f1366613c0d Mon Sep 17 00:00:00 2001
From: Wei Ziran <weiziran125@gmail.com>
Date: Mon, 3 Jun 2024 18:49:11 +0800
Subject: [PATCH] fix ut error

---
 pkg/vm/engine/tae/mergesort/merger.go   | 1 +
 pkg/vm/engine/tae/mergesort/reshaper.go | 1 +
 pkg/vm/engine/tae/mergesort/types.go    | 6 +++++-
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/pkg/vm/engine/tae/mergesort/merger.go b/pkg/vm/engine/tae/mergesort/merger.go
index e473faf397c49..a4d671f01e7a2 100644
--- a/pkg/vm/engine/tae/mergesort/merger.go
+++ b/pkg/vm/engine/tae/mergesort/merger.go
@@ -90,6 +90,7 @@ func newMerger[T any](host MergeTaskHost, lessFunc sort.LessFunc[T], sortKeyPos
 				totalRowCnt:   host.GetTotalRowCnt(),
 				rowSize:       host.GetTotalSize() / host.GetTotalRowCnt(),
 				targetObjSize: host.GetTargetObjSize(),
+				blkPerObj:     host.GetObjectMaxBlocks(),
 			},
 		},
 	}
diff --git a/pkg/vm/engine/tae/mergesort/reshaper.go b/pkg/vm/engine/tae/mergesort/reshaper.go
index 45b056c35d70c..2d4d2d1a2fa6b 100644
--- a/pkg/vm/engine/tae/mergesort/reshaper.go
+++ b/pkg/vm/engine/tae/mergesort/reshaper.go
@@ -53,6 +53,7 @@ func reshape(ctx context.Context, host MergeTaskHost) error {
 			totalRowCnt:   host.GetTotalRowCnt(),
 			rowSize:       host.GetTotalSize() / host.GetTotalRowCnt(),
 			targetObjSize: host.GetTargetObjSize(),
+			blkPerObj:     host.GetObjectMaxBlocks(),
 		},
 	}
 
diff --git a/pkg/vm/engine/tae/mergesort/types.go b/pkg/vm/engine/tae/mergesort/types.go
index 1f293b70585b0..091e3b4358099 100644
--- a/pkg/vm/engine/tae/mergesort/types.go
+++ b/pkg/vm/engine/tae/mergesort/types.go
@@ -63,6 +63,7 @@ func (x *heapSlice[T]) Len() int      { return len(x.s) }
 
 type writerStats struct {
 	totalRowCnt, rowSize, targetObjSize uint32
+	blkPerObj                           uint16
 
 	blkRowCnt, objRowCnt, objBlkCnt int
 	mergedRowCnt, objCnt            int
@@ -70,7 +71,10 @@ type writerStats struct {
 
 func (s *writerStats) needNewObject() bool {
 	if s.targetObjSize == 0 {
-		return s.objBlkCnt == int(options.DefaultBlocksPerObject)
+		if s.blkPerObj == 0 {
+			return s.objBlkCnt == int(options.DefaultBlocksPerObject)
+		}
+		return s.objBlkCnt == int(s.blkPerObj)
 	}
 
 	if uint32(s.objRowCnt)*s.rowSize > s.targetObjSize {