如何在PySpark中编写支持参数替换的SQL查询模板自定义函数？

你可以直接基于DuckDB实现这个需求，DuckDB可以直接对内存中的Pandas DataFrame执行SQL查询，性能比传统的pandasql好很多，不需要做额外的数据导入导出操作。

函数代码 #

import duckdb
import pandas as pd

def func_sql(df: pd.DataFrame, col: str, sql_template: str):
    # 替换SQL模板中的{col}占位符，自动给列名加反引号兼容带特殊字符的列名
    formatted_sql = sql_template.format(col=f"`{col}`")
    # 执行SQL查询
    query_result = duckdb.sql(formatted_sql)
    
    # 单值聚合结果（比如count、sum这类标量结果）直接返回值，其余结果转成DataFrame返回
    raw_res = query_result.fetchall()
    if len(raw_res) == 1 and len(raw_res[0]) == 1:
        return raw_res[0][0]
    return query_result.df()

调用示例 #

# 构造测试DataFrame
df_tbl = pd.DataFrame({
    'age': [18, 20, 22, 18, 25, 20],
    'user_name': ['张三', '李四', '王五', '赵六', '孙七', '周八']
})

# 按示例调用，统计age列去重后的总数
distinct_age_cnt = func_sql(df_tbl, 'age', 'select count(distinct {col}) from df_tbl')
print(distinct_age_cnt) 
# 输出结果：4，对应age列的4个不同取值18、20、22、25

扩展使用说明 #

• 传入的SQL模板里的表名，必须和你环境中存储DataFrame的变量名完全一致，DuckDB会自动匹配内存中对应名称的DataFrame
• 除了去重计数这类单值查询，分组、筛选类查询也可以直接用，比如统计各年龄的人数：

  age_dist = func_sql(df_tbl, 'age', 'select {col}, count(1) as user_cnt from df_tbl group by {col} order by {col}')
  

  调用后会直接返回结构化的分组统计DataFrame，方便后续分析使用
• 如果你的环境已经在用pandasql，只需要把函数里的SQL执行部分替换成pysqldf(formatted_sql)即可，其余逻辑不用改

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