PyAthena ArrowCursor查询ARRAY<STRING>列转字符串的解决方法问询

我用PyAthena的ArrowCursor查询Athena中ARRAY<STRING>类型的列时遇到了类型转换问题：

• 执行SQL：SELECT ARRAY ['a', 'b', 'c'] AS abc
• 用execution_result.description能正确识别列类型为array
• 但调用as_arrow()生成PyArrow表后，该列变成了STRING类型，值为"[a, b, c]"

需要把结果转成Polars或Pandas DataFrame，且abc列保持ARRAY<STRING>类型，不想自己解析字符串转列表，有没有更优方法？

1. 改用默认Cursor直接转换 #

放弃ArrowCursor，用PyAthena的默认Cursor，它会直接把Athena的数组类型返回为Python列表，Polars/Pandas能自动识别为列表类型：

import pyathena
import polars as pl
import pandas as pd

AWS_PARAMETERS = {
    'aws_access_key_id': **,
    'aws_secret_access_key': **,
    'region_name': **,
    's3_staging_dir': **,
    'work_group': **,
}

# 使用默认Cursor
cursor = pyathena.connect(**AWS_PARAMETERS).cursor()
execution_result = cursor.execute("SELECT ARRAY ['a', 'b', 'c'] AS abc")

rows = execution_result.fetchall()
columns = [col[0] for col in execution_result.description]

# 转Polars DataFrame，abc列自动为List[str]
pl_df = pl.DataFrame(rows, schema=columns)
# 转Pandas DataFrame
pd_df = pd.DataFrame(rows, columns=columns)

2. 强制Athena输出Parquet格式 #

在查询中指定输出格式为Parquet，PyArrow可以正确解析Parquet中的数组类型，不会出现字符串化问题：

from pyathena.arrow.cursor import ArrowCursor
import polars as pl

cursor = pyathena.connect(**AWS_PARAMETERS).cursor(ArrowCursor)
# 先设置输出格式为Parquet再执行查询
query = """
SET output_format = 'PARQUET';
SELECT ARRAY ['a', 'b', 'c'] AS abc
"""
execution_result = cursor.execute(query)

arrow_table = execution_result.as_arrow()
# 此时abc列是pyarrow.list_[pyarrow.string()]类型
pl_df = pl.from_arrow(arrow_table)
pd_df = arrow_table.to_pandas()

也可以直接在Athena的WorkGroup中默认配置输出格式为Parquet，不用每次查询都加SET语句。

3. 升级PyAthena版本并使用内置转换方法 #

检查你的PyAthena版本，较新的版本（>=2.0.0）修复了ArrowCursor的类型映射问题，直接用as_pandas()或as_polars()方法就能正确转换数组类型：

execution_result = cursor.execute("SELECT ARRAY ['a', 'b', 'c'] AS abc")
# 转Pandas DataFrame
pd_df = execution_result.as_pandas()
# 转Polars DataFrame
pl_df = execution_result.as_polars()

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