解决 Go 中的十亿行挑战（从 到 s）

不久前，一位朋友告诉我一个挑战，涉及读取 10 亿行文件。我发现这个想法很有趣，但由于当时是大学考试周，所以我最终把它留到以后再看。几个月后，我看到了 Theo 拍摄的有关挑战的视频，并决定仔细观察。

十亿行挑战赛的目标是计算一系列城市的最低、最高和平均温度 - 具体来说，这个列表中有 10 亿个项目，其中每个项目由一个城市的名称组成和温度，每个城市都可以出现多次。最后，程序必须按城市名称的字母顺序显示这些值。

我认为尝试解决挑战会很有趣，即使没有奖励。无论如何，我写了这篇文字来描述我的过程。

第一次尝试：使代码工作

每当我需要解决更复杂的问题时，我的首要目标就是让程序运行。它可能不是最快或最干净的代码，但它是有效的代码。

基本上，我创建了位置结构来表示列表中的每个城市，包含最低和最高温度、温度总和以及城市在列表中出现的次数（最后两个是计算平均值所必需的） 。我知道有一种方法可以直接计算平均值，而无需存储温度数及其总和。但正如我之前提到的，目标是让代码正常工作。

数据列表由城市名称后跟温度组成，并以分号分隔。例如：

Antananarivo;15.6
Iqaluit;-20.7
Dolisie;25.8
Kuopio;-6.8

读取数据最简单的方法是使用 Scan，它允许您一次读取一行。通过该行，您可以将其分为两部分：分号之前和之后的值。要获取温度，您可以使用 strconv.ParseFloat，它将字符串转换为浮点数。第一次实现的完整代码如下：

package main

import (
    "bufio"
    "fmt"
    "math"
    "os"
    "sort"
    "strconv"
    "strings"
)

type Location struct {
    min   float64
    max   float64
    sum   float64
    count int
}

func NewLocation() *Location {
    return &Location{
        min:   math.MaxInt16,
        max:   math.MinInt16,
        sum:   0,
        count: 0,
    }
}

func (loc *Location) Add(temp float64) {
    if temp < loc.min {
        loc.min = temp
    } else if temp > loc.max {
        loc.max = temp
    }

    loc.sum += temp
    loc.count += 1
}

var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")

func main() {
    flag.Parse()
    if *cpuprofile != "" {
        f, err := os.Create(*cpuprofile)
        if err != nil {
            log.Fatal(err)
        }
        pprof.StartCPUProfile(f)
        defer pprof.StopCPUProfile()
    }

    file, _ := os.Open("./measurements.txt")
    defer file.Close()

    m := map[string]*Location{}

    scanner := bufio.NewScanner(file)
    for scanner.Scan() {
        line := scanner.Text()
        name, tempStr, _ := strings.Cut(line, ";")
        temp, _ := strconv.ParseFloat(tempStr, 32)

        loc, ok := m[name]
        if !ok {
            loc = NewLocation()
            m[name] = loc
        }
        loc.Add(temp)
    }

    keys := make([]string, 0, len(m))
    for k := range m {
        keys = append(keys, k)
    }
    sort.Strings(keys)

    for _, name := range keys {
        loc := m[name]
        mean := loc.sum / float64(loc.count)
        fmt.Printf("%s: %.1f/%.1f/%.1f\n", name, loc.min, mean, loc.max)
    }
}

这个更简单的版本运行大约需要 97 秒。

优化字符串到浮点数的转换

分析执行配置文件，我意识到最大的瓶颈之一是 strconv.ParseFloat 函数。基本上，它的总执行时间是 23 秒（约占总时间的 23%）。

这个函数的问题是它是通用的，也就是说，它可以与任何有效的浮点数一起使用。然而，我们的数据具有非常特定的温度格式。请参阅下面的示例：

Antananarivo;15.6
Iqaluit;-20.7
Dolisie;5.8
Kuopio;-6.8

温度格式始终相同：点前一位或两位数字，点后一位数字，并且可能在开头包含一个减号。因此，我们可以创建一个以特定方式转换值的函数，从而优化流程，而无需执行 ParseFloat 的所有通用检查。

func bytesToTemp(b []byte) float64 {
    var v int16
    var isNeg int16 = 1

    for i := 0; i < len(b)-1; i++ {
        char := b[i]
        if char == '-' {
            isNeg = -1
        } else if char == '.' {
            digit := int16(b[i+1] - '0')
            v = v*10 + digit
        } else {
            digit := int16(char - '0')
            v = v*10 + digit
        }
    }

    return float64(v*isNeg) / 10
}

为了以字节格式而不是字符串读取数据，我将扫描器的返回从字符串更改为字节

line := scanner.Bytes()
before, after, _ := bytes.Cut(line, []byte{';'})
name := string(before)
temp := bytesToTemp(after)

这些小改动将执行时间缩短至 75 秒。

读取更大的数据块

使用Scan的最大优点是程序不需要一次将整个文件加载到内存中。相反，它可以让您逐行读取，以性能换取内存。

需要注意的是，一次读取一行和一次加载 14 GB 数据之间存在折衷。这个中间立场是读取块，即内存片段。这样，我们就可以读取 128 MB 的块，而不是一次读取整个文件。

