如何基于Spark DataFrame按二分类目标绘制数值特征分组箱线图

核心实现逻辑 #

分组箱线图的核心是将每个特征的数值，按照二分类目标的两个取值拆分为独立数组，再将数组列表传入boxplot方法即可生成分层对比效果。注意如果Spark数据集量级较大，不要直接全量拉取数据到本地驱动节点，建议先按比例采样避免内存溢出。

完整可运行代码 #

import math
import matplotlib.pyplot as plt
from pyspark.sql.types import NumericType

# 配置项
target_column = "Recommended IND"  # 替换为你的二分类目标列名
subplot_per_row = 3  # 每行展示的子图数量，和原有代码的3列布局对齐
sample_ratio = 1.0   # 数据量超过10万行时建议调小到0.1~0.3，采样绘图

# 1. 自动识别所有数值特征列，排除目标列
numeric_features = [
    field.name for field in final_sdf.schema.fields
    if isinstance(field.dataType, NumericType) and field.name != target_column
]

# 2. 采样后将数据拉取到本地，转成pandas DataFrame方便后续处理
local_df = final_sdf.select(*numeric_features, target_column).sample(fraction=sample_ratio).toPandas()

# 3. 计算子图布局
row_count = math.ceil(len(numeric_features) / subplot_per_row)
fig, axes = plt.subplots(row_count, subplot_per_row, figsize=(5*subplot_per_row, 4*row_count))
axes = axes.flatten()

# 4. 逐特征绘制分组箱线图
for idx, feature in enumerate(numeric_features):
    current_ax = axes[idx]
    # 按二分类取值拆分两组数据
    group_neg = local_df[local_df[target_column] == 0][feature].dropna().values
    group_pos = local_df[local_df[target_column] == 1][feature].dropna().values
    
    current_ax.boxplot(
        [group_neg, group_pos],
        labels=[f"{target_column}=0", f"{target_column}=1"],
        showmeans=True,
        notch=True
    )
    current_ax.set_title(feature, fontsize=15)
    current_ax.set_ylabel("Feature value")

# 隐藏没有用到的空子图
for idx in range(len(numeric_features), len(axes)):
    axes[idx].axis("off")

plt.tight_layout()
plt.show()

对应原有单特征代码的改造版本 #

如果暂时不需要批量生成所有特征的图，只想先实现Age字段的分组箱线图，可以直接修改原有代码如下，不需要引入pandas：

age_lst = final_sdf.select('Age','Recommended IND').collect()
# 按目标值拆分两组数据
age_group0 = [row['Age'] for row in age_lst if row['Recommended IND'] == 0]
age_group1 = [row['Age'] for row in age_lst if row['Recommended IND'] == 1]

fig, ax = plt.subplots()
ax.set_title('Age',fontsize=15)
# 传入两个分组的数组列表，即可生成分组箱线图
ax.boxplot([age_group0, age_group1], labels=['Not Recommended', 'Recommended'], showmeans=True, notch=True)
ax.set_ylabel('years')
plt.show()

关键注意点 #

• 不要在大数据集上直接全量collect()或toPandas()，一定要先做采样，否则会导致驱动节点内存溢出
• boxplot方法接收的分组参数是列表形式，每个元素对应一个分组的数值序列，按顺序展示即可
• 自动识别数值列的逻辑会跳过目标列，不需要手动枚举所有特征名，新增特征时会自动纳入绘图范围

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