如何从Snowflake自动获取带NULL的int16数据并保留Arrow类型至Pandas？

我需要一种自动化方法，通过Snowflake Connector从Snowflake数据库中获取不兼容NumPy数据类型的数据（例如可空整数）。

示例场景：执行查询加载一列带NULL值的int16类型整数列：

query = "SELECT MyIntColumn FROM MyTable"

其中MyIntColumn是带NULL值的int16类型。

执行基础连接与查询代码：

con = snowflake.connector.connect()
cursor = con.cursor()
cursor.execute(query)

尝试两种获取方式都无法保留原类型：

1. 使用fetch_pandas_all()：返回的pd.DataFrame中MyIntColumn的类型为dtype=float，不符合预期。

   df = cursor.fetch_pandas_all()
   

2. 使用fetch_arrow_all()后转Pandas：即便使用Pandas 2.0.2，该列仍然是float类型。

   arrow_data = cursor.fetch_arrow_all()
   df = arrow_data.to_pandas()
   

请问是否存在自动化方法，在获取数据时保留Arrow后端的可空整数类型？

以下几种方法可自动化保留Snowflake中的可空整数类型：

1. 配置Snowflake Connector的Pandas选项 #

在创建连接时指定pandas_options，启用可空数据类型支持，Connector会直接生成带Pandas可空整数类型的DataFrame：

con = snowflake.connector.connect(
    # 填入你的连接参数（account、user、password等）
    pandas_options={"engine": "pyarrow", "use_nullable_dtypes": True}
)
cursor = con.cursor()
cursor.execute(query)
df = cursor.fetch_pandas_all()

2. 从Arrow表转Pandas时指定类型映射 #

如果使用fetch_arrow_all()，可在to_pandas()中通过types_mapper参数，强制将Arrow可空整数类型映射为Pandas对应可空类型：

import pyarrow as pa
from pandas import ArrowDtype

def arrow_to_pandas_types(arrow_type):
    if pa.types.is_integer(arrow_type) and pa.types.is_nullable(arrow_type):
        return ArrowDtype(arrow_type)
    return None

arrow_data = cursor.fetch_arrow_all()
df = arrow_data.to_pandas(types_mapper=arrow_to_pandas_types)

3. 全局启用Pandas Arrow扩展类型 #

在代码开头全局配置Pandas，使其优先使用Arrow兼容的可空类型，避免自动将带NULL的整数列转为float：

import pandas as pd
pd.set_option("future.no_silent_downcasting", True)
pd.set_option("mode.copy_on_write", True)

# 后续执行数据获取与转换
arrow_data = cursor.fetch_arrow_all()
df = arrow_data.to_pandas()

验证类型 #

转换完成后，可通过以下代码确认列类型是否符合预期：

print(df["MyIntColumn"].dtype)
# 预期输出：ArrowDtype(int16) 或 Int16（取决于使用的方法）

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