PyAthena で struct を含む array カラムをクエリすると NULL になる

awspython

PyAthena で array<struct<..., array<...>>> のようなカラムをクエリすると NULL になる。

from pyathena import connect # Version: 3.27.1

conn = connect(region_name="ap-northeast-1", work_group="primary")
cursor = conn.cursor()

cursor.execute("""
SELECT ARRAY[
    CAST(ROW(1, ARRAY['a', 'b']) AS ROW(id INTEGER, tags ARRAY(VARCHAR)))
] AS col
""")

print(cursor.fetchone())  # (None,)

実装を見ると [{id=1, tags=[a, b]}] のような json.loads() でそのままパースできないネストした array は複雑すぎるということでパースせずに None で返している

なお、CAST AS JSON すると [{"id":1,"tags":["a","b"]}] のように返り _to_json() でパースされるので取得できる。

cursor.execute("""
SELECT CAST(
    ARRAY[
        CAST(ROW(1, ARRAY['a', 'b']) AS ROW(id INTEGER, tags ARRAY(VARCHAR)))
    ] AS JSON
) AS col
""")

print(cursor.fetchone())  # ('[{"id":1,"tags":["a","b"]}]',)

また、PandasCursor で unload=True を指定することでも S3 に Parquet として出力したものを読むので問題なく取得できる。

from pyathena.pandas.cursor import PandasCursor

conn = connect(region_name="ap-northeast-1", s3_staging_dir="s3://your-bucket/")
cursor = conn.cursor(PandasCursor, unload=True)

df = cursor.execute("""
SELECT ARRAY[
    CAST(ROW(1, ARRAY['a', 'b']) AS ROW(id INTEGER, tags ARRAY(VARCHAR)))
] AS col
""").as_pandas()

print(df['col'].iloc[0])  # [{'id': 1, 'tags': array(['a', 'b'], dtype=object)}]