golang导出10w+大数据量的csv文件

golang大数据量导出csv

导出文件采用协程+文件指针同时写入一个文件提供写入速度

1. 可以避免数据表数据太大内存爆了
2. 估计前面文件区块大小保证顺序性
3. 采用文件指针可以大大的提高写入速度，起码一倍
4. context控制上下文
5. github.com/zeromicro/go-zero/core/threading 控制并发
6. sync.WaitGroup 等待锁
7. 废话不说了，看代码
8. 核心思想就是计算前面区间大概使用了字节，然后移动文件指针写入自己这部分区间
9. MySQL这块你可以数据库记录一下单行size，然后预估就只要统计这个size字段

安装依赖

go get github.com/zeromicro/go-zero/core/threading

func ExportWhite(req *request.ExportWhiteRequest) (*response.ExportWhiteResp, error) {
	// 打开文件
	filename := fmt.Sprintf("%s/white_%d.csv", global.ExprotPath, time.Now().Unix())
	file, err := os.Create(filename)
	if err != nil {
		panic(err)
	}
	defer file.Close()

	// 第一种简单粗暴的方法 本地耗时 36s 本地机器ThinkPad x13
	// w := csv.NewWriter(file)
	// list, _ := model.GetWhiteListAll(req.PublicWhite, 200000)
	// for _, row := range list {
	// 	context := fmt.Sprintf("%d,%s,%d,%d,%s,%s,%s\n",
	// 		row.ID,
	// 		row.WhiteName,
	// 		row.Status,
	// 		row.CreatorID,
	// 		row.CreatorName,
	// 		row.CreatedAt,
	// 		row.UpdatedAt,
	// 	)
	// 	w.Write([]string{context})
	// 	// 刷新缓冲
	// 	w.Flush()
	// }
	// return nil, nil

	// 时间上一倍多差异
	// 第二种根据文件指针并发插入 本地耗时 18s 本地机器ThinkPad x13
	// 服务器耗时 969ms
	// 获取总量
	totalRows, err := model.GetWhiteTotal(req.PublicWhite)
	if err != nil {
		return nil, err
	}

	// 每一块区域处理数量
	var concurrency uint32 = 1000
	totalLevel := int(math.Ceil(float64(totalRows) / float64(concurrency)))

	// 打开文件并获取文件句柄
	fileHandle, err := os.OpenFile(filename, os.O_RDWR, 0644)
	if err != nil {
		panic(err)
	}
	defer fileHandle.Close()

	// 创建一个父上下文用于协程之间的协作
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 创建等待锁和并发控制
	var wg sync.WaitGroup
	task := threading.NewTaskRunner(100)

	// 指定 UTF-8 编码
	fileHandle.WriteString("\xEF\xBB\xBF") // 写入 UTF-8 文件标识符

	// 写入表头
	header := "ID,白名单名称,状态,操作人ID,操作人名字,创建时间,更新时间\n"
	headerSize := len(header) + len("\xEF\xBB\xBF")
	fileHandle.WriteString(header)

	// 并发写入数据
	for level := 0; level < totalLevel; level++ {
		wg.Add(1)
		level := level

		// 启动协程并发写入
		task.Schedule(func() {
			defer wg.Done()

			// 协程上下文控制
			select {
			case <-ctx.Done():
				fmt.Printf("%d 并发下载遇到错误\r\n", level)
				return
			default:
			}

			// 获取数据
			data, err := model.GetWhiteList(request.GetWhiteListRequest{
				PublicWhite: req.PublicWhite,
				PagingQuery: request.PagingQuery{
					Page:     int(level) + 1,
					PageSize: int(concurrency),
				},
			})
			// fmt.Println(level, data)
			if err != nil {
				fmt.Println(err)
				cancel()
				return
			}
			// 预估数据大小
			thisPtr, err := model.EstimateSizeWhite(req.PublicWhite,
				model.WriteFileEstimateSizeType{
					IsDb:         true,
					MaxId:        uint32(data[0].ID),
					Level:        uint32(level),
					BlockLineNum: uint32(len(data)),
					Concurrency:  concurrency,
				})

			if err != nil {
				fmt.Println(err)
				cancel()
				return
			}

			// 移动指针
			if level == 0 {
				thisPtr = 0
			}
			offset, err := fileHandle.Seek(thisPtr+int64(headerSize), 0)
			if err != nil {
				fmt.Println(err)
				cancel()
				return
			}
			// fmt.Println("offset:", offset-int64(headerSize))
			// yoff := offset
			// 写入数据到文件指针位置
			for _, row := range data {
				context := []byte(fmt.Sprintf("%d,%s,%d,%d,%s,%s,%s\n",
					row.ID,
					row.WhiteName,
					row.Status,
					row.CreatorID,
					row.CreatorName,
					row.CreatedAt,
					row.UpdatedAt,
				))
				_, err := fileHandle.WriteAt(context, offset)
				if err != nil {
					fmt.Println(err)
					cancel()
					return
				}
				offset += int64(len(context))

				// fmt.Println(offset, string(context))

			}
			// fmt.Println("offset end:", yoff-int64(headerSize), offset-int64(headerSize))
		})
	}

	wg.Wait()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}

	// 使用定时任务删除过期文件
	time.AfterFunc(time.Hour, func() {
		os.Remove(filename)
	})
	return &response.ExportWhiteResp{
		Download: fmt.Sprintf("%s/%s", global.CONFIG.Host.ExportUrl, filename),
	}, nil
}

model中预估大小代码

type WriteFileEstimateSizeType struct {
	IsDb         bool
	FieldString  string // DB是true 必传
	MaxId        uint32 // DB是true 必传
	Level        uint32 // DB是false 必传 层级
	BlockLineNum uint32 // DB是false 必传 当前区域内行数数量
	Concurrency  uint32 // DB是false 每次并发写入X行
}
// 单行预估数值
var WhiteRowSize = uint32(len("100105,十万大军2013,1,1,admin,2023-07-10 10:13:35 +0800 CST,2023-07-10 10:13:37 +0800 CST\n"))

// 预估大小
func EstimateSizeWhite(req request.PublicWhite, param WriteFileEstimateSizeType) (int64, error) {
	if param.IsDb {
		// DB 预估大小
		// SELECT SUM(LENGTH(CONCAT_WS('', white_name, creator_id, creator_name,id,`status`,created_at,updated_at))) AS size FROM white_list WHERE id < ?;
		// SELECT SUM(size) AS sizeToal FROM white_list WHERE id id <= ?;
		db := global.DB.Model(WhiteList{})
		db = WhiteListAssembleWhere(db, req)
		// 执行 SQL 查询
		var result struct {
			Size int64
		}
		err := db.Select("SUM(LENGTH(white_name) + ?) as size", WhiteRowSize-16).
			Where("id < ?", param.MaxId).
			Scan(&result).Error
		return result.Size, err
	} else {
		// 单行预估数值
		// rowSize := uint32(lo.RuneLength("十万大军2001,1,admin,100093,0,2023-07-10 10:13:35,2023-07-10 10:13:37"))
		// X行数据大小 = 单行预估数值 * X + 扩容预估值
		estimateSize := (WhiteRowSize * param.Concurrency) + (1000 * param.Level)
		// 当前数据位置 = X行数据大小 * level - 最后一行的多余的预估值
		thisPtr := estimateSize*param.Level - ((param.Concurrency - param.BlockLineNum) * WhiteRowSize)
		// 移动指针
		return int64(thisPtr), nil
	}
}

不使用数据库预估的话不准确会有null写入

各位道友们有什么更优的方案欢迎留言