4 秒處理 10 億行數據! Go 語言的 9 大代碼方案,一個比一個快

編譯 | 核子可樂、凌敏

來自 | InfoQ

2024 年開年,Java “十億行挑戰”(1BRC)火爆外網。該挑戰賽要求開發者編寫一個 Java 程序,從一個包含十億行信息的文本文件中檢索溫度測量值,並計算每個氣象站的最小、平均值和最高溫度。“十億行挑戰” 的目標是爲這項任務創建最快的實現,同時探索現代 Java 的優勢。

這項挑戰聽起來很簡單。但十億行代碼實際是一項龐大的工程,如果以每個數字 3 秒的速度數到 10 億,大約需要 95.1 年!因此該挑戰最大的難度在於,處理文件以在儘可能短的時間內打印輸出:

該挑戰很快在 Hacker News、lobste.rs、Reddit 等社區掀起熱烈討論,不少開發者採用 Rust、Go、C++ 等其他編程語言甚至是數據庫參與挑戰。

日前,從業 20 年的軟件工程師 Ben Hoyt 用 Go 語言參與該挑戰,他一共想出了 9 種解決方案,完成 10 億行數據處理的時間最快只需 4 秒,最慢需要 1 分 45 秒。Ben Hoyt 還給自己提了點限制條件:每種方法都僅使用 Go 標準庫以保證可移植性,不涉及程序集、不涉及 unsafe、不涉及內存映射文件。跟其他作者的發現相比,Ben Hoyt 的解決方案不是最慢的、但也沒能佔據榜首。不過最重要的是,他的解跟其他參賽者的思路都不一樣,這種獨立性可能更具價值。

以下是 Ben Hoyt 用 Go 語言編寫的九種解決方案,每個方案都比前一個速度更快,概要如下:

首先通過幾條基線建立對這個任務的初步認識,看看使用 cat 讀取 13 GB 的數據需要多長時間:

$ time cat measurements.txt >/dev/null
0m1.052s

這裏 Ben Hoyt 一共測試了五次,所以實際上文件是被緩存過的。看起來 Linux 是允許把完整的 13 GB 數據都保留在磁盤緩存中,因爲第一次操作花了接近 6 秒時間,之後速度開始直線上升。

相比之下,對文件實際執行某些操作則要慢得多,wc 幾乎需要整整一分鐘:

$ time wc measurements.txt 
 1000000000  1179173106 13795293380 measurements.txt
0m55.710s

要快速爲這個問題找個簡單的方案,Ben Hoyt 可能會先從 AWK 開始。這種方法使用 Gawk,因爲其中的 asorti 函數可以輕鬆對輸出進行排序。Ben Hoyt 還加上了 -b 選項來使用 “將字符當作字節” 模式,這樣能讓速度更快一些:

$ time gawk -b -f 1brc.awk measurements.txt >measurements.out
7m35.567s

事實證明哪怕是最簡單粗暴的 Go 方法,也能在 7 分鐘左右搞定問題。下面就以此爲基礎摸索答案序列。

Ben Hoyt 首先優化出按序單核版本(方案 1 到 7),之後再對其做並行化調整(方案 8 和 9)。得到的所有結果都是在配備高速 SSD 驅動器和 32 GB 內存的 linux/amd64 筆記本電腦上,配合 Go 1.21.5 版本得到的。

Ben Hoyt 的許多方案和其他參與者給出的大部分最快方案,都會假設有效輸入。例如溫度數值只保留一位小數。如果輸入無效,那麼這幾種方案可能會引發運行時故障或產生錯誤輸出。

方案一:簡單且常見的 Go 代碼

作爲第一種方案,Ben Hoyt 的要求就是簡單且直接,只使用 Go 標準庫中的工具:bufio.Scanner 負責讀取數據行,strings.Cut 通過 “;” 進行分隔,strconv.ParseFloat 用於解析溫度,再加上普通的 Go map 來累積結果。

在方案一中,Ben Hoyt 完整列出首種方案的所有代碼,對於之後的方案就只給出比較有趣的部分:

