Spark SQL Optimizer Agent

Spark SQL Optimizer Agent 是一个以 LLM Agent 为主导的 Spark SQL 优化助手。它读取 SQL 和可选运行时证据，调用工具完成证据收集与校验，再由 LLM 综合理解 SQL、定位问题、生成候选 SQL，并输出最终诊断报告。

默认运行 Session 模式：LLM 负责编排、综合分析、候选 SQL 生成和报告撰写；脚本工具负责静态分析、runtime 归一化、skew 检测、热点验证和 SQL 质量门禁。Workflow 模式仅作为离线兜底。

功能特性

• 分析 Spark SQL 结构，包括 join、filter、aggregation、ORDER BY、COUNT DISTINCT、SELECT * 等常见风险模式。
• 归一化 Spark stage metrics、runtime evidence、event log 和 History Server 数据。
• 诊断 shuffle、spill、scan、write-path、数据倾斜等运行时信号。
• 在证据足够时生成热点 key 验证 SQL 或保守改写候选。
• 让 LLM 综合 SQL、工具输出和运行时证据，生成可审查的候选 SQL。
• 对候选 SQL 运行质量门禁，识别 correctness 和质量风险。
• 由 LLM 生成最终中文诊断报告，并保留 deterministic report 作为兜底。
• 默认使用 LLM 主导的 Session 模式；Workflow 模式仅用于无模型配置、CI 或本地兜底。

快速开始

运行测试，确认环境可用：

PYTHONDONTWRITEBYTECODE=1 python3 -m pytest -p no:cacheprovider

创建配置文件：

cp assets/input-template.yaml config.yaml

示例配置：

task_name: "my-sql-optimization"
sql_file_path: "tests/fixtures/sql_join.sql"
stage_metrics_file: "tests/fixtures/stage_metrics_skew.json"
platform: "generic"
optimization_goal: "reduce_runtime"
rewrite_mode: "conservative"
agent_mode: "session"
planner_mode: "llm"
data_access:
  enabled: false
output_dir: "artifacts/my-sql-optimization"

运行：

python3 scripts/orchestrate.py --config config.yaml

查看报告：

artifacts/my-sql-optimization/final_report.md

LLM 配置

默认的 Session 模式需要配置 LLM planner。

创建 .env：

cp .env.example .env

配置 OpenAI-compatible provider：

SPARK_OPTIMIZER_LLM_PROVIDER=openai
OPENAI_API_KEY=sk-your-key
SPARK_OPTIMIZER_MODEL=gpt-4o

也可以配置命令式模型网关：

SPARK_OPTIMIZER_MODEL_COMMAND="/path/to/model-command"

检查配置：

python3 scripts/agent_config.py --check

如果没有 LLM 配置，可以显式运行 Workflow 兜底模式：

python3 scripts/orchestrate.py --config config.yaml --agent-mode workflow

输入配置

常用字段：

字段	说明
sql_text / sql_file_path	Spark SQL 文本或 SQL 文件路径
explain_text / explain_file_path	可选 Spark EXPLAIN 计划
stage_metrics_file	Spark stage/task metrics JSON
runtime_evidence_file	已整理好的 runtime evidence JSON
event_log_path	Spark event log JSONL
history_server_url + app_id	Spark History Server 数据源
hot_values_file	已知热点值，用于生成热点改写候选
hotspot_table	生成热点 key 验证 SQL 时使用的源表
optimization_goal	优化目标，例如 reduce_runtime 或 fix_skew
rewrite_mode	report-only、conservative 或 assisted
output_dir	产物输出目录

支持的优化目标：

• reduce_runtime
• reduce_cost
• improve_stability
• fix_skew
• reduce_shuffle
• reduce_output_files

运行模式

Session 模式

Session 模式使用 LLM planner，从系统给出的合法候选动作中选择下一步。

python3 scripts/orchestrate.py --config config.yaml --agent-mode session

