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