子类化PySpark DataFrame：如何让父类方法返回子类对象？

你的需求完全可行，下面提供几种解决方案，各有优劣，可根据场景选择：

方案1：通过__getattribute__拦截方法调用 #

利用Python的属性拦截机制，在调用父类方法后，将返回的原生DataFrame包装成子类实例：

from pyspark.sql import DataFrame, col

class EnhancedDataframe(DataFrame):
    def __init__(self, df):
        super().__init__(df._jdf, df.sql_ctx)
    
    def notNullCount(self, col_name):
        return self.filter(col(col_name).isNotNull()).count()
    
    def __getattribute__(self, name):
        attr = super().__getattribute__(name)
        # 仅拦截可调用的方法，排除自定义方法和特殊方法
        if callable(attr) and name not in ['notNullCount', '__init__', '__getattribute__']:
            def wrapper(*args, **kwargs):
                result = attr(*args, **kwargs)
                # 若返回原生DataFrame，包装为子类实例
                if isinstance(result, DataFrame) and not isinstance(result, EnhancedDataframe):
                    return EnhancedDataframe(result)
                return result
            return wrapper
        return attr

注意：需要手动排除不需要包装的方法，避免循环调用或异常，PySpark版本更新可能需要调整过滤逻辑。

方案2：包装器模式（组合代替继承） #

不直接子类化DataFrame，而是创建一个包含原生DataFrame的包装类，通过__getattr__代理父类方法，同时返回包装后的实例：

from pyspark.sql import DataFrame, col

class EnhancedDataframe:
    def __init__(self, df):
        self.df = df
    
    def notNullCount(self, col_name):
        return self.df.filter(col(col_name).isNotNull()).count()
    
    def __getattr__(self, name):
        attr = getattr(self.df, name)
        if callable(attr):
            def wrapper(*args, **kwargs):
                result = attr(*args, **kwargs)
                if isinstance(result, DataFrame):
                    return EnhancedDataframe(result)
                return result
            return wrapper
        return attr
    
    # 代理常用属性，比如schema、columns
    @property
    def schema(self):
        return self.df.schema
    
    @property
    def columns(self):
        return self.df.columns

优势：隔离原生DataFrame的内部实现，兼容性更强，避免子类化带来的潜在冲突。

方案3：利用PySpark原生transform方法 #

无需修改类，将自定义方法写成独立函数，通过transform链式调用：

from pyspark.sql import DataFrame, col

def not_null_count(df: DataFrame, col_name: str) -> int:
    return df.filter(col(col_name).isNotNull()).count()

# 使用方式
# df = spark.read.parquet("data_path")
# df.transform(not_null_count, "target_col")

优势：符合PySpark函数式编程风格，无需维护自定义类，灵活度高。

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