使用Go语言高效合并两个大型排序CSV文件

package main

import (
    "encoding/csv"
    "io"
    "log"
    "os"
)

const outFile = "merged_output.csv" // 定义输出文件路径

func main() {
    // 确保命令行参数正确，需要两个输入文件路径
    if len(os.Args) != 3 {
        log.Panic("\nUsage: go run main.go <file1.csv> <file2.csv>")
    }

    // 打开第一个输入文件
    f1, err := os.Open(os.Args[1])
    if err != nil {
        log.Panicf("\nUnable to open first file %s: %v", os.Args[1], err)
    }
    defer f1.Close() // 确保文件在函数结束时关闭

    // 打开第二个输入文件
    f2, err := os.Open(os.Args[2])
    if err != nil {
        log.Panicf("\nUnable to open second file %s: %v", os.Args[2], err)
    }
    defer f2.Close() // 确保文件在函数结束时关闭

    // 创建输出文件
    w, err := os.Create(outFile)
    if err != nil {
        log.Panicf("\nUnable to create new file %s: %v", outFile, err)
    }
    defer w.Close() // 确保文件在函数结束时关闭

    // 使用csv.NewReader包装文件读取器，以便处理CSV格式
    cr1 := csv.NewReader(f1)
    cr2 := csv.NewReader(f2)

    // 使用csv.NewWriter包装文件写入器，以便写入CSV格式
    cw := csv.NewWriter(w)
    defer cw.Flush() // 确保所有缓冲数据写入文件

    // 初始化：从两个文件各读取第一行
    line1, b1 := readline(cr1)
    if !b1 {
        // 如果文件1为空，则直接拷贝文件2的剩余内容
        log.Println("File 1 is empty or has no CSV lines. Copying File 2 content.")
        copyRemaining(cr2, cw)
        return
    }
    line2, b2 := readline(cr2)
    if !b2 {
        // 如果文件2为空，则直接拷贝文件1的剩余内容
        log.Println("File 2 is empty or has no CSV lines. Copying File 1 content.")
        copyRemaining(cr1, cw)
        return
    }

    // 主合并循环：模拟归并排序的合并步骤
    for {
        // 比较当前两行数据，决定哪一行先写入
        // compare函数需要用户根据实际业务逻辑实现
        if compare(line1, line2) { // 如果line1应该排在line2之前或相等
            writeline(cw, line1) // 写入line1
            line1, b1 = readline(cr1) // 从文件1读取下一行
            if !b1 { // 如果文件1已读完
                copyRemaining(cr2, cw) // 拷贝文件2的剩余内容
                break // 退出循环
            }
        } else { // 如果line2应该排在line1之前
            writeline(cw, line2) // 写入line2
            line2, b2 = readline(cr2) // 从文件2读取下一行
            if !b2 { // 如果文件2已读完
                copyRemaining(cr1, cw) // 拷贝文件1的剩余内容
                break // 退出循环
            }
        }
    }
}

// readline 辅助函数：从CSV读取器中读取一行数据
// 返回[]string表示一行数据，bool表示是否成功读取（false表示EOF或其他错误）
func readline(r *csv.Reader) ([]string, bool) {
    line, err := r.Read()
    if err != nil {
        if err == io.EOF {
            return nil, false // 到达文件末尾
        }
        log.Panicf("\nError reading file: %v", err) // 其他读取错误
    }
    return line, true // 成功读取
}

// writeline 辅助函数：将一行数据写入CSV写入器
func writeline(w *csv.Writer, line []string) {
    err := w.Write(line)
    if err != nil {
        log.Panicf("\nError writing file: %v", err)
    }
}

// copyRemaining 辅助函数：将剩余文件内容直接拷贝到输出文件
func copyRemaining(r *csv.Reader, w *csv.Writer) {
    for {
        line, ok := readline(r)
        if !ok {
            break // 文件已读完
        }
        writeline(w, line)
    }
}

// compare 关键函数：用于比较两行CSV数据
// 用户需要根据实际的排序键和排序规则来实现此函数。
// 如果line1应排在line2之前或与line2相等，则返回true；否则返回false。
// 示例：假设CSV的第一列是排序键（字符串类型）
func compare(line1, line2 []string) bool {
    // 假设排序键是CSV的第一列
    key1 := line1[0]
    key2 := line2[0]

    // 根据键进行比较
    // 这里假设是字符串字典序比较。如果键是数字或其他类型，需要进行相应的转换和比较。
    return key1 <= key2
}