Question

我有一个使用Postgresql copy statement创建CSV文件的python函数。我需要在此电子表格中添加一个名为“ UAL”的新列，并在第一行中输入一个示例值，例如30,000，但无需编辑复制语句。这是当前代码：

copy_sql = 'COPY (
  SELECT
    e.name AS "Employee Name",
    e.title AS "Job Title"
    e.gross AS "Total Pay",
    e.total AS "Total Pay & Benefits",
    e.year AS "Year",
    e.notes AS "Notes",
    j.name AS "Agency",
    e.status AS "Status"

  FROM employee_employee e
  INNER JOIN jurisdiction_jurisdiction j on e.jurisdiction_id = j.id
  WHERE

      e.year = 2011 AND
      j.id = 4479

  ORDER BY "Agency" ASC, "Total Pay & Benefits" DESC
  )'

with open(path, 'w') as csvfile:
    self.cursor.copy_expert(copy_sql, csvfile)

我想要做的是使用csv.writer之类的内容来添加内容，如下所示：

    with open(path, 'w') as csvfile:
        self.cursor.copy_expert(copy_sql, csvfile)
        writer = csv.writer(csvfile)
        writer.writerow('test123')

但这会将文本添加到最后一行。我也不确定如何添加新的标题列。有什么建议吗？

Answer 1

添加标头很容易：在调用copy_expert之前写标头。

with open(path, 'w') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["my","super","header"])
    self.cursor.copy_expert(copy_sql, csvfile)

但是添加列不能不重新读取文件并在每一行上添加信息而完成，因此上述解决方案并没有太大帮助。

如果文件不太大且无法容纳在内存中，则可以将sql输出写入“假”文件中：

import io
fakefile = io.StringIO()
self.cursor.copy_expert(copy_sql, fakefile)

现在快退文件并将其解析为csv，写回时添加额外的列

import csv
fakefile.seek(0)
with open(path, 'w', newline="") as csvfile:
    writer = csv.writer(csvfile)
    reader = csv.reader(fakefile)  # works if copy_expert uses "," as separator, else change it
    writer.writerow(["my","super","header","UAL"])
    for row in reader:
       writer.writerow(row+[30000])

或代替内循环：

    writer.writerows(row+[30000] for row in reader)

如果文件太大，则将其写入临时文件，然后以相同的方式进行操作（性能较低）

将数据添加到CSV文件的新列和第一行

1 个答案: