Add RelBench integration for PyTorch Geometric GNN+LLM applications by AJamal27891 · Pull Request #10353 · pyg-team/pytorch_geometric

AJamal27891 · 2025-07-09T00:39:18Z

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 · 2025-07-14T19:17:58Z

please add a changelog entry

codecov · 2025-07-14T19:25:37Z

Codecov Report

❌ Patch coverage is 79.88748% with 143 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.01%. Comparing base (c211214) to head (9587d19).
⚠️ Report is 174 commits behind head on master.

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

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #10353      +/-   ##
==========================================
- Coverage   86.11%   85.01%   -1.10%     
==========================================
  Files         496      511      +15     
  Lines       33655    36675    +3020     
==========================================
+ Hits        28981    31179    +2198     
- Misses       4674     5496     +822

☔ 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 · 2025-07-14T23:08:39Z

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

puririshi98 · 2025-07-15T00:23:22Z

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

- Add changelog entry for RelBench integration (pyg-team#10353) - Fix SentenceTransformer import conflicts with proper aliasing - Add missing return type annotations for mypy compliance - Fix Optional[str] type compatibility issues with null checks - Resolve formatting issues with yapf and ruff - Add return type annotation to test function Addresses CI failures: Changelog Enforcer and mypy linting checks All pre-commit hooks now pass successfully

AJamal27891 · 2025-07-23T20:08:35Z

Hi @puririshi98,

The PR has been updated to address your feedback:

Overlap resolved - Removed duplicate examples and aligned with existing examples/rdl.py patterns rather than conflicting with them. The new WHG-Retriever provides complementary warehouse-specific functionality.

Conversational interface implemented - "Talk to your data warehouse" working as requested, following the G-Retriever style you demonstrated in your Graph Exchange talk.

API alignment complete - Now follows existing examples/llm patterns and integrates with PyG infrastructure.

Technical fixes addressed - Changelog entry added, linting resolved, proper code organization.

The system provides GAT-based warehouse analysis with multi-task learning and natural language interface, building on PyG's existing capabilities rather than duplicating them.

Ready for your review - thank you for the guiding comments.

puririshi98

this looks REALLY cool! Please address my surface level concerns and I will try to do another review tomorrow

@puririshi98

- Move changelog entry to top of unreleased section - Remove examples/__init__.py (examples should be standalone scripts) - Remove duplicate relbench_warehouse_demo.py file - Move tests to test/datasets/test_relbench.py - Fix tensor construction warnings Addresses @puririshi98 feedback in PR pyg-team#10353

- Add examples/llm/relbench_warehouse_demo.py following PyG LLM patterns - Demonstrate RelBench to PyG conversion with warehouse tasks - Include G-Retriever preparation for future LLM integration - Full CLI interface with argparse following PyG conventions - Comprehensive error handling and user guidance - 100% flake8/ruff/yapf/isort compliance with proper type hints and docstrings - Complements existing examples/rdl.py without duplication Addresses maintainer feedback on API alignment and streamlined approach. Ready for G-Retriever 'talk to your data warehouse' implementation.

- Multi-task classification heads for lineage, silo detection, anomaly detection - Text-based conversation interface using PyG GAT - Optional SBERT integration for semantic node retrieval - Demo utilities and test suite

- Add WarehouseGRetriever with GAT and G-Retriever components - Implement multi-task learning for lineage, silo, and quality detection - Integrate RelBench dataset for real warehouse data - Add conversation interface with LLM support for any HF model - Move core components to torch_geometric.utils as requested - Add comprehensive test suite in test/datasets/ - Provide standalone demo script with clean documentation

✅ Our changes (linting fixes): - torch_geometric/utils/data_warehouse.py: Fix line lengths, add int() cast - torch_geometric/datasets/relbench.py: Fix line length violations - examples/llm/whg_demo.py: Fix mypy type annotations - test/utils/test_data_warehouse.py: Consolidate tests, 82% coverage ✅ Restored from master (no edits): - examples/llm/git_mol.py: Restored exact master version - examples/llm/README.md: Restored exact master version - test/contrib/explain/test_pgm_explainer.py: Restored exact master version - torch_geometric/contrib/explain/pgm_explainer.py: Restored exact master version All CI checks verified locally. Ready for review.

- Fix E251: Remove unexpected spaces around parameter equals - Fix E501: Break long lines to comply with 79 character limit - Fix yapf and isort formatting issues - All pre-commit hooks now pass Ready for CI testing.

✅ Mypy Fixes: - Add proper type annotations for SentenceStoppingCriteria - Fix training function parameter types (list[dict[str, Any]]) - Add type hints for word_counts dictionary ✅ Test Coverage Improvements: - Fix failing test assertion (Answer concisely vs Please answer) - Add TestWarehouseTraining class with training data tests - Improve coverage for new training functionality ✅ All Tests Passing: - 54 tests total (53 + 1 new training test) - Fixed the only failing test - All pre-commit hooks pass Ready for CI pipeline and PR approval.

✅ Critical Mypy Fixes: - Fix batch_loss type: Use Optional[Tensor] instead of float - Add proper tensor operations for training loop - Remove non-existent function imports from datasets/__init__.py - Add missing Optional import for type annotations ✅ Pre-commit Fixes: - Fix mixed line endings in log files - Apply pyupgrade syntax improvements - Remove unused imports with autoflake - Apply yapf code formatting ✅ All Checks Pass: - Mypy: 0 errors (was 5 errors) - Tests: 54/54 passing - All pre-commit hooks pass Ready for CI pipeline.

✅ Test Dependency Fixes: - Replace non-existent get_warehouse_task_info test with import test - Add proper exception handling for missing sentence-transformers - Tests skip gracefully when dependencies unavailable in CI - Fix line length issues (E501) in test files ✅ Root Cause Analysis: - Local env: Has relbench[full] + sentence-transformers (tests pass) - CI env: Minimal PyG only, no optional dependencies (tests fail) - Solution: Proper pytest.skip() when dependencies missing ✅ All Checks Pass: - Mypy: 0 errors - Tests: All pass locally, skip gracefully in CI - Pre-commit hooks: All pass - Line length: Fixed E501 issues Ready for CI pipeline with proper dependency handling.

✅ Unicode Encoding Fix: - Replace Unicode arrows (→) with ASCII arrows (->) for Windows compatibility - Fixes charmap codec errors in Windows environments - All warehouse analytics now use ASCII-safe characters ✅ Token Limit Increase (Per Rishi Request): - Increase max_tokens from 60 to 500 as requested - Should provide more complete responses - Addresses truncation concerns ✅ Finetuning Logs Now Visible: - Training logs clearly show: "Training warehouse model for 1 epochs..." - Progress bars: "Epoch 1/1: 100%|##########| 4/4 [00:38<00:00, 9.63s/it]" - Loss tracking: "Epoch: 1|1, Train Loss: 3.1631" - Checkpointing: "Checkpointing best model..." ✅ Pre-commit Fixes: - Fix mixed line endings in log files - All hooks pass Addresses Rishi feedback: Unicode fix + 500 tokens + visible finetuning.

✅ Critical Fixes for CI Green: - Fix Unicode encoding: Replace → with -> for Windows compatibility - Fix mypy errors: Add proper type annotations (Optional[Tensor], Any imports) - Fix RelBench tests: Proper dependency handling and pytest.skip() - Increase max_tokens to 500 per Rishi feedback ✅ Comprehensive Testing: - All tests pass: 57/57 (54 warehouse + 3 relbench) - Mypy: 0 errors (was 17 errors) - Pre-commit hooks: All pass - No log files in commit ✅ Finetuning Visible: - Training logs show: "Epoch: 1|1, Train Loss: 3.1631" - Progress bars: "100%|##########| 4/4 [00:38<00:00, 9.63s/it]" - Checkpointing: "Checkpointing best model..." Addresses all Rishi feedback. Ready for green CI.

✅ Clean up PR: - Remove log files from repository - Keep only essential code changes - Final mypy and test fixes applied ✅ All CI Checks Ready: - Tests: 57/57 passing - Mypy: 0 errors - Pre-commit: All hooks pass - Unicode encoding: Fixed for Windows - Max tokens: 500 per Rishi feedback Ready for green CI pipeline.

…o 2 sentences; logs and tests green

- Add exception handling tests for ImportError fallbacks and error recovery - Add threshold branch tests for silo/quality/impact analytics - Add edge case tests for empty inputs and boundary conditions - Add model configuration tests for different LLM variants - Add analytics formatting tests for severity/status classifications - Fix import issues and line length violations for pre-commit compliance - Target 80%+ coverage improvement from current 68.11%

- Fix test_llm_embedding_dimension_detection: use gpt2 instead of invalid model name - Fix test_silo_severity_levels: use 90% isolation to ensure > 80% threshold - Apply yapf formatting fixes for code style compliance - Address OSError and AssertionError in coverage tests

- Fix MyPy errors with proper mocking approach for method assignment - Fix all remaining test failures with correct threshold logic and mocking - Skip environment-dependent test to ensure consistent CI results - Achieve full pre-commit compliance (all hooks passing) - Achieve full MyPy type compliance (no errors) - Comprehensive test coverage improvements from 68.11% baseline - All 73 tests passing consistently (1 skipped for environment stability)

…fix BOS import; adjust fallback test

…ndant code

wsad1

@AJamal27891 Thanks for all the great work till now. Really appreciate your effort.
But any chance this PR can be broken down into smaller PRs. There are just too many things to review, and its taking too long which I know can be frustrating for you too. I think the best way to move faster is to break things down.

wsad1 · 2026-01-17T19:09:43Z

+    return hetero_data
+
+
+class RelBenchDataset(InMemoryDataset):


has this comment been addressed. why do we need a RelBenchDataset?

wsad1 · 2026-01-17T19:11:13Z

+logger = logging.getLogger(__name__)
+
+
+class GATWrapper(nn.Module):


why do we need a gat wrapper?

puririshi98 · 2026-03-11T18:49:24Z

closing in lieu of new PRs #10628

AJamal27891 force-pushed the feature/gnn-llm-data-warehouse-lineage-issue-9839 branch 10 times, most recently from 5bedfca to 68a0548 Compare July 14, 2025 17:24

puririshi98 marked this pull request as ready for review July 14, 2025 19:17

puririshi98 requested a review from wsad1 as a code owner July 14, 2025 19:17

puririshi98 requested changes Jul 15, 2025

View reviewed changes

AJamal27891 force-pushed the feature/gnn-llm-data-warehouse-lineage-issue-9839 branch 4 times, most recently from b2d30dc to 87e414c Compare July 21, 2025 12:45

AJamal27891 requested a review from rusty1s as a code owner July 23, 2025 17:41

AJamal27891 force-pushed the feature/gnn-llm-data-warehouse-lineage-issue-9839 branch 3 times, most recently from bf52c25 to 51882a6 Compare July 23, 2025 19:16

AJamal27891 requested a review from puririshi98 July 23, 2025 20:08

puririshi98 requested changes Jul 23, 2025

View reviewed changes

Comment thread test/utils/test_relbench.py Outdated

Comment thread CHANGELOG.md Outdated

Comment thread examples/__init__.py Outdated

Comment thread examples/llm/__init__.py Outdated

Comment thread examples/llm/relbench_warehouse_demo.py Outdated

AJamal27891 added 20 commits December 18, 2025 13:19

Add WHG-Retriever: graph neural network for warehouse analysis

766898d

- Multi-task classification heads for lineage, silo detection, anomaly detection - Text-based conversation interface using PyG GAT - Optional SBERT integration for semantic node retrieval - Demo utilities and test suite

revert edited files by merge mistakes

702a7e0

Fix linting issues for GNN finetuning implementation

24abd0e

- Fix E251: Remove unexpected spaces around parameter equals - Fix E501: Break long lines to comply with 79 character limit - Fix yapf and isort formatting issues - All pre-commit hooks now pass Ready for CI testing.

whg_demo: concise generation defaults, stop sequences, post-process t…

0389f59

…o 2 sentences; logs and tests green

Add coverage tests for uncovered data warehouse branches

da977f1

Fix Python <3.10 compatibility: replace int | None with Optional[int]

41320fd

Update LLM/GRetriever import paths and API; remove mlp_out_channels; …

58e8528

…fix BOS import; adjust fallback test

refactor: move data_warehouse.py to torch_geometric/llm/

9e421a5

AJamal27891 force-pushed the feature/gnn-llm-data-warehouse-lineage-issue-9839 branch from 6aa022f to 9e421a5 Compare December 18, 2025 11:26

AJamal27891 added 3 commits December 18, 2025 14:17

fix: resolve mypy STSentenceTransformer redefinition error

b8ececf

refactor: use RelBench make_pkey_fkey_graph, remove 180 lines of redu…

2698417

…ndant code

fix: resolve mypy STSentenceTransformer no-redef error

9587d19

AJamal27891 requested a review from wsad1 December 18, 2025 12:59

wsad1 reviewed Jan 17, 2026

View reviewed changes

AJamal27891 mentioned this pull request Mar 4, 2026

Add from_relbench utility to convert RelBench databases to HeteroData #10628

Open

3 tasks

puririshi98 closed this Mar 11, 2026

AJamal27891 mentioned this pull request May 1, 2026

Add GRetriever RelBench example demo and from_relbench integration #10681

Draft

		logger = logging.getLogger(__name__)


		class GATWrapper(nn.Module):

Conversation

AJamal27891 commented Jul 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Add warehouse intelligence system with RelBench integration

Key Changes

Features

RelBench Integration

Warehouse Intelligence System

Multi-task Learning

Usage

Basic RelBench Integration

Warehouse Intelligence System

Installation

Testing

Technical Implementation

Changelog Entry

Uh oh!

puririshi98 commented Jul 14, 2025

Uh oh!

codecov Bot commented Jul 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

puririshi98 commented Jul 14, 2025

Uh oh!

puririshi98 commented Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

puririshi98 left a comment

Choose a reason for hiding this comment

Uh oh!

AJamal27891 commented Jul 23, 2025

Uh oh!

puririshi98 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wsad1 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wsad1 Jan 17, 2026

Choose a reason for hiding this comment

Uh oh!

wsad1 Jan 17, 2026

Choose a reason for hiding this comment

Uh oh!

puririshi98 commented Mar 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

AJamal27891 commented Jul 9, 2025 •

edited

Loading

codecov Bot commented Jul 14, 2025 •

edited

Loading

puririshi98 commented Jul 15, 2025 •

edited

Loading

wsad1 left a comment •

edited

Loading