Neo4j Python驱动处理大数据查询的正确用法与性能优化

原代码及问题分析 #

原代码在处理小量DataFrame时运行正常，但数据量超过20万行时性能骤降，核心问题及代码缺陷如下：

原代码 #

import os
from neo4j import GraphDatabase
from tqdm import tqdm

driver = GraphDatabase.driver('bolt://localhost:7687',
              auth=(os.getenv('NEO_USERNAME'),     
                    os.getenv('NEO_PASSWORD')))

def create_constraint(tx):
    tx.run("CREATE CONSTRAINT entityIndex IF NOT EXISTS FOR (e:Entity) REQUIRE (e.EntityId) IS UNIQUE")

def create_entity(tx, row):
    tx.run("MERGE (e:Entity {EntityId: $entity_id}) ON CREATE SET n.LastAccess = timestamp() ON MATCH SET n.LastAccess = timestamp()", entity_id=row["entity_id"])
    
with driver.session() as session:
    session.execute_write(create_constraint)

with driver.session() as session:
    for idx, row in tqdm(df.iterrows(), total=len(df)):
        try:
            session.execute_write(create_entity, row)
        except Exception as e:
            print(e)

核心问题 #

• 逐行请求数据库：每一行数据单独发起一次execute_write调用，20万行对应20万次网络往返，大量时间消耗在网络通信和事务初始化上，这是性能暴跌的核心原因。
• Cypher语法错误：MERGE语句定义的节点变量是e，但ON CREATE/ON MATCH中误用了未定义的变量n，会直接导致运行错误。

高效优化方案：批量处理 #

通过批量提交数据减少网络往返次数，结合Cypher的UNWIND语句批量处理多行数据，是大数据量导入场景的最优解。

优化后代码 #

import os
import numpy as np
from neo4j import GraphDatabase
from tqdm import tqdm

driver = GraphDatabase.driver('bolt://localhost:7687',
              auth=(os.getenv('NEO_USERNAME'),     
                    os.getenv('NEO_PASSWORD')))

# 创建唯一约束（仅需执行一次）
with driver.session() as session:
    session.execute_write(lambda tx: tx.run("CREATE CONSTRAINT entityIndex IF NOT EXISTS FOR (e:Entity) REQUIRE (e.EntityId) IS UNIQUE"))

# 定义批量处理的Cypher查询
query = """
UNWIND $data AS row
MERGE (e:Entity {EntityId: row.entity_id})
ON CREATE SET e.LastAccess = timestamp()
ON MATCH SET e.LastAccess = timestamp()
"""

# 拆分DataFrame为多个批次（可根据数据库性能调整batch_size，建议500-2000）
batch_size = 1000
list_df = np.array_split(df, max(1, len(df) // batch_size))

# 批量提交数据
for adf in tqdm(list_df, total=len(list_df)):
    try:
        driver.execute_query(query, data=adf.to_dict(orient="records"), database_="neo4j")
    except Exception as e:
        print(f"批次处理失败: {e}")

优化说明 #

1. 批量拆分数据：将大DataFrame拆分为小批次，平衡内存占用与请求效率，避免单次请求数据量过大导致的超时或内存溢出。
2. UNWIND批量处理：通过UNWIND $data AS row一次性处理整个批次的数据，将网络请求次数从20万次降至几百次，大幅提升效率。
3. 修正语法错误：统一使用节点变量e，解决原代码的变量未定义问题。
4. 简化API调用：使用driver.execute_query（Neo4j Python驱动4.4+支持），无需手动管理会话，代码更简洁高效。

内容的提问来源于stack exchange，提问作者Simone Gabbriellini