Go语言中高效读取大型CSV文件的随机行：水塘抽样法实践

package main

import (
    "encoding/csv"
    "fmt"
    "io"
    "math/rand"
    "os"
    "time"
)

// ReadRandomCSVLines 使用水塘抽样算法从CSV文件中读取指定数量的随机行。
// 它不会将整个文件加载到内存中，适用于处理大型CSV文件。
// filePath: CSV文件路径
// numSamples: 希望抽取的随机行数
// 返回值: 抽取的随机行（[][]string 格式），或错误信息
func ReadRandomCSVLines(filePath string, numSamples int) ([][]string, error) {
    if numSamples <= 0 {
        return nil, fmt.Errorf("numSamples 必须是正整数")
    }

    file, err := os.Open(filePath)
    if err != nil {
        return nil, fmt.Errorf("无法打开文件 %s: %w", filePath, err)
    }
    defer file.Close()

    reader := csv.NewReader(file)
    // 如果CSV文件包含标题行，可以在这里读取并跳过它：
    // _, err = reader.Read()
    // if err != nil && err != io.EOF {
    //  return nil, fmt.Errorf("无法读取CSV标题行: %w", err)
    // }

    // 初始化水塘，预分配容量以减少扩容开销
    reservoir := make([][]string, 0, numSamples)
    var linesRead int64 = 0 // 记录已读取的总行数，使用int64以支持超大文件

    // 初始化随机数生成器，使用当前时间作为种子以确保每次运行结果不同
    r := rand.New(rand.NewSource(time.Now().UnixNano()))

    for {
        record, err := reader.Read() // 读取下一条CSV记录
        if err == io.EOF {
            break // 文件读取完毕
        }
        if err != nil {
            return nil, fmt.Errorf("读取CSV记录失败: %w", err)
        }

        linesRead++ // 已读取行数加一

        if linesRead <= int64(numSamples) {
            // 如果已读取行数小于或等于目标样本数，直接填充水塘
            reservoir = append(reservoir, record)
        } else {
            // 对于后续行，以 numSamples/linesRead 的概率决定是否替换水塘中的一个元素
            // 生成一个 [0, linesRead-1] 之间的随机整数
            j := r.Int63n(linesRead)

            // 如果 j 小于 numSamples，则替换水塘中索引为 j 的元素
            if j < int64(numSamples) {
                reservoir[j] = record
            }
        }
    }

    return reservoir, nil
}

func main() {
    // 创建一个虚拟的CSV文件用于测试
    dummyData := `id,name,city,age
1,Alice,New York,30
2,Bob,London,24
3,Charlie,Paris,35
4,David,Berlin,29
5,Eve,Tokyo,22
6,Frank,Sydney,41
7,Grace,Rome,28
8,Heidi,Madrid,33
9,Ivan,Beijing,27
10,Judy,Moscow,31
11,Kyle,Dubai,26
12,Liam,Toronto,38
13,Mia,Seoul,25
14,Noah,Cairo,32
15,Olivia,Rio,23
16,Peter,Bangkok,36
17,Quinn,Warsaw,21
18,Rachel,Lisbon,29
19,Sam,Dublin,34
20,Tina,Vienna,26
`
    filePath := "large_dummy.csv"
    err := os.WriteFile(filePath, []byte(dummyData), 0644)
    if err != nil {
        fmt.Println("创建虚拟文件失败:", err)
        return
    }
    defer os.Remove(filePath) // 程序结束时清理虚拟文件

    // 示例1: 抽取 5 条随机行
    numSamples1 := 5
    fmt.Printf("--- 尝试从 %s 中抽取 %d 条随机行 ---\n", filePath, numSamples1)
    randomLines1, err := ReadRandomCSVLines(filePath, numSamples1)
    if err != nil {
        fmt.Println("读取随机行失败:", err)
        return
    }
    fmt.Println("抽取的随机行 (示例1):")
    for i, line := range randomLines1 {
        fmt.Printf("%d: %v\n", i+1, line)
    }
    fmt.Printf("实际抽取行数: %d\n\