github.com/matrixorigin/matrixone@v0.7.0/pkg/vm/engine/tae/samples/sample3/main.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 main
    16  
    17  import (
    18  	"context"
    19  	"os"
    20  	"runtime/pprof"
    21  	"sync"
    22  	"time"
    23  
    24  	"github.com/matrixorigin/matrixone/pkg/common/mpool"
    25  	"github.com/matrixorigin/matrixone/pkg/container/batch"
    26  	"github.com/matrixorigin/matrixone/pkg/container/vector"
    27  	"github.com/matrixorigin/matrixone/pkg/logutil"
    28  	"github.com/matrixorigin/matrixone/pkg/vm/engine"
    29  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/catalog"
    30  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/common"
    31  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/containers"
    32  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/db"
    33  	"github.com/matrixorigin/matrixone/pkg/vm/engine/tae/moengine"
    34  	"github.com/panjf2000/ants/v2"
    35  )
    36  
    37  var sampleDir = "/tmp/sample3"
    38  var dbName = "db"
    39  var cpuprofile = "/tmp/sample3/cpuprofile"
    40  var memprofile = "/tmp/sample3/memprofile"
    41  
    42  func init() {
    43  	os.RemoveAll(sampleDir)
    44  }
    45  
    46  func startProfile() {
    47  	f, _ := os.Create(cpuprofile)
    48  	if err := pprof.StartCPUProfile(f); err != nil {
    49  		panic(err)
    50  	}
    51  }
    52  
    53  func stopProfile() {
    54  	pprof.StopCPUProfile()
    55  	memf, _ := os.Create(memprofile)
    56  	defer memf.Close()
    57  	_ = pprof.Lookup("heap").WriteTo(memf, 0)
    58  }
    59  
    60  func main() {
    61  	tae, _ := db.Open(sampleDir, nil)
    62  	defer tae.Close()
    63  	eng := moengine.NewEngine(tae)
    64  
    65  	ctx := context.TODO()
    66  	schema := catalog.MockSchema(13, 12)
    67  	{
    68  		txn, _ := eng.StartTxn(nil)
    69  		txnOperator := moengine.TxnToTxnOperator(txn)
    70  		err := eng.Create(ctx, dbName, txnOperator)
    71  		if err != nil {
    72  			panic(err)
    73  		}
    74  		db, err := eng.Database(ctx, dbName, txnOperator)
    75  		if err != nil {
    76  			panic(err)
    77  		}
    78  		defs, _ := moengine.SchemaToDefs(schema)
    79  		err = db.Create(ctx, schema.Name, defs)
    80  		if err != nil {
    81  			panic(err)
    82  		}
    83  		if err := txn.Commit(); err != nil {
    84  			panic(err)
    85  		}
    86  	}
    87  	batchCnt := uint32(100)
    88  	batchRows := uint32(10000) * 1 / 2 * batchCnt
    89  	logutil.Info(tae.Opts.Catalog.SimplePPString(common.PPL1))
    90  	bat := catalog.MockBatch(schema, int(batchRows))
    91  	newbat := batch.New(true, bat.Attrs)
    92  	newbat.Vecs = containers.CopyToMoVecs(bat.Vecs)
    93  	bats := containers.SplitBatch(newbat, int(batchCnt))
    94  	var wg sync.WaitGroup
    95  	doAppend := func(b *batch.Batch) func() {
    96  		return func() {
    97  			defer wg.Done()
    98  			txn, _ := eng.StartTxn(nil)
    99  			txnOperator := moengine.TxnToTxnOperator(txn)
   100  			// {
   101  			// 	db, _ := txn.GetDatabase(dbName)
   102  			// 	rel, _ := db.GetRelationByName(schema.Name)
   103  			// }
   104  			db, err := eng.Database(ctx, dbName, txnOperator)
   105  			if err != nil {
   106  				panic(err)
   107  			}
   108  			rel, err := db.Relation(ctx, schema.Name)
   109  			if err != nil {
   110  				panic(err)
   111  			}
   112  			if err := rel.Write(ctx, b); err != nil {
   113  				panic(err)
   114  			}
   115  			if err := txn.Commit(); err != nil {
   116  				panic(err)
   117  			}
   118  		}
   119  	}
   120  	p, _ := ants.NewPool(10)
   121  	now := time.Now()
   122  	startProfile()
   123  	for _, b := range bats {
   124  		wg.Add(1)
   125  		_ = p.Submit(doAppend(b))
   126  	}
   127  	wg.Wait()
   128  	stopProfile()
   129  	logutil.Infof("Append takes: %s", time.Since(now))
   130  	{
   131  		txn, _ := eng.StartTxn(nil)
   132  		txnOperator := moengine.TxnToTxnOperator(txn)
   133  		db, err := eng.Database(ctx, dbName, txnOperator)
   134  		if err != nil {
   135  			panic(err)
   136  		}
   137  		rel, err := db.Relation(ctx, schema.Name)
   138  		if err != nil {
   139  			panic(err)
   140  		}
   141  		m := mpool.MustNewZero()
   142  		readProc := func(reader engine.Reader) {
   143  			defer wg.Done()
   144  			for {
   145  				bat, err := reader.Read(ctx, []string{schema.ColDefs[0].Name}, nil, m)
   146  				if err != nil {
   147  					panic(err)
   148  				}
   149  				if bat == nil {
   150  					break
   151  				}
   152  				logutil.Infof("bat rows: %d", vector.Length(bat.Vecs[0]))
   153  			}
   154  		}
   155  
   156  		parallel := 10
   157  		readers, err := rel.NewReader(ctx, parallel, nil, nil)
   158  		if err != nil {
   159  			panic(err)
   160  		}
   161  		for _, reader := range readers {
   162  			wg.Add(1)
   163  			go readProc(reader)
   164  		}
   165  		wg.Wait()
   166  		if err = txn.Commit(); err != nil {
   167  			panic(err)
   168  		}
   169  	}
   170  	logutil.Info(tae.Opts.Catalog.SimplePPString(common.PPL1))
   171  }