Python实现DataFrame Object类型列迭代自定义函数并导出每页2图PDF

问题需求 #

我有一个自定义柱状图函数label_bar，需要完成以下需求：

• 循环遍历DataFrame中的Object类型（分类）列，应用该函数生成图表；
• 将所有生成的图表导出为PDF文件，要求每页展示2个图表。

目前已完成函数编写、测试数据集构造及单图表调用，但批量处理与PDF导出部分存在困难，需要完善这部分代码并添加注释。

完整代码实现 #

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages

# 全局显示设置
pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)
pd.set_option('display.float_format', lambda x: '%.3f' % x)

# 自定义柱状图函数，修改为支持传入子图轴对象，适配批量PDF导出
def label_bar(data, feature, ax, perc=False, n=None):
    """
    顶部带数值/百分比的柱状图
    data: 数据源DataFrame
    feature: 要绘制的DataFrame列名
    ax: 用于绘制的matplotlib子图轴对象
    perc: 是否显示百分比（默认显示计数）
    n: 仅显示前n个高频类别（默认显示所有类别）
    """
    total = len(data[feature])
    count = data[feature].nunique()
    
    # 绘制计数柱状图
    sns.countplot(
        data=data,
        x=feature,
        palette="Paired",
        order=data[feature].value_counts().index[:n],
        ax=ax
    )
    
    # 设置子图属性
    ax.set_xticklabels(ax.get_xticklabels(), rotation=90, fontsize=12)
    ax.set_title(f'Distribution of {feature}', fontsize=14)
    
    # 为每个柱子添加数值/百分比标注
    for p in ax.patches:
        if perc:
            label = "{:.1f}%".format(100 * p.get_height() / total)
        else:
            label = int(p.get_height())
        
        x = p.get_x() + p.get_width() / 2
        y = p.get_height()
        
        ax.annotate(
            label,
            (x, y),
            ha="center",
            va="center",
            size=10,
            xytext=(0, 5),
            textcoords="offset points"
        )

# 生成测试数据集
data = {
    'pet': ['cat', 'leopard', 'dog', 'goat', 'goat','cat','lizard','rabbit','lizard','goat','fish','cat', 'dog', 'dog', 'goat', 'goat','cat','lizard','rabbit','lizard'],
    'vet_cost': [1200, 150, 300, 450, 200,550,750,800,350,150,320,150, 300, 450, 200,550,750,800,350,150],
    'x_region':['R1','R3','R1','R5','R2','R5','R1','R1','R6','R1','R1','R3','R1','R5','R2','R5','R1','R1','R7','R1'],
    'x_class':['T1','T1','T2','T3','T5','T5','T5','T0','T5','T4','T1','T6','T2','T3','T7','T8','T5','T8','T5','T4'],
    'x_diet':['D1','D3','D1','D5','D2','D4','D4','D1','D6','D1','D1','D3','D1','D5','D2','D5','D1','D1','D7','D1'],
    'y_region':['Y1','Y3','Y1','Y5','Y2','Y5','Y1','Y1','Y6','Y1','Y1','Y3','Y1','Y5','Y2','Y5','Y1','Y1','Y7','Y1'],
    'y_class':['Z1','Z1','Z2','Z3','Z5','Z5','Z5','Z0','Z5','Z4','Z1','Z6','Z2','Z3','Z7','Z8','Z5','Z8','Z5','Z4'],
    'y_diet':['A1','A3','A1','A5','A2','A4','A4','A1','A6','A1','A1','A3','A1','A5','A2','A5','A1','A1','A7','A1'],
}
df = pd.DataFrame(data)

# 筛选所有Object类型的分类列
object_columns = df.select_dtypes('object').columns.tolist()

# 初始化PDF文件，指定保存路径
with PdfPages('category_distributions.pdf') as pdf:
    # 按每页2个图表分组处理
    for i in range(0, len(object_columns), 2):
        # 创建每页的figure和子图，2行1列布局
        fig, axes = plt.subplots(nrows=2, ncols=1, figsize=(10, 12))
        
        # 处理当前页的第一个图表
        if i < len(object_columns):
            label_bar(df, object_columns[i], axes[0], perc=True)
        
        # 处理当前页的第二个图表（如果还有剩余列）
        if i + 1 < len(object_columns):
            label_bar(df, object_columns[i+1], axes[1], perc=True)
        
        # 调整子图间距，避免标题和标签重叠
        plt.tight_layout()
        
        # 将当前页保存到PDF
        pdf.savefig(fig)
        plt.close(fig)  # 关闭figure释放内存

print("PDF文件已成功生成：category_distributions.pdf")

关键修改说明 #

1. 函数适配：修改label_bar函数，新增ax参数，让函数可以在指定的子图上绘制，而不是每次创建独立的figure，解决批量生成的问题；移除原函数中的plt.show()，避免弹窗干扰。
2. 批量遍历：通过df.select_dtypes('object')筛选所有分类列，使用步长为2的循环实现每页2个图表的分组。
3. PDF导出：使用PdfPages上下文管理器自动处理PDF的创建和关闭，每页创建2个子图，绘制完成后保存并关闭figure，避免内存泄漏。
4. 布局优化：使用plt.tight_layout()自动调整子图间距，避免图表元素重叠。

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