github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/mergesort/func.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 mergesort
    16  
    17  import (
    18  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    19  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    20  	"github.com/matrixorigin/matrixone/pkg/sort"
    21  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers"
    22  )
    23  
    24  /// merge things
    25  
    26  func SortBlockColumns(
    27  	cols []containers.Vector, pk int, pool *containers.VectorPool,
    28  ) ([]int64, error) {
    29  	pkCol := cols[pk]
    30  	sortedIdx := make([]int64, pkCol.Length())
    31  	for i := 0; i < len(sortedIdx); i++ {
    32  		sortedIdx[i] = int64(i)
    33  	}
    34  	sort.Sort(false, false, true, sortedIdx, pkCol.GetDownstreamVector(), nil)
    35  
    36  	for i := 0; i < len(cols); i++ {
    37  		err := cols[i].GetDownstreamVector().Shuffle(sortedIdx, pool.GetMPool())
    38  		if err != nil {
    39  			return nil, err
    40  		}
    41  	}
    42  	return sortedIdx, nil
    43  }
    44  
    45  func ReshapeBatches(
    46  	batches []*batch.Batch,
    47  	fromLayout, toLayout []uint32,
    48  	vpool DisposableVecPool) ([]*batch.Batch, func()) {
    49  	// just do reshape, keep sortedIdx nil
    50  	ret := make([]*batch.Batch, 0, len(toLayout))
    51  	rfs := make([]func(), 0, len(toLayout))
    52  	for _, layout := range toLayout {
    53  		bat, releaseF := getSimilarBatch(batches[0], int(layout), vpool)
    54  		bat.SetRowCount(int(layout))
    55  		ret = append(ret, bat)
    56  		rfs = append(rfs, releaseF)
    57  	}
    58  	releaseF := func() {
    59  		for _, rf := range rfs {
    60  			rf()
    61  		}
    62  	}
    63  
    64  	fromIdx := 0
    65  	fromOffset := 0
    66  	for i := 0; i < len(toLayout); i++ {
    67  		toOffset := 0
    68  		for toOffset < int(toLayout[i]) {
    69  			// find offset to fill a full block
    70  			fromLeft := int(fromLayout[fromIdx]) - fromOffset
    71  			if fromLeft == 0 {
    72  				fromIdx++
    73  				fromOffset = 0
    74  				fromLeft = int(fromLayout[fromIdx])
    75  			}
    76  			length := 0
    77  			if fromLeft < int(toLayout[i])-toOffset {
    78  				length = fromLeft
    79  			} else {
    80  				length = int(toLayout[i]) - toOffset
    81  			}
    82  
    83  			for vecIdx, vec := range batches[fromIdx].Vecs {
    84  				window, err := vec.Window(fromOffset, fromOffset+length)
    85  				if err != nil {
    86  					panic(err)
    87  				}
    88  				err = vector.GetUnionAllFunction(*vec.GetType(), vpool.GetMPool())(ret[i].Vecs[vecIdx], window)
    89  				if err != nil {
    90  					panic(err)
    91  				}
    92  			}
    93  
    94  			// update offset
    95  			fromOffset += length
    96  			toOffset += length
    97  		}
    98  	}
    99  	return ret, releaseF
   100  }