Go语言并行读取处理超大文本

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111

package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sync"
	"time"
)

func main() {
	start := time.Now()
	lineschan := make(chan []string, 100)
	go ParallelRead("D:\\10E.txt", lineschan, 0, 0)
	var count uint64
	for range lineschan {
		chunklines := <-lineschan
		for range chunklines {
			count++
		}
	}
	fmt.Println(count, time.Since(start))
}

// ChunkRead 从文件的指定偏移读取块并返回块内所有行
func ChunkRead(f *os.File, offset int64, size int, bufPool *sync.Pool) []string {
	buf := bufPool.Get().([]byte)[:size] // 从Pool获取缓冲区并调整大小
	n, err := f.ReadAt(buf, offset)      // 从指定偏移读取数据到缓冲区
	if err != nil && err != io.EOF {
		log.Fatal(err)
	}
	buf = buf[:n] // 调整切片以匹配实际读取的字节数

	var lines []string
	var lineStart int
	for i, b := range buf {
		if b == '\n' {
			line := string(buf[lineStart:i]) // 不包括换行符
			lines = append(lines, line)
			lineStart = i + 1 // 设置下一行的开始
		}
	}

	if lineStart < len(buf) { // 处理块末尾的行
		line := string(buf[lineStart:])
		lines = append(lines, line)
	}
	bufPool.Put(buf) // 读取完成后，将缓冲区放回Pool
	return lines
}

// 并行读取文件并将每个块的行发送到指定通道
// filePath 文件路径
// ChunkSizeMB 分块的大小（单位MB，设置为0时使用默认100MB）,设置过大反而不利，视情调整
// MaxGoroutine 并发读取分块的数量
// linesCh用于接收返回结果的通道。
func ParallelRead(filePath string, linesCh chan<- []string, ChunkSizeMB, MaxGoroutine int) {
	if ChunkSizeMB == 0 {
		ChunkSizeMB = 100
	}
	ChunkSize := ChunkSizeMB * 1024 * 1024
	// 内存复用
	bufPool := sync.Pool{
		New: func() interface{} {
			return make([]byte, 0, ChunkSize)
		},
	}

	if MaxGoroutine == 0 {
		MaxGoroutine = runtime.NumCPU() // 设置为0时使用CPU核心数
	}

	f, err := os.Open(filePath)
	if err != nil {
		log.Fatalf("failed to open file: %v", err)
	}
	defer f.Close()

	info, err := f.Stat()
	if err != nil {
		log.Fatalf("failed to get file info: %v", err)
	}

	wg := sync.WaitGroup{}
	chunkOffsetCh := make(chan int64, MaxGoroutine)

	// 分配工作
	go func() {
		for i := int64(0); i < info.Size(); i += int64(ChunkSize) {
			chunkOffsetCh <- i
		}
		close(chunkOffsetCh)
	}()

	// 启动工作协程
	for i := 0; i < MaxGoroutine; i++ {
		wg.Add(1)
		go func() {
			for chunkOffset := range chunkOffsetCh {
				linesCh <- ChunkRead(f, chunkOffset, ChunkSize, &bufPool)
			}
			wg.Done()
		}()
	}

	// 等待所有解析完成后关闭行通道
	wg.Wait()
	close(linesCh)
}

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242

package main

import (
	"bufio"
	"database/sql"
	"fmt"
	"io"
	"log"
	"math"
	"os"
	"regexp"
	"strings"
	"sync"
	"time"

	_ "github.com/marcboeker/go-duckdb"
)

const (
	chunkSize     = 300000                         // 控制每个块的行数
	maxGoroutines = 100                            // 控制总的协程数
	fileName      = "D:\\3part1.csv" // 待处理文件名
	// 写数据库相关
	sqinnpath = "d:\\sqinn.exe"        // 下载地址 https://github.com/cvilsmeier/sqinn/releases
	dbname    = "E:\\mydb.duckdb" // 数据库文件路径
)

var (
	// 定义一个通道，用于传递要写入数据库的数据
	dataChan = make(chan []string, 10000000) // 缓冲区大小可以根据需要调整
	// db       = *sql.DB
	rePhone = regexp.MustCompile(`"([\d]{11})"`)
	reXm    = regexp.MustCompile(`"([\p{Han}]{2,4})"`)
	reSfz   = regexp.MustCompile(`"([\d]{18})"`)
	reYhk   = regexp.MustCompile(`"([\d]{16})"`)
)

func main() {
	s := time.Now()
	file, err := os.Open(fileName)
	if err != nil {
		fmt.Println("无法读取文件", err)
		return
	}
	defer file.Close() // 在检查错误后关闭

	var wg_Main sync.WaitGroup
	wg_Main.Add(1)
	go processReciveDatas(&wg_Main, dataChan)

	sem := make(chan struct{}, maxGoroutines) // 用于限制总的协程数
	if err := ProcessFileWithChunk(file, chunkSize, sem, dataChan); err != nil {
		fmt.Println("处理文件时发生错误:", err)
	}
	wg_Main.Wait() // 等待数据库写入协程完成

	fmt.Println("\n所用时间 - ", time.Since(s))
}

