github.com/matrixorigin/matrixone@v1.2.0/pkg/sql/colexec/restrict/restrict.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 restrict
    16  
    17  import (
    18  	"bytes"
    19  	"fmt"
    20  
    21  	"github.com/matrixorigin/matrixone/pkg/pb/plan"
    22  	"github.com/matrixorigin/matrixone/pkg/vm"
    23  
    24  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    25  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    26  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    27  
    28  	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
    29  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    30  )
    31  
    32  const argName = "restrict"
    33  
    34  func (arg *Argument) String(buf *bytes.Buffer) {
    35  	buf.WriteString(argName)
    36  	ap := arg
    37  	buf.WriteString(fmt.Sprintf("filter(%s)", ap.E))
    38  }
    39  
    40  func (arg *Argument) Prepare(proc *process.Process) (err error) {
    41  	ap := arg
    42  	ap.ctr = new(container)
    43  
    44  	filterList := colexec.SplitAndExprs([]*plan.Expr{ap.E})
    45  	ap.ctr.executors, err = colexec.NewExpressionExecutorsFromPlanExpressions(proc, filterList)
    46  	return err
    47  }
    48  
    49  func (arg *Argument) Call(proc *process.Process) (vm.CallResult, error) {
    50  	if err, isCancel := vm.CancelCheck(proc); isCancel {
    51  		return vm.CancelResult, err
    52  	}
    53  
    54  	result, err := arg.GetChildren(0).Call(proc)
    55  	if err != nil {
    56  		return result, err
    57  	}
    58  
    59  	anal := proc.GetAnalyze(arg.GetIdx(), arg.GetParallelIdx(), arg.GetParallelMajor())
    60  	anal.Start()
    61  	defer anal.Stop()
    62  
    63  	if result.Batch == nil || result.Batch.IsEmpty() || result.Batch.Last() {
    64  		return result, nil
    65  	}
    66  	if arg.buf != nil {
    67  		proc.PutBatch(arg.buf)
    68  		arg.buf = nil
    69  	}
    70  	arg.buf = result.Batch
    71  
    72  	anal.Input(arg.buf, arg.GetIsFirst())
    73  
    74  	var sels []int64
    75  	for i := range arg.ctr.executors {
    76  		if arg.buf.IsEmpty() {
    77  			break
    78  		}
    79  
    80  		vec, err := arg.ctr.executors[i].Eval(proc, []*batch.Batch{arg.buf})
    81  		if err != nil {
    82  			result.Batch = nil
    83  			return result, err
    84  		}
    85  
    86  		if proc.OperatorOutofMemory(int64(vec.Size())) {
    87  			return result, moerr.NewOOM(proc.Ctx)
    88  		}
    89  		anal.Alloc(int64(vec.Size()))
    90  		if !vec.GetType().IsBoolean() {
    91  			return result, moerr.NewInvalidInput(proc.Ctx, "filter condition is not boolean")
    92  		}
    93  
    94  		bs := vector.GenerateFunctionFixedTypeParameter[bool](vec)
    95  		if vec.IsConst() {
    96  			v, null := bs.GetValue(0)
    97  			if null || !v {
    98  				arg.buf, err = tryDupBatch(proc, arg.buf)
    99  				if err != nil {
   100  					return result, err
   101  				}
   102  				arg.buf.Shrink(nil, false)
   103  			}
   104  		} else {
   105  			if sels == nil {
   106  				sels = proc.Mp().GetSels()
   107  			}
   108  			sels = sels[:0]
   109  
   110  			l := uint64(vec.Length())
   111  			if bs.WithAnyNullValue() {
   112  				for j := uint64(0); j < l; j++ {
   113  					v, null := bs.GetValue(j)
   114  					if !null && v {
   115  						sels = append(sels, int64(j))
   116  					}
   117  				}
   118  			} else {
   119  				for j := uint64(0); j < l; j++ {
   120  					v, _ := bs.GetValue(j)
   121  					if v {
   122  						sels = append(sels, int64(j))
   123  					}
   124  				}
   125  			}
   126  			arg.buf, err = tryDupBatch(proc, arg.buf)
   127  			if err != nil {
   128  				return result, err
   129  			}
   130  			arg.buf.Shrink(sels, false)
   131  		}
   132  	}
   133  
   134  	if sels != nil {
   135  		proc.Mp().PutSels(sels)
   136  	}
   137  
   138  	// bad design here. should compile a pipeline like `-> restrict -> output (just do clean work or memory reuse) -> `
   139  	// but not use the IsEnd flag to do the clean work.
   140  	if arg.IsEnd {
   141  		result.Batch = nil
   142  	} else {
   143  		anal.Output(arg.buf, arg.GetIsLast())
   144  		if arg.buf == result.Batch {
   145  			arg.buf = nil
   146  		} else {
   147  			result.Batch = arg.buf
   148  		}
   149  	}
   150  	return result, nil
   151  }
   152  
   153  func tryDupBatch(proc *process.Process, bat *batch.Batch) (*batch.Batch, error) {
   154  	cnt := bat.GetCnt()
   155  	if cnt == 1 {
   156  		return bat, nil
   157  	}
   158  	newBat, err := bat.Dup(proc.Mp())
   159  	if err != nil {
   160  		return nil, err
   161  	}
   162  	// proc.PutBatch(bat)
   163  	return newBat, nil
   164  }