refactor: multiupdate relation

matrixorigin · Feb 10, 2025 · bb1fc7a · bb1fc7a
1 parent 0cf6f0d
commit bb1fc7a
Showing 1 changed file with 19 additions and 10 deletions.
diff --git a/pkg/sql/colexec/multi_update/multi_update.go b/pkg/sql/colexec/multi_update/multi_update.go
@@ -68,13 +68,22 @@ func (update *MultiUpdate) Prepare(proc *process.Process) error {
 
 	for _, updateCtx := range update.MultiUpdateCtx {
 		info := update.ctr.updateCtxInfos[updateCtx.TableDef.Name]
-		info.Sources = nil
 		if update.Action != UpdateWriteS3 {
-			rel, err := colexec.GetRelAndPartitionRelsByObjRef(proc.Ctx, proc, update.Engine, updateCtx.ObjRef)
-			if err != nil {
-				return err
+			if len(info.Sources) == 0 {
+				info.Sources = nil
+				rel, err := colexec.GetRelAndPartitionRelsByObjRef(proc.Ctx, proc, update.Engine, updateCtx.ObjRef)
+				if err != nil {
+					return err
+				}
+				info.Sources = append(info.Sources, rel)
+			} else {
+				for _, rel := range info.Sources {
+					err := rel.Reset(proc.GetTxnOperator())
+					if err != nil {
+						return err
+					}
+				}
 			}
-			info.Sources = append(info.Sources, rel)
 		}
 	}
 
@@ -377,13 +386,13 @@ func (update *MultiUpdate) resetMultiUpdateCtxs() {
 func (update *MultiUpdate) resetMultiSources(proc *process.Process) error {
 	for _, updateCtx := range update.MultiUpdateCtx {
 		info := update.ctr.updateCtxInfos[updateCtx.TableDef.Name]
-		info.Sources = nil
 		if update.Action != UpdateWriteS3 {
-			rel, err := colexec.GetRelAndPartitionRelsByObjRef(proc.Ctx, proc, update.Engine, updateCtx.ObjRef)
-			if err != nil {
-				return err
+			for _, rel := range info.Sources {
+				err := rel.Reset(proc.GetTxnOperator())
+				if err != nil {
+					return err
+				}
 			}
-			info.Sources = append(info.Sources, rel)
 		}
 	}
 	return nil