PySpark DataFrame行数统计报错及SHAP计算行数不符问题

问题1：count()报错及死锁警告关联分析 #

调用PySpark DataFrame的df.count()时出现以下错误：

WARN PythonRunner: Detected deadlock while completing task 24.0 in stage 4 (TID 28): Attempting to kill Python Worker
...
ERROR Executor: Exception in task 24.0 in stage 4.0 (TID 28)
...
ValueError: Shape of passed values is (4,1), indices imply (4,4)

已知该ValueError由数据形状不匹配导致，但现有解决方案多针对Pandas DataFrame，不清楚PySpark下的处理方式；同时想明确死锁警告是否需要关注，是否和该ValueError有关联。

问题背景：基于Databricks代码计算SHAP值 #

参考Databricks代码实现模型SHAP值计算，代码如下：

explainer = shap.TreeExplainer(model)
shap_columns = ['feature1', 'feature2', 'feature3', 'feature4']

def calculate_shap(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    for X in iterator:
        yield pd.DataFrame(
            explainer.shap_values(np.array(X), check_additivity=False)[0],
            columns=shap_columns,
        )

return_schema = StructType()
for feature in shap_columns:
    return_schema = return_schema.add(StructField(feature, FloatType()))

df = spark_X.mapInPandas(calculate_shap, schema=return_schema)
df.count()

df与spark_X均为PySpark DataFrame类型，df.printSchema()显示4列结构正常。

问题2：修改代码后行数不符 #

将explainer.shap_values(...)部分用[]包裹后报错消失，但返回行数仅18K，与预期的180M不符，询问原因。

解答 #

问题1：错误原因及死锁关联 #

• ValueError根源：
  问题出在explainer.shap_values(np.array(X))[0]的返回形状与shap_columns列数不匹配。处理每个分区的Pandas DataFrameX时，shap_values返回的数组应为(样本数, 特征数)，但如果[0]索引误取了特征维度而非样本维度，会导致传入pd.DataFrame的数组形状变成(4,1)（4个特征，1个样本），而columns=shap_columns指定了4列，Pandas判定形状不匹配，触发错误。

• 死锁警告与ValueError的关联：
  死锁警告是衍生结果——Python Worker进程执行calculate_shap时抛出未捕获的ValueError，导致任务卡住，Spark的PythonRunner检测到任务长时间无响应，判定为死锁并尝试杀死Worker。解决ValueError后，死锁问题会自动消失，无需额外处理。

问题2：行数不符的原因 #

把explainer.shap_values(...)用[]包裹后，相当于给返回的SHAP数组外层多套了一层列表，Pandas会将整个数组当作一行数据处理，而非按样本拆分。原本每个分区的X有N个样本，应生成N行SHAP值，现在被压缩成1行，最终所有分区的行数累加后仅等于分区数（18K），而非预期的180M样本数。

修复方案 #

1. 确认SHAP返回结构：

   • 分类模型：shap_values返回每个类别的SHAP值数组，形状为(类别数, 样本数, 特征数)，需确保[0]取的是目标类别的数组，且形状为(样本数,4)。
   • 回归模型：shap_values直接返回单个数组，形状为(样本数, 特征数)，无需[0]索引。

2. 添加形状验证：
   在calculate_shap函数中添加打印语句，确认SHAP数组形状：

   def calculate_shap(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
       for X in iterator:
           shap_vals = explainer.shap_values(np.array(X), check_additivity=False)
           # 打印形状确认
           print(f"SHAP values shape: {np.array(shap_vals).shape}")
           # 根据模型类型取正确数组
           if isinstance(shap_vals, list):
               # 分类模型取第一个类别
               shap_vals = shap_vals[0]
           yield pd.DataFrame(shap_vals, columns=shap_columns)
   

3. 确保行匹配：
   保证每个分区的SHAP数组行数与输入X的行数一致，这样mapInPandas会将每个样本的SHAP值映射为一行，最终行数与spark_X一致。

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