pandas：遍历数据框，对每一行进行 SQL 查询答案

【问题标题】：pandas: iterate over dataframe, do SQL query for each rowpandas：遍历数据框，对每一行进行 SQL 查询
【发布时间】：2021-09-10 19:18:27
【问题描述】：

我有一个数据框和一个 500 万行的本地 Postgres 数据库。在数据框的每一行中，我想添加一列，它是针对 Postgres 数据库的查询结果。

这就是我现在拥有的：

for index, row in df_tf.iterrows():
    row = dict(row)
    id = row['National ID']
    q = 'select name from companies where company_number=%s'
    cursor.execute(q, [company_number])
    results = cursor.fetchall()
    if len(results):
        row['name'] = result[0][0]
        writer.writerow(row)
    else:
        row['name'] = ''
        writer.writerow(row)

所以我正在遍历行并将结果写入本地 CSV。

有没有一种方法可以让我更简洁地做到这一点，并将结果保存在本地数据框中？

我知道我可以将 Postgres 数据加载到 pandas 中并直接加入，但它相当大而且速度很慢，所以我更喜欢使用 Postgres 查询。

【问题讨论】：

由于您只使用第一个结果，fetchone() 可能有助于整理并更快
df 的平均大小是多少？
@P.Naoum 它不是很大，大约 80,000 行。 Postgres 数据更大，大约 500 万行 / 2.5GB。我可以将它加载到内存中，但我不希望这样做（而且我也对内存太大的数据集的答案感兴趣！）。

标签： python pandas

【解决方案1】：

使用sqlalchemy declarative_base 的方法。

粗略代码：

from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy import String, Integer # noqa
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm import sessionmaker

base = declarative_base()
engine = create_engine(#some stuff)
session = scoped_session(sessionmaker(bind=engine))()

class Companies(base):
    __tablename__ = 'companies'
    name = Column(String)
    company_number = Column(Integer)
    ...
    #other stuff

 @classmethod
 def get_by_company_number(cls, company_number):
    query = session.query(cls).filter(cls.company_number == company_number)
    if query.count() == 0:
        return ''
    else:
        return query.first().name
 
 df_tf['name'] = df_tf['National ID'].apply(Companies.get_by_company_number)
 df_tf.to_csv('filename.csv')

【讨论】：

【解决方案2】：

我认为我的第一眼看起来会像（未经测试）：

import pandas
import psycopg2
import csv
import contextlib

def get_company_name(cursor, company_number):
    query = 'SELECT name FROM companies WHERE company_number=%s;'
    cursor.execute(query, [company_number])
    results = cursor.fetchone()
    return results[0] if results else ''

df_tf = pandas.DataFrame("...")
with contextlib.ExitStack() as ctx:
    connection = ctx.enter_context(psycopg2.connect("..."))
    cursor = ctx.enter_context(connection.cursor())
    file_out = ctx.enter_context(open("results.csv", "w"))
    writer = csv.DictWriter(file_out, fieldnames=["National ID", "Name"])
    writer.writeheader()
    for _, row in df_tf.iterrows():
        row = dict(row)
        row['Name'] = get_company_name(cursor, row['National ID'])
        writer.writerow(row)

根据数据帧中的数据，缓存来自get_company_name() 的结果可能是值得的。我想有更好的答案，但这是我会尝试的。

【讨论】：