如何让函数每次执行时获取最新全局DataFrame用于Spark SQL查询？

核心问题分析 #

你的函数无法获取后续新增的全局DataFrame，主要原因有两个：

1. 函数内部使用dir()（无参数）会返回函数自身的局部变量，而非全局命名空间的变量。
2. 当函数从外部模块导入时，globals()指向的是模块自身的全局空间，而非调用者（如Notebook）的全局空间。

解决方案 #

方案1：同一Notebook内的函数修正

直接遍历当前Notebook的全局命名空间，筛选出所有DataFrame：

import pandas as pd
import pyspark.pandas as ps
import re

def get_all_dataframes():
    # 遍历全局变量，筛选Pandas DataFrame
    return [name for name, var in globals().items() if isinstance(var, pd.core.frame.DataFrame)]

def query_data(query):
    dataframes = get_all_dataframes()
    
    # 提取SQL中的表名
    tables = re.findall(r'FROM\s+(\w+)|JOIN\s+(\w+)', query, re.IGNORECASE)
    tables = [table for sublist in tables for table in sublist if table]

    tables_found = []
    # 替换表名为占位符并验证存在性
    for table in tables:
        if table in dataframes:
            # 使用正则边界避免部分匹配
            query = re.sub(r'\b' + re.escape(table) + r'\b', f'{{{table}}}', query)
            if table not in tables_found:
                tables_found.append(table)
        else:
            raise ValueError(f"Table '{table}' does not exist in memory.")

    # 构建查询参数并转换为PySpark DataFrame
    print(f"Tables detected: {tables_found}")
    query_args = {table: ps.from_pandas(globals()[table]) for table in tables_found}

    def _run_query(**query_args):
        return ps.sql(query, **query_args).to_pandas()

    query_results = _run_query(**query_args)
    print(f"Table created with {query_results.shape[0]:,} rows and {query_results.shape[1]:,} columns.")

    return query_results

方案2：外部模块导入的函数修正

使用inspect模块获取调用者（如Notebook）的全局命名空间，确保能读取最新的变量：

# 保存为sql_utils.py
import inspect
import pandas as pd
import pyspark.pandas as ps
import re

def get_all_dataframes():
    # 获取调用者的全局命名空间
    caller_frame = inspect.currentframe().f_back
    try:
        caller_globals = caller_frame.f_globals
        return [name for name, var in caller_globals.items() if isinstance(var, pd.core.frame.DataFrame)]
    finally:
        # 避免引用循环导致内存泄漏
        del caller_frame

def query_data(query):
    dataframes = get_all_dataframes()
    
    tables = re.findall(r'FROM\s+(\w+)|JOIN\s+(\w+)', query, re.IGNORECASE)
    tables = [table for sublist in tables for table in sublist if table]

    tables_found = []
    for table in tables:
        if table in dataframes:
            query = re.sub(r'\b' + re.escape(table) + r'\b', f'{{{table}}}', query)
            if table not in tables_found:
                tables_found.append(table)
        else:
            raise ValueError(f"Table '{table}' does not exist in memory.")

    print(f"Tables detected: {tables_found}")
    # 再次获取调用者全局空间以访问DataFrame
    caller_frame = inspect.currentframe().f_back
    try:
        caller_globals = caller_frame.f_globals
        query_args = {table: ps.from_pandas(caller_globals[table]) for table in tables_found}
    finally:
        del caller_frame

    def _run_query(**query_args):
        return ps.sql(query, **query_args).to_pandas()

    query_results = _run_query(**query_args)
    print(f"Table created with {query_results.shape[0]:,} rows and {query_results.shape[1]:,} columns.")

    return query_results

在Notebook中使用：

from sql_utils import query_data

# 创建初始DataFrame
df1 = pd.DataFrame({'a': [1,2], 'key': ['x','y']})
df2 = pd.DataFrame({'b': [3,4], 'key': ['x','y']})

# 执行查询
result = query_data("""
SELECT m1.a, m2.b
FROM df1 m1 INNER JOIN df2 m2
ON m1.key = m2.key
ORDER BY m1.a, m2.b
""")

# 新增DataFrame后再次查询
df3 = pd.DataFrame({'c': [5,6], 'key': ['x','y']})
result2 = query_data("""
SELECT m1.a, m3.c
FROM df1 m1 INNER JOIN df3 m3
ON m1.key = m3.key
""")

关键改进点 #

1. 替换dir()为globals().items()：直接遍历全局变量并验证类型，避免读取函数局部变量。
2. 使用inspect获取调用者命名空间：解决外部模块函数无法访问Notebook全局变量的问题。
3. 移除语法错误：修正get_all_dataframes()中多余的=符号。
4. 正则安全处理：添加re.escape()避免表名含特殊字符时的匹配错误。

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