每一步 observe / plan / act / verify 都会记录到 agent_trace.json。

Workflow 模式

Workflow 模式按固定流程运行，不调用 LLM，只作为兜底路径。

python3 scripts/orchestrate.py --config config.yaml --agent-mode workflow

这个模式适合 CI、本地测试和无模型配置的离线分析。

输出产物

运行完成后，output_dir 中会生成：

文件	说明
final_report.md	面向人的 Spark SQL 诊断报告；Session 模式下优先使用 LLM 生成
llm_final_report.md	LLM report writer 生成的最终报告
deterministic_report.md	确定性报告兜底版本，Session 模式有 LLM 报告时保留
llm_report_response.json	LLM report writer 的结构化响应
optimization_plan.json	结构化优化计划，包含 diagnostic schema
routing_plan.json	Agent 路由结果和证据边界
agent_results/*.json	每个 Agent 的原始结果
agent_session.json	Session 模式摘要
agent_trace.json	Session 模式下的 LLM planner trace
runtime_metadata.json	运行时证据覆盖摘要
candidate_*.sql / hotspot_refactored.sql	生成的候选 SQL
llm_optimized_candidate_*.sql	LLM 生成的候选优化 SQL
llm_optimizer_response.json	LLM 对 SQL、问题和候选优化的结构化分析
candidate_rewrite_diff_*.md	原 SQL 和候选 SQL 的 diff

建议阅读顺序：

1. final_report.md
2. optimization_plan.json
3. routing_plan.json
4. agent_trace.json
5. agent_results/*.json

Agent 流程

整体流程：

evidence collection
-> runtime evidence normalization
-> SQL structure analysis
-> runtime diagnosis
-> skew diagnosis
-> hotspot rewrite candidate generation
-> LLM SQL optimization
-> SQL quality gate
-> LLM report generation
-> result merge and fallback rendering

主要 Agent：

Agent	职责
evidence-collector-agent	校验并物化输入证据
runtime-evidence-collector-agent	归一化运行时证据
sql-structure-agent	分析 SQL 和 EXPLAIN 结构
runtime-diagnosis-agent	诊断运行时瓶颈信号
skew-agent	检测数据倾斜并推断候选热点 key
hotspot-rewrite-agent	生成热点 key 验证 SQL 或受控改写候选
llm-sql-optimizer-agent	综合 SQL、工具发现和运行时证据，生成 LLM 优化建议和候选 SQL
sql-quality-gate-agent	检查 SQL correctness 和质量风险
report-writer-agent	基于完整诊断计划和 AgentResult，用 LLM 生成最终中文诊断报告

项目结构

spark-sql-optimizer-agent/
├── agents/                  # Agent spec 和 registry
├── assets/                  # 配置模板
├── references/              # 契约、证据 schema、路由策略和规则
├── spark_optimizer/          # Python 主实现包
│   ├── cli/                  # CLI 入口
│   ├── core/                 # 契约、路由、workflow、证据 helper
│   ├── agents/               # Agent Python 实现
│   ├── planners/             # LLM planner 和 action policy
│   ├── tools/                # 专家工具适配层
│   ├── reporting/            # 结果合并和报告渲染
│   └── llm/                  # 模型网关
├── scripts/                 # 兼容旧命令路径的 wrapper
├── src/                     # SQL、skew、hotspot、quality 专家脚本
└── tests/                   # 回归测试和 fixtures

开发

运行测试：

PYTHONDONTWRITEBYTECODE=1 python3 -m pytest -p no:cacheprovider

常用入口：

python3 scripts/orchestrate.py --config config.yaml
python3 scripts/orchestrate.py --config config.yaml --agent-mode workflow
python3 scripts/orchestrate.py --interactive
python3 scripts/agent_config.py --check

参考文档

• references/agent_contract.md
• references/evidence-schema.md
• references/routing_policy.md
• references/multi_agent_design.md
• references/rules.md