Spark Connect Protobuf 定义协议文件与 Python 桩代码生成全指南【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址: https://gitcode.com/gh_mirrors/sp/sparkSpark Connect 是 Apache Spark 的客户端—服务器client-server架构方案它把 Spark 的核心执行能力封装为远程服务客户端通过标准协议与服务器通信。而这一协议的定义就落在sql/connect/common/src/main/protobuf/目录下的一组.proto文件中。本篇指南以仓库中的 README.md 为骨架系统讲解 Spark Connect 协议的定义方式、.proto文件的构成以及修改协议后如何通过 Docker 镜像或本地 Python 环境重新生成 PySpark 所需的 Python 桩代码stubs帮助你在开发和提交 Spark Connect 相关改动时快速掌握协议修改的完整闭环。阅读完本文你将能够理解sql/connect/common/src/main/protobuf目录下各.proto文件的职责划分与底层生成链路掌握两种官方推荐的代码生成方法Docker 镜像法与本地buf法并理解其原理差异了解生成产物在python/pyspark/sql/connect/proto/下的组织方式与代码风格校验规则。一、目录概览Spark Connect 协议从何而来sql/connect/common/src/main/protobuf/目录下存放了定义 Spark Connect 协议的全部.proto文件它们统一声明了package spark.connect与syntax proto3组织在 spark/connect/ 子目录中文件职责base.proto协议核心Plan、UserContext、各类请求/响应消息如AnalyzePlanRequest、ExecutePlanRequest、FetchErrorDetailsRequestcatalog.protoCatalog 相关操作CurrentDatabase、TableExists、GetTables等commands.proto服务端命令定义如SqlCommand、WriteOperation等common.proto通用类型如StorageLevel、DataType引用、Expression引用expressions.proto表达式Expression的定义relations.proto关系Relation的定义覆盖 DataFrame 逻辑计划types.protoSpark 数据类型DataType、StructType等ml.protoML 相关 API 的消息定义ml_common.protoML 通用类型与参数处理pipelines.protoDeclarative Pipelines 相关定义example_plugins.proto插件示例定义这些.proto文件相互 import形成完整的协议模型。例如 base.proto 顶部的 import 语句就引入了commands.proto、common.proto、expressions.proto、relations.proto、types.proto、ml.proto与pipelines.proto。从源码结构看这一目录同时是 JVM 侧服务端与 Python 侧客户端协议定义的中枢JVM 侧通过option java_package org.apache.spark.connect.proto与option java_multiple_files true生成 Java 类Python 侧则通过buf的 protoc 插件生成*_pb2.py与*_pb2.pyi文件供 PySpark 客户端使用。二、修改 .proto 之后为什么必须重新生成桩代码协议文件是客户端与服务端之间的“契约”。当你在sql/connect/common/src/main/protobuf/下修改了任何.proto文件比如给消息新增字段、为Plan增加新的操作类型必须重新生成python/pyspark/sql/connect/proto/下的 Python 桩代码否则Python 客户端无法感知新协议PySpark 客户端通过python/pyspark/sql/connect/proto/下的*_pb2.py进行消息序列化/反序列化。协议已改而桩代码未更新客户端将无法构造或解析新消息。类型提示缺失*_pb2.pyi文件为 IDE 与mypy提供类型信息不更新会导致类型检查与自动补全失真。CI 校验失败从仓库配置看pyproject.toml中明确将python/pyspark/sql/connect/proto/*排除在ruff的 lint 范围之外原因是这些文件是自动生成的与python/pyspark/sql/streaming/proto/*同理。任何手工改动都会被后续重新生成覆盖因此必须走正式生成流程。官方在 README.md 中给出了两种生成方法Docker 镜像法推荐与本地 Python 环境法。下文逐一展开。三、Method 1Docker 镜像法推荐免安装、可复现Docker 法无需在本地安装任何代码生成工具构建出的镜像自带固定版本的工具链生成结果可复现适合 CI 或多人协作场景。3.1 构建镜像在 Spark 仓库根目录下执行docker build -t connect-cg \ --build-context root. \ dev/spark-test-image/connect-gen-protos/说明-t connect-cg为镜像命名--build-context root.把Spark 仓库根目录作为名为root的构建上下文这样 Dockerfile 可以COPY --fromroot pyproject.toml ./pyproject.toml获取依赖定义构建目标目录是dev/spark-test-image/connect-gen-protos/其中只有一份 Dockerfile。打开这份 Dockerfile 可以看到镜像的具体构成它基于ubuntu:nobleUbuntu 24.04安装python3.12与python3.12-venv从 GitHub Releases 下载固定版本buf当前 Dockerfile 中ARG BUF_VERSION1.66.1并安装到/usr/local/bin/buf创建虚拟环境/opt/spark-venv并设置为默认PATH从root上下文拷贝pyproject.toml随后执行python3.12 -m pip install --group ci_lint安装 CI lint 依赖组——该组见 pyproject.toml 中的ci_lint组包含internal_lint其中锁定了mypy1.19.1、mypy-protobuf3.3.0、ruff0.14.8等版本与internal_ci_connect包含protobuf6.33.5等以/spark为工作目录ENTRYPOINT指向/spark/dev/connect-gen-protos.sh。因此镜像内工具版本与pyproject.toml中的锁定版本保持一致这正是“可复现环境”的来源。3.2 运行镜像生成代码在 Spark 仓库根目录执行docker run --cpus 1 -it --rm -v $(pwd):/spark connect-cg参数含义--cpus 1限制容器使用 1 个 CPU控制构建期资源占用-it分配交互式终端便于观察生成日志--rm容器退出后自动删除-v $(pwd):/spark将当前 Spark 仓库挂载到容器内的/sparkconnect-cg上一步构建的镜像名。容器启动后ENTRYPOINT会在/spark下执行 dev/connect-gen-protos.sh该脚本内部进一步调用./dev/gen-protos.sh connect $即 dev/gen-protos.sh最终把生成的*_pb2.py、*_pb2.pyi等文件写入挂载目录python/pyspark/sql/connect/proto/。由于仓库是以 bind mount 方式挂载容器内写出的文件会直接落盘到本地 checkout退出容器后改动依然保留。四、Method 2本地 Python 环境法无需 Docker不想用 Docker 时可以在本地准备工具链后直接运行生成脚本。4.1 前置条件需要两个基础工具**buf与 buf.gen.yaml 控制行为Python 3.12生成产物依赖较新的 Python 生态。4.2 安装 Python 依赖先查看 pyproject.toml 中锁定的最新版本——具体而言internal_lint组中定义了ruff0.14.8 mypy1.19.1 mypy-protobuf3.3.0README 给出的安装命令以当前pyproject.toml为准pip install mypyversion mypy-protobufversion ruffversion例如pip install mypy1.19.1 mypy-protobuf3.3.0 ruff0.14.8提示版本号务必与pyproject.toml保持一致。mypy-protobuf负责生成*_pb2.pyi类型桩文件ruff负责对生成代码做统一格式化mypy用于类型检查。若版本漂移生成的代码格式或类型注解可能与仓库 CI 期望不符。4.3 执行生成脚本在 Spark 仓库根目录运行./dev/connect-gen-protos.sh脚本将生成的 Python 文件写入python/pyspark/sql/connect/proto/。也可以传入自定义输出目录./dev/connect-gen-protos.sh /tmp/my-proto-output五、源码视角生成链路内部做了什么理解了两种方法后再深入源码看这条生成链路能更好地把握“改协议 → 重新生成”的内在逻辑。5.1 入口脚本与参数dev/connect-gen-protos.sh 是一个极薄的封装if [[ $# -gt 1 ]]; then echo Illegal number of parameters. echo Usage: ./dev/connect-gen-protos.sh [path] exit -1 fi ./dev/gen-protos.sh connect $它只接收 0 或 1 个参数可选的自定义输出路径其余工作全部委托给dev/gen-protos.sh connect。而dev/gen-protos.sh是一个同时服务connect与streaming两个模块的通用生成脚本if [[ $1 connect ]]; then MODULEconnect OUTPUT_PATH${SPARK_HOME}/python/pyspark/sql/connect/proto/ SOURCE_MODULEspark.connect TARGET_MODULEpyspark.sql.connect.proto elif [[ $1 streaming ]]; then MODULEstreaming OUTPUT_PATH${SPARK_HOME}/python/pyspark/sql/streaming/proto/ SOURCE_MODULEorg.apache.spark.sql.execution.streaming TARGET_MODULEpyspark.sql.streaming.proto ...connect 模式与 streaming 模式共用一套处理逻辑只是源.proto目录、输出目录和模块命名不同。connect 模式进入sql/connect/common/src/main目录执行buf generatestreaming 模式则进入sql/core/src/main。这也是为什么两个模块生成的 Python 产物结构完全一致。5.2 插件配置buf.gen.yamlbuf generate的行为由sql/connect/common/src/main/buf.gen.yaml控制。它一次生成多种语言的产物version: v1 plugins: - plugin: buf.build/protocolbuffers/cpp:v33.5 out: gen/proto/cpp - plugin: buf.build/protocolbuffers/csharp:v33.5 out: gen/proto/csharp - plugin: buf.build/protocolbuffers/java:v33.5 out: gen/proto/java - plugin: buf.build/grpc/ruby:v1.76.0 out: gen/proto/ruby - plugin: buf.build/protocolbuffers/ruby:v33.5 out: gen/proto/ruby - plugin: buf.build/protocolbuffers/python:v33.5 out: gen/proto/python - plugin: buf.build/grpc/python:v1.76.0 out: gen/proto/python - name: mypy out: gen/proto/python可以看到C、C#、Java、Ruby、Python 与 gRPC、mypy 桩都在同一轮buf generate中产出。虽然脚本最终只把gen/proto/python下的 Python 文件拷贝到目标模块但其他语言产物也已就绪可供服务端JVM等端侧在构建时使用。此外buf.yaml中启用了默认 lint 规则与FILE级别的 breaking change 检查用于在生成时把住协议兼容性。5.3 生成后的后处理生成并非简单拷贝dev/gen-protos.sh 还会做四类后处理以适配 PySpark 的实际包布局导入路径改写把生成的 Python 文件中的from spark.connect import ...改写为from pyspark.sql.connect.proto import ..._pb2.py与_pb2_grpc.py并把序列化描述符中的模块名从spark.connect改为pyspark.sql.connect.proto.pyi类型桩修正对*.pyi文件替换 import 与模块引用并删除typing_extensions.final装饰器行添加 Apache License 头为每个生成文件头部追加 ASF License 注释删除空 gRPC 文件如果*_grpc.py内容为空行数仅剩 20 行即 License 头等模板行则直接删除避免产生无意义的空模块。随后脚本用ruff format --config $SPARK_HOME/pyproject.toml对生成代码统一格式化最后将gen/proto/python下所有.py*文件拷贝到OUTPUT_PATH并清理临时目录。5.4 生成产物落点connect 模式的默认输出目录为python/pyspark/sql/connect/proto/。当前仓库中该目录下的产物与.proto文件一一对应例如base_pb2.py/base_pb2.pyi← 对应 base.protorelations_pb2.py/relations_pb2.pyi← 对应 relations.protoexpressions_pb2.py/expressions_pb2.pyi← 对应 expressions.protocommands_pb2.py、catalog_pb2.py、types_pb2.py、common_pb2.py、ml_pb2.py、ml_common_pb2.py、pipelines_pb2.py、example_plugins_pb2.py等均配有对应的.pyi类型桩这些文件属于自动生成产物不应手工编辑对协议的任何修改都应回到.proto源文件再通过上述两种方法之一重新生成。六、典型工作流与注意事项综合以上内容一次完整的“协议修改 → 代码重新生成”工作流如下修改sql/connect/common/src/main/protobuf/spark/connect/下的一个或多个.proto文件新增字段、消息或服务方法在仓库根目录执行docker build -t connect-cg --build-context root. dev/spark-test-image/connect-gen-protos/构建工具镜像首次或工具版本变化时执行docker run --cpus 1 -it --rm -v $(pwd):/spark connect-cg完成生成或改用本地法安装buf与mypy/mypy-protobuf/ruff版本以 pyproject.toml 为准再运行./dev/connect-gen-protos.sh检查git status确认python/pyspark/sql/connect/proto/下的*_pb2.py、*_pb2.pyi等文件已同步更新若只需临时验证可传自定义目录例如./dev/connect-gen-protos.sh /tmp/my-proto-output避免污染工作区。注意事项汇总版本锁定本地法务必以 pyproject.toml 中internal_lint组锁定的mypy、mypy-protobuf、ruff版本为准Docker 法由 Dockerfile 自动保证版本一致。Python 版本Docker 镜像内置 Python 3.12本地环境要求 Python 3.12。生成是全覆盖式的dev/gen-protos.sh 会先删除旧的gen临时目录再重新生成并整体拷贝到输出目录输出目录中的旧文件会被新文件覆盖因此不要手工混入改动。协议兼容性buf.yaml 启用了FILE级别的 breaking change 检查修改消息字段时需遵循 protobuf 的兼容性规则例如不要复用已废弃的字段号避免破坏既有客户端。七、结语Spark Connect 的协议定义集中保存在sql/connect/common/src/main/protobuf/中它是客户端与服务端通信契约的唯一事实来源。修改协议后通过 README.md 中介绍的 Docker 镜像法或本地buf法重新生成 Python 桩代码即可让 PySpark 客户端与服务端保持同步。理解背后的生成链路connect-gen-protos.sh→gen-protos.sh→buf generate→ 后处理与拷贝能让你在遇到生成异常或提交协议变更时快速定位问题也便于在 CI 中对齐工具版本、维持生成产物的可复现性。【免费下载链接】sparkApache Spark - A unified analytics engine for large-scale data processing项目地址: https://gitcode.com/gh_mirrors/sp/spark创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考