Psycopg2：使用额外的列值将 CSV 数据复制到 table

Question

我正在使用 psycopg2 中的 copy_expert 方法将数据从 CSV 文件复制到 postgresql table。我有这样的 postgres table：

create table my_table (
  cola          text,
  colb          text,
  colc          text,
  cold          text,
  cole          text,
  colf          text,
  colg          text
)

还有一个包含前五列数据的 CSV，如下所示：

cola,colb,colc,cold,cole
1,foo,a,10,vvv
2,bar,b,20,www
3,baz,c,30,xxx
4,boo,d,40,yyy
5,baa,e,50,zzz

我想复制前五列的 CSV 数据，同时还指定 colf 和 colg 的值（每一行的 colf 和colg)。

我可以像这样将前五列复制到我的 table：

conn = psycopg2.connect('dbname=name user=username')
cur = conn.cursor()
copy_sql = """
  copy my_table (cola, colb, colc, cold, cole)
  from stdin with
    csv
    header
    delimiter as ','
"""
from_csv = '/path/to/data.csv'
with open(from_csv, 'r') as f:
  cur.copy_expert(sql=copy_sql, file=f)
  conn.commit()
  cur.close()

如何使用 python 指定最后两列的值？我知道我可以在 table DDL 中指定默认值，如下所示：

create table my_table (
  cola          text,
  colb          text,
  colc          text,
  cold          text,
  cole          text,
  colf          text default 'foo',
  colg          text default 'bar'
)

但我想使用 python 添加值，因为每个 CSV 上传都会有自己的 colf 和 colg 值，并且这些值由逻辑决定在我的 python 代码中。

Answer 1

看起来有几种方法可以做到这一点，首先将我需要的列添加到数据中，然后上传更新后的数据。

使用 petl 包：

import psycopg2
from petl import fromcsv, addfield, todb

csv_file = '/path/to/data.csv'
table = fromcsv(csv_file)
table = addfield(table, 'colf', 'Some value')
table = addfield(table, 'colg', 'Another value')

conn = psycopg2.connect('dbname=test user=user')
todb(table, conn, 'my_table')

这在小数据上工作正常，但在大数据上非常慢。 psycopg2 copy_from 和 copy_expert 命令似乎运行得更快，因为它们使用了 postgresql 批量复制。通过首先将我的 csv 文件转换为 pandas dataframe:

，我能够使用 copy_from 复制我的数据

import psycopg2
import pandas as pd
from io import StringIO

csv_file = '/path/to/file'
df = pd.read_csv(csv_file)
df['colf'] = 'My value'
df['colg'] = 'Foobar'

为了使用 psycopg2 copy_ 命令，我需要将 dataframe 转换为具有 read() 和 [=26 的类文件对象=] 方法，我可以使用 StringIO:

buf = StringIO()
df.to_csv(buf, header=False, index=False)
buf.pos = 0

注意需要把buffer的pos设置为0，因为pandas.to_csv好像默认把pos设置到末尾了。有关解释，请参阅。

然后我可以复制那个缓冲区对象：

conn = psycopg2.connect('dbname=test user=user')
cur = conn.cursor()
cur.copy_from(buf, 'my_table', sep=',')
conn.commit()
cur.close()

Psycopg2：使用额外的列值将 CSV 数据复制到 table

Psycopg2: copy CSV data to table with extra column values

python

csv

postgresql

copy

psycopg2