如何优化DataProcessing类代码的可读性与运行效率？

核心思路 #

利用类的实例属性存储中间状态，拆分单一职责的方法避免重复传参，同时针对多进程场景设计纯函数式的策略判断逻辑，兼顾代码可读性与运行效率。

具体实现代码 #

import pandas as pd
from multiprocessing import Pool
import os

class DataProcessing:
    def __init__(self, csv_path, constant_df):
        # 初始化必要配置与属性，避免重复传递参数
        self.csv_path = csv_path
        self.constant_df = constant_df
        self.df = None
        self.metrics_df = None

    def _read_csv(self):
        # 优化CSV读取：指定列、数据类型减少内存占用
        self.df = pd.read_csv(
            self.csv_path,
            usecols=["code", "date", "value1", "value2"],  # 只加载需要的列
            dtype={"code": str, "value1": float, "value2": float}
        )

    def _merge_constant_df(self):
        # 拼接常量数据，若constant_df是小表可考虑广播优化
        self.df = pd.merge(self.df, self.constant_df, on="code", how="left")

    def _calculate_metrics(self):
        # 矢量化计算指标，避免循环提升速度
        self.df["metric1"] = self.df["value1"] * self.df["constant_col"]
        self.df["metric2"] = self.df["value2"].rolling(7).mean()
        # 也可单独存储指标结果
        self.metrics_df = self.df[["code", "date", "metric1", "metric2"]]

    @staticmethod
    def _single_code_strategy(code_data):
        # 纯函数式策略判断：仅依赖传入的单code数据，适合多进程
        code, group = code_data
        # 示例策略：判断metric1是否连续3天大于阈值
        group["signal"] = (group["metric1"] > 100).rolling(3).sum() == 3
        return group

    def run_parallel_strategy(self):
        # 按code分组拆分数据
        code_groups = list(self.df.groupby("code"))
        # 进程数设为CPU核心数，避免资源浪费
        process_num = min(os.cpu_count(), len(code_groups))
        
        with Pool(process_num) as pool:
            # 多进程执行策略判断
            result_groups = pool.map(self._single_code_strategy, code_groups)
        
        # 合并结果
        self.df = pd.concat(result_groups, ignore_index=True)

    def full_process(self):
        # 串联所有步骤，对外提供统一调用入口
        self._read_csv()
        self._merge_constant_df()
        self._calculate_metrics()
        self.run_parallel_strategy()
        return self.df

优势说明 #

1. 代码可读性：每个方法仅负责一个步骤，类属性存储中间结果，无需重复传递code等参数，结构清晰易维护。
2. 运行效率：
   • CSV读取阶段通过指定列和数据类型减少内存开销，加快加载速度；
   • 指标计算使用pandas矢量化操作，比循环快几个数量级；
   • 多进程采用纯函数式的单code处理，避免进程间状态共享的问题，充分利用多核CPU。
3. 扩展性：后续新增步骤只需添加对应方法，修改逻辑只需调整单个方法，耦合度低。

额外优化建议 #

• 如果constant_df是静态小数据，可在类初始化时提前设置索引，提升合并效率；
• 多进程时若数据量极大，可考虑用multiprocessing.Pool.imap分批处理，避免内存溢出；
• 策略判断逻辑复杂时，可将其拆分为更小的静态方法，保持单一职责。

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