PySpark MCP

PySpark MCP 服务

SQL 迁移辅助、AWS Glue 作业模板生成和 Spark 代码优化，均以 MCP 服务的形式提供。

这不是实时 Spark 的 pyspark-mcp 软件包。 本项目用于生成 SQL → PySpark / Glue 源代码，以 pyspark-tools 发布。 SemyonSinchenko/pyspark-mcp 会检查正在运行的 SparkSession。此项目仍保留已弃用的 pyspark-mcp 控制台脚本，以便旧配置继续工作；它会打印一条警告，然后启动此服务。

功能

• SQL 方言转译 — 使用 SQLGlot 在 PostgreSQL、Oracle、Redshift、MySQL、Snowflake 和 Spark SQL 之间进行转换
• PySpark DataFrame API 生成 — 根据 SQL 生成 DataFrame API 源文本，并附带优化提示
• AWS Glue 模板 — 作业脚本字符串、DynamicFrame 转换、Data Catalog 定义、S3 布局建议
• 批处理 — 遍历 SQL 文件/目录并输出转换后的模块
• 代码审查与优化 — 基于模式审查现有 PySpark 源代码
• 模式检测 — 查找重复片段并建议实用工具

不支持的功能

• 递归 CTE → 提供等效的 Spark SQL 和指导（PySpark 不原生支持递归 CTE）
• MERGE/PIVOT/CONNECT BY → 转译为 Spark SQL，并提供 DataFrame API 指导
• 所有 SQL 都能完美地 1:1 转译为 DataFrame API — 复杂查询会获得 Spark SQL 和建议
• 它不会启动 SparkSession、提交 Glue 作业或执行 SQL
• optimize(mode="code") 返回的是建议；它不会重写你的代码
• glue_s3 是一种路径启发式判断（不会调用 AWS，也不会测量性能提升）
• 它不能替代 SemyonSinchenko/pyspark-mcp 来查看实时目录/执行计划

为什么不直接调用 sqlglot

SQLGlot 已经可以转译方言。此 MCP 在其核心功能之外增加了三项能力：通过转换测试锁定的 DataFrame API 格式化输出，以及连接、窗口和类型转换映射；Glue 作业样板字符串（书签、DynamicFrame、目录表），让 Agent 可以直接生成文件，而不必自行拼装；还有一个包含 14 个工具的 FastMCP 接口，让大模型选择 convert / mode=sql，而不是自行配置 sqlglot。如果你只需要 sqlglot.transpile(...)，请使用 sqlglot。

快速开始

pip install pyspark-tools
pyspark-tools

无需克隆的替代方式：uvx pyspark-tools。run_server.py 是一个开发便捷脚本，会插入 sys.path 并打印启动横幅。在配置和生产环境中，优先使用 pyspark-tools。

试用

pip install pyspark-tools
python -c "from pathlib import Path; from pyspark_tools.sql_converter import SQLToPySparkConverter as C; from pyspark_tools.consolidated_tools import glue_job; c,s,o=C(),Path('examples'),Path('examples/out'); [(o/f'{n}.py').write_text(c.convert_sql_to_pyspark((s/f'{n}.sql').read_text(), dialect=d).pyspark_code) for n,d in [('postgres_orders','postgres'),('oracle_decode','oracle')]]; (o/'orders_etl_glue.py').write_text(glue_job(mode='template', job_name='orders_etl', sql_query=(s/'postgres_orders.sql').read_text())['template'])"

会写入与 examples/out/ 相同的文件。MCP stdio CLI：pyspark-tools。

示例：SQL → PySpark

SELECT o.customer_id, c.name, SUM(o.amount) AS total
FROM orders o
JOIN customers c ON o.customer_id = c.id
WHERE o.status = 'paid'
GROUP BY o.customer_id, c.name

调用 convert 并设置 mode=sql。捕获的转换器输出（dialect=spark）：

from pyspark.sql import SparkSession
from pyspark.sql.functions import (
    col, lit, when, count, sum as spark_sum, avg, min, max, countDistinct,
    coalesce, concat, datediff, date_add, to_date,
    row_number, rank, lag, lead,
)
from pyspark.sql.window import Window

