PySpark toPandas()在单元测试中触发未关闭套接字ResourceWarning

问题根源 #

这个警告的核心和unittest的测试隔离机制、Spark toPandas()的资源管理逻辑直接相关：

• unittest为保证用例间独立，会在每个测试结束后强制回收内存资源。而toPandas()内部依赖socket连接从Executor拉取数据到Driver，这个socket的正常关闭依赖Python上下文管理器或垃圾回收的延迟处理。当unittest强制回收时，socket还没来得及完成关闭流程，就触发了ResourceWarning。
• 直接运行代码时，Python垃圾回收有足够时间等待socket的关闭逻辑执行，因此不会触发警告。

可行的解决方法 #

1. 显式管理SparkSession生命周期 #

在测试用例的setUp()中初始化SparkSession，tearDown()里显式关闭，确保所有数据传输操作完成后再释放资源：

import unittest
import gc
from pyspark.sql import SparkSession

class TestSparkPandasConversion(unittest.TestCase):
    def setUp(self):
        self.spark = SparkSession.builder \
            .master("local[1]") \
            .appName("TestToPandasWarning") \
            .getOrCreate()

    def tearDown(self):
        # 先关闭SparkSession，主动释放所有网络连接
        self.spark.stop()
        # 强制触发垃圾回收，清理残留对象
        gc.collect()

    def test_to_pandas_no_warning(self):
        test_df = self.spark.createDataFrame([(1, "test"), (2, "demo")], ["id", "content"])
        pd_df = test_df.toPandas()
        # 你的测试断言逻辑
        self.assertEqual(pd_df["id"].tolist(), [1, 2])

2. 启用Arrow优化的Pandas转换 #

Spark的Arrow优化版toPandas()资源管理更严谨，能减少这类socket泄漏问题。初始化SparkSession时添加对应配置：

self.spark = SparkSession.builder \
    .master("local[1]") \
    .appName("TestToPandasArrow") \
    .config("spark.sql.execution.arrow.pyspark.enabled", "true") \
    .getOrCreate()

注意：启用Arrow需要保证Spark（2.3+）和Pandas（0.23+）版本兼容。

3. 手动触发垃圾回收（应急方案） #

如果不想修改Spark配置或会话管理，也可以在toPandas()调用后、测试结束前手动触发GC：

def test_to_pandas_gc(self):
    test_df = self.spark.createDataFrame([(1, "test")], ["id", "content"])
    pd_df = test_df.toPandas()
    # 手动GC清理残留socket对象
    import gc
    gc.collect()
    self.assertEqual(len(pd_df), 1)

验证方法 #

运行测试时加上-W default参数，强制显示所有警告，确认问题是否解决：

python -m unittest -W default your_test_file.py

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