12 11
用Go统计大文件小点滴

今天,团队里面的小伙伴说手头上有一些几个G以上的文件,需要统计文件里面的数据(文件里面的数据很简单都是Key:Value这样的组织形式)。然后咱俩随便抽了一个5G多的文件,用很普通的方式写了几行代码,流程是 读文件-->扫文件-->识别KV-->统计

然后程序跑了一趟,发现竟然要5分多钟。十分惊讶下,我想,明明很简单的逻辑的代码为什么要跑这么慢。然后我看了一下初版的代码,发现引起性能问题的几个点,分别是:

  • strconv的频繁使用
  • map[string]uint64形式的对象频繁调用
  • 为了识别KV,遍历文件时使用了 split 函数

下面是最开始的代码:


package main

import (
    "bufio"
    "bytes"
    "fmt"
    "log"
    "os"
    "strconv"
    "time"
)

var splitB = []byte(":")

var statsMap = make(map[string]int64, 4096)

func readAndHandleDataFile(filepath string) {
    f, err := os.Open(filepath)
    if err != nil {
        return
    }
    defer func() {
        f.Close()
    }()
    s := bufio.NewScanner(f)

    for s.Scan() {
        if b := s.Bytes(); b != nil {
            bs := bytes.Split(b, splitB)
            v, _ := strconv.ParseInt(string(bs[1]), 10, 0)
            statsMap[string(bs[0])] += v
        }
    }
    log.Printf("read  %s completed !\n", filepath)
}

func main() {
    args := os.Args
    if len(args) < 2 {
        panic("please tell me the data file")
    }
    dataFile := args[1]

    start := time.Now()
    if dataFile != "" {
        readAndHandleDataFile(dataFile)
    }
    elapsed := time.Now().Sub(start)
    log.Printf("total elapsed time: %f seconds", elapsed.Seconds())

    for k, v := range statsMap {
        fmt.Printf("[%s]:%d\n", k, v)
    }
}

为什么使用strconv的原因很简单,我们代码中需要把字符串转化为数值进行统计,于是直接使用了strconv.atoi的方法。

map[string]uint64 这样的结构主要是作为本地的统计信息集合,用于每次分析KV的时候对数值重新设置。

split 函数的目的更明确了,因为需要拿到Key和Value,当时脑袋第一反应就是使用了strings.split函数了

以上这三个点其实就是引起大文件耗时超过5分钟的直接原因了。

于是我对代码进行了小小的改造,如下:


package main

import (
    "bufio"
    "bytes"
    "fmt"
    "log"
    "os"
    "time"
)

var statsMap = make(map[uint64]uint64, 256)
var nameMap = make(map[uint64]string, 256)

//自定义哈希函数
func hashBytes(data []byte) uint64 {
    var h uint64 = 14695981039346656037
    for _, c := range data {
        h = (h ^ uint64(c)) * 1024
    }

    if _, ok := nameMap[h]; !ok {
        nameMap[h] = string(data)
    }
    return h
}

//把字符数组转化为无符号整型
func parsebyteToUint64(b []byte) (n uint64) {
    for i := 0; i < len(b); i++ {
        var v byte
        d := b[i]
        v = d - '0'
        n *= uint64(10)
        n1 := n + uint64(v)
        n = n1
    }
    return n
}

//处理数据文件
func readAndHandleDataFile(filepath string) {
    f, err := os.Open(filepath)
    if err != nil {
        return
    }
    defer f.Close()
    s := bufio.NewScanner(f)
    for s.Scan() {
        if b := s.Bytes(); b != nil {
            idx := bytes.IndexByte(b, ':') //分隔符所在索引位置
            hashVal := hashBytes(b[0:idx]) //计算哈希值
            statsMap[hashVal] += parsebyteToUint64(b[idx+1:])
        }
    }
    log.Printf("read %s completed !\n", filepath)
}

func main() {
    args := os.Args
    if len(args) < 2 {
        panic("please tell me the data file")
    }
    dataFile := args[1]

    start := time.Now()
    log.Println("read start")
    if dataFile != "" {
        readAndHandleDataFile(dataFile)
    }
    elapsed := time.Now().Sub(start)
    log.Printf("total elapsed time: %f seconds", elapsed.Seconds())

    for k, v := range statsMap {
        fmt.Printf("[%s]:%d\n", nameMap[k], v)
    }
}

然后再跑了一趟,耗时从原来的300多秒降到只需70秒左右,性能得到较大的提升。

最后的总结:

  • 对性能敏感的部分,尽量使用[]byte而不是直接操作字符串
  • 频繁操作map的时候,可以事先对key进行自定义快速的hash,而不直接使用string作为key(当然要考虑具体场景)
  • 类型转换在密集型场景下是代价很大的