如何在Pyathena的SQL参数中传入多个月份值？

问题场景 #

在AWS SageMaker的Jupyter Notebook中，通过Pyathena连接AWS Athena表生成DataFrame时，传入单个月份参数（如202206）可正常执行SQL查询，但传入月份列表（如[202206,202203]）时触发报错。原代码示例如下：

from pandasql import sqldf
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor

curr_month = 202206
prev_month = 202203

conn = connect(work_group='prod_user',
               region_name= 'eu-west-1')

sql_query = ('''
select 
year_month,
id,
field_1,
field_2,
field_3
from
mart.table_xyz
WHERE
year_month = (%(param)s)
;
''')

cursor = connect(work_group='prod_user',
               region_name= 'eu-west-1',
                 cursor_class=PandasCursor).cursor()
user_df = cursor.execute(sql_query,{"param": curr_month}).as_pandas()

尝试传入列表的报错代码：

cursor = connect(work_group='prod_user',
               region_name= 'eu-west-1',
                 cursor_class=PandasCursor).cursor()
user_df = cursor.execute(sql_query,{"param": [curr_month,prev_month]}).as_pandas()

解决方法 #

核心是调整SQL的条件逻辑，并利用Pyathena对列表参数的原生支持：

1. 修改SQL的WHERE子句：将等于判断=改为IN，保留命名参数占位符以适配列表输入
2. 直接传入列表参数：Pyathena会自动将列表转换为SQL标准的多值格式（如(202206,202203)）

修改后的完整代码：

from pandasql import sqldf
from pyathena import connect
from pyathena.pandas.cursor import PandasCursor

curr_month = 202206
prev_month = 202203

# 调整WHERE条件为IN适配列表参数
sql_query = ('''
select 
year_month,
id,
field_1,
field_2,
field_3
from
mart.table_xyz
WHERE
year_month IN %(param)s
;
''')

cursor = connect(work_group='prod_user',
               region_name= 'eu-west-1',
                 cursor_class=PandasCursor).cursor()
# 直接传入月份列表参数
user_df = cursor.execute(sql_query,{"param": [curr_month,prev_month]}).as_pandas()

说明 #

该方式无需手动拼接SQL字符串，既避免了SQL注入风险，也能适配任意长度的月份列表，Pyathena会自动完成参数的格式转换。

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