如何加速需传入列名与值的df.apply()？PostgreSQL批量插入优化

我正在优化一个upsert方法，该方法接收Pandas DataFrame并执行查询将数据写入PostgreSQL数据库。当前性能瓶颈在于循环调用cast函数，把Python数据类型转换为对应PostgreSQL类型，生成插入用的vals字符串。

现有代码 #

生成vals字符串的代码 #

cdef str vals = ''
df['row_statement'] = '('
cdef str c
for c in cols:
    df.row_statement += df[c].apply(lambda x: self.cast(c, x)).astype(str) + ', '

df.row_statement = df.row_statement.str[:-2] + ')'
df.row_statement += ','
vals = df.row_statement.str.cat()
vals = vals[1:-2]
df = df.drop('row_statement', axis=1)

cast函数实现 #

def cast(self, column_name, value, emptyList=False):
    cdef n = sqlalchemy.sql.null()
    if not isinstance(value, list):
        if pd.isnull(value) or value == '':
            return n
    elif emptyList and not value:
        return n
    cdef str typ = self._data_types[column_name]
    return self._cast_map[typ](value)

_cast_map示例 #

'text': lambda x: f"'{x}'",
'date': lambda x: f"'{x}'::date",
'timestamp without time zone': lambda x: f"'{x}'::timestamp",

最终查询字符串 #

qry = f"""
            INSERT INTO {self.schema}.{self.table} ({','.join(df.columns.tolist())})
            VALUES ({vals})
            {self._conflict if self._conflict else ''}
    """

尝试过的优化方案（效果更差） #

将DataFrame转为NumPy数组遍历（已Cython优化）：

cdef int r
cdef int c
cdef int c_max = df_arr.shape[1]
cdef int r_max = df_arr.shape[0]

for r in range(r_max):
    vals += '('
    for c in range(c_max):
        vals += str(self.cast(index[c], df_arr[r, c])) + ','
    vals = vals[:-1]
    vals += '),'
vals = vals[1:-2]

疑问 #

1. 我知道向量化比apply更快，但cast函数需要列名和值，这种场景能否实现向量化？
2. 还有哪些方法可以加速该函数？
3. 是否有比lambda字典更优的类型转换方式？

1. 实现向量化转换的可行方案 #

可以针对不同列类型批量处理——因为每一列的类型是固定的，不需要逐元素传入列名，直接按列批量转换：

• 遍历每一列，根据列类型获取对应的转换函数
• 对整列数据做向量化处理，避免逐行apply的开销

示例代码：

vals_list = []
for col in cols:
    col_type = self._data_types[col]
    cast_func = self._cast_map[col_type]
    # 先处理空值，再对非空值应用转换函数
    col_series = df[col].apply(lambda x: cast_func(x) if pd.notnull(x) and x != '' else 'NULL')
    vals_list.append(col_series)

# 按行拼接成( val1, val2, ... )格式
row_statements = vals_list[0].str.cat([s.add(', ') for s in vals_list[1:]], sep='')
row_statements = '(' + row_statements + ')'
# 一次性拼接所有行
vals = ', '.join(row_statements.tolist())

这种方式把逐元素的apply缩小到列级别，利用Pandas的向量化字符串操作减少循环开销，同时避免了原代码中多次字符串拼接的低效操作。

2. 加速函数的其他方法 #

（1）避免字符串拼接的性能损耗 #

原代码中多次对vals或row_statement做+=拼接，Python字符串是不可变对象，每次拼接都会生成新对象，数据量大时性能极差。改用列表收集所有行的字符串，最后一次join是最优方案，如上面示例中的', '.join(row_statements.tolist())。

（2）Cython优化cast函数 #

当前cast是Python函数，即使外层用Cython，调用Python函数的开销依然存在。可以把cast逻辑直接移到Cython代码中，避免跨层调用的开销：

cdef str cy_cast(str col_type, object value):
    if pd.isnull(value) or value == '':
        return 'NULL'
    if col_type == 'text':
        return f"'{value}'"
    elif col_type == 'date':
        return f"'{value}'::date"
    elif col_type == 'timestamp without time zone':
        return f"'{value}'::timestamp"
    # 补充其他类型的处理逻辑

处理列时直接调用这个Cython函数，能显著降低调用开销。

（3）使用psycopg2的execute_values #

完全绕开手动生成vals字符串的逻辑，用psycopg2内置的批量插入工具——它会自动处理类型转换，性能远高于手动拼接SQL：

from psycopg2.extras import execute_values

# 构造基础SQL
columns = ','.join(cols)
insert_sql = f"INSERT INTO {self.schema}.{self.table} ({columns}) VALUES %s {self._conflict if self._conflict else ''}"
# 把DataFrame转为元组列表
data = [tuple(row) for row in df.to_numpy()]
# 执行批量插入
execute_values(self.cursor, insert_sql, data)
self.connection.commit()

这种方式不需要手动处理类型转换，psycopg2会自动把Python类型映射到PostgreSQL类型，同时底层用二进制协议传输，性能提升非常明显。

3. 替代lambda字典的类型转换方式 #

（1）预先编译转换模板 #

用字符串模板替代lambda，避免lambda函数调用的开销：

_cast_templates = {
    'text': "'{}'",
    'date': "'{}'::date",
    'timestamp without time zone': "'{}'::timestamp"
}

def cast(self, column_name, value):
    if pd.isnull(value) or value == '':
        return 'NULL'
    template = self._cast_templates[self._data_types[column_name]]
    return template.format(value)

模板格式化的开销比lambda略低，尤其是在批量调用时。

（2）利用Pandas的to_sql #

如果可以接受使用to_sql，可以直接指定method='multi'或者配合psycopg2的execute_values，Pandas会自动处理类型转换：

df.to_sql(
    name=self.table,
    con=self.engine,
    schema=self.schema,
    if_exists='append',
    method='multi',
    index=False
)

如果需要upsert，可以结合SQLAlchemy的on_conflict_do_update，无需手动拼接SQL。

内容的提问来源于stack exchange，提问作者DZ14