Go语言中从大型CSV文件随机读取行的策略：蓄水池抽样详解

reader := csv.NewReader(file)
lines, err := reader.ReadAll() // 潜在的内存和性能瓶颈
// 然后从 lines 中随机选择

package main

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

// ReadRandomCSVRecordsReservoir 从指定CSV文件中使用蓄水池抽样算法读取 k 个随机记录。
// 返回一个包含 k 个 [][]string 类型的记录切片。
func ReadRandomCSVRecordsReservoir(filePath string, k int) ([][]string, error) {
    file, err := os.Open(filePath)
    if err != nil {
        return nil, fmt.Errorf("打开文件失败: %w", err)
    }
    defer file.Close()

    reader := csv.NewReader(file)
    // reservoir 用于存放随机抽取的 k 个记录。
    // 预分配容量以避免多次扩容。
    reservoir := make([][]string, 0, k)

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

    recordCount := 0 // 记录当前已处理的CSV记录总数
    for {
        record, err := reader.Read() // 读取单个CSV记录
        if err == io.EOF {
            break // 文件读取完毕
        }
        if err != nil {
            return nil, fmt.Errorf("读取CSV记录错误: %w", err)
        }

        recordCount++ // 增加已处理记录计数

        if recordCount <= k {
            // 前 k 个记录直接填充蓄水池
            reservoir = append(reservoir, record)
        } else {
            // 对于第 k+1 个及之后的记录
            // 生成一个 [0, recordCount-1] 之间的随机整数 j
            j := r.Intn(recordCount) 
            // 如果 j 小于 k，则用当前记录替换蓄水池中的第 j 个记录
            if j < k {
                reservoir[j] = record
            }
        }
    }

    // 如果文件中的总记录数少于 k，则返回所有记录
    if recordCount < k {
        return reservoir, nil
    }

    return reservoir, nil
}

func main() {
    // 创建一个大型虚拟CSV文件用于测试
    dummyFileName := "large_data.csv"
    createDummyCSVFile(dummyFileName, 1000000) // 创建包含100万行的文件

    k := 10 // 需要抽取的随机记录数量

    fmt.Printf("正在从 %s 中抽取 %d 个随机CSV记录...\n", dummyFileName, k)
    selectedRecords, err := ReadRandomCSVRecordsReservoir(dummyFileName, k)
    if err != nil {
        log.Fatalf("抽取随机CSV记录失败: %v", err)
    }

    fmt.Println("抽取的记录:")
    for i, record := range selectedRecords {
        fmt.Printf("%d: %v\n", i+1, record)
    }

    // 清理虚拟文件
    os.Remove(dummyFileName)
}

// createDummyCSVFile 是一个辅助函数，用于生成一个指定行数的虚拟CSV文件。
func createDummyCSVFile(filename string, numLines int) {
    file, err := os.Create(filename)
    if err != nil {
        log.Fatalf("创建虚拟CSV文件失败: %v", err)
    }
    defer file.Close()

    writer := csv.NewWriter(file)
    defer writer.Flush() // 确保所有缓冲数据写入文件

    // 写入CSV头部
    writer.Write([]string{"ID", "Name", "Description", "Value"})

    for i := 0; i < numLines; i++ {
        record := []string{
            strconv.Itoa(i + 1),
            fmt.Sprintf("Item_%d", i+1),
            // 包含逗号和换行的描述，以测试CSV解析器的鲁棒性
            fmt.Sprintf("这是第 %d 项的描述，其中可能包含逗号, 甚至换行符。\n描述的第二行。", i+1),
            fmt.Sprintf("%.2f", float64(i)*1.23),
        }
        err := writer.Write(record)
        if err != nil {
            log.Fatalf("写入虚拟CSV记录失败: %v", err)
        }
    }
    fmt.Printf("已创建虚拟CSV文件 %s，包含 %d 条记录 (含头部).\n", filename, numLines+1)
}