github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/colexec/mergeblock/mergeblock.go (about)

     1  // Copyright 2021 Matrix Origin
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //	http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  package mergeblock
    15  
    16  import (
    17  	"bytes"
    18  
    19  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    20  	"github.com/matrixorigin/matrixone/pkg/vm"
    21  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    22  )
    23  
    24  const argName = "merge_block"
    25  
    26  func (arg *Argument) String(buf *bytes.Buffer) {
    27  	buf.WriteString(argName)
    28  	buf.WriteString(": MergeS3BlocksMetaLoc ")
    29  }
    30  
    31  func (arg *Argument) Prepare(proc *process.Process) error {
    32  	ap := arg
    33  	ap.container = new(Container)
    34  	ap.container.mp = make(map[int]*batch.Batch)
    35  	ap.container.mp2 = make(map[int][]*batch.Batch)
    36  	return nil
    37  }
    38  
    39  func (arg *Argument) Call(proc *process.Process) (vm.CallResult, error) {
    40  	if err, isCancel := vm.CancelCheck(proc); isCancel {
    41  		return vm.CancelResult, err
    42  	}
    43  
    44  	var err error
    45  	ap := arg
    46  	result, err := arg.GetChildren(0).Call(proc)
    47  	if err != nil {
    48  		return result, err
    49  	}
    50  
    51  	anal := proc.GetAnalyze(arg.GetIdx(), arg.GetParallelIdx(), arg.GetParallelMajor())
    52  	anal.Start()
    53  	defer anal.Stop()
    54  
    55  	if result.Batch == nil {
    56  		result.Status = vm.ExecStop
    57  		return result, nil
    58  	}
    59  	if result.Batch.IsEmpty() {
    60  		result.Batch = batch.EmptyBatch
    61  		return result, nil
    62  	}
    63  	bat := result.Batch
    64  	if err := ap.Split(proc, bat); err != nil {
    65  		return result, err
    66  	}
    67  
    68  	// If the target is a partition table
    69  	if len(ap.PartitionSources) > 0 {
    70  		// 'i' aligns with partition number
    71  		for i := range ap.PartitionSources {
    72  			if ap.container.mp[i].RowCount() > 0 {
    73  				// batches in mp will be deeply copied into txn's workspace.
    74  				if err = ap.PartitionSources[i].Write(proc.Ctx, ap.container.mp[i]); err != nil {
    75  					return result, err
    76  				}
    77  			}
    78  
    79  			for _, bat := range ap.container.mp2[i] {
    80  				// batches in mp2 will be deeply copied into txn's workspace.
    81  				if err = ap.PartitionSources[i].Write(proc.Ctx, bat); err != nil {
    82  					return result, err
    83  				}
    84  
    85  			}
    86  			ap.container.mp2[i] = ap.container.mp2[i][:0]
    87  		}
    88  	} else {
    89  		// handle origin/main table.
    90  		if ap.container.mp[0].RowCount() > 0 {
    91  			//batches in mp will be deeply copied into txn's workspace.
    92  			if err = ap.Tbl.Write(proc.Ctx, ap.container.mp[0]); err != nil {
    93  				return result, err
    94  			}
    95  		}
    96  
    97  		for _, bat := range ap.container.mp2[0] {
    98  			//batches in mp2 will be deeply copied into txn's workspace.
    99  			if err = ap.Tbl.Write(proc.Ctx, bat); err != nil {
   100  				return result, err
   101  			}
   102  		}
   103  		ap.container.mp2[0] = ap.container.mp2[0][:0]
   104  	}
   105  
   106  	return result, nil
   107  }