github.com/matrixorigin/matrixone@v1.2.0/pkg/vm/engine/tae/mergesort/func_test.go (about)

     1  // Copyright 2024 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  	"runtime"
    19  	"testing"
    20  
    21  	"github.com/matrixorigin/matrixone/pkg/common/mpool"
    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/vm/engine/tae/common"
    26  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers"
    27  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/testutils/mocks"
    28  	"github.com/stretchr/testify/require"
    29  )
    30  
    31  type testPool struct {
    32  	pool *containers.VectorPool
    33  }
    34  
    35  func (p *testPool) GetVector(typ *types.Type) (ret *vector.Vector, release func()) {
    36  	v := p.pool.GetVector(typ)
    37  	return v.GetDownstreamVector(), v.Close
    38  }
    39  
    40  func (p *testPool) GetMPool() *mpool.MPool {
    41  	return p.pool.GetMPool()
    42  }
    43  
    44  func TestReshapeBatches(t *testing.T) {
    45  	pool := &testPool{pool: mocks.GetTestVectorPool()}
    46  	batchSize := 5
    47  	fromLayout := []uint32{2000, 4000, 6000}
    48  	toLayout := []uint32{8192, 3808}
    49  
    50  	batches := make([]*batch.Batch, len(fromLayout))
    51  	for i := range fromLayout {
    52  		batches[i] = batch.NewWithSize(batchSize)
    53  		for j := 0; j < batchSize; j++ {
    54  			vec := containers.MakeVector(types.T_int32.ToType(), common.DefaultAllocator)
    55  			for k := uint32(0); k < fromLayout[i]; k++ {
    56  				vec.Append(int32(k), false)
    57  			}
    58  			batches[i].Vecs[j] = vec.GetDownstreamVector()
    59  			batches[i].SetRowCount(vec.Length())
    60  		}
    61  	}
    62  
    63  	var m1, m2 runtime.MemStats
    64  	runtime.GC()
    65  	runtime.ReadMemStats(&m1)
    66  	retBatch, releaseF := ReshapeBatches(batches, fromLayout, toLayout, pool)
    67  	runtime.ReadMemStats(&m2)
    68  	t.Log("total:", m2.TotalAlloc-m1.TotalAlloc)
    69  	t.Log("mallocs:", m2.Mallocs-m1.Mallocs)
    70  
    71  	for i, v := range retBatch {
    72  		require.Equal(t, toLayout[i], uint32(v.RowCount()))
    73  	}
    74  	releaseF()
    75  }