github.com/ttys3/gojieba-bleve@v1.0.2/bleve_test.go (about)

     1  package jbleve
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"log"
     7  	"os"
     8  	"runtime"
     9  	"runtime/pprof"
    10  	"testing"
    11  
    12  	"github.com/blevesearch/bleve"
    13  	"github.com/yanyiwu/gojieba"
    14  )
    15  
    16  func Example() {
    17  	INDEX_DIR := "gojieba.bleve"
    18  	messages := []struct {
    19  		Id   string
    20  		Body string
    21  	}{
    22  		{
    23  			Id:   "1",
    24  			Body: "你好",
    25  		},
    26  		{
    27  			Id:   "2",
    28  			Body: "交代",
    29  		},
    30  		{
    31  			Id:   "3",
    32  			Body: "长江大桥",
    33  		},
    34  	}
    35  
    36  	indexMapping := bleve.NewIndexMapping()
    37  	os.RemoveAll(INDEX_DIR)
    38  	// clean index when example finished
    39  	defer os.RemoveAll(INDEX_DIR)
    40  
    41  	err := indexMapping.AddCustomTokenizer("gojieba",
    42  		map[string]interface{}{
    43  			"dictpath":     gojieba.DICT_PATH,
    44  			"hmmpath":      gojieba.HMM_PATH,
    45  			"userdictpath": gojieba.USER_DICT_PATH,
    46  			"idf":          gojieba.IDF_PATH,
    47  			"stop_words":   gojieba.STOP_WORDS_PATH,
    48  			"type":         "gojieba",
    49  		},
    50  	)
    51  	if err != nil {
    52  		panic(err)
    53  	}
    54  	err = indexMapping.AddCustomAnalyzer("gojieba",
    55  		map[string]interface{}{
    56  			"type":      "gojieba",
    57  			"tokenizer": "gojieba",
    58  		},
    59  	)
    60  	if err != nil {
    61  		panic(err)
    62  	}
    63  	indexMapping.DefaultAnalyzer = "gojieba"
    64  
    65  	index, err := bleve.New(INDEX_DIR, indexMapping)
    66  	if err != nil {
    67  		panic(err)
    68  	}
    69  	for _, msg := range messages {
    70  		if err := index.Index(msg.Id, msg); err != nil {
    71  			panic(err)
    72  		}
    73  	}
    74  
    75  	querys := []string{
    76  		"你好世界",
    77  		"亲口交代",
    78  		"长江",
    79  	}
    80  
    81  	for _, q := range querys {
    82  		req := bleve.NewSearchRequest(bleve.NewQueryStringQuery(q))
    83  		req.Highlight = bleve.NewHighlight()
    84  		res, err := index.Search(req)
    85  		if err != nil {
    86  			panic(err)
    87  		}
    88  		fmt.Println(prettify(res))
    89  	}
    90  
    91  	//cleanup cgo allocated heap memory
    92  	if jieba, ok := (index.Mapping().AnalyzerNamed("gojieba").Tokenizer).(*JiebaTokenizer); !ok {
    93  		panic("jieba.Free() failed")
    94  	} else {
    95  		jieba.Free()
    96  	}
    97  	index.Close()
    98  	// Output:
    99  	// [{"id":"1","score":0.27650412875470115}]
   100  	// [{"id":"2","score":0.27650412875470115}]
   101  	// [{"id":"3","score":0.7027325540540822}]
   102  }
   103  
   104  func BenchmarkExample(b *testing.B) {
   105  	// CPU Profile
   106  	cpuProfile, err := os.Create("cpu.prof")
   107  	if err != nil {
   108  		log.Fatal(err)
   109  	}
   110  	pprof.StartCPUProfile(cpuProfile)
   111  	defer pprof.StopCPUProfile()
   112  
   113  	for i := 0; i < 300; i++ {
   114  		Example()
   115  	}
   116  
   117  	// Memory Profile
   118  	f, err := os.Create("mem.prof")
   119  	if err != nil {
   120  		log.Fatal("could not create memory profile: ", err)
   121  	}
   122  	defer f.Close()
   123  	runtime.GC() // get up-to-date statistics
   124  	if err := pprof.WriteHeapProfile(f); err != nil {
   125  		log.Fatal("could not write memory profile: ", err)
   126  	}
   127  }
   128  
   129  func prettify(res *bleve.SearchResult) string {
   130  	type Result struct {
   131  		Id    string  `json:"id"`
   132  		Score float64 `json:"score"`
   133  	}
   134  	results := []Result{}
   135  	for _, item := range res.Hits {
   136  		results = append(results, Result{item.ID, item.Score})
   137  	}
   138  	b, err := json.Marshal(results)
   139  	if err != nil {
   140  		panic(err)
   141  	}
   142  	return string(b)
   143  }