如何在函数中获取传入的PySpark DataFrame实例的变量名？

核心思路 #

通过Python的inspect模块抓取函数调用处的源码上下文，再用正则或AST解析提取传入的变量名——因为PySpark DataFrame不可变且没法统一添加name属性，这种从调用上下文逆向解析的方式是最可行的替代方案。

实现方案（正则解析版） #

import inspect
import re
from pyspark.sql import DataFrame

def display(df: DataFrame):
    # 获取调用display函数的上层栈帧
    caller_frame = inspect.currentframe().f_back
    # 提取调用行的源码文本
    caller_line = inspect.getsource(caller_frame).strip()
    
    # 用正则匹配display的参数部分
    match_result = re.match(r'display\((.*?)\)', caller_line)
    if match_result:
        df_name = match_result.group(1).strip()
        # 这里替换成你的实际展示逻辑
        print(f"当前展示DataFrame: {df_name}")
        df.show()
    else:
        # 处理参数不是直接变量名的情况（比如表达式）
        print("无法识别DataFrame变量名")
        df.show()

# 测试示例
df_on_step_42 = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "val"])
display(df_on_step_42)

更鲁棒的AST解析版 #

如果要支持复杂的源码格式（比如参数跨行），可以用AST解析代替正则：

import inspect
import ast
from pyspark.sql import DataFrame

def display(df: DataFrame):
    caller_frame = inspect.currentframe().f_back
    caller_line = inspect.getsource(caller_frame).strip()
    try:
        # 把源码解析成抽象语法树
        tree = ast.parse(caller_line)
        call_node = tree.body[0].value
        arg_node = call_node.args[0]
        # 判断参数是否是直接变量名
        if isinstance(arg_node, ast.Name):
            df_name = arg_node.id
            print(f"当前展示DataFrame: {df_name}")
            df.show()
        else:
            print("传入的不是直接变量名（比如表达式）")
            df.show()
    except Exception as e:
        print(f"解析变量名失败: {str(e)}")
        df.show()

注意事项 #

• 仅当调用display时传入直接变量名才有效，如果传入的是表达式（比如display(df.filter("id > 1"))），无法解析出有意义的名称，必须做好降级处理。
• 依赖调用处的源码格式，如果调用时参数拆成多行，正则版本需要调整规则，AST版本兼容性更好。

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