func r1(inputPath string, output io.Writer) error {
    type stats struct {
        min, max, sum float64
        count         int64
    }
    f, err := os.Open(inputPath)
    if err != nil {
        return err
    }
    defer f.Close()
    stationStats := make(map[string]stats)
    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        line := scanner.Text()
        station, tempStr, hasSemi := strings.Cut(line, ";")
        if !hasSemi {
            continue
        }
        temp, err := strconv.ParseFloat(tempStr, 64)
        if err != nil {
            return err
        }
        s, ok := stationStats[station]
        if !ok {
            s.min = temp
            s.max = temp
            s.sum = temp
            s.count = 1
        } else {
            s.min = min(s.min, temp)
            s.max = max(s.max, temp)
            s.sum += temp
            s.count++
        }
        stationStats[station] = s
    }
    stations := make([]string, 0, len(stationStats))
    for station := range stationStats {
        stations = append(stations, station)
    }
    sort.Strings(stations)
    fmt.Fprint(output, "{")
    for i, station := range stations {
        if i > 0 {
            fmt.Fprint(output, ", ")
        }
        s := stationStats[station]
        mean := s.sum / float64(s.count)
        fmt.Fprintf(output, "%s=%.1f/%.1f/%.1f", station, s.min, mean, s.max)
    }
    fmt.Fprint(output, "}\n")
    return nil
}

這種最基本的方案能夠在 1 分 45 秒內完成 10 億行數據的處理。相較於 AWK 方案的 7 分鐘,這明顯是有了質的飛躍。

方案二:帶指針值的 map

Ben Hoyt 之前開發過一款單詞計數程序,當時就發現實際執行的哈希處理比理論需要的數量要多得多。對於每一行,我們都會對字符串執行兩次哈希處理:第一次用於從 map 中獲取值,第二次則是更新該 map。

爲了避免這種情況,我們可以使用 map[string]*stats(指針值)並更新指向的 struct,而不再使用 map[string]stats 並更新哈希表本體。

這裏,Ben Hoyt 首先想到用 Go 分析器來做確認。只需幾行代碼,即可將 CPU 分析添加到 Go 程序當中。

$ ./go-1brc -cpuprofile=cpu.prof -revision=1 measurements-10000000.txt >measurements-10000000.out
Processed 131.6MB in 965.888929ms
$ go tool pprof -http=: cpu.prof
...

通過在精簡後的 1000 萬行輸入文件上運行,這些命令爲方案一生成了以下概覽:

Map 操作佔用了整整 30% 的時間:其中 12.24% 用於分配,17.35% 用於查找。所以只要使用指針值,我們應該就能消除大部分 map 分配時間。

順帶一提,這張概覽圖還顯示出其餘時間的具體用途:

總的來說,方案二其實就是對 map 操作做一點小小調整:

stationStats := make(map[string]*stats)
scanner := bufio.NewScanner(f)
for scanner.Scan() {
    // ...
    s := stationStats[station]
    if s == nil {
        stationStats[station] = &stats{
            min:   temp,
            max:   temp,
            sum:   temp,
            count: 1,
        }
    } else {
        s.min = min(s.min, temp)
        s.max = max(s.max, temp)
        s.sum += temp
        s.count++
    }
}

在 map 中存在氣象站的常見情況下,我們現在只須執行一次 map 操作 s := stationStats[station],也就是說對氣象站名稱進行哈希處理並訪問哈希表的過程只需執行一次。即在氣象站已存在於 map 內的情況(在 10 億行數據中佔多數比例),我們會更新現有指向 struct。

雖然這對性能的提升不是太大,但也有一定效果:在 map 中使用指針值,可以將整個處理時間由 1 分 45 秒縮短至 1 分 31 秒。

方案三:去掉 strconv.ParseFloat

第三種方案相對比較硬核:用自定義代碼來取代 strconv.ParseFloat 進行溫度解析。標準庫函數會處理大量我們並不需要支持的極端溫度輸入情況,畢竟實際數據格式就是 1.2 或 34.5 這類 2 到 3 位數字(有些前面再多個負號)。

另外,strconv.ParseFloat 會接受一條字符串參數。現在我們不需要調用該參數,因此可以直接從 Scanner.Bytes 使用字節切片,而非藉助 Scanner.Text 進行字符串的分配和複製。

方法如下:

negative := false
index := 0
if tempBytes[index] == '-' {
    index++
    negative = true
}
temp := float64(tempBytes[index] - '0') // parse first digit
index++
if tempBytes[index] != '.' {
    temp = temp*10 + float64(tempBytes[index]-'0') // parse optional second digit
    index++
}
index++ // skip '.'
temp += float64(tempBytes[index]-'0') / 10 // parse decimal digit
if negative {
    temp = -temp
}

