惯性聚合 高效追踪和阅读你感兴趣的博客、新闻、科技资讯
阅读原文 在惯性聚合中打开

推荐订阅源

C
Check Point Blog
美团技术团队
Jina AI
Jina AI
人人都是产品经理
人人都是产品经理
The Cloudflare Blog
V
Visual Studio Blog
Google DeepMind News
Google DeepMind News
Hugging Face - Blog
Hugging Face - Blog
云风的 BLOG
云风的 BLOG
有赞技术团队
有赞技术团队
T
The Blog of Author Tim Ferriss
WordPress大学
WordPress大学
月光博客
月光博客
宝玉的分享
宝玉的分享
小众软件
小众软件
MongoDB | Blog
MongoDB | Blog
Apple Machine Learning Research
Apple Machine Learning Research
A
About on SuperTechFans
J
Java Code Geeks
博客园_首页
Cyber Security Advisories - MS-ISAC
Cyber Security Advisories - MS-ISAC
N
Netflix TechBlog - Medium
Vercel News
Vercel News
博客园 - 聂微东

Jiajun的技术笔记

你好,2026! TiDB 源码阅读(六):TiDB Coprocessor 源码解析 性能优化的核心思想 TiDB 源码阅读(五):索引 TiDB 源码阅读(四):AST、逻辑计划、物理计划 CockroachDB Serverless Architecture podman 无故退出 Cursor Control-L (CTRL-L) Keyboard Shortcuts in Terminal Replace docker with podman Using xmonad with xfce4 A RC script for freebsd frpc 自己动手写一个k8s controller AI 会取代你的(编程)岗位吗? 自建DERP服务器提升Tailscale连接速度(使用Nginx转发) 自动升级Docker容器 再读《程序员修炼之道-从小工到专家》 让浏览器下载文件 再读《软件随想录》/《黑客与画家》/《软技能》 HTTP 压力测试中的 Coordinated Omission 2的补码 编程语言中的 context 是什么? flutter macOS 构建出错 Flatpak 使用小记 Golang CAS 操作是怎么实现的 PostgreSQL 当MQ来使用 Clash 结合 工作VPN 的网络设计 使用 PostgreSQL 搭建 JuiceFS PostgreSQL 配置优化和日志分析 有GitHub Copilot?那就可以搭建你的ChatGPT4服务 窗口函数的使用(以PG为例)
用Go导入大型CSV到PostgreSQL
Jiajun Huang · 2021-12-11 · via Jiajun的技术笔记

最近我想试试 PostgreSQL,素闻美名,一直没有尝试过。从网上下载了一个超大的CSV,解压后达18G,一般的文件编辑器 直接打不开,简单的方案是直接用 PostgreSQL 提供的 \copy 命令,或者 COPY 语句,但是这个文件无法使用,因为 其中有几行都是坏数据。

如果是MySQL的话,可以使用 LOAD FILE IGNORE... 来忽略错误,但是PostgreSQL没有这个选项,所以我只能选择用Go自己 来导入。

吐个槽,MySQL用IGNORE之后,连数据错误也会忽略,导致我导入数据之后,才发现 int 不够表示CSV里的数据字段,导入的 很多数据直接变成了 2 ** 31 -1 也就是 2147483647 了,白等了一个小时。

对于大型文件,如果没有足够的内存,也确实是很难处理,我们采取的基本策略就是分块处理,为了提高吞吐量,我们要做 批量提交和并发处理,为了处理异常数据,我们要能dump出有问题的那一块数据,以便处理之后我们再次导入, 由于每一个块是相对较小的,dump出来之后,我们是可以直接用文本编辑器处理问题行的,此外由于涉及到 string和bytes 的转换,我们需要避免频繁申请内存,可以使用上 黑科技

// b2s converts byte slice to a string without memory allocation.
// See https://groups.google.com/forum/#!msg/Golang-Nuts/ENgbUzYvCuU/90yGx7GUAgAJ .
//
// Note it may break if string and/or slice header will change
// in the future go versions.
func b2s(b []byte) string {
	/* #nosec G103 */
	return *(*string)(unsafe.Pointer(&b))
}

