Python调用Julia传递大型数据结构时如何降低开销？

在数值模拟场景中，Python调用Julia处理含百万级元素的字典时，出现Julia内部计算仅0.6s，但整体耗时达6s的情况，核心瓶颈在于Python与Julia之间的字典序列化/反序列化开销。以下是针对该问题的低开销解决方案，基于你的代码示例优化。

核心问题根源 #

默认情况下，PyJulia会将Python字典完整序列化为Julia的Dict类型，处理完成后再将Julia的Dict反序列化为Python字典。对于百万级元素，这种逐元素的类型转换和数据拷贝是耗时的主要原因。

解决方案：拆分键值对为数组传递 #

利用数组的零拷贝传递特性（通过PyArray直接访问Python端数据），将字典拆分为键数组和值数组传递给Julia，处理后返回结果数组，最后在Python端重建字典。这种方式避免了整个字典的序列化开销。

优化后的代码示例 #

Python端（main.py）

from time import time
import numpy as np
import julia

# 初始化Julia环境
jl = julia.Julia(compiled_modules=False)
from julia import Main
Main.include("main.jl")

# 生成大型字典
n = 1_000_000
d = {i: str(i) for i in range(n)}

# 将字典拆分为键、值数组（numpy数组支持零拷贝传递）
keys = np.array(list(d.keys()), dtype=np.int64)
values = np.array(list(d.values()), dtype=np.str_)

# 调用优化后的Julia函数
t1 = time()
res_keys, res_values = Main.func_opt(keys, values)
# 在Python端重建字典
res = dict(zip(res_keys, res_values))
t2 = time()

print(f"Elapsed overall :: {t2-t1} s")

Julia端（main.jl）

using PyCall

function func_opt(keys::PyArray{Int64}, values::PyArray{String})
    t = @elapsed begin
        # 预分配结果数组容量，避免动态扩容开销
        d2_keys = Int64[]
        d2_vals = String[]
        sizehint!(d2_keys, length(keys) ÷ 2)
        sizehint!(d2_vals, length(keys) ÷ 2)
        
        # 直接遍历PyArray（零拷贝访问Python端数据）
        for (k, v) in zip(keys, values)
            if iseven(k)
                push!(d2_keys, k)
                push!(d2_vals, "0")
            end
        end
    end
    println("In Julia body elapsed:  ", t)
    
    # 返回结果数组，避免字典序列化
    return d2_keys, d2_vals
end

额外优化建议 #

• 预编译Julia函数：将函数放入模块并预编译，减少首次调用的启动开销。例如在Julia模块中添加@precompile_setup。
• 避免不必要的数据转换：如果Python端可以直接使用数组而非字典，可省去最后重建字典的步骤。
• 使用更高效的序列化库：若仍需传递字典，可尝试使用MsgPack或Arrow等高性能序列化库替代默认的序列化机制，但数组拆分方法通常效果更显著。

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