golang解析csv並匯入elastic search

這是一個建立於 的文章，其中的資訊可能已經有所發展或是發生改變。

今天我們一起來寫一個從csv檔案將資料匯入elastic search的小程式

準備工作

在gopath的src檔案夾下建立csv2es檔案夾，並建立main.go檔案。下載一些csv檔案備用

解析命令列參數

首先我們需要使用flag package解析命令列參數,代碼如下

func main() {  // 解析命令列輸入  host := flag.String("host", "http://localhost:9200", "host, e.g. http://localhost:9200")  file := flag.String("file", "", "file path")  esIndex := flag.String("index", "", "elastic search index")  esType := flag.String("type", "", "elastic search type")  flag.Parse()  if *file == "" {    fmt.Println("please set which csv file you want to import clearly")    return  }  if *esIndex == "" {    fmt.Println("please set elastic search index")    return  }  if *esType == "" {    fmt.Println("please set elastic search type")    return  }...}

主要解析的參數有es的地址，待匯入的檔案的路徑，匯入到的es的index和type。運行go build，產生可執行檔csv2go,執行

csv2go -h-file string        file path (default "")  -host string        host, e.g. http://localhost:9200 (default "http://localhost:9200")  -index string        elastic search index (default "")  -type string        elastic search type (default "")

串連es

elastic這個開源項目可以協助我們串連elastic

  // 串連es  ctx := context.Background()  client, err := elastic.NewClient(    elastic.SetURL(*host),    elastic.SetSniff(false))  if err != nil {    panic(err)  }  // 檢查index是否存在，如果不存在則建立index  exists, err := client.IndexExists(*esIndex).Do(ctx)  if err != nil {    panic(err)  }  if !exists {    createIndex, err := client.CreateIndex(*esIndex).Do(ctx)    if err != nil {      panic(err)    }  }

解析csv並匯入(index)到elastic search

這裡需要注意幾點。第一，Mac上會存在\r結尾的檔案的問題，所以我們使用macreader這個包對io.Reader包了一層，有興趣的同學可以看我之前的文章《mac上的檔案有毒》，第二，我們預設csv檔案的第一行為column name，後面各行都是合法的記錄。

  // 解析csv  f, _ := os.Open(*file)  r := csv.NewReader(macreader.New(bufio.NewReader(f)))  keys, err := r.Read()  for {                                            //1    record, err := r.Read()    if err == io.EOF {      break    }    m := make(map[string]string)    for i, key := range keys {      m[key] = record[i]    }    jsonStr, err := json.Marshal(m)    if err != nil {      panic(err)    }    put1, err := client.Index().            Index(*esIndex).            Type(*esType).            BodyString(string(jsonStr)).            Do(ctx)    if err != nil {      // Handle error      panic(err)    }    fmt.Printf("Indexed tweet %s to index %s, type %s\n", put1.Id, put1.Index, put1.Type)  }                                           //2

ok,一個基本的將csv中的資料匯入elastic search的程式完成了，我們來測試一下效能吧。在上面的（1）行代碼前面加上

start := time.Now().Unix()

在上面的（2）行代碼後面加上

end := time.Now().Unix()fmt.Println(end, start)

測試了一下跑了一個61567條記錄的檔案，一共跑了36分鐘。夠我睡一個午覺了...

提升效率

elastic search有一個bulk api,可以將一些操作合并起來，同時傳遞給elastic search處理並返回

  // 建立一個Bulk  bulkRequest := client.Bulk()  for {    ...    // 為每一條記錄產生一個IndexRequest並加入Bulk    req := elastic.NewBulkIndexRequest().Index(*esIndex).Type(*esType).Doc(string(jsonStr))    bulkRequest.Add(req)  }  // 一次性完成請求  bulkResponse, err := bulkRequest.Do(ctx)  if err != nil {  }  indexed := bulkResponse.Indexed()  fmt.Println("向es匯入了",len(indexed),"條資料")

最佳化過後，插入相同的6萬多條記錄只需要幾秒鐘。cool ^_^

本文的代碼已經開源在github,歡迎使用或者提出意見。