用 Petgraph + PyO3 取代 NetworkX:Python 圖運算的 Rust 加速
用 Petgraph + PyO3 取代 NetworkX:Python 圖運算的 Rust 加速
當 NetworkX 成為效能瓶頸時,用 Rust Petgraph 提速 10-100 倍。本文分享如何用 PyO3 + maturin 將 Rust 圖庫編譯成 Python 原生模組,並保持 API 相容。
為什麼要換 NetworkX
| 情境 | NetworkX | 需求 |
|---|---|---|
| 小型圖 (<1000 nodes) | ✅ 足夠 | 不需換 |
| 大型圖 (>10K nodes) | 🐢 慢 | 考慮換 |
| 即時查詢 (<10ms) | ❌ 難達成 | 需要換 |
| 批量圖遍歷 | 🐢 瓶頸 | 建議換 |
我們的 GraphRAG 系統需要在每次查詢時做 2-hop 遍歷,NetworkX 約需 50ms,Petgraph 可壓到 <1ms。
技術選型
為什麼選 Petgraph + PyO3
| 方案 | 優點 | 缺點 |
|---|---|---|
| Petgraph + PyO3 | 最快、原生 Python 整合 | 需寫 Rust |
| graph-tool | 快、已有 Python binding | 安裝複雜 |
| igraph | 中等速度 | C 底層、API 不友善 |
| rustworkx | Rust 速度、Python API | 功能較少 |
選擇 Petgraph 的原因:
- Rust 生態最成熟的圖庫
- PyO3 讓 Python 整合無痛
- 可完全客製化資料結構
專案結構
py-kb/
├── rust_graph/ # Rust 模組
│ ├── Cargo.toml # Rust 依賴
│ ├── pyproject.toml # maturin 設定
│ └── src/lib.rs # CareerGraph 實作
└── src/career_kb/graph/
└── knowledge_graph.py # Python 橋接層
Step 1: Cargo.toml 設定
[package]
name = "career_graph"
version = "0.1.0"
edition = "2021"
[lib]
name = "career_graph"
crate-type = ["cdylib"] # 動態連結庫,給 Python 用
[dependencies]
pyo3 = { version = "0.23", features = ["extension-module"] }
petgraph = "0.7"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
[profile.release]
opt-level = 3
lto = true # Link-time optimization
關鍵設定說明
| 設定 | 說明 |
|---|---|
crate-type = ["cdylib"] |
編譯成 Python 可載入的動態連結庫 |
pyo3/extension-module |
告訴 PyO3 這是 Python 模組 |
lto = true |
啟用跨 crate 優化,進一步提速 |
Step 2: pyproject.toml (maturin 設定)
[build-system]
requires = ["maturin>=1.0,<2.0"]
build-backend = "maturin"
[project]
name = "career_graph"
version = "0.1.0"
requires-python = ">=3.11"
[tool.maturin]
features = ["pyo3/extension-module"]
module-name = "career_graph"
Step 3: Rust 實作
資料結構
use petgraph::graph::{DiGraph, NodeIndex};
use std::collections::HashMap;
// 節點資料
#[derive(Clone, Debug)]
struct NodeData {
name: String,
node_type: String, // skill, company, project
category: Option<String>,
}
// 邊資料
#[derive(Clone, Debug)]
struct EdgeData {
relation: String, // implies, relatedTo
}
// 主結構:用 HashMap 做 name -> NodeIndex 映射
#[pyclass]
pub struct CareerGraph {
graph: DiGraph<NodeData, EdgeData>,
node_indices: HashMap<String, NodeIndex>,
}
PyO3 綁定
use pyo3::prelude::*;
#[pymethods]
impl CareerGraph {
// Python: graph = CareerGraph()
#[new]
pub fn new() -> Self {
CareerGraph {
graph: DiGraph::new(),
node_indices: HashMap::new(),
}
}
// Python: count = graph.load_from_json("path.json")
pub fn load_from_json(&mut self, path: &str) -> PyResult<usize> {
// 讀取 JSON 並建立圖
// ... 詳細實作省略
}
// Python: neighbors = graph.get_neighbors("LangChain", 2)
pub fn get_neighbors(&self, node: &str, hops: usize) -> Vec<String> {
// BFS 遍歷
// ... 詳細實作省略
}
// Python: count = graph.node_count()
pub fn node_count(&self) -> usize {
self.graph.node_count()
}
}
// 模組定義
#[pymodule]
fn career_graph(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<CareerGraph>()?;
Ok(())
}
BFS 遍歷實作
use petgraph::Direction;
use petgraph::visit::EdgeRef; // 重要:edge.target() 需要這個 trait
use std::collections::{HashSet, VecDeque};
pub fn get_neighbors(&self, node: &str, hops: usize) -> Vec<String> {
let start_idx = match self.node_indices.get(node) {
Some(idx) => *idx,
None => return Vec::new(),
};
let mut visited: HashSet<NodeIndex> = HashSet::new();
let mut queue: VecDeque<(NodeIndex, usize)> = VecDeque::new();
visited.insert(start_idx);
queue.push_back((start_idx, 0));
while let Some((current, depth)) = queue.pop_front() {
if depth >= hops {
continue;
}
// 雙向遍歷 (outgoing + incoming)
for neighbor in self.graph.neighbors_directed(current, Direction::Outgoing) {
if !visited.contains(&neighbor) {
visited.insert(neighbor);
queue.push_back((neighbor, depth + 1));
}
}
for neighbor in self.graph.neighbors_directed(current, Direction::Incoming) {
if !visited.contains(&neighbor) {
visited.insert(neighbor);
queue.push_back((neighbor, depth + 1));
}
}
}
// 移除起點,轉成名稱列表
visited.remove(&start_idx);
visited
.iter()
.filter_map(|idx| self.graph.node_weight(*idx).map(|n| n.name.clone()))
.collect()
}
JSON 同步方法
/// 取得所有邊 (source, target, relation)
pub fn get_all_edges(&self) -> Vec<(String, String, String)> {
self.graph
.edge_references()
.filter_map(|edge| {
let source = self.graph.node_weight(edge.source())?;
let target = self.graph.node_weight(edge.target())?;
Some((
source.name.clone(),
target.name.clone(),
edge.weight().relation.clone(),
))
})
.collect()
}
/// 直接存成 skill-graph.json 格式
pub fn save_to_json(&self, path: &str) -> PyResult<()> {
let json = self.to_skill_graph_json();
fs::write(path, json)
.map_err(|e| PyErr::new::<pyo3::exceptions::PyIOError, _>(e.to_string()))
}
Step 4: 編譯與安裝
# 安裝 Rust (若尚未安裝)
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
source $HOME/.cargo/env
# 安裝 maturin
uv pip install maturin
# 編譯並安裝到 venv
cd py-kb
uv run maturin develop --manifest-path rust_graph/Cargo.toml --release
成功輸出:
🔗 Found pyo3 bindings
🐍 Found CPython 3.12
📦 Built wheel for CPython 3.12
✏️ Setting installed package as editable
🛠 Installed career_graph-0.1.0
Step 5: Python 橋接層
# 嘗試載入 Rust 模組,失敗則 fallback 到 NetworkX
try:
from career_graph import CareerGraph as RustCareerGraph
USE_RUST_BACKEND = True
except ImportError:
USE_RUST_BACKEND = False
import networkx as nx
class CareerKnowledgeGraph:
def __init__(self):
self._use_rust = USE_RUST_BACKEND
if self._use_rust:
self._rust_graph = RustCareerGraph()
else:
import networkx as nx
self.graph = nx.DiGraph()
def load_from_skill_graph(self) -> int:
if self._use_rust:
return self._rust_graph.load_from_json(str(self.skill_graph_path))
# NetworkX 實作...
def node_count(self) -> int:
if self._use_rust:
return self._rust_graph.node_count()
return self.graph.number_of_nodes()
為什麼用 Fallback 設計
| 情境 | 行為 |
|---|---|
| Rust 模組已編譯 | 使用 Petgraph(快) |
| Rust 模組未編譯 | 使用 NetworkX(慢但能用) |
| CI/CD 環境 | 可選擇跳過 Rust 編譯 |
這樣開發者不需強制安裝 Rust 也能使用專案。
效能對比
import time
from career_graph import CareerGraph
g = CareerGraph()
g.load_from_json("data/skill-graph.json")
# 測試 2-hop 遍歷
start = time.perf_counter()
for _ in range(1000):
g.get_neighbors("LangChain", 2)
elapsed = time.perf_counter() - start
print(f"1000 次遍歷: {elapsed*1000:.2f}ms")
print(f"平均每次: {elapsed:.4f}ms")
| Backend | 1000 次遍歷 | 平均每次 |
|---|---|---|
| NetworkX | ~50,000ms | ~50ms |
| Petgraph | ~500ms | ~0.5ms |
| 加速比 | 100x |
常見問題
1. linker cc not found
# 安裝 C 編譯工具
sudo apt-get install build-essential
2. ModuleNotFoundError: No module named 'career_graph'
Rust 模組未編譯或未安裝到正確的 venv:
cd py-kb
uv run maturin develop --manifest-path rust_graph/Cargo.toml --release
3. edge.target() 找不到
需要導入 EdgeRef trait:
use petgraph::visit::EdgeRef; // 加這行
4. Pyo3 signature 警告
// 加上 signature 屬性
#[pyo3(signature = (name, entity_type, category=None))]
pub fn add_entity(&mut self, name: &str, entity_type: &str, category: Option<&str>) -> bool {
何時該用這個方案
| ✅ 適合 | ❌ 不適合 |
|---|---|
| 圖遍歷是熱點 | 一次性分析腳本 |
| 需要即時回應 | 離線批次處理 |
| 圖結構固定 | 頻繁修改圖結構 |
| 長期維護專案 | 快速原型開發 |
總結
| 元件 | 技術 |
|---|---|
| 圖庫 | Rust Petgraph |
| Python 綁定 | PyO3 |
| 建構工具 | maturin |
| Fallback | NetworkX |
用 Rust 加速 Python 的核心路徑,同時保持 Python 的開發便利性。
Career Knowledge Base 是一個本地優先的履歷知識庫系統,使用 Python + Rust + LanceDB + LangChain 建構。
- ← Previous
手刻 GraphRAG:用 NetworkX + LangChain 建構知識圖譜增強檢索 - Next →
GraphRAG 全局檢索:社區分群 + 社區摘要 + Map-Reduce 聚合(LangChain 版)