github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/colexec/product/product.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  
    15  package product
    16  
    17  import (
    18  	"bytes"
    19  
    20  	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    23  	"github.com/matrixorigin/matrixone/pkg/vm"
    24  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    25  )
    26  
    27  const argName = "product"
    28  
    29  func (arg *Argument) String(buf *bytes.Buffer) {
    30  	buf.WriteString(argName)
    31  	buf.WriteString(": cross join ")
    32  }
    33  
    34  func (arg *Argument) Prepare(proc *process.Process) error {
    35  	ap := arg
    36  	ap.ctr = new(container)
    37  	ap.ctr.InitReceiver(proc, false)
    38  	return nil
    39  }
    40  
    41  func (arg *Argument) Call(proc *process.Process) (vm.CallResult, error) {
    42  	if err, isCancel := vm.CancelCheck(proc); isCancel {
    43  		return vm.CancelResult, err
    44  	}
    45  
    46  	anal := proc.GetAnalyze(arg.GetIdx(), arg.GetParallelIdx(), arg.GetParallelMajor())
    47  	anal.Start()
    48  	defer anal.Stop()
    49  	ap := arg
    50  	ctr := ap.ctr
    51  	result := vm.NewCallResult()
    52  	var err error
    53  	for {
    54  		switch ctr.state {
    55  		case Build:
    56  			if err := ctr.build(proc, anal); err != nil {
    57  				return result, err
    58  			}
    59  			ctr.state = Probe
    60  
    61  		case Probe:
    62  			if ctr.inBat != nil {
    63  				if err := ctr.probe(ap, proc, anal, arg.GetIsLast(), &result); err != nil {
    64  					return result, err
    65  				}
    66  				return result, nil
    67  			}
    68  			ctr.inBat, _, err = ctr.ReceiveFromSingleReg(0, anal)
    69  			if err != nil {
    70  				return result, err
    71  			}
    72  
    73  			if ctr.inBat == nil {
    74  				ctr.state = End
    75  				continue
    76  			}
    77  			if ctr.inBat.IsEmpty() {
    78  				proc.PutBatch(ctr.inBat)
    79  				ctr.inBat = nil
    80  				continue
    81  			}
    82  			if ctr.bat == nil {
    83  				proc.PutBatch(ctr.inBat)
    84  				ctr.inBat = nil
    85  				continue
    86  			}
    87  			anal.Input(ctr.inBat, arg.GetIsFirst())
    88  			if err := ctr.probe(ap, proc, anal, arg.GetIsLast(), &result); err != nil {
    89  				return result, err
    90  			}
    91  			return result, nil
    92  
    93  		default:
    94  			result.Batch = nil
    95  			result.Status = vm.ExecStop
    96  			return result, nil
    97  		}
    98  	}
    99  }
   100  
   101  func (ctr *container) build(proc *process.Process, anal process.Analyze) error {
   102  	for {
   103  		bat, _, err := ctr.ReceiveFromSingleReg(1, anal)
   104  		if err != nil {
   105  			return err
   106  		}
   107  		if bat == nil {
   108  			break
   109  		}
   110  		ctr.bat, err = ctr.bat.AppendWithCopy(proc.Ctx, proc.Mp(), bat)
   111  		if err != nil {
   112  			return err
   113  		}
   114  		proc.PutBatch(bat)
   115  	}
   116  	return nil
   117  }
   118  
   119  func (ctr *container) probe(ap *Argument, proc *process.Process, anal process.Analyze, isLast bool, result *vm.CallResult) error {
   120  	if ctr.rbat != nil {
   121  		proc.PutBatch(ctr.rbat)
   122  		ctr.rbat = nil
   123  	}
   124  	ctr.rbat = batch.NewWithSize(len(ap.Result))
   125  	for i, rp := range ap.Result {
   126  		if rp.Rel == 0 {
   127  			ctr.rbat.Vecs[i] = proc.GetVector(*ctr.inBat.Vecs[rp.Pos].GetType())
   128  		} else {
   129  			ctr.rbat.Vecs[i] = proc.GetVector(*ctr.bat.Vecs[rp.Pos].GetType())
   130  		}
   131  	}
   132  	count := ctr.inBat.RowCount()
   133  	count2 := ctr.bat.RowCount()
   134  	var i, j int
   135  	for j = ctr.probeIdx; j < count2; j++ {
   136  		for i = 0; i < count; i++ {
   137  			for k, rp := range ap.Result {
   138  				if rp.Rel == 0 {
   139  					if err := ctr.rbat.Vecs[k].UnionOne(ctr.inBat.Vecs[rp.Pos], int64(i), proc.Mp()); err != nil {
   140  						return err
   141  					}
   142  				} else {
   143  					if err := ctr.rbat.Vecs[k].UnionOne(ctr.bat.Vecs[rp.Pos], int64(j), proc.Mp()); err != nil {
   144  						return err
   145  					}
   146  				}
   147  			}
   148  		}
   149  		if ctr.rbat.Vecs[0].Length() >= colexec.DefaultBatchSize {
   150  			anal.Output(ctr.rbat, isLast)
   151  			result.Batch = ctr.rbat
   152  			ctr.rbat.SetRowCount(ctr.rbat.Vecs[0].Length())
   153  			ctr.probeIdx = j + 1
   154  			return nil
   155  		}
   156  	}
   157  	// ctr.rbat.AddRowCount(count * count2)
   158  	ctr.probeIdx = 0
   159  	ctr.rbat.SetRowCount(ctr.rbat.Vecs[0].Length())
   160  	anal.Output(ctr.rbat, isLast)
   161  	result.Batch = ctr.rbat
   162  
   163  	proc.PutBatch(ctr.inBat)
   164  	ctr.inBat = nil
   165  	return nil
   166  }