buf := make([]byte, chunkSize)
reader := bufio.NewReader(file)
var leftData []byte
for {
    n, err := reader.Read(buf)
    if err != nil {
        if err == io.EOF {
            break
        }
        panic(err)
    }

    chunk := append(leftData, buf[:n]...)
    lastIndex := bytes.LastIndex(chunk, []byte{'\n'})
    leftData = chunk[lastIndex+1:]

    lines := bytes.Split(chunk[:lastIndex], []byte{'\n'})

    for _, line := range lines {
        before, after, _ := bytes.Cut(line, []byte{';'})
        name := string(before)
        temp := bytesToTemp(after)

        loc, ok := m[name]
        if !ok {
            loc = NewLocation()
            m[name] = loc
        }
        loc.Add(temp)
    }
}

结果，执行时间下降到了 70 秒。比以前好多了，但仍有改进的空间。

更改数据类型

事实上，整个挑战都围绕着带小数位的数字。然而，处理浮点始终是一个巨大的挑战（参见 IEEE-754）。既然如此，为什么不用整数来表示温度呢？

type Location struct {
    min   int16
    max   int16
    sum   int32
    count int32
}

如之前所定义，温度始终由最多三位数字表示。因此，除去逗号，值可以在-999和999之间变化，因此int16足以表示它们。对于求和和计数，int32 绰绰有余，因为该类型的范围在 -2147483648 到 2147483647 之间。

Dado que agora esperamos um valor inteiro de 16 bits para a temperatura, precisamos modificar a função bytesToTemp. Para isso, mudamos o retorno para int16 e removemos a divisão no final. Assim, a função vai sempre vai retornar um número inteiro.

func bytesToTemp(b []byte) int16 {
    var v int16
    var isNeg int16 = 1

    for i := 0; i < len(b)-1; i++ {
        char := b[i]
        if char == '-' {
            isNeg = -1
        } else if char == '.' {
            digit := int16(b[i+1] - '0')
            v = v*10 + digit
        } else {
            digit := int16(char - '0')
            v = v*10 + digit
        }
    }

    return v * isNeg
}

Para finalizar, modifiquei a função Add para aceitar os valores inteiros e ajustei o print para dividir os valores antes de mostrá-los na tela. Com isso, o tempo caiu três segundos, indo para 60 segundos. Não é muito, mas uma vitória é uma vitória.

Melhorando a Performance da Conversão de Bytes para String

Novamente analisando o profile, vi que tinha uma certa função chamada slicebytetostring que custava 13,5 segundos de tempo de execução. Analisando, descobri que essa função é a responsável por converter um conjunto de bytes em uma string (o próprio nome da função deixa claro isso). No caso, essa é a função chamada internamente quando se usa a função string(bytes).

Em Go, assim como na maioria das linguagens, strings são imutáveis, o que significa que não podem ser modificadas após serem criadas (normalmente, quando se faz isso, uma nova string é criada). Por outro lado, listas são mutáveis. Ou seja, quando se faz uma conversão de uma lista de bytes para string, é preciso criar uma cópia da lista para garantir que a string não mude se a lista mudar.

Para evitar o custo adicional de alocação de memória nessas conversões, podemos utilizar a biblioteca unsafe para realizar a conversão de bytes para string (Nota: ela é chamada de unsafe por um motivo).

name := unsafe.String(unsafe.SliceData(before), len(before))

Diferente do caso anterior, a função acima reutiliza os bytes passados para gerar a string. O problema disso é que, se a lista original mudar, a string resultante também será afetada. Embora possamos garantir que isso não ocorrerá neste contexto específico, em aplicações maiores e mais complexas, o uso de unsafe pode se tornar bem inseguro.

Com essa mudança, reduzimos o tempo de execução para 51 segundos. Nada mal.

Reimplementando bytes.Cut

Lembra que eu mencionei que as temperaturas sempre tinham formatos específicos? Então, segundo o profile da execução, que separa a linha em duas partes (nome da cidade e temperatura), custa 5.38 segundos para rodar. E refizermos ela na mão?

Para separar os dois valores, precisamos encontrar onde está o ";". Como a gente já sabe, os valores da temperatura podem ter entre três e cinco caracteres. Assim, precisamos verificar se o caractere anterior aos dígitos é um ";". Simples, não?

idx := 0
if line[len(line)-4] == ';' {
    idx = len(line) - 4
} else if line[len(line)-5] == ';' {
    idx = len(line) - 5
} else {
    idx = len(line) - 6
}

before := line[:idx]
after := line[idx+1:]

Com isso, o tempo de execução foi para 46 segundos, cerca de 5 segundos a menos que antes (quem diria, não é?).

Paralelismo para acelerar o processamento

Todo esse tempo, nosso objetivo foi tornar o código o mais rápido possível em um núcleo. Mudando coisas aqui e ali, diminuímos o tempo de 97 segundos para 46 segundos. Claro, ainda daria para melhorar o tempo sem ter que lidar com paralelismo, mas a vida é curta demais para se preocupar com isso, não é?

