Databricks PySpark执行Join时报DataSourcePartitioning异常求助

这个错误是Spark 3.2.1（Databricks 10.4 LTS内置版本）的已知bug，触发场景为：

• 输入DataFrame来自V2数据源（如Delta Lake），带有DataSourcePartitioning类型的分区元数据
• 自适应查询执行（AQE，默认开启）处于启用状态
• 执行Join操作时，AQE的EnsureRequirements优化器无法识别DataSourcePartitioning类型，导致抛出IllegalStateException

1. 临时禁用AQE或相关优化 #

在执行目标函数前，设置Spark配置关闭AQE，或者仅关闭引发问题的分区调整逻辑：

# 方案1：全局禁用AQE（最简单直接）
spark.conf.set("spark.sql.adaptive.enabled", "false")

# 方案2：仅禁用AQE的分区合并和倾斜Join优化（保留其他AQE功能）
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "false")
spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "false")

执行完函数后可恢复配置：

spark.conf.set("spark.sql.adaptive.enabled", "true")

2. 清除输入DataFrame的数据源分区元数据 #

通过将DataFrame转换为RDD再转回DataFrame，彻底清除数据源自带的分区信息，避免优化器识别到DataSourcePartitioning：
修改函数开头，对输入DataFrame做转换：

from pyspark.sql import Row, functions as F
from pyspark.sql.types import DataFrame

def expand_category(input_df: DataFrame, category_list: list, column_list: list) -> DataFrame:
    '''
    @input: 待扩展的DataFrame
    @category_list: 类别列表
    @column_list: 标识记录唯一性的列
    '''
    # 清除数据源分区元数据
    input_df = input_df.rdd.toDF()
    
    input_device_id_df = input_df.select(F.col('device_id')).distinct()
    category_list = [{'category': i} for i in category_list]
    category_df = spark.createDataFrame(Row(**x) for x in category_list)
    cross_join_df = input_device_id_df.crossJoin(category_df)
    output_df_with_null = cross_join_df.join(input_df, on = column_list)
    output_df = output_df_with_null.na.fill(0)
    return output_df

3. 重构函数逻辑，优化Join方式 #

改用先聚合原数据，再与全量组合左关联的方式，减少优化器处理复杂分区的概率：

from pyspark.sql import functions as F
from pyspark.sql.types import DataFrame, StringType

def expand_category(input_df: DataFrame, category_list: list, column_list: list) -> DataFrame:
    '''
    @input: 待扩展的DataFrame
    @category_list: 类别列表
    @column_list: 标识记录唯一性的列
    '''
    # 生成所有device_id与category的全量组合
    full_category_combinations = input_df.select("device_id").distinct() \
        .crossJoin(spark.createDataFrame(category_list, StringType()).toDF("category"))
    
    # 对原数据按唯一键聚合，确保每个device_id+category只有一条记录
    aggregated_input = input_df.groupBy(column_list).agg(
        F.sum("active_photos").alias("active_photos"),
        F.sum("attentiont_milliseconds").alias("attentiont_milliseconds"),
        F.sum("dwell_milliseconds").alias("dwell_milliseconds"),
        F.sum("dwell_ots").alias("dwell_ots"),
        F.sum("ots").alias("ots"),
        F.sum("total_photos").alias("total_photos"),
        F.sum("views").alias("views")
    )
    
    # 左关联后填充缺失值为0
    output_df = full_category_combinations.join(aggregated_input, on=column_list, how="left") \
        .na.fill(0)
    
    return output_df

该bug在Spark 3.3.x及以上版本已修复，若条件允许，可升级Databricks Runtime至11.3 LTS或更高版本，从根本上解决问题。

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