github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/colexec/loopjoin/join.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 loopjoin
    16  
    17  import (
    18  	"bytes"
    19  	"time"
    20  
    21  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    22  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    23  	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
    24  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    25  )
    26  
    27  func String(_ any, buf *bytes.Buffer) {
    28  	buf.WriteString(" loop join ")
    29  }
    30  
    31  func Prepare(proc *process.Process, arg any) error {
    32  	ap := arg.(*Argument)
    33  	ap.ctr = new(container)
    34  	return nil
    35  }
    36  
    37  func Call(idx int, proc *process.Process, arg any, isFirst bool, isLast bool) (bool, error) {
    38  	anal := proc.GetAnalyze(idx)
    39  	anal.Start()
    40  	defer anal.Stop()
    41  	ap := arg.(*Argument)
    42  	ctr := ap.ctr
    43  	for {
    44  		switch ctr.state {
    45  		case Build:
    46  			if err := ctr.build(ap, proc, anal); err != nil {
    47  				return false, err
    48  			}
    49  			ctr.state = Probe
    50  
    51  		case Probe:
    52  			start := time.Now()
    53  			bat := <-proc.Reg.MergeReceivers[0].Ch
    54  			anal.WaitStop(start)
    55  
    56  			if bat == nil {
    57  				ctr.state = End
    58  				continue
    59  			}
    60  			if bat.Length() == 0 {
    61  				continue
    62  			}
    63  			if ctr.bat == nil || ctr.bat.Length() == 0 {
    64  				bat.Clean(proc.Mp())
    65  				continue
    66  			}
    67  			err := ctr.probe(bat, ap, proc, anal, isFirst, isLast)
    68  			bat.Clean(proc.Mp())
    69  			return false, err
    70  
    71  		default:
    72  			ap.Free(proc, false)
    73  			proc.SetInputBatch(nil)
    74  			return true, nil
    75  		}
    76  	}
    77  }
    78  
    79  func (ctr *container) build(ap *Argument, proc *process.Process, anal process.Analyze) error {
    80  	start := time.Now()
    81  	bat := <-proc.Reg.MergeReceivers[1].Ch
    82  	anal.WaitStop(start)
    83  
    84  	if bat != nil {
    85  		ctr.bat = bat
    86  	}
    87  	return nil
    88  }
    89  
    90  func (ctr *container) probe(bat *batch.Batch, ap *Argument, proc *process.Process, anal process.Analyze, isFirst bool, isLast bool) error {
    91  	anal.Input(bat, isFirst)
    92  	rbat := batch.NewWithSize(len(ap.Result))
    93  	rbat.Zs = proc.Mp().GetSels()
    94  	for i, rp := range ap.Result {
    95  		if rp.Rel == 0 {
    96  			rbat.Vecs[i] = vector.New(bat.Vecs[rp.Pos].Typ)
    97  		} else {
    98  			rbat.Vecs[i] = vector.New(ctr.bat.Vecs[rp.Pos].Typ)
    99  		}
   100  	}
   101  	count := bat.Length()
   102  	for i := 0; i < count; i++ {
   103  		vec, err := colexec.JoinFilterEvalExpr(bat, ctr.bat, i, proc, ap.Cond)
   104  		if err != nil {
   105  			rbat.Clean(proc.Mp())
   106  			return err
   107  		}
   108  		bs := vector.MustTCols[bool](vec)
   109  		if len(bs) == 1 {
   110  			if bs[0] {
   111  				for j := 0; j < len(ctr.bat.Zs); j++ {
   112  					for k, rp := range ap.Result {
   113  						if rp.Rel == 0 {
   114  							if err := vector.UnionOne(rbat.Vecs[k], bat.Vecs[rp.Pos], int64(i), proc.Mp()); err != nil {
   115  								vec.Free(proc.Mp())
   116  								rbat.Clean(proc.Mp())
   117  								return err
   118  							}
   119  						} else {
   120  							if err := vector.UnionOne(rbat.Vecs[k], ctr.bat.Vecs[rp.Pos], int64(j), proc.Mp()); err != nil {
   121  								vec.Free(proc.Mp())
   122  								rbat.Clean(proc.Mp())
   123  								return err
   124  							}
   125  						}
   126  					}
   127  					rbat.Zs = append(rbat.Zs, ctr.bat.Zs[j])
   128  				}
   129  			}
   130  		} else {
   131  			for j, b := range bs {
   132  				if b {
   133  					for k, rp := range ap.Result {
   134  						if rp.Rel == 0 {
   135  							if err := vector.UnionOne(rbat.Vecs[k], bat.Vecs[rp.Pos], int64(i), proc.Mp()); err != nil {
   136  								rbat.Clean(proc.Mp())
   137  								vec.Free(proc.Mp())
   138  								return err
   139  							}
   140  						} else {
   141  							if err := vector.UnionOne(rbat.Vecs[k], ctr.bat.Vecs[rp.Pos], int64(j), proc.Mp()); err != nil {
   142  								rbat.Clean(proc.Mp())
   143  								vec.Free(proc.Mp())
   144  								return err
   145  							}
   146  						}
   147  					}
   148  					rbat.Zs = append(rbat.Zs, ctr.bat.Zs[j])
   149  				}
   150  			}
   151  		}
   152  		vec.Free(proc.Mp())
   153  	}
   154  	rbat.ExpandNulls()
   155  	anal.Output(rbat, isLast)
   156  	proc.SetInputBatch(rbat)
   157  	return nil
   158  }