// 打开文件以块处理
func ProcessFileWithChunk(f *os.File, chunkSize int, sem chan struct{}, dataChan chan<- []string) error {
	linesPool := sync.Pool{New: func() interface{} {
		lines := make([]byte, 250*1024) // 多行最大占用字节大小
		return &lines
	}}

	r := bufio.NewReader(f)
	var wg_ProcessFileWithChunk sync.WaitGroup

	for {
		bufPtr := linesPool.Get().(*[]byte)
		buf := *bufPtr
		n, err := r.Read(buf)
		buf = buf[:n]

		if n == 0 {
			if err != nil {
				fmt.Println(err)
				break
			}
			if err == io.EOF {
				break
			}
			return err
		}

		nextUntillNewline, err := r.ReadBytes('\n')
		if err != io.EOF {
			buf = append(buf, nextUntillNewline...)
		}

		wg_ProcessFileWithChunk.Add(1)
		sem <- struct{}{}
		go func(b []byte) {
			defer wg_ProcessFileWithChunk.Done()
			ProcessChunk(b, chunkSize, dataChan)
			<-sem
			linesPool.Put(bufPtr)
		}(buf)
	}

	wg_ProcessFileWithChunk.Wait()
	close(dataChan) // 然后关闭dataChan
	return nil
}


// 处理块
func ProcessChunk(chunk []byte, chunkSize int, dataChan chan<- []string) {
	logs := string(chunk)
	chunkLines := strings.Split(logs, "\n")

	n := len(chunkLines)
	noOfThread := n / chunkSize
	if n%chunkSize != 0 {
		noOfThread++
	}

	var wg_ProcessChunk sync.WaitGroup
	for i := 0; i < noOfThread; i++ {
		start := i * chunkSize
		end := int(math.Min(float64((i+1)*chunkSize), float64(n)))
		wg_ProcessChunk.Add(1)
		go func(s int, e int) {
			defer wg_ProcessChunk.Done()
			ProcessLines(chunkLines[s:e], dataChan)
		}(start, end)
	}

	wg_ProcessChunk.Wait()
}


// 数据生产者-处理行
func ProcessLines(lines []string, dataChan chan<- []string) {
	// var wg_ProcessLines sync.WaitGroup

	for _, line := range lines {
		// wg_ProcessLines.Add(1) // 递增计数
		// go func(ln string) {   // 使用ln捕获当前迭代的line值
		// defer wg_ProcessLines.Done()
		matches := processSingeLine(line) // 应使用ln，而不是外部的line变量
		dataChan <- matches
		// }(line) // 这里传递line作为ln到协程
	}

	// wg_ProcessLines.Wait() // 等待所有协程完成
}

// 数据消费者
func processReciveDatas(wg *sync.WaitGroup, dataChan <-chan []string) {
	defer wg.Done()
	// 启动 sqinn
	sq := sqinn.MustLaunch(sqinn.Options{
		SqinnPath: sqinnpath,
	})
	defer sq.Terminate()
	// 打开数据库
	sq.MustOpen(dbname)
	defer sq.Close()
	// 创建表
	sq.MustExecOne(`
	CREATE TABLE IF NOT EXISTS your_table (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		xm TEXT,
		sfz TEXT,
		phone TEXT,
		yhk TEXT
	)`)

	// 启动事务
	sq.MustExecOne("BEGIN")

	var params []any // 用于存储所有记录的所有字段值
	var count int    // 记录数计数器

	// 从dataChan中读取数据，每个元素是一条记录的字段值切片
	for record := range dataChan {
		// 遍历每条记录的每个字段值
		for _, fieldValue := range record {
			// fmt.Println(fieldValue)
			params = append(params, fieldValue) // 将字段值添加到params切片中
		}
		count++ // 增加记录数计数器

		// 如果达到一定数量的记录，执行插入操作
		if count == 100000 { // 假设每次批量处理100000条记录
			// 执行插入操作
			ttt := params[0]
			fmt.Println(ttt)
			sq.MustExec("INSERT INTO your_table (xm, sfz, phone, yhk) VALUES (?, ?, ?, ?)", count, 4, params)
			sq.MustExecOne("COMMIT")
			// 重置params切片和计数器，准备下一批数据
			params = []any{}
			count = 0
		}
	}

	// 检查是否还有未处理的数据（最后一批不足100000条记录）
	if count > 0 {
		// 执行插入操作
		sq.MustExec("INSERT INTO your_table (xm, sfz, phone, yhk) VALUES (?, ?, ?, ?)", count, 4, params)
		sq.MustExecOne("COMMIT")
	}
	fmt.Println("processReciveDatas运行完毕！")
}


// 单行逻辑
func processSingeLine(line string) []string {
	// 使用正则表达式从line中提取值
	phoneMatches := rePhone.FindStringSubmatch(line)
	xmMatches := reXm.FindStringSubmatch(line)
	sfzMatches := reSfz.FindStringSubmatch(line)
	yhkMatches := reYhk.FindStringSubmatch(line)

	// 初始化一个字符串切片来存储匹配到的值
	var matches []string

	// 将匹配到的值添加到切片中
	if len(phoneMatches) > 1 {
		matches = append(matches, phoneMatches[1])
	} else {
		matches = append(matches, "")
	}
	if len(xmMatches) > 1 {
		matches = append(matches, xmMatches[1])
	} else {
		matches = append(matches, "")
	}
	if len(sfzMatches) > 1 {
		matches = append(matches, sfzMatches[1])
	} else {
		matches = append(matches, "")
	}
	if len(yhkMatches) > 1 {
		matches = append(matches, yhkMatches[1])
	} else {
		matches = append(matches, "")
	}
	return matches
}