Para poder rodar o código em vários núcleos, decidi usar a estrutura de canais nativa do Go. Além disso, também criei um grupo de espera que vai indicar quando o processamento dos dados terminaram.

Vale destacar que workers, nesse caso, é uma constante que define quantas goroutines serão criadas para processar os dados. No meu caso, são 12, visto que eu tenho um processador com 6 núcleos e 12 threads.

chunkChan := make(chan []byte, workers)
var wg sync.WaitGroup
wg.Add(workers)

O próximo passo foi criar as goroutines que serão responsável por receber os dados do canal e processá-los. A lógica de processamento dos dados é semelhante ao modelo single thread.

for i := 0; i < workers; i++ {
    go func() {
        for chunk := range chunkChan {
            lines := bytes.Split(chunk, []byte{'\n'})
            for _, line := range lines {
                before, after := parseLine(line)
                name := unsafe.String(unsafe.SliceData(before), len(before))
                temp := bytesToTemp(after)

                loc, ok := m[name]
                if !ok {
                    loc = NewLocation()
                    m[name] = loc
                }
                loc.Add(temp)
            }
        }
        wg.Done()
    }()
}

Por fim, o código responsável por ler os dados do disco e enviá-los ao canal:

for {
    n, err := reader.Read(buf)
    if err != nil {
        if err == io.EOF {
            break
        }
        panic(err)
    }

    chunk := append(leftData, buf[:n]...)
    lastIndex := bytes.LastIndex(chunk, []byte{'\n'})
    leftData = chunk[lastIndex+1:]
    chunkChan <- chunk[:lastIndex]
}
close(chunkChan)
wg.Wait()

Vale ressaltar que os mapas em Go não são thread-safe. Isso significa que acessar ou alterar dados no mesmo mapa de forma concorrente pode levar a problemas de consistência ou erros. Embora não tenha observado problemas durante meus testes, vale a pena tratar esse problema.

Uma das maneiras de resolver esse problema seria criar um mecanismo de trava para o mapa, permitindo que somente uma goroutine consiga utilizá-lo por vez. Isso, claro, poderia tornar a execução um pouco mais lenta.

A segunda alternativa consiste em criar um mapa para cada uma das goroutines, de modo que não vai existir concorrência entre elas. Por fim, os mapas são colocados em um novo canal e os valores do mapa principal calculados a partir deles. Essa solução ainda vai ter um custo, mas vai ser menor que a anterior.

close(chunkChan)

go func() {
    wg.Wait()
    close(mapChan)
}()

keys := make([]string, 0, 825)

m := map[string]*Location{}
for lm := range mapChan {
    for lk, lLoc := range lm {
        loc, ok := m[lk]
        if !ok {
            keys = append(keys, lk)
            m[lk] = lLoc
            continue
        }

        if lLoc.min < loc.min {
            loc.min = lLoc.min
        }
        if lLoc.max > loc.max {
            loc.max = lLoc.max
        }
        loc.sum += lLoc.sum
        loc.count += lLoc.count
    }
}

Além disso, como o processamento passou a ser distribuído entre diferentes núcleos, diminui o tamanho do chunk de 128 MB para 2 MB. Cheguei nesse número testando vários valores, tendo entre 1 MB e 5 MB os melhores resultando. Na média, 2 MB obteve o melhor desempenho.

Enfim, o nosso tempo de processamento caiu de 46 segundos para... 12 segundos.

Otimizando a quebra de linhas no chunk

Todas as vezes que eu analisava o profile, a função bytes.Split chamava a minha atenção. O tempo de execução dela era de 16 segundos (tempo total, considerando todas as goroutines), o que parece justo, visto que ela é responsável por quebrar um chunk em linhas. No entanto, parecia um trabalho dobrado, dado que ela primeiro quebra as linhas para, em seguida, as linhas serem lidas uma a uma. Por que não fazer ambos ao mesmo tempo?

Minha abordagem foi percorrer o chunk e verificar se o byte atual correspondia a um \n. Dessa forma, consegui percorrer todas as linhas ao mesmo tempo em que as quebrava, processando em seguida.

start := 0
start := 0
for end, b := range chunk {
    if b != '\n' {
        continue
    }
    before, after := parseLine(chunk[start:end])
    // ...
    start = end + 1
}

Com essa simples mudança, o tempo de execução caiu para aproximadamente 9 segundos.

Executed in    8.45 secs    fish           external
   usr time   58.47 secs  152.00 micros   58.47 secs
   sys time    4.52 secs  136.00 micros    4.52 secs

Próximos passos

Atualmente, o maior gargalo da aplicação é o mapa. Somando todas as operações de leitura e escrita, são 32 segundos (de longe, o tempo mais alto). Talvez criar um algoritmo de hash mais eficiente resolva? Fica como ideia para o futuro.

No mais, conseguimos diminuir o tempo de 1 minuto e 40 segundos para quase 8 segundos, sem usar qualquer biblioteca externa. Além disso, tentando fazer a aplicação ficar cada vez mais rápida, me fez aprender muita coisa.

以上是解决 Go 中的十亿行挑战（从 到 s）的详细内容。更多信息请关注PHP中文网其他相关文章！