github.com/matrixorigin/matrixone@v0.7.0/pkg/sql/colexec/mergeblock/mergeblock_test.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  package mergeblock
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"reflect"
    20  	"testing"
    21  
    22  	"github.com/matrixorigin/matrixone/pkg/catalog"
    23  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    24  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    25  	"github.com/matrixorigin/matrixone/pkg/testutil"
    26  	"github.com/matrixorigin/matrixone/pkg/vm/engine"
    27  	"github.com/stretchr/testify/require"
    28  )
    29  
    30  type mockRelation struct {
    31  	engine.Relation
    32  	result *batch.Batch
    33  }
    34  
    35  func (e *mockRelation) Write(_ context.Context, b *batch.Batch) error {
    36  	e.result = b
    37  	return nil
    38  }
    39  
    40  func TestMergeBlock(t *testing.T) {
    41  	proc := testutil.NewProc()
    42  	proc.Ctx = context.TODO()
    43  	batch1 := &batch.Batch{
    44  		Attrs: []string{catalog.BlockMeta_TableIdx_Insert, catalog.BlockMeta_MetaLoc},
    45  		Vecs: []*vector.Vector{
    46  			testutil.MakeUint16Vector([]uint16{0, 1, 2}, nil),
    47  			testutil.MakeVarcharVector([]string{
    48  				"a.seg:magic:15",
    49  				"b.seg:magic:15",
    50  				"c.seg:magic:15"}, nil),
    51  		},
    52  		Zs: []int64{1, 1, 1},
    53  	}
    54  	argument1 := Argument{
    55  		Tbl:          &mockRelation{},
    56  		Unique_tbls:  []engine.Relation{&mockRelation{}, &mockRelation{}},
    57  		AffectedRows: 0,
    58  		notFreeBatch: true,
    59  	}
    60  	proc.Reg.InputBatch = batch1
    61  	Prepare(proc, &argument1)
    62  	_, err := Call(0, proc, &argument1, false, false)
    63  	require.NoError(t, err)
    64  	require.Equal(t, uint64(15), argument1.AffectedRows)
    65  	// Check Tbl
    66  	{
    67  		result := argument1.Tbl.(*mockRelation).result
    68  		// check attr names
    69  		require.True(t, reflect.DeepEqual(
    70  			[]string{catalog.BlockMeta_MetaLoc},
    71  			result.Attrs,
    72  		))
    73  		// check vector
    74  		require.Equal(t, 1, len(result.Vecs))
    75  		for i, vec := range result.Vecs {
    76  			require.Equal(t, 1, vector.Length(vec), fmt.Sprintf("column number: %d", i))
    77  		}
    78  	}
    79  	// Check UniqueTables
    80  	for j := range argument1.Unique_tbls {
    81  		tbl := argument1.Unique_tbls[j]
    82  		result := tbl.(*mockRelation).result
    83  		// check attr names
    84  		require.True(t, reflect.DeepEqual(
    85  			[]string{catalog.BlockMeta_MetaLoc},
    86  			result.Attrs,
    87  		))
    88  		// check vector
    89  		require.Equal(t, 1, len(result.Vecs))
    90  		for i, vec := range result.Vecs {
    91  			require.Equal(t, 1, vector.Length(vec), fmt.Sprintf("column number: %d", i))
    92  		}
    93  	}
    94  	argument1.Free(proc, false)
    95  	require.Equal(t, int64(0), proc.GetMPool().CurrNB())
    96  }