SnowPark中DataFrame分组聚合：NULL作通配符的实现问题

原代码存在的问题 #

1. 列名引用错误：DataFrame schema定义为小写的a/b/c/amount，但代码中使用了大写的df.A/df.B/df.C，会触发列不存在的异常。
2. 聚合函数错误：需求是对Amount列求和，代码中误用了F.avg()，应替换为F.sum()。
3. 自连接逻辑缺陷：原连接条件的方向错误，应该是左表行的非NULL列必须匹配右表行的对应列，而非左表为NULL时匹配所有；同时distinct后的左表会导致分组标识重复计算。

正确实现方案 #

方式一：自连接+精准过滤 #

这种方式通过自连接匹配所有符合条件的行，再分组聚合：

import snowflake.snowpark.functions as F

# 初始化DataFrame
df = session.create_dataframe([
      ("BI", 2,    "c", 2),
      ("MO", 3,    "c", 4),
      (None, 2,    "c", 9),
      ("BI", 2,    "c", 2),
      ("MO", 2,    "c", 4),
      ("MO", 3,    "c", 2),
      ("MO", None, "d", 4),
      ("MO", 3,    "d", 2),
      ("MO", 2,    "d", 3)],
      schema=["a", "b", "c", "amount"]
    )

# 构建正确的连接条件：右表行必须在左表行的非NULL列上完全匹配
join_expr = (
    F.when(F.col("left.a").is_not_null(), F.col("left.a") == F.col("right.a")).otherwise(F.lit(True)) &
    F.when(F.col("left.b").is_not_null(), F.col("left.b") == F.col("right.b")).otherwise(F.lit(True)) &
    F.when(F.col("left.c").is_not_null(), F.col("left.c") == F.col("right.c")).otherwise(F.lit(True))
)

# 执行自连接、分组求和、去重
result = (
    df.alias("left")
    .join(df.alias("right"), join_expr)
    .groupBy("left.a", "left.b", "left.c")
    .agg(F.sum("right.amount").alias("Amount_sum"))
    .distinct()
)

result.show()

方式二：窗口函数（更高效） #

通过构造动态分区键，用窗口函数实现聚合，避免自连接的性能开销：

import snowflake.snowpark.functions as F
from snowflake.snowpark.window import Window

df = session.create_dataframe([
      ("BI", 2,    "c", 2),
      ("MO", 3,    "c", 4),
      (None, 2,    "c", 9),
      ("BI", 2,    "c", 2),
      ("MO", 2,    "c", 4),
      ("MO", 3,    "c", 2),
      ("MO", None, "d", 4),
      ("MO", 3,    "d", 2),
      ("MO", 2,    "d", 3)],
      schema=["a", "b", "c", "amount"]
    )

# 生成动态分区键：非NULL列用原值，NULL列用固定占位符统一分组
partition_keys = [
    F.when(F.col("a").is_not_null(), F.col("a")).otherwise(F.lit("__IGNORE__")),
    F.when(F.col("b").is_not_null(), F.col("b")).otherwise(F.lit("__IGNORE__")),
    F.when(F.col("c").is_not_null(), F.col("c")).otherwise(F.lit("__IGNORE__"))
]

window_spec = Window.partitionBy(*partition_keys)

# 计算分区内的求和，保留原始列并去重
result = (
    df.withColumn("Amount_sum", F.sum("amount").over(window_spec))
    .select("a", "b", "c", "Amount_sum")
    .distinct()
)

result.show()

验证结果 #

两种方案都会生成符合预期的输出：

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