如何将Polars DataFrame保存到Snowflake且不丢失数据类型信息？

问题概述 #

使用Polars的write_database（ADBC引擎）将DataFrame持久化到Snowflake，再用read_database_uri读取时，出现严重类型丢失：

• 原i64类型列变为f64
• pl.List、pl.Struct类型列转为pl.String
• pl.Categorical类型列转为pl.String

测试用DataFrame代码：

df_for_testing = pl.DataFrame({
    'A1': [1,2,3],
    'A2': [2,8,7],
    'B1':[[2.2,4.4,6.6],[12,16,14],[152,257,252]],
    'B2':[[1,1],[1,2],[2,1]],
    'C1':['Hello', 'World','!'],
    'C2':['The Test', 'Was For', 'The Takers']
}).with_columns([
    pl.struct(['A1','A2']).alias('As'),
    pl.struct(['B1','B2']).alias('Bs'),
]).with_columns(
    D1 = pl.col('C1').cast(pl.Categorical), 
    D2 = pl.col('C2').cast(pl.Categorical)
)

可行解决方案 #

1. 使用Parquet+Snowpipe（推荐） #

Parquet格式原生支持复杂类型存储，Snowflake通过Snowpipe读取Parquet文件时能准确映射类型，避免ADBC直接写入的类型丢失问题：

• 步骤1：将Polars DataFrame导出为Parquet文件（保留完整类型信息）

  df_for_testing.write_parquet("s3://your-bucket/path/to/file.parquet")  # 或其他云存储服务
  

• 步骤2：在Snowflake中创建外部阶段（Stage），指向存储Parquet文件的位置
• 步骤3：创建Snowpipe管道，自动或手动加载Parquet文件到Snowflake表
• 步骤4：读取时直接从Snowflake读取，Polars能通过ADBC正确解析大部分复杂类型

  df_new = pl.read_database_uri('SELECT * FROM test_table', uri=con_str, engine='adbc')
  

  注：Snowflake会将Parquet的List映射为ARRAY、Struct映射为OBJECT，Polars读取时会自动转换回对应类型；Categorical类型在Parquet中以STRING+元数据存储，Snowflake读取后仍为STRING，需手动转回Categorical。

2. 改进类型转换函数 #

你的match_to_outbound_schema函数可优化，解决Categorical类型转换的类别丢失问题，以及更严谨的嵌套类型解析：

def match_to_outbound_schema(df: pl.DataFrame, outbound_schema: dict, verbose: bool = False) -> pl.DataFrame:
    assert df.schema.keys() == outbound_schema.keys(), "列名不匹配"

    for column, target_dtype in outbound_schema.items():
        current_dtype = df.schema[column]
        if current_dtype == target_dtype:
            if verbose:
                print(f'{column} 类型匹配，无需处理')
            continue

        # 数值类型转换
        if current_dtype.is_numeric() and target_dtype.is_numeric():
            if verbose:
                print(f'{column}: {current_dtype} -> {target_dtype}（数值转换）')
            df = df.with_columns(pl.col(column).cast(target_dtype))

        # 嵌套类型（List/Struct）从JSON字符串解析
        elif current_dtype.is_(pl.String()) and target_dtype.is_nested():
            if verbose:
                print(f'{column}: {current_dtype} -> {target_dtype}（嵌套类型解析）')
            df = df.with_columns(pl.col(column).str.json_decode().cast(target_dtype))

        # 时间类型转换
        elif current_dtype.is_temporal() and target_dtype.is_temporal():
            if verbose:
                print(f'{column}: {current_dtype} -> {target_dtype}（时间类型转换）')
            df = df.with_columns(pl.col(column).cast(target_dtype))

        # Categorical类型转换（保留原类别集合）
        elif current_dtype.is_(pl.String()) and isinstance(target_dtype, pl.Categorical):
            if verbose:
                print(f'{column}: {current_dtype} -> {target_dtype}（Categorical转换）')
            # 复用原Categorical的类别集合
            cat_categories = target_dtype.categories
            df = df.with_columns(
                pl.col(column).cast(pl.Categorical).cat.set_categories(cat_categories)
            )

    # 最终强制转换并校验
    df = df.cast(outbound_schema)
    assert df.schema == outbound_schema, "类型转换不完全匹配"
    return df

# 使用示例
df_from_snowflake = pl.read_database_uri('SELECT * FROM test_table', uri=con_str, engine='adbc')
df_restored = match_to_outbound_schema(df_from_snowflake, df_for_testing.schema, verbose=True)

3. 手动创建Snowflake表Schema后写入 #

避免ADBC自动推断Schema导致的类型偏差，先手动在Snowflake创建与Polars类型对应的表：

• Polars与Snowflake类型映射参考：

  Polars类型	Snowflake类型
  pl.Int64	NUMBER(38, 0)
  pl.Float64	FLOAT
  pl.List	ARRAY
  pl.Struct	OBJECT
  pl.Categorical	STRING

• 步骤1：在Snowflake中执行建表语句：

  CREATE OR REPLACE TABLE test_table (
      A1 NUMBER(38,0),
      A2 NUMBER(38,0),
      B1 ARRAY,
      B2 ARRAY,
      C1 STRING,
      C2 STRING,
      As OBJECT,
      Bs OBJECT,
      D1 STRING,
      D2 STRING
  );
  

• 步骤2：使用Polars写入时指定if_table_exists='append'（避免重建表）：

  df_for_testing.write_database(
      table_name='test_table',
      if_table_exists='append',
      connection=con_str,
      engine='adbc'
  )
  

• 步骤3：读取后用改进的转换函数转回原类型。

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