github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/colexec/loopsemi/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 loopsemi
    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(" ⨝ ")
    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, pos := range ap.Result {
    95  		rbat.Vecs[i] = vector.New(bat.Vecs[pos].Typ)
    96  	}
    97  	count := bat.Length()
    98  	for i := 0; i < count; i++ {
    99  		vec, err := colexec.JoinFilterEvalExpr(bat, ctr.bat, i, proc, ap.Cond)
   100  		if err != nil {
   101  			rbat.Clean(proc.Mp())
   102  			return err
   103  		}
   104  		bs := vector.MustTCols[bool](vec)
   105  		for _, b := range bs {
   106  			if b {
   107  				for k, pos := range ap.Result {
   108  					if err := vector.UnionOne(rbat.Vecs[k], bat.Vecs[pos], int64(i), proc.Mp()); err != nil {
   109  						vec.Free(proc.Mp())
   110  						rbat.Clean(proc.Mp())
   111  						return err
   112  					}
   113  				}
   114  				rbat.Zs = append(rbat.Zs, bat.Zs[i])
   115  				break
   116  			}
   117  		}
   118  		vec.Free(proc.Mp())
   119  	}
   120  	rbat.ExpandNulls()
   121  	anal.Output(rbat, isLast)
   122  	proc.SetInputBatch(rbat)
   123  	return nil
   124  }