不太直觀,但也不至於特別難理解。方案三的處理時長從 1 分 31 秒成功縮短到了 1 分鐘以內:55.8 秒。

方案四:定點整數

曾幾何時,浮點指令的執行速度要比整數指令慢得多。現如今速度差距仍然存在,只是沒那麼誇張。但如果可以,把浮點轉換成整數還是會提高性能。

對於這個問題,每項溫度都有一個小數位,因此可以輕鬆用定點整數進行表示。例如,我們可以將 34.5 表示爲整數 345,然後在最終輸出結果之前再將其轉換回浮點數。

所以方案四跟方案三基本相同,只是將 stats struct 字段調整如下:

type stats struct {
    min, max, count int32
    sum             int64
}

之後在輸出結果時,再將數字除以 10:

mean := float64(s.sum) / float64(s.count) / 10
fmt.Fprintf(output, "%s=%.1f/%.1f/%.1f",
    station, float64(s.min)/10, mean, float64(s.max)/10)

這裏 Ben Hoyt 用 32 位整數表示最低和最高溫度,因爲最高溫度很可能達到約 500(即 50 攝氏度)。也可以使用 int16,但從過往的開發經驗來看,現代 64 位 CPU 在處理 16 位整數時速度要比 32 位整數更慢。在具體測試中,二者似乎沒有什麼可感知的差異,但 Ben Hoyt 還是優先選擇了 32 位。

使用整數之後,運行時間從 55.8 秒縮短到了 51.0 秒,效果也算顯著。

方案五:去掉 bytes.Cut

爲了推衍出方案五,Ben Hoyt 先爲方案四生成了一份概覽圖:

看來繼續優化是越來越困難了。很明顯,map 操作在其中占主導地位,轉爲自定義哈希表和去掉 bufio.Scanner 也非易事。所以這裏我們先試着去掉 bytes.Cut。

Ben Hoyt 想到一個簡單的辦法來節約時間。以原始文件中的一行數據爲例:

New Orleans;11.7

直接從後往前查找 “;” 來解析溫度,其速度會比直接掃描完整氣象站名稱來查找 “;” 更快。下面這段不怎麼優雅的代碼就是幹這個的:

end := len(line)
tenths := int32(line[end-1] - '0')
ones := int32(line[end-3] - '0') // line[end-2] is '.'
var temp int32
var semicolon int
if line[end-4] == ';' {          // positive N.N temperature
    temp = ones*10 + tenths
    semicolon = end - 4
} else if line[end-4] == '-' {   // negative -N.N temperature
    temp = -(ones*10 + tenths)
    semicolon = end - 5
} else {
    tens := int32(line[end-4] - '0')
    if line[end-5] == ';' {      // positive NN.N temperature
        temp = tens*100 + ones*10 + tenths
        semicolon = end - 5
    } else {                     // negative -NN.N temperature
        temp = -(tens*100 + ones*10 + tenths)
        semicolon = end - 6
    }
}
station := line[:semicolon]

迴避掉 bytes.Cut 之後,運行時間從 51.0 秒縮短到了 46.0 秒,又是一場小小的勝利。

方案六:去掉 bufio.Scanner

現在要嘗試去掉 bufio.Scanner 了。可以想到,要想查找每行的末尾,掃描器必須查看所有字節並尋找換行符。接下來,就是處理大量字節來解析溫度並找到 “;”。因此,我們可以嘗試把這些步驟整合起來,避免使用 bufio.Scanner。

在方案六中,我們分配了一個 1 MB 的緩衝區來讀取大塊文件,查找塊中的最後一個換行符來確保不會把單行截斷,之後再處理這些單個塊。具體代碼如下:

