Windows下Databricks Connect中PySpark UDF无法加载PyNaCl模块

verify_signature_udf = udf(lambda x: verify_signature(x, public_keys), BooleanType())

data_signed = data.withColumn(
"is_signature_valid", verify_signature_udf("state_values")
)

PyNaCl已在本地通过Databricks Connect安装，但执行器上未安装，报错：

File "/databricks/spark/python/pyspark/cloudpickle/cloudpickle.py", line 679, in subimport
import(name)
ModuleNotFoundError: No module named 'nacl'

按照Spark官方Python打包文档尝试配置虚拟环境：
```python
import os
os.environ['PYSPARK_PYTHON'] = "./environment/bin/python"
spark = SparkSession.builder.config(
    "spark.archives",
    "pyspark_venv.tar.gz#environment").getOrCreate()

问题依旧。将nacl包从tar.gz提取为zip文件，通过spark.sparkContext.addPyFile(path="nacl.zip")加载，又出现新错误：

File "/local_disk0/spark-xxx8db3a-5436-4ce8-8ff5-19eaeb4397b4/executor-xxxb7a74-4e1b-40bf-aae2-fc3553155f91/spark-xxx70cb9-482d-42a9-901a-c36f66a42a19/isolatedSparkFiles/0e10cb02-db69-4d63-b7ea-6c2b415fb5d9/nacl.zip/nacl/bindings/crypto_aead.py", line 17, in <module>
    from nacl._sodium import ffi, lib
ModuleNotFoundError: No module named 'nacl._sodium'

注：zip文件包含以下sodium组件，tar.gz无额外相关内容：

./nacl/bindings/sodium_core.py
./nacl/bindings/__pycache__/sodium_core.cpython-39.pyc
./nacl/_sodium.pyd

尝试将导入语句移到for循环中，本地Databricks Connect无报错，但执行器仍报错：

def verify_signature(msg, keys):
    for key in keys:
       # 避免本地Databricks Connect报错
       from nacl import bindings as c  
       c.crypto_sign_ed25519ph_update(...)
       ...

verify_signature_udf = udf(lambda x: verify_signature(x, public_keys), BooleanType())

data_signed = data.withColumn(
    "is_signature_valid", verify_signature_udf(my_keys)("state_values")
)

询问可行的解决办法，包括dbx是否适用，以及非UDF的实现方式。

解决办法 #

1. 正确构建Linux环境下的虚拟环境归档 #

PyNaCl包含平台相关的C扩展（Windows下是.pyd，Linux下是.so），Windows打包的依赖无法在Linux执行器上运行，必须在Linux环境构建虚拟环境：

• 执行以下命令构建并打包虚拟环境：

  python -m venv pyspark_venv
  source pyspark_venv/bin/activate
  pip install pynacl
  deactivate
  tar -czf pyspark_venv.tar.gz pyspark_venv
  

• 将pyspark_venv.tar.gz上传到DBFS或Databricks工作区
• 在Databricks Connect代码中配置：

  import os
  os.environ['PYSPARK_PYTHON'] = "./environment/bin/python"
  spark = SparkSession.builder.config(
      "spark.archives", "dbfs:/path/to/pyspark_venv.tar.gz#environment"
  ).getOrCreate()
  

2. 直接在Databricks集群安装PyNaCl #

如果使用固定集群，通过UI直接安装库：

• 进入集群页面 → 库 → 安装新库
• 选择PyPI源，输入pynacl完成安装
• 重启集群后，Databricks Connect连接该集群时，执行器会自动加载PyNaCl依赖，无需手动打包

3. 使用dbx管理依赖（可行） #

dbx可简化Databricks项目的依赖部署：

• 初始化dbx项目，在requirements.txt中添加pynacl
• 配置deploy.json指定依赖安装方式（虚拟环境或直接PyPI安装）
• 运行dbx deploy部署项目，dbx会自动处理执行器上的依赖安装
• 运行UDF时依赖会被正确加载，无需手动处理归档或addPyFile

4. 非UDF替代方案 #

Pandas UDF

利用Pandas批量处理数据，同时确保依赖正确加载（需先解决依赖问题）：

from pyspark.sql.functions import pandas_udf
import pandas as pd
from nacl import bindings as c

@pandas_udf('boolean')
def verify_signature_pandas_udf(msg_series: pd.Series) -> pd.Series:
    def verify_single(msg):
        result = False
        for key in public_keys:
            # 签名验证逻辑
            c.crypto_sign_ed25519ph_update(...)
            # 根据验证结果更新result
        return result
    return msg_series.apply(verify_single)

data_signed = data.withColumn("is_signature_valid", verify_signature_pandas_udf("state_values"))

注意事项

Spark SQL原生函数无法直接实现PyNaCl的加密验证逻辑，仅当验证逻辑可简化为SQL操作时才适用，此场景下不推荐。

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