Architecture & Integration

System Overview

We built a modular integration layer that connects Neo4j to PyTorch Geometric, allowing GNNs to be trained and run inference using data stored in a graph database. The architecture consists of three main components:

1. Graph Store & Feature Store: Interfaces that fetch graph structure and node features from Neo4j
2. Samplers: Custom implementations of neighbor sampling and GraphSAINT sampling that query Neo4j
3. Java UDP Plugin: A custom Neo4j plugin that enables in-database computation

Neo4j-PyG Integration

PyTorch Geometric uses a GraphStore and FeatureStore abstraction to decouple data storage from model code. We implemented these interfaces to fetch data from Neo4j on-demand:

• Graph Store: Stores the graph topology (edges) and provides sampling methods
• Feature Store: Retrieves node features (embeddings, labels) by node ID

This design allows the same GNN model code to work with different data sources, including in-memory PyG datasets, Neo4j, or other databases.

Why a Custom Plugin?

The initial Cypher-based approach for sampling neighborhoods was slow. Profiling revealed that Cypher’s iterative pattern matching materializes the k-hop neighborhood in RAM, resulting in many database hits per query:

The scatter plot shows that operators like Expand and Filter produce thousands of rows while hitting the database tens of thousands of times. This motivated building the Java UDP plugin to move sampling logic into the database itself, avoiding repeated round-trips.

The Java UDP Plugin

One of our key contributions was a custom Java User Defined Procedure (UDP) plugin for Neo4j. This plugin, located in the neo4j-gcn-plugin directory, provides:

• In-database inference: Run a trained GNN model directly inside Neo4j using Cypher queries
• Optimized sampling: Faster neighborhood sampling with compact data representation
• Feature fetching: Efficient retrieval of node features in byte-array format

The plugin moves computation into the database, minimizing data transfer between Neo4j and the Python training process. This is critical for performance when working with large graphs.

Data Representation

We experimented with different feature storage formats:

• embedding_bytes: Features stored as byte arrays, fastest for training
• embedding_bytes_floats: Features stored as floats, required for in-database Cypher inference but ~2x slower for training

The choice of data format has a significant impact on performance, and the optimal choice depends on the use case.

Datasets

We benchmarked our approach on multiple datasets:

• Cora: Small citation network (2,708 nodes)
• ogbn-arxiv: Medium citation network (169,343 nodes)
• ogbn-products: Large product co-purchase graph (2,449,029 nodes)
• ogbn-papers100M: Very large citation network (111,059,956 nodes)
• Coauthor Physics: Large co-authorship network (539,289 nodes)

This range allowed us to evaluate how different strategies scale with graph size.