如何在Snowpark中用Snowpark或Python UDF重映射分组字典值生成新列

方法1：用when()+otherwise()构建条件分支 #

这是最贴合SQL CASE WHEN逻辑的写法，适合映射规则简单的场景：

from snowflake.snowpark.functions import when, col

# 初始化原DataFrame
df = session.create_dataframe([("Apple",), ("Banana",), ("Potato",), ("Onion",)], schema=["Item"])

# 生成Category列
df_with_category = df.with_column(
    "Category",
    when(col("Item").isin(["Apple", "Banana"]), "Fruit")
    .when(col("Item").isin(["Potato", "Onion"]), "Vegetable")
    .otherwise("Unknown")  # 可选，处理未匹配的兜底情况
)

df_with_category.show()

方法2：映射字典+自定义UDF #

如果映射条目较多或规则复杂，用字典存储映射关系再配合UDF更高效：

from snowflake.snowpark.functions import udf
from snowflake.snowpark.types import StringType

# 定义映射规则字典
item_category_map = {
    "Apple": "Fruit",
    "Banana": "Fruit",
    "Potato": "Vegetable",
    "Onion": "Vegetable"
}

# 注册自定义函数
@udf(return_type=StringType())
def map_category(item):
    return item_category_map.get(item, "Unknown")

# 应用UDF生成新列
df_with_category = df.with_column("Category", map_category(col("Item")))

df_with_category.show()

方法3：通过Join关联映射表 #

如果映射规则需要频繁维护，可以把映射关系单独存为DataFrame或Snowflake表，用Join实现关联：

# 创建映射规则DataFrame
mapping_df = session.create_dataframe(
    [("Apple", "Fruit"), ("Banana", "Fruit"), ("Potato", "Vegetable"), ("Onion", "Vegetable")],
    schema=["Item", "Category"]
)

# 左关联生成新列
df_with_category = df.join(mapping_df, on="Item", how="left")

df_with_category.show()

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