Ian Chou's Blog

用 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 的原因:

  1. Rust 生態最成熟的圖庫
  2. PyO3 讓 Python 整合無痛
  3. 可完全客製化資料結構

專案結構

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 建構。