如何找出Spark中所有已持久化的DataFrame并存储为列表？

现有一段Spark代码如下：

df = spark.sql("""
select * from abc
""")

df.persist()

df2 = spark.sql("""
select * from def
""")

df2.persist()

df3 = spark.sql("""
select * from mno""")

需要找出所有调用过persist()的DataFrame变量，并生成包含这些变量名的列表，期望输出为：

l1 = [df, df2]

可以根据需求选择静态代码分析或运行时追踪两种方式：

1. 静态代码分析（直接处理代码文本） #

如果只是需要从代码字符串中提取目标变量名，用正则表达式匹配persist()调用的模式即可：

import re

# 原始代码文本
code_content = """
df = spark.sql("""
select * from abc
""")

df.persist()

df2 = spark.sql("""
select * from def
""")

df2.persist()

df3 = spark.sql("""
select * from mno""")
"""

# 匹配变量名.persist()的格式，兼容前后可能的空格
matched_vars = re.findall(r'(\w+)\s*\.\s*persist\s*\(\s*\)', code_content)
# 去重同时保持变量出现的原有顺序
unique_dfs = list(dict.fromkeys(matched_vars))
# 生成目标列表代码
print(f"l1 = [{', '.join(unique_dfs)}]")

执行后会直接输出l1 = [df, df2]，正则表达式能兼容变量名与方法、方法与括号间的空格，覆盖多数代码格式场景。

2. 运行时动态收集 #

如果需要在代码执行过程中实时记录调用过persist()的DataFrame，可以通过包装DataFrame的persist方法实现：

# 初始化存储列表
persisted_dfs = []

# 获取原始的persist方法
from pyspark.sql.dataframe import DataFrame
original_persist = DataFrame.persist

# 定义包装后的persist方法，调用时将DataFrame实例加入列表
def tracked_persist(self, storageLevel=None):
    persisted_dfs.append(self)
    if storageLevel:
        return original_persist(self, storageLevel)
    return original_persist(self)

# 替换原persist方法
DataFrame.persist = tracked_persist

# 执行你的Spark代码
df = spark.sql("""select * from abc""")
df.persist()

df2 = spark.sql("""select * from def""")
df2.persist()

df3 = spark.sql("""select * from mno""")

# 此时persisted_dfs就是包含所有调用过persist的DataFrame实例的列表
l1 = persisted_dfs

这种方式能在程序运行过程中自动收集目标DataFrame，适合需要在代码内部直接使用该列表的场景。

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