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