Add RelBench integration for PyTorch Geometric GNN+LLM applications

[AJamal27891]

Add warehouse intelligence system with RelBench integration

Closes #9839

This PR implements a warehouse intelligence system for PyTorch Geometric, providing RelBench dataset integration and graph-based warehouse analysis capabilities using G-Retriever architecture for multi-task learning on data lineage, silo detection, and quality assessment tasks.

Key Changes

New Files:

• torch_geometric/datasets/relbench.py - RelBench to HeteroData conversion utilities
• torch_geometric/utils/data_warehouse.py - Warehouse intelligence with G-Retriever integration
• examples/llm/whg_demo.py - Warehouse intelligence demonstration
• test/datasets/test_relbench.py - Comprehensive RelBench functionality tests (9 tests)
• test/utils/test_data_warehouse.py - Warehouse intelligence system tests (13 tests)

Updated Files:

• torch_geometric/datasets/__init__.py - Export RelBench utilities
• pyproject.toml - Optional dependency groups for relbench and whg

Features

RelBench Integration

• create_relbench_hetero_data() - Convert RelBench datasets to PyG HeteroData
• RelBenchDataset - PyG dataset wrapper for RelBench data
• HeuristicLabeler - Generate warehouse task labels from graph structure
• RelBenchProcessor - Process RelBench data with semantic embeddings

Warehouse Intelligence System

• WarehouseGRetriever - G-Retriever architecture for warehouse analysis
• WarehouseTaskHead - Multi-task prediction (lineage, silo, quality)
• WarehouseConversationSystem - Natural language interface for warehouse queries
• SimpleWarehouseModel - Lightweight model for basic warehouse operations

Multi-task Learning

• Lineage prediction - Trace data flow and dependencies
• Silo detection - Identify isolated data components
• Quality assessment - Detect anomalies and data quality issues

Usage

Basic RelBench Integration

from torch_geometric.datasets.relbench import create_relbench_hetero_data

# Convert RelBench dataset to PyG format
hetero_data = create_relbench_hetero_data(
    dataset_name='rel-f1',
    sample_size=100,
    create_lineage_labels=True,
    create_silo_labels=True,
    create_anomaly_labels=True
)

Warehouse Intelligence System

from torch_geometric.utils.data_warehouse import create_warehouse_demo

# Create warehouse conversation system
warehouse_system = create_warehouse_demo()

# Query warehouse intelligence
result = warehouse_system.process_query(
    "What is the data lineage in this warehouse?", 
    graph_data
)
print(result['answer'])

Installation

# RelBench integration
pip install torch-geometric[relbench]

# Warehouse intelligence system
pip install torch-geometric[whg]

# Both features
pip install torch-geometric[relbench,whg]

Testing

Test Coverage: 22 tests across 2 files, all passing ✅

• test/datasets/test_relbench.py - RelBench integration tests (9 tests)
• test/utils/test_data_warehouse.py - Warehouse intelligence tests (13 tests)

Includes comprehensive coverage of core functionality, edge cases, and error handling.

Technical Implementation

The system integrates G-Retriever architecture with RelBench datasets to provide warehouse intelligence capabilities. Key technical features:

• Semantic embeddings using sentence transformers for text-based node features
• Multi-task learning with shared GNN backbone and task-specific heads
• Heuristic labeling for automatic warehouse task label generation
• LLM integration with fallback to traditional GNN approaches
• Modular design supporting both standalone and integrated usage

Changelog Entry

Added to CHANGELOG.md under version 2.7.0:

