类方法中使用PySpark applyInPandas时SparkSession缺失问题排查

错误根源 #

报错AttributeError: 'NoneType' object has no attribute '_jvm'的核心原因是：

• Spark的Worker进程启动时，会自动导入你的业务模块（这里是utils.py），但Worker进程不会自动初始化SparkContext/SparkSession。
• 你的utils.py在**模块顶层（文件初始化阶段）**直接调用了pyspark.sql.functions.col()，这个函数依赖Driver端的SparkContext，Worker进程中没有可用的上下文，因此抛出NoneType错误。
• Databricks环境默认会给Worker进程自动初始化Spark上下文，所以不会触发该问题；但CI/CD的VM环境无此逻辑，导致报错。

修复方案（优先保留applyInPandas） #

1. 重构模块级Spark代码 #

将utils.py中模块顶层的Spark函数调用，移到函数内部延迟执行：

# 原错误写法（utils.py）
from pyspark.sql import functions as f
presenceDict = {
    'CCH w/ coffee machine': ((f.col('source') == 'customer_md') & (f.col('has_coffee_machine') == 'YES')),
    # 其他键值对
}

# 修改后写法
from pyspark.sql import functions as f

def get_presence_dict():
    return {
        'CCH w/ coffee machine': ((f.col('source') == 'customer_md') & (f.col('has_coffee_machine') == 'YES')),
        # 其他键值对
    }

这样只有当Driver端调用get_presence_dict()时才会执行f.col()，避免Worker进程初始化时触发错误。

2. 确保applyInPandas的处理函数纯Pandas化 #

验证_MyPandasImputerFunction及其依赖的所有代码，仅使用Pandas API，绝对不能包含任何Spark相关的函数或初始化逻辑——因为该函数是在Worker进程的Pandas环境中执行的，与Driver端的Spark上下文完全隔离。

用pandas_udf实现分组DataFrame处理 #

如果需要切换到pandas_udf，可使用GROUPED_MAP类型，逻辑与applyInPandas完全一致，代码示例如下：

from pyspark.sql.session import SparkSession
from pyspark.sql import functions as F
from pyspark.sql import types as t
import pandas as pd

class MyClass():
    def transform(self, spark_df):
        self.spark = SparkSession.builder.getOrCreate()
        final_df = self._MySparkMethod(spark_df)
        return final_df  # 原代码遗漏返回，需补充

    @F.pandas_udf(
        returnType=t.StructType([
            t.StructField('col_1', t.StringType()),
            t.StructField('col_2', t.DoubleType())
        ]),
        functionType=F.PandasUDFType.GROUPED_MAP
    )
    def _MyPandasImputerFunction(pdf: pd.DataFrame) -> pd.DataFrame:
        # 纯Pandas填充操作示例
        pdf['col_2'] = pdf['col_2'].fillna(pdf['col_2'].median())
        return pdf

    def _MySparkMethod(self, spark_df):
        return spark_df.groupby("col_1").apply(self._MyPandasImputerFunction)

注意：

• GROUPED_MAP类型的pandas_udf输入为分组后的Pandas DataFrame，输出也必须是Pandas DataFrame，返回schema需与输出结构严格匹配。
• 同样要保证该UDF内部及依赖模块无Spark API调用。

额外验证步骤 #

1. 在CI/CD环境中，提前在Driver端完成SparkSession初始化，再执行类的transform方法。
2. 单独测试utils.py的导入，确认模块初始化时不会触发Spark相关报错。

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