自定义Pandas类的方法链式调用实现及结果异常问题

为啥会踩这俩坑 #

1. 一开始没法链式调用：因为你的return_agg和return_sample方法返回的是聚合结果或者抽样后的DataFrame，不是MyClass的实例本身，自然没法像Pandas那样.到底。
2. 抽样后聚合结果不对：后来改链式调用时，估计只是随便返回了self，但return_sample没更新类内部存的DataFrame，导致return_agg还是用最开始的原始数据算，根本没用到抽样后的结果。

正确实现思路 #

核心就是：类内部要存一份当前状态的DataFrame，修改数据的方法更新这个内部数据后返回自身，输出结果的方法按需返回结果或者实例。

完整可运行代码 #

import pandas as pd
import numpy as np

class MyClass:
    def __init__(self, df: pd.DataFrame):
        # 内部存一份当前数据，复制是为了不污染用户传入的原始df
        self.current_df = df.copy()

    def return_sample(self, sample_size: int = 100, random_seed: int = 42):
        # 抽样后更新内部的current_df，然后返回self支持链式调用
        self.current_df = self.current_df.sample(n=sample_size, random_state=random_seed)
        return self

    def return_agg(self, agg_columns: list, agg_method: str = 'mean'):
        # 基于当前内部的抽样后数据做聚合，直接返回结果
        return self.current_df[agg_columns].agg(agg_method)

# 测试一下
if __name__ == '__main__':
    # 造点测试数据
    raw_df = pd.DataFrame({
        'value1': np.random.randn(1000),
        'value2': np.random.randint(0, 20, 1000),
        'group': np.random.choice(['X', 'Y', 'Z'], 1000)
    })

    # 直接聚合的结果
    direct_result = MyClass(raw_df).return_agg(['value1', 'value2'])
    print("直接聚合结果:\n", direct_result)

    # 先抽样再聚合的结果
    sample_result = MyClass(raw_df).return_sample(sample_size=200).return_agg(['value1', 'value2'])
    print("\n抽样后聚合结果:\n", sample_result)

关键细节说清楚 #

• 内部数据隔离：__init__里用df.copy()是怕你改内部数据时把用户传的原始df也改了，避免不必要的bug。
• 链式调用的核心：return_sample必须返回self，这样调用完它之后，还能接着调用return_agg。
• 灵活调整返回值：如果return_agg之后还想链式调用其他方法，也可以让它返回self，同时把聚合结果存在类的属性里，比如：

  def return_agg(self, agg_columns: list, agg_method: str = 'mean'):
      self.agg_result = self.current_df[agg_columns].agg(agg_method)
      return self
  

  调用时就这么用：MyClass(raw_df).return_sample(200).return_agg(['value1']).agg_result

避坑提醒 #

• 别在方法里直接返回处理后的DataFrame，除非你就是不想链式调用，否则一返回df就断了链式。
• 千万别忘了更新类内部的current_df，不然所有方法都是基于最开始的原始数据在跑。
• 如果需要保留原始数据，还可以多存一个self.original_df，方便后续对比或者恢复。

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