// s2b converts string to a byte slice without memory allocation.
//
// Note it may break if string and/or slice header will change
// in the future go versions.
func s2b(s string) (b []byte) {
	/* #nosec G103 */
	bh := (*reflect.SliceHeader)(unsafe.Pointer(&b))
	/* #nosec G103 */
	sh := (*reflect.StringHeader)(unsafe.Pointer(&s))
	bh.Data = sh.Data
	bh.Cap = sh.Len
	bh.Len = sh.Len
	return b
}

以及 sync.Pool,最开始我使用的是 bufio.Scanner,但是没想到超过了它的限制,报了 bufio.Scanner: token too long, 翻了一下代码如下:

var (
    ErrTooLong         = errors.New("bufio.Scanner: token too long")
)

//...

if len(s.buf) >= s.maxTokenSize || len(s.buf) > maxInt/2 {
    s.setErr(ErrTooLong)
    return false
}

//...

const (
	// MaxScanTokenSize is the maximum size used to buffer a token
	// unless the user provides an explicit buffer with Scanner.Buffer.
	// The actual maximum token size may be smaller as the buffer
	// may need to include, for instance, a newline.
	MaxScanTokenSize = 64 * 1024
)

可以自己把buffer调大,不过我选择自己逐行读取并且处理,毕竟咱也不知道这么大的数据里,最长的那行到底有多长。

最后代码如下:

package main

import (
	"bufio"
	"io"
	"io/ioutil"
	"log"
	"os"
	"strings"
	"sync"

	"github.com/jmoiron/sqlx"
	"github.com/lib/pq"
)

const size uint64 = 10000

var (
	tokens          = make(chan bool, 50)
	stringSlicePool = sync.Pool{
		New: func() interface{} {
			cache := make([]string, size)
			return cache[:0]
		},
	}
)

// 用wrapper避免参数是 []string 时,是值拷贝的问题
type wrapper struct {
	lines []string
}

func dumpData(w *wrapper) {
	file, err := ioutil.TempFile("./dumps/", "damage")
	if err != nil {
		log.Printf("failed to open temp file")
		return
	}
	defer file.Close()

	for _, line := range w.lines {
		file.WriteString(line)
	}
}

func newWrapper() *wrapper {
	lines := stringSlicePool.Get().([]string)
	return &wrapper{lines: lines}
}

func deleteWrapper(w *wrapper) {
	w.lines = w.lines[:0]
	stringSlicePool.Put(w.lines)
}

func writeData(wg *sync.WaitGroup, db *sqlx.DB, w *wrapper) {
	wg.Add(1)

	token := <-tokens // 并发控制

	tx := db.MustBegin()
	stmt, err := tx.Prepare(pq.CopyIn("表名", "字段1", "字段2" /*字段3...*/))
	if err != nil {
		log.Printf("failed to prepare: %s", err)
		goto done
	}

	if len(w.lines) == 0 {
		goto done
	}

	for _, line := range w.lines {
		data := strings.Split(line, "\t") // 此处是一个频繁内存申请的点
		if len(data) < 2 {
			log.Printf("ignore %s", line)
			continue
		}

		stmt.Exec(data[len(data)-2], data[len(data)-1][:len(data[1])-1])
	}
	stmt.Close()
	if err = tx.Commit(); err != nil {
		log.Printf("failed to commit: %s", err)
		dumpData(w)
		goto done
	}

	log.Printf("saving %d lines", len(w.lines))

done:
	tokens <- token
	deleteWrapper(w)
	wg.Done()
}

func main() {
	var wg sync.WaitGroup

	db, err := sqlx.Connect("postgres", "user=postgres dbname=数据库名 sslmode=disable password=密码")
	if err != nil {
		log.Fatalln(err)
	}
	log.Printf("%v, %s", db, err)

	file, err := os.Open("./to_import.csv")
	if err != nil {
		log.Fatal(err)
	}
	defer file.Close()

	for i := 0; i < cap(tokens); i++ {
		tokens <- true
	}

	reader := bufio.NewReader(file)
	cache := newWrapper()

	// optionally, resize scanner's capacity for lines over 64K, see next example
	var i uint64 = 0
	reader.ReadString('\n')
	for {
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}
		cache.lines = append(cache.lines, line)
		i += 1

		if i%size == 0 {
			oldCache := cache
			go writeData(&wg, db, oldCache)
			cache = newWrapper()
		}
	}

	log.Printf("wait wg...")
	wg.Wait()
	log.Printf("done...")
}