github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/colexec/order/order.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 order
    16  
    17  import (
    18  	"bytes"
    19  
    20  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    21  	"github.com/matrixorigin/matrixone/pkg/container/nulls"
    22  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    23  	"github.com/matrixorigin/matrixone/pkg/partition"
    24  	"github.com/matrixorigin/matrixone/pkg/pb/plan"
    25  	"github.com/matrixorigin/matrixone/pkg/sort"
    26  	"github.com/matrixorigin/matrixone/pkg/sql/colexec"
    27  	"github.com/matrixorigin/matrixone/pkg/vm/process"
    28  )
    29  
    30  func String(arg any, buf *bytes.Buffer) {
    31  	ap := arg.(*Argument)
    32  	buf.WriteString("τ([")
    33  	for i, f := range ap.Fs {
    34  		if i > 0 {
    35  			buf.WriteString(", ")
    36  		}
    37  		buf.WriteString(f.String())
    38  	}
    39  	buf.WriteString("])")
    40  }
    41  
    42  func Prepare(_ *process.Process, arg any) error {
    43  	ap := arg.(*Argument)
    44  	ap.ctr = new(container)
    45  	{
    46  		ap.ctr.desc = make([]bool, len(ap.Fs))
    47  		ap.ctr.nullsLast = make([]bool, len(ap.Fs))
    48  		ap.ctr.vecs = make([]evalVector, len(ap.Fs))
    49  		for i, f := range ap.Fs {
    50  			ap.ctr.desc[i] = f.Flag&plan.OrderBySpec_DESC != 0
    51  			if f.Flag&plan.OrderBySpec_NULLS_FIRST != 0 {
    52  				ap.ctr.nullsLast[i] = false
    53  			} else if f.Flag&plan.OrderBySpec_NULLS_LAST != 0 {
    54  				ap.ctr.nullsLast[i] = true
    55  			} else {
    56  				ap.ctr.nullsLast[i] = ap.ctr.desc[i]
    57  			}
    58  		}
    59  	}
    60  	return nil
    61  }
    62  
    63  func Call(idx int, proc *process.Process, arg any, isFirst bool, isLast bool) (bool, error) {
    64  	anal := proc.GetAnalyze(idx)
    65  	anal.Start()
    66  	defer anal.Stop()
    67  
    68  	bat := proc.InputBatch()
    69  	ap := arg.(*Argument)
    70  	if bat == nil {
    71  		ap.Free(proc, false)
    72  		return true, nil
    73  	}
    74  	if bat.Length() == 0 {
    75  		return false, nil
    76  	}
    77  	end, err := ap.ctr.process(ap, bat, proc)
    78  	if err != nil {
    79  		ap.Free(proc, true)
    80  		return false, err
    81  	}
    82  	return end, nil
    83  }
    84  
    85  func (ctr *container) process(ap *Argument, bat *batch.Batch, proc *process.Process) (bool, error) {
    86  	for i := 0; i < bat.VectorCount(); i++ {
    87  		vec := bat.GetVector(int32(i))
    88  		if vec.IsOriginal() {
    89  			nvec, err := vector.Dup(bat.Vecs[i], proc.Mp())
    90  			if err != nil {
    91  				return false, err
    92  			}
    93  			bat.SetVector(int32(i), nvec)
    94  
    95  		}
    96  	}
    97  	for i, f := range ap.Fs {
    98  		vec, err := colexec.EvalExpr(bat, proc, f.Expr)
    99  		if err != nil {
   100  			return false, err
   101  		}
   102  		ctr.vecs[i].vec = vec
   103  		ctr.vecs[i].needFree = true
   104  		for j := range bat.Vecs {
   105  			if bat.Vecs[j] == vec {
   106  				ctr.vecs[i].needFree = false
   107  				break
   108  			}
   109  		}
   110  	}
   111  	defer ctr.cleanEvalVectors(proc.Mp())
   112  	ovec := ctr.vecs[0].vec
   113  	var strCol []string
   114  
   115  	sels := make([]int64, len(bat.Zs))
   116  	for i := 0; i < len(bat.Zs); i++ {
   117  		sels[i] = int64(i)
   118  	}
   119  
   120  	nullCnt := nulls.Length(ovec.Nsp)
   121  	// skip sort for all nulls
   122  	if nullCnt < ovec.Length() {
   123  		if ovec.Typ.IsString() {
   124  			strCol = vector.GetStrVectorValues(ovec)
   125  		} else {
   126  			strCol = nil
   127  		}
   128  		sort.Sort(ctr.desc[0], ctr.nullsLast[0], nullCnt > 0, sels, ovec, strCol)
   129  	}
   130  	if len(ctr.vecs) == 1 {
   131  		if err := bat.Shuffle(sels, proc.Mp()); err != nil {
   132  			panic(err)
   133  		}
   134  		return false, nil
   135  	}
   136  	ps := make([]int64, 0, 16)
   137  	ds := make([]bool, len(sels))
   138  	for i, j := 1, len(ctr.vecs); i < j; i++ {
   139  		desc := ctr.desc[i]
   140  		nullsLast := ctr.nullsLast[i]
   141  		ps = partition.Partition(sels, ds, ps, ovec)
   142  		vec := ctr.vecs[i].vec
   143  		nullCnt = nulls.Length(vec.Nsp)
   144  		// skip sort for all nulls
   145  		if nullCnt < vec.Length() {
   146  			if vec.Typ.IsString() {
   147  				strCol = vector.GetStrVectorValues(vec)
   148  			} else {
   149  				strCol = nil
   150  			}
   151  			for i, j := 0, len(ps); i < j; i++ {
   152  				if i == j-1 {
   153  					sort.Sort(desc, nullsLast, nullCnt > 0, sels[ps[i]:], vec, strCol)
   154  				} else {
   155  					sort.Sort(desc, nullsLast, nullCnt > 0, sels[ps[i]:ps[i+1]], vec, strCol)
   156  				}
   157  			}
   158  		}
   159  		ovec = vec
   160  	}
   161  	if err := bat.Shuffle(sels, proc.Mp()); err != nil {
   162  		panic(err)
   163  	}
   164  	return false, nil
   165  }