Azure Databricks中如何让multiprocessing map_async返回含日志的字典？

核心思路是让每个子进程的任务函数直接返回包含所有所需信息的字典，利用map_async自动收集所有子进程的返回结果，完全规避全局变量共享的问题——这是Databricks multiprocessing环境下最可靠的方案。

修改后的代码示例 #

import multiprocessing
import time
import os
from datetime import datetime
from neo4j import GraphDatabase

def singleq_test(query):
    # 每个子进程独立初始化Neo4j驱动，避免跨进程资源冲突
    driver = GraphDatabase.driver("bolt://your-neo4j-host:7687", auth=("neo4j-user", "neo4j-pass"))
    
    # 收集当前任务的元数据
    pid = os.getpid()
    task_start = datetime.now()
    runtimes = []
    last_query_result = None

    # 执行查询并统计运行时（这里按5次循环计算平均，可按需调整）
    for _ in range(5):
        start = time.perf_counter()
        with driver.session() as session:
            last_query_result = session.run(query).data()
        end = time.perf_counter()
        runtimes.append(end - start)
    
    avg_runtime = sum(runtimes) / len(runtimes)
    task_end = datetime.now()

    driver.close()

    # 返回完整的结果字典
    return {
        "pid": pid,
        "task_start_time": task_start.strftime("%Y-%m-%d %H:%M:%S.%f"),
        "task_end_time": task_end.strftime("%Y-%m-%d %H:%M:%S.%f"),
        "query": query,
        "query_response": last_query_result,
        "average_runtime_seconds": round(avg_runtime, 6),
        "all_runtimes": runtimes
    }

def qlist_test(query_list):
    # 根据Databricks集群CPU核数调整进程池大小，避免资源过载
    with multiprocessing.Pool(processes=multiprocessing.cpu_count()//2) as pool:
        async_result = pool.map_async(singleq_test, query_list)
        # 阻塞等待所有任务完成，获取结果列表
        full_results = async_result.get()
    
    return full_results

# 测试调用（Databricks中可直接运行，无需__main__判断也能执行，但加上更规范）
if __name__ == "__main__":
    sample_queries = [
        "MATCH (u:User) RETURN u.name LIMIT 10",
        "MATCH (p:Product)-[:SOLD_TO]->(u:User) RETURN count(*) AS sales_count"
    ]
    results = qlist_test(sample_queries)
    # 可将结果写入Databricks Delta表或打印输出
    for res in results:
        print(f"PID {res['pid']} 平均耗时: {res['average_runtime_seconds']}s")

关键说明 #

1. 规避全局变量问题：每个子进程的任务数据（PID、时间、查询结果）都在进程内生成并打包返回，map_async会自动将所有子进程的字典结果收集成一个列表，无需跨进程共享状态。
2. 独立Neo4j驱动：Databricks的进程隔离环境下，跨进程共享Neo4j驱动会导致连接泄漏或异常，每个子进程单独创建/关闭驱动是最安全的做法。
3. 灵活调整参数：
   • 可以修改查询循环次数来调整平均运行时的计算精度；
   • 如果查询结果过大，可只返回结果行数或关键统计值，减少内存占用；
   • 进程池大小建议设为集群CPU核数的一半，避免抢占Spark任务的资源。

注意事项 #

• 确保Neo4j的连接信息（地址、账号）在子进程中可访问，推荐通过Databricks Secrets管理敏感信息，避免硬编码；
• 如果需要将结果持久化，可直接把full_results列表转换为DataFrame，写入Delta Lake或其他存储。

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