Psycopg2：將CSV數據複製到帶有額外列值的表格

Question

我在psycopg2中使用copy_expert方法將數據從CSV文件複製到postgresql表格。我有Postgres的表是這樣的：

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

和含有類似數據這樣的第5列的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）。

我可以複製前五列到我的表是這樣的：

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的最後兩列指定值？我知道我可以在表中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') 

這適用於小型數據不錯，但它在大型數據慢得令人難以置信。 psycopg2copy_from和copy_expert由於使用postgresql批量複製，因此命令似乎要快得多。我能夠首先把我的CSV文件導入到複製使用copy_from我的數據pandasdataframe：

import psycopg2 
import pandas as pd 
from StringIO import StringIO 

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

爲了利用psycopg2copy_的命令，我需要將dataframe轉換成類文件對象與read()和readline()方法，我可以不使用StringIO：

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

請注意，您需要設置緩衝區設置爲0的pos，因爲pandas.to_csv似乎默認設置pos到最後。有關說明，請參閱this SO answer。

然後我可以複製緩衝對象：

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