buf := make([]byte, 1024*1024)
readStart := 0
for {
    n, err := f.Read(buf[readStart:])
    if err != nil && err != io.EOF {
        return err
    }
    if readStart+n == 0 {
        break
    }
    chunk := buf[:readStart+n]
    newline := bytes.LastIndexByte(chunk, '\n')
    if newline < 0 {
        break
    }
    remaining := chunk[newline+1:]
    chunk = chunk[:newline+1]
    for {
        station, after, hasSemi := bytes.Cut(chunk, []byte(";"))
        // ... from here, same temperature processing as r4 ...

去掉 bufio.Scanner 並進行自主掃描之後,處理時間從 46.0 秒縮短到了 41.3 秒。效果不算太好,但至少感知得到。

方案七:自定義哈希表

方案七是這次探索中的真正核心。我們會自行建立一個自定義哈希表,而不再使用 Go map。這樣做有兩大優點:

  1. 我們可以在查找 “;” 時對氣象站名稱進行哈希處理,從而避免對字節的二次處理。

  2. 我們可以將哈希表中的每個鍵存儲爲字節切片,從而避免將各個鍵轉換爲 string(將在每一行上分配和複製)。

在 Go 中自定義哈希表並不複雜,只需使用帶有線性探測的 FNV-1a 哈希算法即可。如果發生衝突,則使用下一空槽。

爲了簡單起見,Ben Hoyt 預先分配了大量哈希桶(這裏共用到 10 萬個)以避免編寫邏輯來調整表的大小。但如果表的佔用比例超過了一半,代碼還是會出問題。通過測試,Ben Hoyt 發現引發哈希衝突的幾率大概是 2%。

爲了解決這次的問題,Ben Hoyt 還添加了一堆新代碼,包括哈希表設置、哈希本體以及表探測與插入:

// The hash table structure:
type item struct {
    key  []byte
    stat *stats
}
items := make([]item, 100000) // hash buckets, linearly probed
size := 0                     // number of active items in items slice
buf := make([]byte, 1024*1024)
readStart := 0
for {
    // ... same chunking as r6 ...
    for {
        const (
            // FNV-1 64-bit constants from hash/fnv.
            offset64 = 14695981039346656037
            prime64  = 1099511628211
        )
        // Hash the station name and look for ';'.
        var station, after []byte
        hash := uint64(offset64)
        i := 0
        for ; i < len(chunk); i++ {
            c := chunk[i]
            if c == ';' {
                station = chunk[:i]
                after = chunk[i+1:]
                break
            }
            hash ^= uint64(c) // FNV-1a is XOR then *
            hash *= prime64
        }
        if i == len(chunk) {
            break
        }
        // ... same temperature parsing as r6 ...
        // Go to correct bucket in hash table.
        hashIndex := int(hash & uint64(len(items)-1))
        for {
            if items[hashIndex].key == nil {
                // Found empty slot, add new item (copying key).
                key := make([]byte, len(station))
                copy(key, station)
                items[hashIndex] = item{
                    key: key,
                    stat: &stats{
                        min:   temp,
                        max:   temp,
                        sum:   int64(temp),
                        count: 1,
                    },
                }
                size++
                if size > len(items)/2 {
                    panic("too many items in hash table")
                }
                break
            }
            if bytes.Equal(items[hashIndex].key, station) {
                // Found matching slot, add to existing stats.
                s := items[hashIndex].stat
                s.min = min(s.min, temp)
                s.max = max(s.max, temp)
                s.sum += int64(temp)
                s.count++
                break
            }
            // Slot already holds another key, try next slot (linear probe).
            hashIndex++
            if hashIndex >= len(items) {
                hashIndex = 0
            }
        }
    }
    readStart = copy(buf, remaining)
}

這部分代碼帶來了巨大回報:自定義哈希表將處理時長從 41.3 秒縮短到了 25.8 秒。

方案八:並行處理各塊

在方案八中,Ben Hoyt 想引入一些並行性。但爲了控制變量,他打算繼續沿用方案 1 的代碼,畢竟更簡單且常見。方案一中保留了 bufio.Scanner 和 strconv.ParseFloat,這裏姑且直接將其並行化。在並行化成功之後,Ben Hoyt 再嘗試引入之後幾種方案的優化手段,雙管齊下的結果就是最終的方案九。

對這類 Map-Reduce 問題進行並行化並不困難:把文件拆分成大小相似的多個塊(每個 CPU 核心對應一個塊)、啓動一個線程(在 Go 中叫作 goroutine)來處理各個塊,最後把結果合併起來即可。

所以總體來看,代碼表示如下所示:

// Determine non-overlapping parts for file split (each part has offset and size).
parts, err := splitFile(inputPath, maxGoroutines)
if err != nil {
    return err
}
// Start a goroutine to process each part, returning results on a channel.
resultsCh := make(chan map[string]r8Stats)
for _, part := range parts {
    go r8ProcessPart(inputPath, part.offset, part.size, resultsCh)
}
// Wait for the results to come back in and aggregate them.
totals := make(map[string]r8Stats)
for i := 0; i < len(parts); i++ {
    result := <-resultsCh
    for station, s := range result {
        ts, ok := totals[station]
        if !ok {
            totals[station] = r8Stats{
                min:   s.min,
                max:   s.max,
                sum:   s.sum,
                count: s.count,
            }
            continue
        }
        ts.min = min(ts.min, s.min)
        ts.max = max(ts.max, s.max)
        ts.sum += s.sum
        ts.count += s.count
        totals[station] = ts
    }
}

由於 splitFile 函數有點繁瑣,所以這裏沒有使用。它負責查看文件的大小,除以我們指定的拆分塊數,然後查找每一塊,在末尾讀取 100 個字節並查找最後一個換行符,藉此確保每個塊在結尾都保留了整行(未將原始數據行截斷)。

r8ProcessPart 函數與方案 1 基本相同,但它會首先查找各個塊的偏移並將長度限制爲塊大小之內(使用 io.LimitedReader)。完成後,它會發回自己的 stats map:

func r8ProcessPart(inputPath string, fileOffset, fileSize int64,
                   resultsCh chan map[string]r8Stats) {
    file, err := os.Open(inputPath)
    if err != nil {
        panic(err)
    }
    defer file.Close()
    _, err = file.Seek(fileOffset, io.SeekStart)
    if err != nil {
        panic(err)
    }
    f := io.LimitedReader{R: file, N: fileSize}
    stationStats := make(map[string]r8Stats)
    scanner := bufio.NewScanner(&f)
    for scanner.Scan() {
        // ... same processing as r1 ...
    }
    resultsCh <- stationStats
}

相較於方案一,並行處理的性能表現出巨大優勢,成功將時間從 1 分 45 秒縮短到了 24.3 秒。相比之下,之前的 “優化但非並行” 版本(即方案七)需要耗費 25.8 秒。也就是說並行化比優化的性能增強效果更好,而且也簡單得多。

方案九:優化加並行

在方案九,也就是最終答案中,我們簡單將之前從方案一到七的所有優化方法,跟方案八中的並行化結合起來。

這裏 Ben Hoyt 使用了方案八中的 splitFile 函數,其餘代碼則直接從方案七處複製而來,所以這裏就不再贅述了。從結果來看,最終方案將處理時長從 24.3 秒進一步縮短到 3.99 秒。

有趣的是,由於所有實際處理現在都在單一大函數 r9ProcessPart 中進行,因此概覽圖就沒什麼用了。整個過程如下所示:

如大家所見,有 82% 的時間都花在了 r9ProcessPart 上,其中 bytes.Equal 佔用了 13%,文件讀取則佔用了餘下的 5%。

如果想要進一步分析,我們就得進一步下探到源視圖的層次。下面來看內部循環:

但這份報告還是有讓人迷惑的地方。爲什麼 if items[hashIndex].key == nil 顯示消費了 5.01 秒,但調用 bytes.Equal 則只用了 390 毫秒?難道說切片查找就是要比函數調用快得多?Ben Hoyt 自己也不太理解,歡迎各位 Go 性能大神在評論區中答疑解惑。

總而言之,這裏肯定還有更大的性能優化空間,但 4 秒之內處理 10 億行數據,也就是每秒 2.5 億行,這對不少開發者來說已經相當夠用了。

寫在最後

也許有人會問,折騰這些有意義嗎?

對於大多數日常編程任務,使用最簡單、最常見的代碼往往纔是王道。哪怕是面對超過 10 億行的溫度統計數據,如果只需要獲得一次答案,那麼 1 分 45 秒也絕非不可接受。

但如果我們正在構建數據處理管線,並且用以上種種方法把代碼執行速度提高了 4 倍、甚至是 26 倍,那麼不僅用戶體驗會更好,也能節約下大量計算成本。換言之,系統負載水平更低,且計算成本很可能只是原先的 1/4 甚至 1/26!

或者,如果大家正在構建像 GraalVM 這樣的運行時,或者像 Ben Hoyt 的 GoAWK 這種解釋器,那這樣的性能差異將極爲重要:解釋器的速度越快,一切用戶程序的運行速度都將同步提升。

哪怕退一萬步,單純嘗試讓代碼充分發揮機器性能本身也是種既有益、也有趣的嘗試,不是嗎?

參考鏈接

https://benhoyt.com/writings/go-1brc/

https://www.morling.dev/blog/one-billion-row-challenge/

本文由 Readfog 進行 AMP 轉碼,版權歸原作者所有。
來源https://mp.weixin.qq.com/s/_tP6bJYTSHGSZWpNZYdlEg