优化 Go 中对 CSV 的写入

一分耕耘，一分收获！既然打开了这篇文章《优化 Go 中对 CSV 的写入》，就坚持看下去吧！文中内容包含等等知识点...希望你能在阅读本文后，能真真实实学到知识或者帮你解决心中的疑惑，也欢迎大佬或者新人朋友们多留言评论，多给建议！谢谢！

以下代码片段验证电话号码并将详细信息写入 csv。

func Parse(phone Input, output *PhoneNumber) error {
    var n PhoneNumber
    num, _ := phonenumbers.Parse(phone.Number, phone.Prefix)
    n.PhoneNumber = phonenumbers.Format(num, phonenumbers.E164)
    n.CountryCode = num.GetCountryCode()
    n.PhoneType = phonenumbers.GetNumberType(num)
    n.NetworkName, _ = phonenumbers.GetCarrierForNumber(num, "EN")
    n.Region = phonenumbers.GetRegionCodeForNumber(num)
    *output = n
    return nil
}

func createFile(path string) {
    // detect if file exists
    var _, err = os.Stat(path)

    // create file if not exists
    if os.IsNotExist(err) {
        var file, err = os.Create(path)
        if err != nil {
            return
        }
        defer file.Close()
    }
}

func worker(ctx context.Context, dst chan string, src chan []string) {
    for {
        select {
        case dataArray, ok := <-src: // you must check for readable state of the channel.
            if !ok {
                return
            }
            go processNumber(dataArray[0])
        case <-ctx.Done(): // if the context is cancelled, quit.
            return
        }
    }
}

func processNumber(number string) {
    num, e := phonenumbers.Parse(number, "")
    if e != nil {
        return
    }
    region := phonenumbers.GetRegionCodeForNumber(num)
    carrier, _ := phonenumbers.GetCarrierForNumber(num, "EN")
    path := "sample_all.csv"
    createFile(path)
    var csvFile, _ = os.OpenFile(path, os.O_APPEND|os.O_WRONLY, os.ModeAppend)
    csvwriter := csv.NewWriter(csvFile)
    _ = csvwriter.Write([]string{phonenumbers.Format(num, phonenumbers.E164), fmt.Sprintf("%v", num.GetCountryCode()), fmt.Sprintf("%v", phonenumbers.GetNumberType(num)), carrier, region})

    defer csvFile.Close()

    csvwriter.Flush()
}

func ParseFile(phone Input, output *PhoneNumber) error {

    // create a context
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    // that cancels at ctrl+C
    go onSignal(os.Interrupt, cancel)

    numberOfWorkers := 2

    start := time.Now()

    csvfile, err := os.Open(phone.File)
    if err != nil {
        log.Fatal(err)
    }
    defer csvfile.Close()

    reader := csv.NewReader(csvfile)

    // create the pair of input/output channels for the controller=>workers com.
    src := make(chan []string)
    out := make(chan string)

    // use a waitgroup to manage synchronization
    var wg sync.WaitGroup

    // declare the workers
    for i := 0; i < numberOfWorkers; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            worker(ctx, out, src)
        }()
    }

    // read the csv and write it to src
    go func() {
        for {
            record, err := reader.Read()
            if err == io.EOF {
                break
            } else if err != nil {
                log.Fatal(err)
            }
            src <- record // you might select on ctx.Done().
        }
        close(src) // close src to signal workers that no more job are incoming.
    }()

    // wait for worker group to finish and close out
    go func() {
        wg.Wait()  // wait for writers to quit.
        close(out) // when you close(out) it breaks the below loop.
    }()

    // drain the output
    for res := range out {
        fmt.Println(res)
    }

    fmt.Printf("\n%2fs", time.Since(start).Seconds())
    return nil
}

在 processnumber 函数中，如果我跳过写入 csv，验证数字的过程会完成 6 秒，但在 csv 上一次写入一条记录会将时间消耗延长至 15 秒。

如何优化代码？ 我可以将记录分块并以块的形式写入，而不是一次写入一行吗？

解决方案

直接在工作协程中工作，而不是为每个任务触发协程。

打开文件输出文件一次。刷新输出文件一次。

func worker(ctx context.Context, dst chan []string, src chan []string) {
    for {
        select {
        case dataArray, ok := <-src: // you must check for readable state of the channel.
            if !ok {
                return
            }
            dst <- processNumber(dataArray[0])
        case <-ctx.Done(): // if the context is cancelled, quit.
            return
        }
    }
}

func processNumber(number string) []string {
    num, e := phonenumbers.Parse(number, "")
    if e != nil {
        return
    }
    region := phonenumbers.GetRegionCodeForNumber(num)
    carrier, _ := phonenumbers.GetCarrierForNumber(num, "EN")
    return []string{phonenumbers.Format(num, phonenumbers.E164), fmt.Sprintf("%v", num.GetCountryCode()), fmt.Sprintf("%v", phonenumbers.GetNumberType(num)), carrier, region}
}

func ParseFile(phone Input, output *PhoneNumber) error {

    // create a context
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()
    // that cancels at ctrl+C
    go onSignal(os.Interrupt, cancel)

    numberOfWorkers := 2

    start := time.Now()

    csvfile, err := os.Open(phone.File)
    if err != nil {
        log.Fatal(err)
    }
    defer csvfile.Close()

    reader := csv.NewReader(csvfile)

    // create the pair of input/output channels for the controller=>workers com.
    src := make(chan []string)
    out := make(chan string)

    // use a waitgroup to manage synchronization
    var wg sync.WaitGroup

    // declare the workers
    for i := 0; i < numberOfWorkers; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            worker(ctx, out, src)
        }()
    }

    // read the csv and write it to src
    go func() {
        for {
            record, err := reader.Read()
            if err == io.EOF {
                break
            } else if err != nil {
                log.Fatal(err)
            }
            src <- record // you might select on ctx.Done().
        }
        close(src) // close src to signal workers that no more job are incoming.
    }()

    // wait for worker group to finish and close out
    go func() {
        wg.Wait()  // wait for writers to quit.
        close(out) // when you close(out) it breaks the below loop.
    }()

    path := "sample_all.csv"

    file, err := os.Create(path)
    if err != nil {
        return err
    }
    defer file.Close()
    csvwriter := csv.NewWriter(csvFile)

    // drain the output
    for res := range out {
        csvwriter.Write(res)
    }
    csvwriter.Flush()

    fmt.Printf("\n%2fs", time.Since(start).Seconds())
    return nil
}

今天带大家了解了的相关知识，希望对你有所帮助；关于Golang的技术知识我们会一点点深入介绍，欢迎大家关注golang学习网公众号，一起学习编程~