Neo4j Python驱动高效写入实践

本文详解如何通过批量 UNWIND 查询与合理会话管理，显著提升 Neo4j Python 驱动在数十万级数据写入场景下的性能，避免逐行执行导致的严重延迟。

在使用 Neo4j Python 官方驱动（neo4j==5.20+）处理大规模数据（如 >20 万条记录）时，若沿用 session.execute_write() 逐行调用 Cypher（例如对 DataFrame 每行调用一次 MERGE），性能会急剧下降——这并非驱动缺陷，而是因网络往返开销、事务开销和单次查询解析成本叠加所致。根本优化路径是：减少请求次数 + 提升单次查询吞吐量。

核心方案是采用 UNWIND + 批量参数化写入。UNWIND 将传入的列表参数展开为行流，配合 CREATE/MERGE 实现“一次请求、多行写入”。配合合理的批大小（通常 5,000–20,000），可将写入速度提升 10–100 倍。

以下为生产就绪的推荐实现：

from neo4j import GraphDatabase
import pandas as pd
from tqdm import tqdm

# 初始化驱动（建议复用全局实例）
driver = GraphDatabase.driver(
    "bolt://localhost:7687",
    auth=(os.getenv("NEO_USERNAME"), os.getenv("NEO_PASSWORD"))
)

# ✅ 正确的约束创建（仅需执行一次，建议独立运行）
def create_constraint():
    with driver.session(database="neo4j") as session:
        session.run("CREATE CONSTRAINT entityIndex IF NOT EXISTS ON (e:Entity) ASSERT e.EntityId IS UNIQUE")

# ✅ 高效批量写入：使用 UNWIND + MERGE
BATCH_SIZE = 10_000
query = """
UNWIND $rows AS row
MERGE (e:Entity {EntityId: row.entity_id})
ON CREATE SET e.LastAccess = timestamp()
ON MATCH SET e.LastAccess = timestamp()
"""

def bulk_upsert_entities(df: pd.DataFrame):
    # 转为字典列表（列名需与 Cypher 中 row.xxx 严格一致）
    records = df[["entity_id"]].to_dict(orient="records")

    for i in tqdm(range(0, len(records), BATCH_SIZE), desc="Uploading batches"):
        batch = records[i : i + BATCH_SIZE]
        try:
            # execute_query 是 v5.0+ 推荐的顶层方法，自动管理会话与事务
            driver.execute_query(
                query,
                rows=batch,
                database_="neo4j"  # 注意下划线命名（非 database）
            )
        except Exception as e:
            print(f"Batch {i//BATCH_SIZE} failed: {e}")
            raise

# 使用示例
# bulk_upsert_entities(df)

⚠️ 关键注意事项：

• 不要在循环中新建 Session：每个 with driver.session() 都有连接开销；execute_query() 内部已优化会话复用。
• database_ 参数名含下划线：这是 Python 驱动的保留关键字规避写法（非 typo），务必使用 database_ 而非 database。
• MERGE 中变量名需匹配：Cypher 中 row.entity_id 必须与 df[["entity_id"]] 列名完全一致（区分大小写）。
• 错误处理粒度：按批捕获异常，而非单行——单行失败不应中断整个批次，可记录失败批次后重试或排查数据质量。
• 索引/约束先行：确保 :Entity(EntityId) 约束已存在（如上 create_constraint），否则 MERGE 性能将退化为全表扫描。

? 进阶提示：对于关系批量创建，同样使用 UNWIND，但需先确保起点/终点节点已存在（或用 MATCH + MERGE 组合）。例如：

UNWIND $rels AS rel
MATCH (a:Entity {EntityId: rel.start_id})
MATCH (b:Entity {EntityId: rel.end_id})
MERGE (a)-[r:RELATED_TO]->(b)
SET r.weight = rel.weight

综上，告别逐行 execute_write，拥抱 UNWIND 批量模式——这是 Neo4j 官方文档明确推荐的大数据写入范式，也是生产环境保障吞吐与稳定性的基石。