fix ut error

matrixorigin · Jun 3, 2024 · 47b5b59 · 47b5b59
1 parent e6bb769
commit 47b5b59
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git 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
@@ -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
@@ -63,14 +63,18 @@ 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
 }
 
 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 {