Snowflake中DBSCAN聚类UDF结果异常问题排查与解决问询

问题描述 #

在Snowflake中通过Pandas UDF运行DBSCAN聚类时，结果与本地执行差异极大：本地全量计算得到24个簇，Snowflake中UDF却输出71个簇。排查时发现，用于输出处理行号的测试UDF返回了重复行号，且最大值远小于表总行数，怀疑是并行化机制导致的问题。

附用户原代码（存在笔误，return resul应为return pd.Series(DBSCAN_cluster.labels_)）：

@funcs.pandas_udf(name='DBSCAN_TEST', is_permanent=True,
                stage_location='@UDF', replace=True,
                packages=['scikit-learn==1.0.2', 'pandas', 'numpy'])
def DBSCAN_TEST(data_x: types.PandasDataFrame[float, float]) -> types.PandasSeries[float]:
    data_x.columns = [features]
    DBSCAN_cluster = DBSCAN(eps=2.5, min_samples=4)
    DBSCAN_cluster.fit(data_x)
    return resul  # 此处存在笔误，应为return pd.Series(DBSCAN_cluster.labels_)

原因分析 #

Snowflake的普通Pandas UDF采用分片并行处理机制：系统会将数据集拆分为多个分片，每个UDF实例仅处理其中一个分片，而非全量数据。而DBSCAN是全局聚类算法，依赖所有数据点的距离关系划分簇，分片独立聚类必然会生成大量小簇，与本地全量计算结果完全不符。行号重复、最大值偏小的现象也验证了这一点——每个UDF实例处理的是局部数据，行号是分片内的编号，而非全局连续编号。

正确实现方式 #

要在Snowflake中正确运行DBSCAN这类全局聚类算法，需避免普通Pandas UDF的分片处理，采用以下两种方案：

方案1：全局聚合UDF（适用于中小数据集） #

将全量数据收集到单个节点完成聚类，再将簇标签映射回原数据集：

from snowflake.snowpark import functions as funcs
from snowflake.snowpark.types import StructType, StructField, FloatType, IntegerType
from sklearn.cluster import DBSCAN
import pandas as pd

# 定义全局聚合UDF，接收全量数据点数组，返回对应簇标签列表
@funcs.udf(name='DBSCAN_GLOBAL', is_permanent=True, stage_location='@UDF', replace=True,
           packages=['scikit-learn==1.0.2', 'pandas', 'numpy'],
           input_types=[StructType([StructField('X', FloatType()), StructField('Y', FloatType())])],
           return_type=IntegerType())
def dbscan_global(data_points):
    df = pd.DataFrame(data_points, columns=['X', 'Y'])
    dbscan = DBSCAN(eps=2.5, min_samples=4)
    labels = dbscan.fit_predict(df)
    return labels.tolist()

# 执行流程：收集全量数据→计算标签→关联回原数据
df = session.table('your_table')
# 聚合所有数据点为数组
agg_df = df.agg(funcs.collect_list(funcs.struct(df['X'], df['Y'])).alias('ALL_POINTS'))
# 生成簇标签并展开
label_df = agg_df.select(funcs.flatten(dbscan_global('ALL_POINTS')).alias('CLUSTER_LABEL'))
# 通过行号关联，保证标签与原数据顺序一致
result_df = df.with_column('ROW_ID', funcs.row_number().over(funcs.order_by('X', 'Y'))) \
              .join(label_df.with_column('ROW_ID', funcs.row_number().over(funcs.order_by(funcs.lit(1)))), on='ROW_ID')

方案2：Snowpark ML分布式聚类（推荐大数据量） #

使用Snowpark ML提供的分布式DBSCAN实现，无需手动处理分片逻辑：

from snowflake.ml.cluster import DBSCAN
from snowflake.snowpark.session import Session

# 初始化Snowpark会话
session = Session.builder.configs(your_config).create()
df = session.table('your_table')

# 分布式执行DBSCAN聚类
dbscan = DBSCAN(eps=2.5, min_samples=4, input_cols=['X', 'Y'], output_col='CLUSTER_LABEL')
model = dbscan.fit(df)
result_df = model.transform(df)

注意事项 #

• 禁止用普通Pandas UDF处理DBSCAN、K-Means这类依赖全局数据的算法，此类算法需要基于全量数据计算距离矩阵或全局统计量。
• 若数据集过大，单节点处理性能不足，优先选择Snowpark ML的分布式实现，或先对数据做降维预处理后再聚类。

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