- Added RelBench integration with data warehouse lineage tasks (#10353)

This entry covers the complete warehouse intelligence system implementation including RelBench integration, G-Retriever architecture, multi-task learning capabilities, and comprehensive testing.

[puririshi98]

please add a changelog entry

[codecov]

Codecov Report

❌ Patch coverage is 79.91573% with 143 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.83%. Comparing base (c211214) to head (3e486c3).
⚠️ Report is 94 commits behind head on master.

Files with missing lines	Patch %	Lines
torch_geometric/utils/data_warehouse.py	79.88%	143 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #10353      +/-   ##
==========================================
- Coverage   86.11%   85.83%   -0.28%     
==========================================
  Files         496      503       +7     
  Lines       33655    35840    +2185     
==========================================
+ Hits        28981    30765    +1784     
- Misses       4674     5075     +401     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[puririshi98]

please also address linting, will do a deep review this week

[puririshi98]

there is alot of overlap with this existing example: https://github.com/pyg-team/pytorch_geometric/blob/master/examples/rdl.py
I was under the impression that you were going to integrate this into a pipeline similar to G-retriever (see https://github.com/pyg-team/pytorch_geometric/tree/master/examples/llm) where you could "talk to your data warehouse" (since G-retriever style GNN+LLM enables "talk to your graph"). please align your API's with the existing examples/rdl from the core contributors and let me know when you have a working "talk to your data warehouse" example.
See my talk here if you'd like more details about "talk to your graph"
https://www.devreal.ai/graph-exchange-may-2025/

[puririshi98]

see above comment

[puririshi98]

can you run this with https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pyg/ 25.05 and make sure everything runs. Please attach a log of that run so that I can review

[AJamal27891]

I spun everything up in nvcr.io/nvidia/pyg:25.05-py3 (inside the container, no local mods) and the demo runs end-to-end:

"docker exec -w /opt/pyg/pytorch_geometric 7ddc1ab2fd5f python examples/llm/whg_demo.py"
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
PyG SentenceTransformer requires tokenized input
/opt/pyg/pytorch_geometric/examples/llm/../../torch_geometric/nn/nlp/llm.py:97: UserWarning: LLM is being used on CPU, which may be slow
  warnings.warn("LLM is being used on CPU, which may be slow")
/opt/pyg/pytorch_geometric/examples/llm/../../torch_geometric/utils/data_warehouse.py:655: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/native/Scalar.cpp:22.)
  confidence = probs.max(dim=-1).values.mean().item()
Warehouse Intelligence Demo
==============================

Step 1: Loading RelBench data
Loading Database object from /root/.cache/relbench/rel-f1/db...
Done in 0.06 seconds.
Loaded graph with 9 node types
Node types: ['standings', 'constructor_standings', 'circuits', 'constructor_results', 'qualifying', 'drivers', 'races', 'results', 'constructors']
Converted to homogeneous: 450 nodes, 236 edges

Step 2: Creating warehouse conversation system
Setting up 'TinyLlama/TinyLlama-1.1B-Chat-v0.1' with configuration: {'revision': 'main'}
Warehouse conversation system ready

Step 3: Running warehouse intelligence queries

--- Query 1: What is the data lineage in this warehouse? ---
Answer:  The data lineage in this warehouse can be traced as follows:

Direct connections: The raw data sources for this warehouse are spread across multiple data centers and data sources. The data lineage traces back to these sources.

Staged data: The intermediate processing steps involved in data transfo...

Quantitative Analysis: Unknown lineage detectedacross 450 entities (confidence: 0.206)
Query type: lineage
Confidence: N/A

--- Query 2: Are there any data silos? ---
Answer:  The data warehouse is composed of 450 entities, 236 relationships, and sparsely connected connections. The graph structure suggests a hierarchical structure with 450 nodes and 236 relationships. The data sources are isolated or poorly connected, with no bridges connecting different data domains. Cl...

Quantitative Analysis: Analytics: 0 isolated silosout of 450 entities (0.2% isolation rate)
Query type: silo
Confidence: N/A

--- Query 3: What is the data quality status? ---
Answer:  The data in this warehouse is in a highly inconsistent and unstructured format. The data is not cleaned, validated, or formatted in a consistent manner. The data is also not structured in a way that allows for easy analysis and querying.

The data is also not regularly updated or maintained. This m...

Quantitative Analysis: Analytics: Quality score 0.623(GOOD overall status)
Query type: quality
Confidence: N/A

--- Query 4: Analyze the impact of changes in this warehouse ---
Answer:  The warehouse is a large-scale relational database with multiple entity types. The graph structure is comprised of 450 entities, 236 relationships. The connectivity is sparse, with an average degree of 1.0. The impact analysis will focus on the downstream effects of data changes, considering high-i...

Quantitative Analysis: Analytics: 0 high-impact entitiesdetected (LOW risk)
Query type: impact
Confidence: N/A

Step 4: Conversation History
------------------------------
1. Q: Are there any data silos?...
   A:  The data warehouse is composed of 450 entities, 236 relationships, and sparsely...
2. Q: What is the data quality status?...
   A:  The data in this warehouse is in a highly inconsistent and unstructured format....
3. Q: Analyze the impact of changes in this warehouse...
   A:  The warehouse is a large-scale relational database with multiple entity types. ...

Demo completed. Processed 4 queries total.

Features demonstrated:
- RelBench data integration
- Multi-task warehouse intelligence
- Natural language query processing
- Lineage, silo, and quality analysis

The demo runs out of the box no modifications just the command python examples/llm/whg_demo.py
the output shows the integration between rel-bench and the data warehouse

[puririshi98]

curious why does this cut off:
--- Query 1: What is the data lineage in this warehouse? ---

Answer:  The data lineage in this warehouse can be traced as follows:

Direct connections: The raw data sources for this warehouse are spread across multiple data centers and data sources. The data lineage traces back to these sources.

Staged data: The intermediate processing steps involved in data transfo...

Quantitative Analysis: Unknown lineage detectedacross 450 entities (confidence: 0.206)

i notice alot of sentences are truncated. please help me understand why or share the full logs. otherwise this looks really cool

[puririshi98]

can you please make CI green as well?

[puririshi98]

please make sure CI green and you add fineutning to the example

[AJamal27891]

please make sure CI green and you add fineutning to the example

whg_phi3_train_log_concise_clean.txt

this file contains the full log with 1 epoch training example using phi3

[wsad1]

Thanks a lot for adding this.
Left some initial comments. Will continue to review since this is a large PR.

[wsad1]

Suggested change

	def test_relbench_imports() -> None:
	"""Test RelBench module imports."""
	try:
	from torch_geometric.datasets.relbench import (
	RelBenchDataset,
	RelBenchProcessor,
	create_relbench_hetero_data,
	)
	# Test that classes can be imported
	assert RelBenchDataset is not None
	assert RelBenchProcessor is not None
	assert create_relbench_hetero_data is not None
	except ImportError:
	pytest.skip("RelBench not available")
	from torch_geometric.datasets.relbench import (
	RelBenchDataset,
	RelBenchProcessor,
	create_relbench_hetero_data,
	)

No need to check this, the test will only be run on a version that has the relbench.py file.

[wsad1]

Same as above, no need to put these imports in a Try catch block.

[wsad1]

Suggested change

	try:
	from torch_geometric.nn.models import GRetriever
	from torch_geometric.nn.nlp import LLM
	HAS_GRETRIEVER = True
	except ImportError:
	HAS_GRETRIEVER = False
	class GRetriever: # type: ignore
	def __init__(self, *args: Any, **kwargs: Any) -> None:
	raise ImportError("GRetriever requires PyG with LLM support")
	class LLM: # type: ignore
	def __init__(self, *args: Any, **kwargs: Any) -> None:
	raise ImportError("LLM requires PyG with LLM support")
	from torch_geometric.nn.models import GRetriever
	from torch_geometric.nn.nlp import LLM
	HAS_GRETRIEVER = True

Why do we need this try catch block. GRetriever, is part of torch_geometric and it was added before data_warehouse.py, so no need to check for this.
We only need to have such try catch blocks for optional external dependencies.

[wsad1]

Could you add more details about these labels.
Is there a paper one could look at to understand how to use these labels. What is the shape of external_labels.