# Generated from SPARK SQL
spark = SparkSession.builder.appName('SQLToPySpark').getOrCreate()

# Load table: customers
customers_df = spark.table('customers')
# Load table: orders
orders_df = spark.table('orders')

# Main query
result_df = (orders_df.alias('o')
    .join(customers_df.alias('c'), (col('o.customer_id') == col('c.id')), 'inner')
    .filter((col('o.status') == lit('paid')))
    .groupBy(col('o.customer_id'), col('c.name'))
    .select(col('o.customer_id'), col('c.name'), (spark_sum(col('o.amount'))).alias('total')))

实际输出取决于方言检测和回退逻辑；tests/test_sql_conversion_fixes.py 中的转换测试锁定了重要结构。转换器可通过 style="notebook" 启用 Notebook 风格的 import * / show()。

MCP 配置

Claude Desktop

macOS：~/Library/Application Support/Claude/claude_desktop_config.json

Linux：~/.config/Claude/claude_desktop_config.json

{
  "mcpServers": {
    "pyspark": {
      "command": "pyspark-tools",
      "args": []
    }
  }
}

Hermes Agent

添加到 ~/.hermes/config.yaml：

mcp:
  servers:
    pyspark:
      command: pyspark-tools
      enabled_tools: all

Docker

此镜像仅支持 stdio（FastMCP 通过 stdin/stdout 通信）。它不会在端口 8000 上启动 HTTP 服务。docker compose up 用于本地测试，不是可进行健康检查的 Web 服务。

docker compose --profile test run --rm pyspark-tools-test

工具

三个主要工具。其他十一个路由器在此次小版本中仍会注册，但已弃用——优先使用 convert、glue_job 和 review。

convert — SQL → PySpark（包括 mode=batch_dir）

convert(mode="sql", sql_query="SELECT id FROM users", dialect="postgres")
convert(mode="batch_dir", directory_path="etl/", output_dir="out")

glue_job — Glue 5.0 作业模板字符串

glue_job(mode="template", job_name="orders_etl", sql_query="SELECT * FROM orders")

review — 代码审查、模式、重复项

review(mode="code", code="df = spark.table('t')\ndf.collect()")

旧版 / 已弃用： analyze、optimize、glue_schema、glue_s3、 glue_data、refactor、search、context、batch_status、s3_source、 analytics。仍可调用；不要向新的 Agent 宣传。

安全

此 MCP 可以读取本地文件（SQL、TXT、PDF）；如果安装了 [aws] 扩展，还可以使用主机的默认 AWS 凭证列出 / 读取 S3。文件工具仅允许访问进程工作目录下的路径（或显式指定的 base_path / FileHandler(base_directory=...)）。这并非沙箱。

请在受限的操作系统账户下运行服务。不要将其指向存放机密的目录。除非你打算通过 s3_source / glue_s3 读取 S3，否则不要附加具有写入权限的 AWS 凭证。可选扩展：

pip install "pyspark-tools[aws]"    # boto3 for S3/Glue catalog helpers
pip install "pyspark-tools[spark]"  # pyspark — not required at runtime; generated code only

开发

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

# Test
pytest tests/ -v --cov=pyspark_tools

# Format
black pyspark_tools tests
isort pyspark_tools tests

# Lint
flake8 pyspark_tools tests

需要 Python 3.11+（与 CI 矩阵一致）。

架构

pyspark_tools/
├── server.py              # FastMCP server + helper implementations
├── consolidated_tools.py  # 14 @app.tool() routers
├── sql_converter.py       # SQLGlot-based transpilation + DataFrame API generation
├── aws_glue_integration.py # Glue job templates, DynamicFrame, Data Catalog
├── advanced_optimizer.py  # Performance analysis + optimization suggestions
├── batch_processor.py     # Concurrent file processing
├── code_reviewer.py       # PySpark code review patterns
├── duplicate_detector.py  # Code deduplication
├── data_source_analyzer.py # Data source analysis (optional boto3)
└── file_utils.py          # File I/O with allow-root checks

许可证

MIT — 详见 LICENSE。

---

mcp-name: io.github.AnnasMazhar/pyspark-mcp