fix hashbuild / shufflehashbuild / indexbuild / fuzzyfilter reset bug. (#19690)

m-schen · web-flow · commit a42cff11e875 · 2024-10-30T16:37:56.000+08:00
diff --git a/pkg/container/pSpool/sender.go b/pkg/container/pSpool/sender.go
@@ -106,7 +106,9 @@ func (ps *PipelineSpool) ReceiveBatch(idx int) (data *batch.Batch, info error) {
 
 	next := ps.rs[idx].popNextIndex()
 	if ps.shardPool[next].dataContent == nil {
-		ps.csDoneSignal <- struct{}{}
+		defer func() {
+			ps.csDoneSignal <- struct{}{}
+		}()
 	}
 	return ps.shardPool[next].dataContent, ps.shardPool[next].errContent
 }
diff --git a/pkg/sql/colexec/fuzzyfilter/types.go b/pkg/sql/colexec/fuzzyfilter/types.go
@@ -107,7 +107,9 @@ func (fuzzyFilter *FuzzyFilter) getProbeIdx() int {
 }
 
 func (fuzzyFilter *FuzzyFilter) Reset(proc *process.Process, pipelineFailed bool, err error) {
-	message.FinalizeRuntimeFilter(fuzzyFilter.RuntimeFilterSpec, pipelineFailed, err, proc.GetMessageBoard())
+	runtimeSucceed := fuzzyFilter.ctr.state > HandleRuntimeFilter
+
+	message.FinalizeRuntimeFilter(fuzzyFilter.RuntimeFilterSpec, runtimeSucceed, proc.GetMessageBoard())
 	ctr := &fuzzyFilter.ctr
 	ctr.state = Build
 	ctr.collisionCnt = 0
diff --git a/pkg/sql/colexec/hashbuild/build.go b/pkg/sql/colexec/hashbuild/build.go
@@ -91,7 +91,9 @@ func (hashBuild *HashBuild) Call(proc *process.Process) (vm.CallResult, error) {
 				panic("wrong joinmap message tag!")
 			}
 			message.SendMessage(message.JoinMapMsg{JoinMapPtr: jm, Tag: ap.JoinMapTag}, proc.GetMessageBoard())
+			ctr.state = SendSucceed
 
+		case SendSucceed:
 			result.Batch = nil
 			result.Status = vm.ExecStop
 			analyzer.Output(result.Batch)
diff --git a/pkg/sql/colexec/hashbuild/types.go b/pkg/sql/colexec/hashbuild/types.go
@@ -30,6 +30,7 @@ const (
 	BuildHashMap = iota
 	HandleRuntimeFilter
 	SendJoinMap
+	SendSucceed
 )
 
 type container struct {
@@ -83,21 +84,15 @@ func (hashBuild *HashBuild) Release() {
 }
 
 func (hashBuild *HashBuild) Reset(proc *process.Process, pipelineFailed bool, err error) {
+	runtimeSucceed := hashBuild.ctr.state > HandleRuntimeFilter
+	mapSucceed := hashBuild.ctr.state == SendSucceed
+
+	hashBuild.ctr.hashmapBuilder.Reset(proc, !mapSucceed)
 	hashBuild.ctr.state = BuildHashMap
 	hashBuild.ctr.runtimeFilterIn = false
-	message.FinalizeRuntimeFilter(hashBuild.RuntimeFilterSpec, pipelineFailed, err, proc.GetMessageBoard())
-	message.FinalizeJoinMapMessage(proc.GetMessageBoard(), hashBuild.JoinMapTag, false, 0, pipelineFailed, err)
-	if pipelineFailed || err != nil {
-		hashBuild.ctr.hashmapBuilder.FreeWithError(proc)
-	} else {
-		hashBuild.ctr.hashmapBuilder.Reset(proc)
-	}
+	message.FinalizeRuntimeFilter(hashBuild.RuntimeFilterSpec, runtimeSucceed, proc.GetMessageBoard())
+	message.FinalizeJoinMapMessage(proc.GetMessageBoard(), hashBuild.JoinMapTag, false, 0, mapSucceed)
 }
 func (hashBuild *HashBuild) Free(proc *process.Process, pipelineFailed bool, err error) {
-
-	if pipelineFailed || err != nil {
-		hashBuild.ctr.hashmapBuilder.FreeWithError(proc)
-	} else {
-		hashBuild.ctr.hashmapBuilder.Free(proc)
-	}
+	hashBuild.ctr.hashmapBuilder.Free(proc)
 }
diff --git a/pkg/sql/colexec/hashmap_util/hashmap_util.go b/pkg/sql/colexec/hashmap_util/hashmap_util.go
@@ -70,6 +70,7 @@ func (hb *HashmapBuilder) Prepare(Conditions []*plan.Expr, proc *process.Process
 		hb.vecs = make([][]*vector.Vector, 0)
 		hb.executor = make([]colexec.ExpressionExecutor, len(Conditions))
 		hb.keyWidth = 0
+		hb.InputBatchRowCount = 0
 		for i, expr := range Conditions {
 			if _, ok := Conditions[i].Expr.(*pbplan.Expr_Col); !ok {
 				hb.needDupVec = true
@@ -90,24 +91,25 @@ func (hb *HashmapBuilder) Prepare(Conditions []*plan.Expr, proc *process.Process
 	return nil
 }
 
-func (hb *HashmapBuilder) Reset(proc *process.Process) {
+func (hb *HashmapBuilder) Reset(proc *process.Process, hashTableHasNotSent bool) {
+	if hashTableHasNotSent || hb.InputBatchRowCount == 0 {
+		hb.FreeHashMapAndBatches(proc)
+	}
+
 	if hb.needDupVec {
 		for i := range hb.vecs {
 			for j := range hb.vecs[i] {
 				hb.vecs[i][j].Free(proc.Mp())
 			}
 		}
 	}
-	if hb.InputBatchRowCount == 0 {
-		hb.FreeHashMapAndBatches(proc)
-	}
 	hb.InputBatchRowCount = 0
 	hb.Batches.Reset()
 	hb.IntHashMap = nil
 	hb.StrHashMap = nil
 	hb.vecs = nil
 	for i := range hb.UniqueJoinKeys {
-		hb.UniqueJoinKeys[i].CleanOnlyData()
+		hb.UniqueJoinKeys[i].Free(proc.Mp())
 	}
 	hb.UniqueJoinKeys = nil
 	hb.MultiSels.Free()
@@ -137,8 +139,6 @@ func (hb *HashmapBuilder) Free(proc *process.Process) {
 	hb.UniqueJoinKeys = nil
 }
 
-// hashmap and batches are owned by probe operators
-// build operator can only call this when error occurs, or inputbatch rowcount is 0
 func (hb *HashmapBuilder) FreeHashMapAndBatches(proc *process.Process) {
 	if hb.IntHashMap != nil {
 		hb.IntHashMap.Free()
@@ -151,28 +151,6 @@ func (hb *HashmapBuilder) FreeHashMapAndBatches(proc *process.Process) {
 	hb.Batches.Clean(proc.Mp())
 }
 
-func (hb *HashmapBuilder) FreeWithError(proc *process.Process) {
-	hb.needDupVec = false
-	hb.FreeHashMapAndBatches(proc)
-	hb.MultiSels.Free()
-	for i := range hb.executor {
-		if hb.executor[i] != nil {
-			hb.executor[i].Free()
-		}
-	}
-	hb.executor = nil
-	for i := range hb.vecs {
-		for j := range hb.vecs[i] {
-			hb.vecs[i][j].Free(proc.Mp())
-		}
-	}
-	hb.vecs = nil
-	for i := range hb.UniqueJoinKeys {
-		hb.UniqueJoinKeys[i].Free(proc.Mp())
-	}
-	hb.UniqueJoinKeys = nil
-}
-
 func (hb *HashmapBuilder) evalJoinCondition(proc *process.Process) error {
 	for idx1 := range hb.Batches.Buf {
 		tmpVes := make([]*vector.Vector, len(hb.executor))
diff --git a/pkg/sql/colexec/hashmap_util/hashmap_util_test.go b/pkg/sql/colexec/hashmap_util/hashmap_util_test.go
@@ -57,7 +57,8 @@ func TestBuildHashMap(t *testing.T) {
 	require.NoError(t, err)
 	require.Less(t, int64(0), hb.GetSize())
 	require.Less(t, uint64(0), hb.GetGroupCount())
-	hb.FreeWithError(proc)
+	hb.Reset(proc, true)
+	hb.Free(proc)
 	require.Equal(t, int64(0), proc.Mp().CurrNB())
 }
 
diff --git a/pkg/sql/colexec/indexbuild/types.go b/pkg/sql/colexec/indexbuild/types.go
@@ -74,7 +74,9 @@ func (indexBuild *IndexBuild) Release() {
 }
 
 func (indexBuild *IndexBuild) Reset(proc *process.Process, pipelineFailed bool, err error) {
-	message.FinalizeRuntimeFilter(indexBuild.RuntimeFilterSpec, pipelineFailed, err, proc.GetMessageBoard())
+	runtimeSucceed := indexBuild.ctr.state > HandleRuntimeFilter
+
+	message.FinalizeRuntimeFilter(indexBuild.RuntimeFilterSpec, runtimeSucceed, proc.GetMessageBoard())
 	indexBuild.ctr.state = ReceiveBatch
 	if indexBuild.ctr.buf != nil {
 		indexBuild.ctr.buf.CleanOnlyData()
diff --git a/pkg/sql/colexec/shufflebuild/build.go b/pkg/sql/colexec/shufflebuild/build.go
@@ -98,7 +98,9 @@ func (shuffleBuild *ShuffleBuild) Call(proc *process.Process) (vm.CallResult, er
 				jm.IncRef(1)
 			}
 			message.SendMessage(message.JoinMapMsg{JoinMapPtr: jm, IsShuffle: true, ShuffleIdx: ap.ShuffleIdx, Tag: ap.JoinMapTag}, proc.GetMessageBoard())
+			ctr.state = SendSucceed
 
+		case SendSucceed:
 			result.Batch = nil
 			result.Status = vm.ExecStop
 			analyzer.Output(result.Batch)
diff --git a/pkg/sql/colexec/shufflebuild/types.go b/pkg/sql/colexec/shufflebuild/types.go
@@ -30,6 +30,7 @@ const (
 	ReceiveBatch = iota
 	BuildHashMap
 	SendJoinMap
+	SendSucceed
 )
 
 type container struct {
@@ -81,22 +82,15 @@ func (shuffleBuild *ShuffleBuild) Release() {
 }
 
 func (shuffleBuild *ShuffleBuild) Reset(proc *process.Process, pipelineFailed bool, err error) {
-	shuffleBuild.ctr.state = ReceiveBatch
-	message.FinalizeRuntimeFilter(shuffleBuild.RuntimeFilterSpec, pipelineFailed, err, proc.GetMessageBoard())
-	message.FinalizeJoinMapMessage(proc.GetMessageBoard(), shuffleBuild.JoinMapTag, true, shuffleBuild.ShuffleIdx, pipelineFailed, err)
+	runtimeSucceed := shuffleBuild.ctr.state > ReceiveBatch
+	mapSucceed := shuffleBuild.ctr.state == SendSucceed
 
-	if pipelineFailed || err != nil {
-		shuffleBuild.ctr.hashmapBuilder.FreeWithError(proc)
-	} else {
-		shuffleBuild.ctr.hashmapBuilder.Reset(proc)
-	}
+	shuffleBuild.ctr.hashmapBuilder.Reset(proc, !mapSucceed)
+	shuffleBuild.ctr.state = ReceiveBatch
+	message.FinalizeRuntimeFilter(shuffleBuild.RuntimeFilterSpec, runtimeSucceed, proc.GetMessageBoard())
+	message.FinalizeJoinMapMessage(proc.GetMessageBoard(), shuffleBuild.JoinMapTag, true, shuffleBuild.ShuffleIdx, mapSucceed)
 }
 
 func (shuffleBuild *ShuffleBuild) Free(proc *process.Process, pipelineFailed bool, err error) {
-
-	if pipelineFailed || err != nil {
-		shuffleBuild.ctr.hashmapBuilder.FreeWithError(proc)
-	} else {
-		shuffleBuild.ctr.hashmapBuilder.Free(proc)
-	}
+	shuffleBuild.ctr.hashmapBuilder.Free(proc)
 }
diff --git a/pkg/vm/message/joinMapMsg.go b/pkg/vm/message/joinMapMsg.go
@@ -259,8 +259,8 @@ func ReceiveJoinMap(tag int32, isShuffle bool, shuffleIdx int32, mb *MessageBoar
 	}
 }
 
-func FinalizeJoinMapMessage(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, pipelineFailed bool, err error) {
-	if pipelineFailed || err != nil {
+func FinalizeJoinMapMessage(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, sendMapSucceed bool) {
+	if !sendMapSucceed {
 		SendMessage(JoinMapMsg{JoinMapPtr: nil, IsShuffle: isShuffle, ShuffleIdx: shuffleIdx, Tag: tag}, mb)
 	}
 }
diff --git a/pkg/vm/message/runtimeFilterMsg.go b/pkg/vm/message/runtimeFilterMsg.go
@@ -77,9 +77,9 @@ func SendRuntimeFilter(rt RuntimeFilterMessage, m *plan.RuntimeFilterSpec, mb *M
 	}
 }
 
-func FinalizeRuntimeFilter(m *plan.RuntimeFilterSpec, pipelineFailed bool, err error, mb *MessageBoard) {
+func FinalizeRuntimeFilter(m *plan.RuntimeFilterSpec, sendSucceed bool, mb *MessageBoard) {
 	if m != nil {
-		if pipelineFailed || err != nil {
+		if !sendSucceed {
 			var runtimeFilter RuntimeFilterMessage
 			runtimeFilter.Tag = m.Tag
 			runtimeFilter.Typ = RuntimeFilter_DROP

Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,9 @@ func (ps PipelineSpool) ReceiveBatch(idx int) (data batch.Batch, info error) {`
`106`	`106`
`107`	`107`	`next := ps.rs[idx].popNextIndex()`
`108`	`108`	`if ps.shardPool[next].dataContent == nil {`
`109`		`- ps.csDoneSignal <- struct{}{}`
	`109`	`+ defer func() {`
	`110`	`+ ps.csDoneSignal <- struct{}{}`
	`111`	`+ }()`
`110`	`112`	`}`
`111`	`113`	`return ps.shardPool[next].dataContent, ps.shardPool[next].errContent`
`112`	`114`	`}`
Original file line number	Diff line number	Diff line change
`@@ -91,7 +91,9 @@ func (hashBuild HashBuild) Call(proc process.Process) (vm.CallResult, error) {`
`91`	`91`	`panic("wrong joinmap message tag!")`
`92`	`92`	`}`
`93`	`93`	`message.SendMessage(message.JoinMapMsg{JoinMapPtr: jm, Tag: ap.JoinMapTag}, proc.GetMessageBoard())`
	`94`	`+ ctr.state = SendSucceed`
`94`	`95`
	`96`	`+ case SendSucceed:`
`95`	`97`	`result.Batch = nil`
`96`	`98`	`result.Status = vm.ExecStop`
`97`	`99`	`analyzer.Output(result.Batch)`
Original file line number	Diff line number	Diff line change
`@@ -57,7 +57,8 @@ func TestBuildHashMap(t *testing.T) {`
`57`	`57`	`require.NoError(t, err)`
`58`	`58`	`require.Less(t, int64(0), hb.GetSize())`
`59`	`59`	`require.Less(t, uint64(0), hb.GetGroupCount())`
`60`		`- hb.FreeWithError(proc)`
	`60`	`+ hb.Reset(proc, true)`
	`61`	`+ hb.Free(proc)`
`61`	`62`	`require.Equal(t, int64(0), proc.Mp().CurrNB())`
`62`	`63`	`}`
`63`	`64`
Original file line number	Diff line number	Diff line change
`@@ -98,7 +98,9 @@ func (shuffleBuild ShuffleBuild) Call(proc process.Process) (vm.CallResult, er`
`98`	`98`	`jm.IncRef(1)`
`99`	`99`	`}`
`100`	`100`	`message.SendMessage(message.JoinMapMsg{JoinMapPtr: jm, IsShuffle: true, ShuffleIdx: ap.ShuffleIdx, Tag: ap.JoinMapTag}, proc.GetMessageBoard())`
	`101`	`+ ctr.state = SendSucceed`
`101`	`102`
	`103`	`+ case SendSucceed:`
`102`	`104`	`result.Batch = nil`
`103`	`105`	`result.Status = vm.ExecStop`
`104`	`106`	`analyzer.Output(result.Batch)`
Original file line number	Diff line number	Diff line change
`@@ -259,8 +259,8 @@ func ReceiveJoinMap(tag int32, isShuffle bool, shuffleIdx int32, mb *MessageBoar`
`259`	`259`	`}`
`260`	`260`	`}`
`261`	`261`
`262`		`-func FinalizeJoinMapMessage(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, pipelineFailed bool, err error) {`
`263`		`- if pipelineFailed \|\| err != nil {`
	`262`	`+func FinalizeJoinMapMessage(mb *MessageBoard, tag int32, isShuffle bool, shuffleIdx int32, sendMapSucceed bool) {`
	`263`	`+ if !sendMapSucceed {`
`264`	`264`	`SendMessage(JoinMapMsg{JoinMapPtr: nil, IsShuffle: isShuffle, ShuffleIdx: shuffleIdx, Tag: tag}, mb)`
`265`	`265`	`}`
`266`	`266`	`}`