diff --git a/Cargo.lock b/Cargo.lock
index a8d52a1b0..196c10b23 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2603,6 +2603,16 @@ dependencies = [
"num-traits",
]
+[[package]]
+name = "page_size"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da"
+dependencies = [
+ "libc",
+ "winapi",
+]
+
[[package]]
name = "papergrid"
version = "0.12.0"
@@ -3673,6 +3683,57 @@ dependencies = [
"uuid",
]
+[[package]]
+name = "ruvector-gnn"
+version = "0.1.1"
+dependencies = [
+ "anyhow",
+ "criterion",
+ "dashmap",
+ "libc",
+ "memmap2",
+ "napi",
+ "napi-derive",
+ "ndarray 0.16.1",
+ "page_size",
+ "parking_lot 0.12.5",
+ "proptest",
+ "rand 0.8.5",
+ "rand_distr",
+ "rayon",
+ "ruvector-core",
+ "serde",
+ "serde_json",
+ "tempfile",
+ "thiserror 2.0.17",
+]
+
+[[package]]
+name = "ruvector-gnn-node"
+version = "0.1.1"
+dependencies = [
+ "napi",
+ "napi-build",
+ "napi-derive",
+ "ruvector-gnn",
+ "serde_json",
+]
+
+[[package]]
+name = "ruvector-gnn-wasm"
+version = "0.1.1"
+dependencies = [
+ "console_error_panic_hook",
+ "getrandom 0.2.16",
+ "getrandom 0.3.4",
+ "js-sys",
+ "ruvector-gnn",
+ "serde",
+ "serde-wasm-bindgen",
+ "wasm-bindgen",
+ "wasm-bindgen-test",
+]
+
[[package]]
name = "ruvector-graph"
version = "0.1.1"
@@ -3764,6 +3825,7 @@ dependencies = [
"parking_lot 0.12.5",
"regex",
"ruvector-core",
+ "ruvector-graph",
"serde",
"serde-wasm-bindgen",
"serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index 5599a123e..fe6a99c84 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,6 +23,9 @@ members = [
"crates/ruvector-graph",
"crates/ruvector-graph-node",
"crates/ruvector-graph-wasm",
+ "crates/ruvector-gnn",
+ "crates/ruvector-gnn-node",
+ "crates/ruvector-gnn-wasm",
]
resolver = "2"
diff --git a/README.md b/README.md
index 2e858c188..cd5faa252 100644
--- a/README.md
+++ b/README.md
@@ -1,468 +1,277 @@
-# Ruvector
+# RuVector
[](https://opensource.org/licenses/MIT)
[](https://www.rust-lang.org)
[](https://github.com/ruvnet/ruvector)
-[](./docs/TECHNICAL_PLAN.md)
-[](./docs/TECHNICAL_PLAN.md)
-[](./docs/IMPLEMENTATION_SUMMARY.md)
-[](https://github.com/ruvnet/ruvector)
-[](https://github.com/ruvnet/ruvector)
-[](https://www.npmjs.com/package/ruvector)
-[](https://discord.gg/ruvnet)
-[](https://twitter.com/ruvnet)
-**Next-generation vector database built in Rust for extreme performance and universal deployment.**
+**The index is the neural network.** A high-performance vector database with built-in Graph Neural Networks, Neo4j-compatible hypergraph storage, and adaptive tensor compression.
-> Transform your AI applications with **sub-millisecond vector search** that scales from edge devices to **500M+ concurrent global streams**. Built by [rUv](https://ruv.io) and the open-source community at [GitHub/ruvnet](https://github.com/ruvnet).
+## What is RuVector?
-## π Why Ruvector?
+RuVector combines three powerful concepts into one unified system:
-In the age of AI, **vector similarity search is the foundation** of modern applicationsβfrom RAG systems to recommendation engines. But existing solutions force you to choose between **performance**, **scale**, or **portability**.
-
-**Ruvector eliminates that compromise.**
-
-### The rUv Advantage
-
-Developed by **[rUv](https://ruv.io)**βpioneers in agentic AI systems and high-performance distributed computingβRuvector brings enterprise-grade vector search to everyone. Whether you're building the next AI startup or scaling to billions of users, Ruvector adapts to your needs.
-
-π **Learn more**: [ruv.io](https://ruv.io) | [GitHub](https://github.com/ruvnet/ruvector)
-
-### Built for the Modern AI Stack
-
-- β‘ **Blazing Fast**: <0.5ms p50 latency with HNSW indexing and SIMD optimizations
-- π **Globally Scalable**: Deploy to 500M+ concurrent streams across 15 regions with auto-scaling
-- π― **Universal Deployment**: Run anywhereβNative Rust, Node.js, WebAssembly, browsers, edge devices
-- π° **Cost Optimized**: 60% cost reduction through intelligent caching and batching strategies
-- π§ **AI-Native**: Built specifically for embeddings, RAG, semantic search, and agent memory
-- π **Open Source**: MIT licensed, community-driven, production-ready
-
-## π Features
-
-### Core Capabilities
-
-- **Sub-Millisecond Queries**: <0.5ms p50 local latency with state-of-the-art HNSW indexing
-- **Memory Efficient**: 4-32x compression with advanced quantization techniques
-- **High Recall**: 95%+ accuracy with HNSW + Product Quantization
-- **Zero Dependencies**: Pure Rust implementation with minimal external dependencies
-- **Production Ready**: Battle-tested algorithms with comprehensive benchmarks
-- **AgenticDB Compatible**: Drop-in replacement with familiar API patterns
-
-### Global Cloud Scale β¨
-
-- **500M+ Concurrent Streams**: Baseline capacity with burst to 25B for major events
-- **15 Global Regions**: Multi-region deployment with automatic failover
-- **<10ms Global Latency**: p50 worldwide with multi-level caching
-- **99.99% Availability**: Enterprise SLA with redundancy and health monitoring
-- **Adaptive Auto-Scaling**: Predictive + reactive scaling for traffic spikes
-- **60% Cost Savings**: Optimized infrastructure reducing costs from $2.75M to $1.74M/month
-
-### Universal Platform Support
-
-| Platform | Status | Package | Use Case |
-|----------|--------|---------|----------|
-| **Rust Native** | β
Ready | `cargo add ruvector-core` | Servers, microservices, CLI tools |
-| **Node.js** | β
Ready | `npm install ruvector` | APIs, serverless, backend apps |
-| **WebAssembly** | β
Ready | `npm install ruvector-wasm` | Browsers, edge computing, offline |
-| **Cloud Run** | β
Ready | Docker + Terraform | Global scale, 500M+ streams |
-
-## π Performance Benchmarks
-
-### Local Performance (Single Instance)
+1. **Vector Database** β Sub-millisecond HNSW search with 95%+ recall
+2. **Graph Neural Network** β The HNSW topology becomes a trainable GNN
+3. **Hypergraph Storage** β Neo4j-compatible Cypher queries with N-ary relationships
```
-Metric Ruvector Pinecone Qdrant ChromaDB
-ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
-Query Latency (p50) <0.5ms ~2ms ~1ms ~50ms
-Throughput (QPS) 50K+ ~10K ~20K ~1K
-Memory (1M vectors) ~800MB ~2GB ~1.5GB ~3GB
-Recall @ k=10 95%+ 93% 94% 85%
-Browser Support β
β β β
-Offline Capable β
β β
β
+βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
+β RuVector Stack β
+βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
+β Query API β Cypher Parser β Differentiable Search β
+βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
+β GNN Layers β Message Passing β Multi-Head Attention β
+βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
+β HNSW Index β Vector Storage β Tensor Compression (2-32x) β
+βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ€
+β Rust Core β WASM Bindings β Node.js (napi-rs) β
+βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
```
-### Global Cloud Performance (500M Streams)
+## Features
-```
-Metric Value Details
-ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ
-Concurrent Streams 500M baseline Burst to 25B (50x)
-Global Latency (p50) <10ms Multi-region + CDN
-Availability 99.99% SLA 15 regions, auto-failover
-Cost per Stream/Month $0.0035 60% optimized ($1.74M total)
-Regions 15 global Americas, EMEA, APAC
-Throughput per Region 100K+ QPS Adaptive batching
-```
+| Feature | Description |
+|---------|-------------|
+| **GNN on HNSW** | Graph neural network layers that treat the index topology as a trainable graph |
+| **Cypher Queries** | Neo4j-compatible query language with `MATCH`, `WHERE`, `RETURN`, `CREATE` |
+| **Hyperedges** | N-ary relationships connecting multiple nodes (not just pairs) |
+| **Adaptive Compression** | 5-tier tensor compression: f32 β f16 β PQ8 β PQ4 β Binary (2-32x) |
+| **Differentiable Search** | Soft attention over candidates with gradient flow for end-to-end training |
+| **WASM Support** | Full browser support with WebAssembly bindings |
+| **Memory-Mapped Training** | Efficient gradient accumulation on memory-mapped embeddings |
-## β‘ Quick Start
+## Quick Start
### Installation
-**Rust:**
```bash
-cargo add ruvector-core
+# Rust
+cargo add ruvector-graph
+
+# Node.js
+npm install ruvector-gnn-node
+
+# Browser (WASM)
+npm install ruvector-gnn-wasm
```
-**Node.js:**
-```bash
-npm install ruvector
-```
+### Basic Usage
-**WebAssembly:**
-```bash
-npm install ruvector-wasm
-```
-
-### Usage Examples
-
-**Rust:**
+**Cypher Queries (Neo4j-compatible):**
```rust
-use ruvector_core::{VectorDB, Config};
+use ruvector_graph::{GraphDB, cypher::CypherExecutor};
-// Create database
-let db = VectorDB::new(Config::default())?;
+let db = GraphDB::new();
+let executor = CypherExecutor::new(&db);
-// Insert vectors
-db.insert("doc1", vec![0.1, 0.2, 0.3])?;
-db.insert("doc2", vec![0.4, 0.5, 0.6])?;
+// Create nodes and relationships
+executor.execute("CREATE (a:Person {name: 'Alice'})-[:KNOWS]->(b:Person {name: 'Bob'})")?;
-// Search similar vectors
-let results = db.search(vec![0.1, 0.2, 0.3], 10)?;
-for (id, score) in results {
- println!("{}: {}", id, score);
-}
+// Query with pattern matching
+let results = executor.execute("MATCH (p:Person)-[:KNOWS]->(friend) RETURN p.name, friend.name")?;
```
-**Node.js:**
-```javascript
-const { VectorDB } = require('ruvector');
+**GNN Forward Pass:**
+```rust
+use ruvector_gnn::{RuvectorLayer, differentiable_search};
-// Create database
-const db = new VectorDB();
+// Create GNN layer
+let layer = RuvectorLayer::new(128, 256, 4, 0.1); // input, hidden, heads, dropout
-// Insert vectors
-await db.insert('doc1', [0.1, 0.2, 0.3]);
-await db.insert('doc2', [0.4, 0.5, 0.6]);
+// Forward pass with neighbor aggregation
+let output = layer.forward(&node_embedding, &neighbor_embeddings, &edge_weights);
-// Search similar vectors
-const results = await db.search([0.1, 0.2, 0.3], 10);
-results.forEach(({ id, score }) => {
- console.log(`${id}: ${score}`);
-});
+// Differentiable search (soft k-NN)
+let (indices, weights) = differentiable_search(&query, &candidates, 10, 0.07);
```
-**WebAssembly (Browser):**
-```javascript
-import init, { VectorDB } from 'ruvector-wasm';
+**Tensor Compression:**
+```rust
+use ruvector_gnn::TensorCompress;
+
+let compressor = TensorCompress::new();
+
+// Adaptive compression based on access frequency
+let compressed = compressor.compress(&embedding, 0.5)?; // Warm data β f16
+let restored = compressor.decompress(&compressed)?;
+```
+
+### Browser Usage (WASM)
+
+```javascript
+import init, { JsRuvectorLayer, JsTensorCompress, differentiableSearch } from 'ruvector-gnn-wasm';
-// Initialize WASM module
await init();
-// Create database (runs entirely in browser!)
-const db = new VectorDB();
+// GNN layer
+const layer = new JsRuvectorLayer(128, 256, 4, 0.1);
+const output = layer.forward(nodeEmbedding, neighbors, weights);
-// Insert and search
-db.insert('doc1', new Float32Array([0.1, 0.2, 0.3]));
-const results = db.search(new Float32Array([0.1, 0.2, 0.3]), 10);
+// Compression
+const compressor = new JsTensorCompress();
+const compressed = compressor.compress(embedding, 0.5);
```
-### Global Cloud Deployment
+## Architecture
-Deploy Ruvector to handle 500M+ concurrent streams worldwide:
+### Crate Structure
+
+```
+crates/
+βββ ruvector-core/ # Vector database core (HNSW, storage)
+βββ ruvector-graph/ # Neo4j-compatible hypergraph + Cypher
+βββ ruvector-gnn/ # GNN layers, compression, training
+βββ ruvector-gnn-wasm/ # WebAssembly bindings
+βββ ruvector-gnn-node/ # Node.js bindings (napi-rs)
+βββ ruvector-graph-wasm/ # Graph WASM bindings
+βββ ruvector-graph-node/ # Graph Node.js bindings
+```
+
+### Compression Tiers
+
+| Tier | Access Freq | Format | Compression | Use Case |
+|------|------------|--------|-------------|----------|
+| Hot | >80% | f32 | 1x | Active queries |
+| Warm | 40-80% | f16 | 2x | Recent data |
+| Cool | 10-40% | PQ8 | ~8x | Older data |
+| Cold | 1-10% | PQ4 | ~16x | Archived |
+| Archive | <1% | Binary | ~32x | Rarely accessed |
+
+### GNN Message Passing
+
+The RuvectorLayer implements attention-based message passing on the HNSW graph:
+
+```
+h_new = LayerNorm(h + GRU(h, Attention(W_msg(neighbors), edge_weights)))
+```
+
+1. **Message**: Transform neighbor embeddings with learned weights
+2. **Aggregate**: Multi-head attention over messages, weighted by edge similarity
+3. **Update**: GRU cell combines current state with aggregated messages
+4. **Normalize**: Layer normalization with residual connection
+
+## Tutorial
+
+### 1. Creating a Knowledge Graph
+
+```rust
+use ruvector_graph::{GraphDB, NodeBuilder, EdgeBuilder};
+
+let db = GraphDB::new();
+
+// Create nodes
+let alice = NodeBuilder::new("alice")
+ .label("Person")
+ .property("name", "Alice")
+ .property("age", 30)
+ .build();
+
+let knows_rust = NodeBuilder::new("rust")
+ .label("Skill")
+ .property("name", "Rust")
+ .build();
+
+db.create_node(alice)?;
+db.create_node(knows_rust)?;
+
+// Create relationship
+let edge = EdgeBuilder::new("e1", "alice", "rust")
+ .edge_type("KNOWS")
+ .property("level", "expert")
+ .build();
+
+db.create_edge(edge)?;
+```
+
+### 2. Semantic Search with GNN Enhancement
+
+```rust
+use ruvector_gnn::{RuvectorLayer, RuvectorQuery, QueryMode};
+
+// Initialize GNN layer
+let gnn = RuvectorLayer::new(384, 512, 8, 0.1);
+
+// Build query
+let query = RuvectorQuery::neural_search(query_embedding, 10, 2)
+ .with_temperature(0.07);
+
+// Search with GNN-enhanced representations
+let enhanced = gnn.forward(&query.vector.unwrap(), &neighbor_embs, &weights);
+```
+
+### 3. Training with InfoNCE Loss
+
+```rust
+use ruvector_gnn::training::{info_nce_loss, sgd_step, TrainConfig};
+
+let config = TrainConfig::default();
+
+// Compute contrastive loss
+let loss = info_nce_loss(
+ &anchor_embedding,
+ &positive_embeddings,
+ &negative_embeddings,
+ config.temperature
+);
+
+// Update embeddings
+sgd_step(&mut embedding, &gradient, config.learning_rate);
+```
+
+## Documentation
+
+| Topic | Link |
+|-------|------|
+| **Getting Started** | [docs/guide/GETTING_STARTED.md](./docs/guide/GETTING_STARTED.md) |
+| **Cypher Query Language** | [docs/api/CYPHER_REFERENCE.md](./docs/api/CYPHER_REFERENCE.md) |
+| **GNN Architecture** | [docs/gnn-layer-implementation.md](./docs/gnn-layer-implementation.md) |
+| **Compression Guide** | [docs/optimization/COMPRESSION.md](./docs/optimization/COMPRESSION.md) |
+| **WASM Bindings** | [crates/ruvector-gnn-wasm/README.md](./crates/ruvector-gnn-wasm/README.md) |
+| **Node.js Bindings** | [crates/ruvector-gnn-node/README.md](./crates/ruvector-gnn-node/README.md) |
+| **API Reference** | [docs/api/](./docs/api/) |
+| **Performance Tuning** | [docs/optimization/PERFORMANCE_TUNING_GUIDE.md](./docs/optimization/PERFORMANCE_TUNING_GUIDE.md) |
+
+## Performance
+
+```
+Query Latency (p50) <0.5ms HNSW with SIMD
+GNN Forward Pass ~1ms Per-node with neighbors
+Compression (PQ8) ~8x Memory reduction
+Recall @ k=10 95%+ High accuracy search
+Browser (WASM) ~2ms Full functionality
+```
+
+## Building from Source
```bash
-# 1. Clone repository
+# Clone
git clone https://github.com/ruvnet/ruvector.git
cd ruvector
-# 2. Deploy infrastructure (Terraform)
-cd src/burst-scaling/terraform
-terraform init && terraform apply
-
-# 3. Deploy Cloud Run services (multi-region)
-cd ../cloud-run
-gcloud builds submit --config=cloudbuild.yaml
-
-# 4. Initialize agentic coordination
-cd ../agentic-integration
-npm install && npm run swarm:init
-
-# 5. Run validation tests
-cd ../../benchmarks
-npm run test:quick
-```
-
-**Deployment Time**: 4-6 hours for full global infrastructure
-**Cost**: $1.74M/month (500M streams, optimized)
-
-See [Deployment Guide](./docs/cloud-architecture/DEPLOYMENT_GUIDE.md) for complete instructions.
-
-## π― Use Cases
-
-### Local & Edge Computing
-
-- **RAG Systems**: Fast vector retrieval for Large Language Models with <0.5ms latency
-- **Semantic Search**: AI-powered similarity search for documents, images, and code
-- **Recommender Systems**: Real-time personalized recommendations on edge devices
-- **Agent Memory**: Reflexion memory and skill libraries for autonomous AI agents
-- **Code Search**: Find similar code patterns across repositories instantly
-- **Offline AI**: Run powerful vector search entirely in the browser (WebAssembly)
-
-### Global Cloud Scale
-
-- **Streaming Platforms**: 500M+ concurrent learners with real-time recommendations
-- **Live Events**: Handle 50x traffic spikes (World Cup: 25B concurrent streams)
-- **Multi-Region AI**: Global vector search with <10ms latency anywhere
-- **Enterprise RAG**: Planet-scale retrieval for distributed AI applications
-- **Real-Time Analytics**: Process billions of similarity queries per day
-- **E-Commerce**: Product recommendations at massive scale with auto-scaling
-
-## ποΈ Architecture
-
-### Project Structure
-
-Ruvector is organized as a Rust workspace with specialized crates:
-
-```
-ruvector/
-βββ crates/
-β βββ ruvector-core/ # Core vector database engine (Rust)
-β βββ ruvector-node/ # Node.js bindings via NAPI-RS
-β βββ ruvector-wasm/ # WebAssembly bindings (browser)
-β βββ ruvector-cli/ # Command-line interface
-β βββ ruvector-bench/ # Performance benchmarks
-β βββ router-core/ # Neural routing and inference
-β βββ router-cli/ # Router command-line tools
-β βββ router-ffi/ # Foreign function interface
-β βββ router-wasm/ # Router WebAssembly bindings
-βββ src/
-β βββ burst-scaling/ # Auto-scaling for traffic spikes
-β βββ cloud-run/ # Google Cloud Run deployment
-β βββ agentic-integration/ # AI agent coordination
-βββ benchmarks/ # Load testing and scenarios
-βββ docs/ # Comprehensive documentation
-```
-
-### Core Technologies
-
-- **HNSW Indexing**: Hierarchical Navigable Small World for fast approximate nearest neighbor search
-- **Product Quantization**: Memory-efficient vector compression (4-32x reduction)
-- **SIMD Optimizations**: Hardware-accelerated vector operations via simsimd
-- **Zero-Copy I/O**: Memory-mapped files for efficient data access
-- **Google Cloud Run**: Multi-region serverless deployment with auto-scaling
-- **Adaptive Batching**: Intelligent request batching for 70% latency reduction
-
-## π Documentation
-
-### Getting Started
-
-- **[Quick Start Guide](./docs/guide/GETTING_STARTED.md)** - Get up and running in 5 minutes
-- **[Installation Guide](./docs/guide/INSTALLATION.md)** - Detailed setup for all platforms
-- **[Basic Tutorial](./docs/guide/BASIC_TUTORIAL.md)** - Step-by-step vector search tutorial
-- **[AgenticDB Quick Start](./docs/getting-started/AGENTICDB_QUICKSTART.md)** - Migration from AgenticDB
-
-### API Documentation
-
-- **[Rust API Reference](./docs/api/RUST_API.md)** - Complete Rust API documentation
-- **[Node.js API Reference](./docs/api/NODEJS_API.md)** - JavaScript/TypeScript API
-- **[WebAssembly API](./docs/getting-started/wasm-api.md)** - Browser and edge usage
-- **[AgenticDB API](./docs/getting-started/AGENTICDB_API.md)** - AgenticDB compatibility layer
-
-### Advanced Topics
-
-- **[Advanced Features](./docs/guide/ADVANCED_FEATURES.md)** - Quantization, indexing, optimization
-- **[Performance Tuning](./docs/optimization/PERFORMANCE_TUNING_GUIDE.md)** - Optimize for your workload
-- **[Optimization Guide](./docs/getting-started/OPTIMIZATION_QUICK_START.md)** - Best practices
-- **[Build Optimization](./docs/optimization/BUILD_OPTIMIZATION.md)** - Compile-time optimizations
-
-### Cloud Deployment
-
-- **[Implementation Summary](./docs/IMPLEMENTATION_SUMMARY.md)** - Complete overview of global deployment
-- **[Architecture Overview](./docs/cloud-architecture/architecture-overview.md)** - 15-region global design
-- **[Deployment Guide](./docs/cloud-architecture/DEPLOYMENT_GUIDE.md)** - Step-by-step setup (4-6 hours)
-- **[Scaling Strategy](./docs/cloud-architecture/scaling-strategy.md)** - Auto-scaling & burst handling
-- **[Performance Optimization](./docs/cloud-architecture/PERFORMANCE_OPTIMIZATION_GUIDE.md)** - 70% latency reduction
-- **[Cost Optimization](./src/cloud-run/COST_OPTIMIZATIONS.md)** - 60% cost savings ($3.66M/year)
-- **[Load Testing](./benchmarks/LOAD_TEST_SCENARIOS.md)** - World Cup and burst scenarios
-
-### Development
-
-- **[Contributing Guidelines](./docs/development/CONTRIBUTING.md)** - How to contribute
-- **[Development Guide](./docs/development/MIGRATION.md)** - Development setup
-- **[Benchmarking Guide](./docs/benchmarks/BENCHMARKING_GUIDE.md)** - Run performance tests
-- **[Technical Plan](./docs/TECHNICAL_PLAN.md)** - Architecture and design decisions
-
-### Complete Index
-
-- **[Documentation Index](./docs/README.md)** - Complete documentation organization
-- **[Changelog](./CHANGELOG.md)** - Version history and updates
-
-## π¨ Building from Source
-
-### Prerequisites
-
-- **Rust**: 1.77 or higher
-- **Node.js**: 18.0 or higher (for Node.js/WASM builds)
-- **wasm-pack**: For WebAssembly builds
-
-### Build Commands
-
-```bash
-# Build all Rust crates (release mode)
+# Build all crates
cargo build --release
# Run tests
cargo test --workspace
-# Run benchmarks
-cargo bench --workspace
-
-# Build Node.js bindings
-cd crates/ruvector-node
-npm install && npm run build
-
-# Build WebAssembly
-cd crates/ruvector-wasm
-wasm-pack build --target web
-
-# Run CLI
-cargo run -p ruvector-cli -- --help
+# Build WASM
+cargo build --package ruvector-gnn-wasm --target wasm32-unknown-unknown
```
-### Development Workflow
+## Contributing
-```bash
-# Format code
-cargo fmt --all
+We welcome contributions! See [CONTRIBUTING.md](./docs/development/CONTRIBUTING.md).
-# Lint code
-cargo clippy --workspace -- -D warnings
+## License
-# Type check
-cargo check --workspace
-
-# Run specific tests
-cargo test -p ruvector-core
-
-# Run benchmarks with specific features
-cargo bench -p ruvector-bench --features simd
-```
-
-## π€ Contributing
-
-We welcome contributions from the community! Ruvector is built by developers, for developers.
-
-### How to Contribute
-
-1. **Fork** the repository at [github.com/ruvnet/ruvector](https://github.com/ruvnet/ruvector)
-2. **Create** a feature branch (`git checkout -b feature/amazing-feature`)
-3. **Commit** your changes (`git commit -m 'Add amazing feature'`)
-4. **Push** to the branch (`git push origin feature/amazing-feature`)
-5. **Open** a Pull Request
-
-### Contribution Areas
-
-- π **Bug Fixes**: Help us squash bugs and improve stability
-- β¨ **New Features**: Add new capabilities and integrations
-- π **Documentation**: Improve guides, tutorials, and API docs
-- π§ͺ **Testing**: Add test coverage and benchmarks
-- π **Translations**: Translate documentation to other languages
-- π‘ **Ideas**: Propose new features and improvements
-
-See [Contributing Guidelines](./docs/development/CONTRIBUTING.md) for detailed instructions.
-
-## π Community & Support
-
-### Connect with Us
-
-- **GitHub**: [github.com/ruvnet/ruvector](https://github.com/ruvnet/ruvector) - Star β and follow for updates
-- **Discord**: [Join our community](https://discord.gg/ruvnet) - Chat with developers and users
-- **Twitter**: [@ruvnet](https://twitter.com/ruvnet) - Follow for announcements and tips
-- **Website**: [ruv.io](https://ruv.io) - Learn about rUv's AI platform and tools
-- **Issues**: [GitHub Issues](https://github.com/ruvnet/ruvector/issues) - Report bugs and request features
-- **Discussions**: [GitHub Discussions](https://github.com/ruvnet/ruvector/discussions) - Ask questions and share ideas
-
-### Enterprise Support
-
-Need enterprise support, custom development, or consulting services?
-
-π§ Contact us at [enterprise@ruv.io](mailto:enterprise@ruv.io)
-
-## π Comparison with Alternatives
-
-| Feature | Ruvector | Pinecone | Qdrant | ChromaDB | Milvus |
-|---------|----------|----------|--------|----------|--------|
-| **Language** | Rust | ? | Rust | Python | C++/Go |
-| **Local Latency (p50)** | <0.5ms | ~2ms | ~1ms | ~50ms | ~5ms |
-| **Global Scale** | 500M+ β¨ | Limited | Limited | No | Limited |
-| **Browser Support** | β
WASM | β | β | β | β |
-| **Offline Capable** | β
| β | β
| β
| β
|
-| **NPM Package** | β
| β
| β | β
| β |
-| **Native Binary** | β
| β | β
| β | β
|
-| **Burst Capacity** | 50x β¨ | Unknown | Unknown | No | Unknown |
-| **Open Source** | β
MIT | β | β
Apache | β
Apache | β
Apache |
-| **Cost (500M)** | $1.74M/mo | $$$$ | $$$ | Self-host | Self-host |
-| **Edge Deployment** | β
| β | Partial | Partial | β |
-
-## π― Latest Updates
-
-### v0.1.0 - Global Streaming Optimization β¨
-
-Complete implementation for massive-scale deployment:
-
-- β
**Architecture**: 15-region global topology with 99.99% SLA
-- β
**Cloud Run Service**: HTTP/2 + WebSocket with adaptive batching (70% latency improvement)
-- β
**Agentic Coordination**: Distributed agent swarm with auto-scaling (6 files, 3,550 lines)
-- β
**Burst Scaling**: Predictive + reactive scaling for 50x spikes (11 files, 4,844 lines)
-- β
**Benchmarking**: Comprehensive test suite supporting 25B concurrent (13 files, 4,582 lines)
-- β
**Cost Optimization**: 60% reduction through caching/batching ($3.66M/year savings)
-- β
**Query Optimization**: 5x throughput increase, 70% latency reduction
-- β
**Production-Ready**: 45+ files, 28,000+ lines of tested code
-
-**Deployment Time**: 4-6 hours for full global infrastructure
-**Cost**: $2.75M/month baseline β **$1.74M with optimizations (60% savings)**
-**Capacity**: 500M concurrent β 25B burst (50x for major events)
-
-See [Implementation Summary](./docs/IMPLEMENTATION_SUMMARY.md) for complete details.
-
-## π License
-
-**MIT License** - see [LICENSE](./LICENSE) for details.
-
-Free to use for commercial and personal projects. We believe in open source.
-
-## π Acknowledgments
-
-Built with battle-tested algorithms and technologies:
-
-- **HNSW**: Hierarchical Navigable Small World graphs
-- **Product Quantization**: Efficient vector compression
-- **simsimd**: SIMD-accelerated similarity computations
-- **Google Cloud Run**: Serverless multi-region deployment
-- **Advanced Caching**: Multi-level caching strategies
-- **Community Contributors**: Thank you to all our contributors! π
-
-### Special Thanks
-
-- The Rust community for incredible tooling and ecosystem
-- Contributors to HNSW, quantization research, and SIMD libraries
-- Our users and beta testers for valuable feedback
-- The [rUv](https://ruv.io) team for making this possible
+MIT License - see [LICENSE](./LICENSE).
---
-**Built by [rUv](https://ruv.io) β’ Open Source on [GitHub](https://github.com/ruvnet/ruvector) β’ Production Ready**
+**Built by [rUv](https://ruv.io)** β’ [GitHub](https://github.com/ruvnet/ruvector) β’ [Documentation](./docs/)
-[](https://github.com/ruvnet/ruvector)
-[](https://twitter.com/ruvnet)
-[](https://discord.gg/ruvnet)
-
-**Status**: Production Ready | **Version**: 0.1.0 | **Scale**: Local to 500M+ concurrent
-
-**Ready for**: World Cup (25B concurrent) β’ Olympics β’ Product Launches β’ Streaming Platforms
-
-[Get Started](./docs/guide/GETTING_STARTED.md) β’ [Documentation](./docs/README.md) β’ [API Reference](./docs/api/RUST_API.md) β’ [Contributing](./docs/development/CONTRIBUTING.md)
+*"The index is a sparse neural network whose topology encodes learned similarity."*
diff --git a/crates/ruvector-core/Cargo.toml b/crates/ruvector-core/Cargo.toml
index 945ebe7c5..89b2c2bc2 100644
--- a/crates/ruvector-core/Cargo.toml
+++ b/crates/ruvector-core/Cargo.toml
@@ -40,7 +40,7 @@ once_cell = { workspace = true }
# Time and UUID
chrono = { workspace = true }
-uuid = { workspace = true, optional = true }
+uuid = { workspace = true, features = ["v4"] }
[dev-dependencies]
criterion = { workspace = true }
@@ -70,12 +70,12 @@ name = "comprehensive_bench"
harness = false
[features]
-default = ["simd", "uuid-support", "storage", "hnsw"]
-uuid-support = ["uuid"]
+default = ["simd", "storage", "hnsw"]
simd = []
storage = ["redb", "memmap2"] # File-based storage (not available in WASM)
hnsw = ["hnsw_rs"] # HNSW indexing (not available in WASM due to mmap dependency)
memory-only = [] # Pure in-memory storage for WASM
+uuid-support = [] # Deprecated: uuid is now always included
[lib]
crate-type = ["rlib"]
diff --git a/crates/ruvector-core/src/error.rs b/crates/ruvector-core/src/error.rs
index 079ea2362..f317a187e 100644
--- a/crates/ruvector-core/src/error.rs
+++ b/crates/ruvector-core/src/error.rs
@@ -54,36 +54,42 @@ pub enum RuvectorError {
Internal(String),
}
+#[cfg(feature = "storage")]
impl From for RuvectorError {
fn from(err: redb::Error) -> Self {
RuvectorError::DatabaseError(err.to_string())
}
}
+#[cfg(feature = "storage")]
impl From for RuvectorError {
fn from(err: redb::DatabaseError) -> Self {
RuvectorError::DatabaseError(err.to_string())
}
}
+#[cfg(feature = "storage")]
impl From for RuvectorError {
fn from(err: redb::StorageError) -> Self {
RuvectorError::DatabaseError(err.to_string())
}
}
+#[cfg(feature = "storage")]
impl From for RuvectorError {
fn from(err: redb::TableError) -> Self {
RuvectorError::DatabaseError(err.to_string())
}
}
+#[cfg(feature = "storage")]
impl From for RuvectorError {
fn from(err: redb::TransactionError) -> Self {
RuvectorError::DatabaseError(err.to_string())
}
}
+#[cfg(feature = "storage")]
impl From for RuvectorError {
fn from(err: redb::CommitError) -> Self {
RuvectorError::DatabaseError(err.to_string())
diff --git a/crates/ruvector-core/src/index.rs b/crates/ruvector-core/src/index.rs
index a3c1ec8a9..34374e39d 100644
--- a/crates/ruvector-core/src/index.rs
+++ b/crates/ruvector-core/src/index.rs
@@ -1,5 +1,6 @@
//! Index structures for efficient vector search
+#[cfg(feature = "hnsw")]
pub mod hnsw;
pub mod flat;
diff --git a/crates/ruvector-core/src/storage_memory.rs b/crates/ruvector-core/src/storage_memory.rs
index 1faaef782..63a9a72f0 100644
--- a/crates/ruvector-core/src/storage_memory.rs
+++ b/crates/ruvector-core/src/storage_memory.rs
@@ -50,7 +50,9 @@ impl MemoryStorage {
// Insert metadata if present
if let Some(metadata) = &entry.metadata {
- self.metadata.insert(id.clone(), metadata.clone());
+ self.metadata.insert(id.clone(), serde_json::Value::Object(
+ metadata.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
+ ));
}
Ok(id)
@@ -73,7 +75,9 @@ impl MemoryStorage {
self.vectors.insert(id.clone(), entry.vector.clone());
if let Some(metadata) = &entry.metadata {
- self.metadata.insert(id.clone(), metadata.clone());
+ self.metadata.insert(id.clone(), serde_json::Value::Object(
+ metadata.iter().map(|(k, v)| (k.clone(), v.clone())).collect()
+ ));
}
ids.push(id);
@@ -86,7 +90,13 @@ impl MemoryStorage {
pub fn get(&self, id: &str) -> Result> {
if let Some(vector_ref) = self.vectors.get(id) {
let vector = vector_ref.clone();
- let metadata = self.metadata.get(id).map(|m| m.clone());
+ let metadata = self.metadata.get(id).and_then(|m| {
+ if let serde_json::Value::Object(map) = m.value() {
+ Some(map.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
+ } else {
+ None
+ }
+ });
Ok(Some(VectorEntry {
id: Some(id.to_string()),
diff --git a/crates/ruvector-gnn-node/.github/workflows/build.yml b/crates/ruvector-gnn-node/.github/workflows/build.yml
new file mode 100644
index 000000000..901779135
--- /dev/null
+++ b/crates/ruvector-gnn-node/.github/workflows/build.yml
@@ -0,0 +1,110 @@
+name: Build and Test
+
+on:
+ push:
+ branches: [main, develop]
+ pull_request:
+ branches: [main, develop]
+
+jobs:
+ build:
+ strategy:
+ fail-fast: false
+ matrix:
+ settings:
+ - host: macos-latest
+ target: x86_64-apple-darwin
+ build: npm run build
+ - host: macos-latest
+ target: aarch64-apple-darwin
+ build: npm run build -- --target aarch64-apple-darwin
+ - host: ubuntu-latest
+ target: x86_64-unknown-linux-gnu
+ build: npm run build
+ - host: ubuntu-latest
+ target: x86_64-unknown-linux-musl
+ build: npm run build -- --target x86_64-unknown-linux-musl
+ - host: ubuntu-latest
+ target: aarch64-unknown-linux-gnu
+ build: npm run build -- --target aarch64-unknown-linux-gnu
+ - host: ubuntu-latest
+ target: aarch64-unknown-linux-musl
+ build: npm run build -- --target aarch64-unknown-linux-musl
+ - host: windows-latest
+ target: x86_64-pc-windows-msvc
+ build: npm run build
+
+ name: Build ${{ matrix.settings.target }}
+ runs-on: ${{ matrix.settings.host }}
+
+ steps:
+ - uses: actions/checkout@v4
+
+ - name: Setup Node.js
+ uses: actions/setup-node@v4
+ with:
+ node-version: 18
+ cache: npm
+ cache-dependency-path: crates/ruvector-gnn-node/package.json
+
+ - name: Install Rust
+ uses: dtolnay/rust-toolchain@stable
+ with:
+ targets: ${{ matrix.settings.target }}
+
+ - name: Cache cargo
+ uses: actions/cache@v4
+ with:
+ path: |
+ ~/.cargo/registry/index/
+ ~/.cargo/registry/cache/
+ ~/.cargo/git/db/
+ target/
+ key: ${{ runner.os }}-cargo-${{ matrix.settings.target }}-${{ hashFiles('**/Cargo.lock') }}
+
+ - name: Install dependencies
+ working-directory: crates/ruvector-gnn-node
+ run: npm install
+
+ - name: Build
+ working-directory: crates/ruvector-gnn-node
+ run: ${{ matrix.settings.build }}
+
+ - name: Test (non-cross compile only)
+ if: matrix.settings.host == 'ubuntu-latest' && matrix.settings.target == 'x86_64-unknown-linux-gnu'
+ working-directory: crates/ruvector-gnn-node
+ run: npm test
+
+ - name: Upload artifacts
+ uses: actions/upload-artifact@v4
+ with:
+ name: bindings-${{ matrix.settings.target }}
+ path: crates/ruvector-gnn-node/*.node
+ if-no-files-found: error
+
+ test:
+ name: Test Node.js bindings
+ runs-on: ubuntu-latest
+ needs: build
+
+ steps:
+ - uses: actions/checkout@v4
+
+ - name: Setup Node.js
+ uses: actions/setup-node@v4
+ with:
+ node-version: 18
+
+ - name: Download artifacts
+ uses: actions/download-artifact@v4
+ with:
+ name: bindings-x86_64-unknown-linux-gnu
+ path: crates/ruvector-gnn-node
+
+ - name: Install dependencies
+ working-directory: crates/ruvector-gnn-node
+ run: npm install --ignore-scripts
+
+ - name: Run tests
+ working-directory: crates/ruvector-gnn-node
+ run: npm test
diff --git a/crates/ruvector-gnn-node/.npmignore b/crates/ruvector-gnn-node/.npmignore
new file mode 100644
index 000000000..9d2874528
--- /dev/null
+++ b/crates/ruvector-gnn-node/.npmignore
@@ -0,0 +1,13 @@
+target/
+Cargo.lock
+.cargo/
+*.node
+*.iml
+.idea/
+.vscode/
+.DS_Store
+*.swp
+*.swo
+*~
+.#*
+\#*#
diff --git a/crates/ruvector-gnn-node/Cargo.toml b/crates/ruvector-gnn-node/Cargo.toml
new file mode 100644
index 000000000..dc0550256
--- /dev/null
+++ b/crates/ruvector-gnn-node/Cargo.toml
@@ -0,0 +1,25 @@
+[package]
+name = "ruvector-gnn-node"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+authors.workspace = true
+repository.workspace = true
+description = "Node.js bindings for Ruvector GNN via NAPI-RS"
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+napi = { workspace = true }
+napi-derive = { workspace = true }
+ruvector-gnn = { path = "../ruvector-gnn", default-features = false }
+serde_json = { workspace = true }
+
+[build-dependencies]
+napi-build = "2"
+
+[profile.release]
+lto = true
+strip = true
diff --git a/crates/ruvector-gnn-node/README.md b/crates/ruvector-gnn-node/README.md
new file mode 100644
index 000000000..74a2b9736
--- /dev/null
+++ b/crates/ruvector-gnn-node/README.md
@@ -0,0 +1,249 @@
+# @ruvector/gnn - Graph Neural Network Node.js Bindings
+
+High-performance Graph Neural Network (GNN) capabilities for Ruvector, powered by Rust and NAPI-RS.
+
+## Features
+
+- **GNN Layers**: Multi-head attention, layer normalization, GRU cells
+- **Tensor Compression**: Adaptive compression with 5 levels (None, Half, PQ8, PQ4, Binary)
+- **Differentiable Search**: Soft attention-based search with temperature scaling
+- **Hierarchical Processing**: Multi-layer GNN forward pass
+- **Zero-copy**: Efficient data transfer between JavaScript and Rust
+- **TypeScript Support**: Full type definitions included
+
+## Installation
+
+```bash
+npm install @ruvector/gnn
+```
+
+## Quick Start
+
+### Creating a GNN Layer
+
+```javascript
+const { RuvectorLayer } = require('@ruvector/gnn');
+
+// Create a GNN layer with:
+// - Input dimension: 128
+// - Hidden dimension: 256
+// - Attention heads: 4
+// - Dropout rate: 0.1
+const layer = new RuvectorLayer(128, 256, 4, 0.1);
+
+// Forward pass
+const nodeEmbedding = new Array(128).fill(0).map(() => Math.random());
+const neighborEmbeddings = [
+ new Array(128).fill(0).map(() => Math.random()),
+ new Array(128).fill(0).map(() => Math.random()),
+];
+const edgeWeights = [0.7, 0.3];
+
+const output = layer.forward(nodeEmbedding, neighborEmbeddings, edgeWeights);
+console.log('Output dimension:', output.length); // 256
+```
+
+### Tensor Compression
+
+```javascript
+const { TensorCompress, getCompressionLevel } = require('@ruvector/gnn');
+
+const compressor = new TensorCompress();
+const embedding = new Array(128).fill(0).map(() => Math.random());
+
+// Adaptive compression based on access frequency
+const accessFreq = 0.5; // 50% access rate
+console.log('Selected level:', getCompressionLevel(accessFreq)); // "half"
+
+const compressed = compressor.compress(embedding, accessFreq);
+const decompressed = compressor.decompress(compressed);
+
+console.log('Original size:', embedding.length);
+console.log('Compression ratio:', compressed.length / JSON.stringify(embedding).length);
+
+// Explicit compression level
+const level = {
+ level_type: 'pq8',
+ subvectors: 8,
+ centroids: 16
+};
+const compressedPQ = compressor.compressWithLevel(embedding, level);
+```
+
+### Differentiable Search
+
+```javascript
+const { differentiableSearch } = require('@ruvector/gnn');
+
+const query = [1.0, 0.0, 0.0];
+const candidates = [
+ [1.0, 0.0, 0.0], // Perfect match
+ [0.9, 0.1, 0.0], // Close match
+ [0.0, 1.0, 0.0], // Orthogonal
+];
+
+const result = differentiableSearch(query, candidates, 2, 1.0);
+console.log('Top-2 indices:', result.indices); // [0, 1]
+console.log('Soft weights:', result.weights); // [0.x, 0.y]
+```
+
+### Hierarchical Forward Pass
+
+```javascript
+const { hierarchicalForward, RuvectorLayer } = require('@ruvector/gnn');
+
+const query = [1.0, 0.0];
+
+// Layer embeddings (organized by HNSW layers)
+const layerEmbeddings = [
+ [[1.0, 0.0], [0.0, 1.0]], // Layer 0 embeddings
+];
+
+// Create and serialize GNN layers
+const layer1 = new RuvectorLayer(2, 2, 1, 0.0);
+const layers = [layer1.toJson()];
+
+// Hierarchical processing
+const result = hierarchicalForward(query, layerEmbeddings, layers);
+console.log('Final embedding:', result);
+```
+
+## API Reference
+
+### RuvectorLayer
+
+#### Constructor
+
+```typescript
+new RuvectorLayer(
+ inputDim: number,
+ hiddenDim: number,
+ heads: number,
+ dropout: number
+): RuvectorLayer
+```
+
+#### Methods
+
+- `forward(nodeEmbedding: number[], neighborEmbeddings: number[][], edgeWeights: number[]): number[]`
+- `toJson(): string` - Serialize layer to JSON
+- `fromJson(json: string): RuvectorLayer` - Deserialize layer from JSON
+
+### TensorCompress
+
+#### Constructor
+
+```typescript
+new TensorCompress(): TensorCompress
+```
+
+#### Methods
+
+- `compress(embedding: number[], accessFreq: number): string` - Adaptive compression
+- `compressWithLevel(embedding: number[], level: CompressionLevelConfig): string` - Explicit level
+- `decompress(compressedJson: string): number[]` - Decompress tensor
+
+#### CompressionLevelConfig
+
+```typescript
+interface CompressionLevelConfig {
+ level_type: 'none' | 'half' | 'pq8' | 'pq4' | 'binary';
+ scale?: number; // For 'half'
+ subvectors?: number; // For 'pq8', 'pq4'
+ centroids?: number; // For 'pq8'
+ outlier_threshold?: number; // For 'pq4'
+ threshold?: number; // For 'binary'
+}
+```
+
+### Search Functions
+
+#### differentiableSearch
+
+```typescript
+function differentiableSearch(
+ query: number[],
+ candidateEmbeddings: number[][],
+ k: number,
+ temperature: number
+): { indices: number[], weights: number[] }
+```
+
+#### hierarchicalForward
+
+```typescript
+function hierarchicalForward(
+ query: number[],
+ layerEmbeddings: number[][][],
+ gnnLayersJson: string[]
+): number[]
+```
+
+### Utility Functions
+
+#### getCompressionLevel
+
+```typescript
+function getCompressionLevel(accessFreq: number): string
+```
+
+Returns the compression level that would be selected for the given access frequency:
+- `accessFreq > 0.8`: "none" (hot data)
+- `accessFreq > 0.4`: "half" (warm data)
+- `accessFreq > 0.1`: "pq8" (cool data)
+- `accessFreq > 0.01`: "pq4" (cold data)
+- `accessFreq <= 0.01`: "binary" (archive)
+
+## Compression Levels
+
+### None
+Full precision, no compression. Best for frequently accessed data.
+
+### Half Precision
+~50% space savings with minimal quality loss. Good for warm data.
+
+### PQ8 (8-bit Product Quantization)
+~8x compression using 8-bit codes. Suitable for cool data.
+
+### PQ4 (4-bit Product Quantization)
+~16x compression with outlier handling. For cold data.
+
+### Binary
+~32x compression, values become +1/-1. For archival data.
+
+## Performance
+
+- **Zero-copy operations** where possible
+- **SIMD optimizations** for vector operations
+- **Parallel processing** with Rayon
+- **Native performance** with Rust backend
+
+## Building from Source
+
+```bash
+# Install dependencies
+npm install
+
+# Build debug
+npm run build:debug
+
+# Build release
+npm run build
+
+# Run tests
+npm test
+```
+
+## License
+
+MIT - See LICENSE file for details
+
+## Contributing
+
+Contributions are welcome! Please see the main Ruvector repository for guidelines.
+
+## Links
+
+- [GitHub Repository](https://github.com/ruvnet/ruvector)
+- [Documentation](https://docs.ruvector.io)
+- [Issues](https://github.com/ruvnet/ruvector/issues)
diff --git a/crates/ruvector-gnn-node/build.rs b/crates/ruvector-gnn-node/build.rs
new file mode 100644
index 000000000..9fc236788
--- /dev/null
+++ b/crates/ruvector-gnn-node/build.rs
@@ -0,0 +1,5 @@
+extern crate napi_build;
+
+fn main() {
+ napi_build::setup();
+}
diff --git a/crates/ruvector-gnn-node/examples/basic.js b/crates/ruvector-gnn-node/examples/basic.js
new file mode 100644
index 000000000..381f59927
--- /dev/null
+++ b/crates/ruvector-gnn-node/examples/basic.js
@@ -0,0 +1,132 @@
+// Example: Basic usage of Ruvector GNN Node.js bindings
+
+const {
+ RuvectorLayer,
+ TensorCompress,
+ differentiableSearch,
+ hierarchicalForward,
+ getCompressionLevel,
+ init
+} = require('../index.js');
+
+console.log(init());
+console.log('');
+
+// ==================== Example 1: GNN Layer ====================
+console.log('=== Example 1: GNN Layer ===');
+
+const layer = new RuvectorLayer(4, 8, 2, 0.1);
+console.log('Created GNN layer (input_dim: 4, hidden_dim: 8, heads: 2, dropout: 0.1)');
+
+const nodeEmbedding = [1.0, 2.0, 3.0, 4.0];
+const neighborEmbeddings = [
+ [0.5, 1.0, 1.5, 2.0],
+ [2.0, 3.0, 4.0, 5.0],
+];
+const edgeWeights = [0.3, 0.7];
+
+const output = layer.forward(nodeEmbedding, neighborEmbeddings, edgeWeights);
+console.log('Input embedding:', nodeEmbedding);
+console.log('Output embedding (length):', output.length);
+console.log('Output embedding (first 4 values):', output.slice(0, 4).map(x => x.toFixed(4)));
+console.log('');
+
+// ==================== Example 2: Tensor Compression ====================
+console.log('=== Example 2: Tensor Compression ===');
+
+const compressor = new TensorCompress();
+const embedding = Array.from({ length: 64 }, (_, i) => Math.sin(i * 0.1));
+
+// Test different access frequencies
+const frequencies = [0.9, 0.5, 0.2, 0.05, 0.001];
+
+frequencies.forEach(freq => {
+ const level = getCompressionLevel(freq);
+ const compressed = compressor.compress(embedding, freq);
+ const decompressed = compressor.decompress(compressed);
+
+ const originalSize = JSON.stringify(embedding).length;
+ const compressedSize = compressed.length;
+ const ratio = (compressedSize / originalSize * 100).toFixed(1);
+
+ console.log(`Frequency: ${freq.toFixed(3)} | Level: ${level.padEnd(6)} | Size: ${ratio}% | Error: ${calculateMSE(embedding, decompressed).toFixed(6)}`);
+});
+console.log('');
+
+// ==================== Example 3: Differentiable Search ====================
+console.log('=== Example 3: Differentiable Search ===');
+
+const query = [1.0, 0.0, 0.0];
+const candidates = [
+ [1.0, 0.0, 0.0], // Perfect match
+ [0.9, 0.1, 0.0], // Close match
+ [0.7, 0.3, 0.0], // Medium match
+ [0.0, 1.0, 0.0], // Orthogonal
+ [0.0, 0.0, 1.0], // Orthogonal
+];
+
+console.log('Query:', query);
+console.log('Number of candidates:', candidates.length);
+
+const result = differentiableSearch(query, candidates, 3, 1.0);
+console.log('Top-3 indices:', result.indices);
+console.log('Soft weights:', result.weights.map(w => w.toFixed(4)));
+console.log('Weights sum:', result.weights.reduce((a, b) => a + b, 0).toFixed(4));
+console.log('');
+
+// ==================== Example 4: Hierarchical Forward ====================
+console.log('=== Example 4: Hierarchical Forward ===');
+
+const query2 = [1.0, 0.0];
+const layerEmbeddings = [
+ [
+ [1.0, 0.0],
+ [0.0, 1.0],
+ [0.7, 0.7],
+ ],
+];
+
+const layer1 = new RuvectorLayer(2, 2, 1, 0.0);
+const layers = [layer1.toJson()];
+
+const finalEmbedding = hierarchicalForward(query2, layerEmbeddings, layers);
+console.log('Query:', query2);
+console.log('Final embedding:', finalEmbedding.map(x => x.toFixed(4)));
+console.log('');
+
+// ==================== Example 5: Layer Serialization ====================
+console.log('=== Example 5: Layer Serialization ===');
+
+const originalLayer = new RuvectorLayer(8, 16, 4, 0.2);
+const serialized = originalLayer.toJson();
+const deserialized = RuvectorLayer.fromJson(serialized);
+
+console.log('Original layer created (8 -> 16, heads: 4, dropout: 0.2)');
+console.log('Serialized size:', serialized.length, 'bytes');
+console.log('Successfully deserialized');
+
+// Test that deserialized layer works
+const testInput = Array.from({ length: 8 }, () => Math.random());
+const testNeighbors = [Array.from({ length: 8 }, () => Math.random())];
+const testWeights = [1.0];
+
+const output1 = originalLayer.forward(testInput, testNeighbors, testWeights);
+const output2 = deserialized.forward(testInput, testNeighbors, testWeights);
+
+console.log('Original output matches deserialized:', arraysEqual(output1, output2, 1e-6));
+console.log('');
+
+// ==================== Helper Functions ====================
+
+function calculateMSE(a, b) {
+ if (a.length !== b.length) return Infinity;
+ const sum = a.reduce((acc, val, i) => acc + Math.pow(val - b[i], 2), 0);
+ return sum / a.length;
+}
+
+function arraysEqual(a, b, epsilon = 1e-10) {
+ if (a.length !== b.length) return false;
+ return a.every((val, i) => Math.abs(val - b[i]) < epsilon);
+}
+
+console.log('All examples completed successfully!');
diff --git a/crates/ruvector-gnn-node/package.json b/crates/ruvector-gnn-node/package.json
new file mode 100644
index 000000000..640f95ced
--- /dev/null
+++ b/crates/ruvector-gnn-node/package.json
@@ -0,0 +1,53 @@
+{
+ "name": "@ruvector/gnn",
+ "version": "0.1.1",
+ "description": "Graph Neural Network capabilities for Ruvector - Node.js bindings",
+ "main": "index.js",
+ "types": "index.d.ts",
+ "napi": {
+ "name": "ruvector-gnn",
+ "triples": {
+ "defaults": true,
+ "additional": [
+ "x86_64-unknown-linux-musl",
+ "aarch64-unknown-linux-gnu",
+ "aarch64-unknown-linux-musl",
+ "aarch64-apple-darwin",
+ "x86_64-pc-windows-msvc"
+ ]
+ }
+ },
+ "scripts": {
+ "artifacts": "napi artifacts",
+ "build": "napi build --platform --release",
+ "build:debug": "napi build --platform",
+ "prepublishOnly": "napi prepublish -t npm",
+ "test": "node --test test/*.test.js",
+ "version": "napi version"
+ },
+ "keywords": [
+ "ruvector",
+ "gnn",
+ "graph-neural-network",
+ "machine-learning",
+ "vector-database",
+ "hnsw",
+ "napi-rs"
+ ],
+ "author": "Ruvector Team",
+ "license": "MIT",
+ "repository": {
+ "type": "git",
+ "url": "https://github.com/ruvnet/ruvector"
+ },
+ "devDependencies": {
+ "@napi-rs/cli": "^2.16.0"
+ },
+ "engines": {
+ "node": ">= 10"
+ },
+ "publishConfig": {
+ "registry": "https://registry.npmjs.org/",
+ "access": "public"
+ }
+}
diff --git a/crates/ruvector-gnn-node/src/lib.rs b/crates/ruvector-gnn-node/src/lib.rs
new file mode 100644
index 000000000..84591f427
--- /dev/null
+++ b/crates/ruvector-gnn-node/src/lib.rs
@@ -0,0 +1,400 @@
+//! Node.js bindings for Ruvector GNN via NAPI-RS
+//!
+//! This module provides JavaScript bindings for the Ruvector GNN library,
+//! enabling graph neural network operations, tensor compression, and
+//! differentiable search in Node.js applications.
+
+#![deny(clippy::all)]
+
+use napi::bindgen_prelude::*;
+use napi_derive::napi;
+use ruvector_gnn::{
+ compress::{CompressionLevel as RustCompressionLevel, CompressedTensor as RustCompressedTensor, TensorCompress as RustTensorCompress},
+ layer::RuvectorLayer as RustRuvectorLayer,
+ search::{differentiable_search as rust_differentiable_search, hierarchical_forward as rust_hierarchical_forward},
+};
+
+// ==================== RuvectorLayer Bindings ====================
+
+/// Graph Neural Network layer for HNSW topology
+#[napi]
+pub struct RuvectorLayer {
+ inner: RustRuvectorLayer,
+}
+
+#[napi]
+impl RuvectorLayer {
+ /// Create a new Ruvector GNN layer
+ ///
+ /// # Arguments
+ /// * `input_dim` - Dimension of input node embeddings
+ /// * `hidden_dim` - Dimension of hidden representations
+ /// * `heads` - Number of attention heads
+ /// * `dropout` - Dropout rate (0.0 to 1.0)
+ ///
+ /// # Example
+ /// ```javascript
+ /// const layer = new RuvectorLayer(128, 256, 4, 0.1);
+ /// ```
+ #[napi(constructor)]
+ pub fn new(input_dim: u32, hidden_dim: u32, heads: u32, dropout: f64) -> Result {
+ if dropout < 0.0 || dropout > 1.0 {
+ return Err(Error::new(
+ Status::InvalidArg,
+ "Dropout must be between 0.0 and 1.0".to_string(),
+ ));
+ }
+
+ Ok(Self {
+ inner: RustRuvectorLayer::new(
+ input_dim as usize,
+ hidden_dim as usize,
+ heads as usize,
+ dropout as f32,
+ ),
+ })
+ }
+
+ /// Forward pass through the GNN layer
+ ///
+ /// # Arguments
+ /// * `node_embedding` - Current node's embedding
+ /// * `neighbor_embeddings` - Embeddings of neighbor nodes
+ /// * `edge_weights` - Weights of edges to neighbors
+ ///
+ /// # Returns
+ /// Updated node embedding
+ ///
+ /// # Example
+ /// ```javascript
+ /// const node = [1.0, 2.0, 3.0, 4.0];
+ /// const neighbors = [[0.5, 1.0, 1.5, 2.0], [2.0, 3.0, 4.0, 5.0]];
+ /// const weights = [0.3, 0.7];
+ /// const output = layer.forward(node, neighbors, weights);
+ /// ```
+ #[napi]
+ pub fn forward(
+ &self,
+ node_embedding: Vec,
+ neighbor_embeddings: Vec>,
+ edge_weights: Vec,
+ ) -> Result> {
+ // Convert f64 to f32
+ let node_f32: Vec = node_embedding.iter().map(|&x| x as f32).collect();
+ let neighbors_f32: Vec> = neighbor_embeddings
+ .iter()
+ .map(|v| v.iter().map(|&x| x as f32).collect())
+ .collect();
+ let weights_f32: Vec = edge_weights.iter().map(|&x| x as f32).collect();
+
+ let result = self.inner.forward(&node_f32, &neighbors_f32, &weights_f32);
+
+ // Convert back to f64
+ Ok(result.iter().map(|&x| x as f64).collect())
+ }
+
+ /// Serialize the layer to JSON
+ #[napi]
+ pub fn to_json(&self) -> Result {
+ serde_json::to_string(&self.inner)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Serialization error: {}", e)))
+ }
+
+ /// Deserialize the layer from JSON
+ #[napi(factory)]
+ pub fn from_json(json: String) -> Result {
+ let inner: RustRuvectorLayer = serde_json::from_str(&json)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Deserialization error: {}", e)))?;
+ Ok(Self { inner })
+ }
+}
+
+// ==================== TensorCompress Bindings ====================
+
+/// Compression level for tensor compression
+#[napi(object)]
+pub struct CompressionLevelConfig {
+ /// Type of compression: "none", "half", "pq8", "pq4", "binary"
+ pub level_type: String,
+ /// Scale factor (for "half" compression)
+ pub scale: Option,
+ /// Number of subvectors (for PQ compression)
+ pub subvectors: Option,
+ /// Number of centroids (for PQ8)
+ pub centroids: Option,
+ /// Outlier threshold (for PQ4)
+ pub outlier_threshold: Option,
+ /// Binary threshold (for binary compression)
+ pub threshold: Option,
+}
+
+impl CompressionLevelConfig {
+ fn to_rust(&self) -> Result {
+ match self.level_type.as_str() {
+ "none" => Ok(RustCompressionLevel::None),
+ "half" => Ok(RustCompressionLevel::Half {
+ scale: self.scale.unwrap_or(1.0) as f32,
+ }),
+ "pq8" => Ok(RustCompressionLevel::PQ8 {
+ subvectors: self.subvectors.unwrap_or(8) as u8,
+ centroids: self.centroids.unwrap_or(16) as u8,
+ }),
+ "pq4" => Ok(RustCompressionLevel::PQ4 {
+ subvectors: self.subvectors.unwrap_or(8) as u8,
+ outlier_threshold: self.outlier_threshold.unwrap_or(3.0) as f32,
+ }),
+ "binary" => Ok(RustCompressionLevel::Binary {
+ threshold: self.threshold.unwrap_or(0.0) as f32,
+ }),
+ _ => Err(Error::new(
+ Status::InvalidArg,
+ format!("Invalid compression level: {}", self.level_type),
+ )),
+ }
+ }
+}
+
+/// Tensor compressor with adaptive level selection
+#[napi]
+pub struct TensorCompress {
+ inner: RustTensorCompress,
+}
+
+#[napi]
+impl TensorCompress {
+ /// Create a new tensor compressor
+ ///
+ /// # Example
+ /// ```javascript
+ /// const compressor = new TensorCompress();
+ /// ```
+ #[napi(constructor)]
+ pub fn new() -> Self {
+ Self {
+ inner: RustTensorCompress::new(),
+ }
+ }
+
+ /// Compress an embedding based on access frequency
+ ///
+ /// # Arguments
+ /// * `embedding` - The input embedding vector
+ /// * `access_freq` - Access frequency in range [0.0, 1.0]
+ ///
+ /// # Returns
+ /// Compressed tensor as JSON string
+ ///
+ /// # Example
+ /// ```javascript
+ /// const embedding = [1.0, 2.0, 3.0, 4.0];
+ /// const compressed = compressor.compress(embedding, 0.5);
+ /// ```
+ #[napi]
+ pub fn compress(&self, embedding: Vec, access_freq: f64) -> Result {
+ let embedding_f32: Vec = embedding.iter().map(|&x| x as f32).collect();
+
+ let compressed = self.inner
+ .compress(&embedding_f32, access_freq as f32)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Compression error: {}", e)))?;
+
+ serde_json::to_string(&compressed)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Serialization error: {}", e)))
+ }
+
+ /// Compress with explicit compression level
+ ///
+ /// # Arguments
+ /// * `embedding` - The input embedding vector
+ /// * `level` - Compression level configuration
+ ///
+ /// # Returns
+ /// Compressed tensor as JSON string
+ ///
+ /// # Example
+ /// ```javascript
+ /// const embedding = [1.0, 2.0, 3.0, 4.0];
+ /// const level = { level_type: "half", scale: 1.0 };
+ /// const compressed = compressor.compressWithLevel(embedding, level);
+ /// ```
+ #[napi]
+ pub fn compress_with_level(
+ &self,
+ embedding: Vec,
+ level: CompressionLevelConfig,
+ ) -> Result {
+ let embedding_f32: Vec = embedding.iter().map(|&x| x as f32).collect();
+ let rust_level = level.to_rust()?;
+
+ let compressed = self.inner
+ .compress_with_level(&embedding_f32, &rust_level)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Compression error: {}", e)))?;
+
+ serde_json::to_string(&compressed)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Serialization error: {}", e)))
+ }
+
+ /// Decompress a compressed tensor
+ ///
+ /// # Arguments
+ /// * `compressed_json` - Compressed tensor as JSON string
+ ///
+ /// # Returns
+ /// Decompressed embedding vector
+ ///
+ /// # Example
+ /// ```javascript
+ /// const decompressed = compressor.decompress(compressed);
+ /// ```
+ #[napi]
+ pub fn decompress(&self, compressed_json: String) -> Result> {
+ let compressed: RustCompressedTensor = serde_json::from_str(&compressed_json)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Deserialization error: {}", e)))?;
+
+ let result = self.inner
+ .decompress(&compressed)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Decompression error: {}", e)))?;
+
+ Ok(result.iter().map(|&x| x as f64).collect())
+ }
+}
+
+// ==================== Search Functions ====================
+
+/// Result from differentiable search
+#[napi(object)]
+pub struct SearchResult {
+ /// Indices of top-k candidates
+ pub indices: Vec,
+ /// Soft weights for top-k candidates
+ pub weights: Vec,
+}
+
+/// Differentiable search using soft attention mechanism
+///
+/// # Arguments
+/// * `query` - The query vector
+/// * `candidate_embeddings` - List of candidate embedding vectors
+/// * `k` - Number of top results to return
+/// * `temperature` - Temperature for softmax (lower = sharper, higher = smoother)
+///
+/// # Returns
+/// Search result with indices and soft weights
+///
+/// # Example
+/// ```javascript
+/// const query = [1.0, 0.0, 0.0];
+/// const candidates = [[1.0, 0.0, 0.0], [0.9, 0.1, 0.0], [0.0, 1.0, 0.0]];
+/// const result = differentiableSearch(query, candidates, 2, 1.0);
+/// console.log(result.indices); // [0, 1]
+/// console.log(result.weights); // [0.x, 0.y]
+/// ```
+#[napi]
+pub fn differentiable_search(
+ query: Vec,
+ candidate_embeddings: Vec>,
+ k: u32,
+ temperature: f64,
+) -> Result {
+ let query_f32: Vec = query.iter().map(|&x| x as f32).collect();
+ let candidates_f32: Vec> = candidate_embeddings
+ .iter()
+ .map(|v| v.iter().map(|&x| x as f32).collect())
+ .collect();
+
+ let (indices, weights) = rust_differentiable_search(
+ &query_f32,
+ &candidates_f32,
+ k as usize,
+ temperature as f32,
+ );
+
+ Ok(SearchResult {
+ indices: indices.iter().map(|&i| i as u32).collect(),
+ weights: weights.iter().map(|&w| w as f64).collect(),
+ })
+}
+
+/// Hierarchical forward pass through GNN layers
+///
+/// # Arguments
+/// * `query` - The query vector
+/// * `layer_embeddings` - Embeddings organized by layer
+/// * `gnn_layers_json` - JSON array of serialized GNN layers
+///
+/// # Returns
+/// Final embedding after hierarchical processing
+///
+/// # Example
+/// ```javascript
+/// const query = [1.0, 0.0];
+/// const layerEmbeddings = [[[1.0, 0.0], [0.0, 1.0]]];
+/// const layer1 = new RuvectorLayer(2, 2, 1, 0.0);
+/// const layers = [layer1.toJson()];
+/// const result = hierarchicalForward(query, layerEmbeddings, layers);
+/// ```
+#[napi]
+pub fn hierarchical_forward(
+ query: Vec,
+ layer_embeddings: Vec>>,
+ gnn_layers_json: Vec,
+) -> Result> {
+ let query_f32: Vec = query.iter().map(|&x| x as f32).collect();
+
+ let embeddings_f32: Vec>> = layer_embeddings
+ .iter()
+ .map(|layer| {
+ layer
+ .iter()
+ .map(|v| v.iter().map(|&x| x as f32).collect())
+ .collect()
+ })
+ .collect();
+
+ let gnn_layers: Vec = gnn_layers_json
+ .iter()
+ .map(|json| {
+ serde_json::from_str(json)
+ .map_err(|e| Error::new(Status::GenericFailure, format!("Layer deserialization error: {}", e)))
+ })
+ .collect::>>()?;
+
+ let result = rust_hierarchical_forward(&query_f32, &embeddings_f32, &gnn_layers);
+
+ Ok(result.iter().map(|&x| x as f64).collect())
+}
+
+// ==================== Helper Functions ====================
+
+/// Get the compression level that would be selected for a given access frequency
+///
+/// # Arguments
+/// * `access_freq` - Access frequency in range [0.0, 1.0]
+///
+/// # Returns
+/// String describing the compression level: "none", "half", "pq8", "pq4", or "binary"
+///
+/// # Example
+/// ```javascript
+/// const level = getCompressionLevel(0.9); // "none" (hot data)
+/// const level2 = getCompressionLevel(0.5); // "half" (warm data)
+/// ```
+#[napi]
+pub fn get_compression_level(access_freq: f64) -> String {
+ if access_freq > 0.8 {
+ "none".to_string()
+ } else if access_freq > 0.4 {
+ "half".to_string()
+ } else if access_freq > 0.1 {
+ "pq8".to_string()
+ } else if access_freq > 0.01 {
+ "pq4".to_string()
+ } else {
+ "binary".to_string()
+ }
+}
+
+/// Module initialization
+#[napi]
+pub fn init() -> String {
+ "Ruvector GNN Node.js bindings initialized".to_string()
+}
diff --git a/crates/ruvector-gnn-node/test/basic.test.js b/crates/ruvector-gnn-node/test/basic.test.js
new file mode 100644
index 000000000..8cea55bd9
--- /dev/null
+++ b/crates/ruvector-gnn-node/test/basic.test.js
@@ -0,0 +1,204 @@
+// Basic tests for Ruvector GNN Node.js bindings
+
+const { test } = require('node:test');
+const assert = require('node:assert');
+
+const {
+ RuvectorLayer,
+ TensorCompress,
+ differentiableSearch,
+ hierarchicalForward,
+ getCompressionLevel,
+ init
+} = require('../index.js');
+
+test('initialization', () => {
+ const result = init();
+ assert.strictEqual(typeof result, 'string');
+ assert.ok(result.includes('initialized'));
+});
+
+test('RuvectorLayer creation', () => {
+ const layer = new RuvectorLayer(4, 8, 2, 0.1);
+ assert.ok(layer instanceof RuvectorLayer);
+});
+
+test('RuvectorLayer forward pass', () => {
+ const layer = new RuvectorLayer(4, 8, 2, 0.1);
+ const node = [1.0, 2.0, 3.0, 4.0];
+ const neighbors = [[0.5, 1.0, 1.5, 2.0], [2.0, 3.0, 4.0, 5.0]];
+ const weights = [0.3, 0.7];
+
+ const output = layer.forward(node, neighbors, weights);
+ assert.strictEqual(output.length, 8);
+ assert.ok(output.every(x => typeof x === 'number'));
+});
+
+test('RuvectorLayer forward with no neighbors', () => {
+ const layer = new RuvectorLayer(4, 8, 2, 0.1);
+ const node = [1.0, 2.0, 3.0, 4.0];
+ const neighbors = [];
+ const weights = [];
+
+ const output = layer.forward(node, neighbors, weights);
+ assert.strictEqual(output.length, 8);
+});
+
+test('RuvectorLayer serialization', () => {
+ const layer = new RuvectorLayer(4, 8, 2, 0.1);
+ const json = layer.toJson();
+ assert.strictEqual(typeof json, 'string');
+ assert.ok(json.length > 0);
+});
+
+test('RuvectorLayer deserialization', () => {
+ const layer1 = new RuvectorLayer(4, 8, 2, 0.1);
+ const json = layer1.toJson();
+ const layer2 = RuvectorLayer.fromJson(json);
+
+ assert.ok(layer2 instanceof RuvectorLayer);
+
+ // Test that they produce same output
+ const node = [1.0, 2.0, 3.0, 4.0];
+ const neighbors = [[0.5, 1.0, 1.5, 2.0]];
+ const weights = [1.0];
+
+ const output1 = layer1.forward(node, neighbors, weights);
+ const output2 = layer2.forward(node, neighbors, weights);
+
+ assert.strictEqual(output1.length, output2.length);
+ output1.forEach((val, i) => {
+ assert.ok(Math.abs(val - output2[i]) < 1e-6);
+ });
+});
+
+test('TensorCompress creation', () => {
+ const compressor = new TensorCompress();
+ assert.ok(compressor instanceof TensorCompress);
+});
+
+test('TensorCompress adaptive compression', () => {
+ const compressor = new TensorCompress();
+ const embedding = [1.0, 2.0, 3.0, 4.0];
+
+ const compressed = compressor.compress(embedding, 0.5);
+ assert.strictEqual(typeof compressed, 'string');
+ assert.ok(compressed.length > 0);
+});
+
+test('TensorCompress round-trip', () => {
+ const compressor = new TensorCompress();
+ const embedding = [1.0, 2.0, 3.0, 4.0];
+
+ const compressed = compressor.compress(embedding, 1.0); // No compression
+ const decompressed = compressor.decompress(compressed);
+
+ assert.strictEqual(decompressed.length, embedding.length);
+ decompressed.forEach((val, i) => {
+ assert.ok(Math.abs(val - embedding[i]) < 1e-6);
+ });
+});
+
+test('TensorCompress with explicit level', () => {
+ const compressor = new TensorCompress();
+ const embedding = Array.from({ length: 64 }, (_, i) => i * 0.1);
+
+ const level = {
+ level_type: 'half',
+ scale: 1.0
+ };
+
+ const compressed = compressor.compressWithLevel(embedding, level);
+ const decompressed = compressor.decompress(compressed);
+
+ assert.strictEqual(decompressed.length, embedding.length);
+});
+
+test('getCompressionLevel', () => {
+ assert.strictEqual(getCompressionLevel(0.9), 'none');
+ assert.strictEqual(getCompressionLevel(0.5), 'half');
+ assert.strictEqual(getCompressionLevel(0.2), 'pq8');
+ assert.strictEqual(getCompressionLevel(0.05), 'pq4');
+ assert.strictEqual(getCompressionLevel(0.001), 'binary');
+});
+
+test('differentiableSearch', () => {
+ const query = [1.0, 0.0, 0.0];
+ const candidates = [
+ [1.0, 0.0, 0.0],
+ [0.9, 0.1, 0.0],
+ [0.0, 1.0, 0.0],
+ ];
+
+ const result = differentiableSearch(query, candidates, 2, 1.0);
+
+ assert.ok(Array.isArray(result.indices));
+ assert.ok(Array.isArray(result.weights));
+ assert.strictEqual(result.indices.length, 2);
+ assert.strictEqual(result.weights.length, 2);
+
+ // First result should be perfect match
+ assert.strictEqual(result.indices[0], 0);
+
+ // Weights should be valid probabilities
+ result.weights.forEach(w => {
+ assert.ok(w >= 0 && w <= 1);
+ });
+});
+
+test('differentiableSearch with empty candidates', () => {
+ const query = [1.0, 0.0, 0.0];
+ const candidates = [];
+
+ const result = differentiableSearch(query, candidates, 2, 1.0);
+
+ assert.strictEqual(result.indices.length, 0);
+ assert.strictEqual(result.weights.length, 0);
+});
+
+test('hierarchicalForward', () => {
+ const query = [1.0, 0.0];
+ const layerEmbeddings = [
+ [[1.0, 0.0], [0.0, 1.0]],
+ ];
+
+ const layer = new RuvectorLayer(2, 2, 1, 0.0);
+ const layers = [layer.toJson()];
+
+ const result = hierarchicalForward(query, layerEmbeddings, layers);
+
+ assert.ok(Array.isArray(result));
+ assert.strictEqual(result.length, 2);
+ assert.ok(result.every(x => typeof x === 'number'));
+});
+
+test('invalid dropout rate throws error', () => {
+ assert.throws(() => {
+ new RuvectorLayer(4, 8, 2, 1.5); // dropout > 1.0
+ });
+
+ assert.throws(() => {
+ new RuvectorLayer(4, 8, 2, -0.1); // dropout < 0.0
+ });
+});
+
+test('compression with empty embedding throws error', () => {
+ const compressor = new TensorCompress();
+ assert.throws(() => {
+ compressor.compress([], 0.5);
+ });
+});
+
+test('compression levels produce different sizes', () => {
+ const compressor = new TensorCompress();
+ const embedding = Array.from({ length: 64 }, (_, i) => Math.sin(i * 0.1));
+
+ const none = compressor.compress(embedding, 1.0); // No compression
+ const half = compressor.compress(embedding, 0.5); // Half precision
+ const binary = compressor.compress(embedding, 0.001); // Binary
+
+ // Binary should be smallest
+ assert.ok(binary.length < half.length);
+ // None should be largest (or close to half)
+ assert.ok(none.length >= half.length * 0.8);
+});
diff --git a/crates/ruvector-gnn-wasm/Cargo.toml b/crates/ruvector-gnn-wasm/Cargo.toml
new file mode 100644
index 000000000..386fb8f32
--- /dev/null
+++ b/crates/ruvector-gnn-wasm/Cargo.toml
@@ -0,0 +1,48 @@
+[package]
+name = "ruvector-gnn-wasm"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+authors.workspace = true
+repository.workspace = true
+description = "WebAssembly bindings for RuVector GNN with tensor compression and differentiable search"
+
+[lib]
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+ruvector-gnn = { path = "../ruvector-gnn", default-features = false, features = ["wasm"] }
+
+# WASM
+wasm-bindgen = { workspace = true }
+js-sys = { workspace = true }
+getrandom = { workspace = true }
+
+# Serialization
+serde = { workspace = true }
+serde-wasm-bindgen = "0.6"
+
+# Utils
+console_error_panic_hook = { version = "0.1", optional = true }
+
+[dev-dependencies]
+wasm-bindgen-test = "0.3"
+
+[features]
+default = []
+console_error_panic_hook = ["dep:console_error_panic_hook"]
+
+# Ensure getrandom uses wasm_js/js feature for WASM
+[target.'cfg(target_arch = "wasm32")'.dependencies]
+getrandom = { workspace = true, features = ["wasm_js"] }
+getrandom02 = { package = "getrandom", version = "0.2", features = ["js"] }
+
+[profile.release]
+opt-level = "z"
+lto = true
+codegen-units = 1
+panic = "abort"
+
+[profile.release.package."*"]
+opt-level = "z"
diff --git a/crates/ruvector-gnn-wasm/README.md b/crates/ruvector-gnn-wasm/README.md
new file mode 100644
index 000000000..e5e693609
--- /dev/null
+++ b/crates/ruvector-gnn-wasm/README.md
@@ -0,0 +1,190 @@
+# RuVector GNN WASM
+
+WebAssembly bindings for RuVector Graph Neural Network operations.
+
+## Features
+
+- **GNN Layer Operations**: Multi-head attention, GRU updates, layer normalization
+- **Tensor Compression**: Adaptive compression based on access frequency
+- **Differentiable Search**: Soft attention-based similarity search
+- **Hierarchical Forward**: Multi-layer GNN processing
+
+## Installation
+
+```bash
+npm install ruvector-gnn-wasm
+```
+
+## Usage
+
+### Initialize
+
+```typescript
+import init, {
+ JsRuvectorLayer,
+ JsTensorCompress,
+ differentiableSearch,
+ SearchConfig
+} from 'ruvector-gnn-wasm';
+
+await init();
+```
+
+### GNN Layer
+
+```typescript
+// Create a GNN layer
+const layer = new JsRuvectorLayer(
+ 4, // input dimension
+ 8, // hidden dimension
+ 2, // number of attention heads
+ 0.1 // dropout rate
+);
+
+// Forward pass
+const nodeEmbedding = new Float32Array([1.0, 2.0, 3.0, 4.0]);
+const neighbors = [
+ new Float32Array([0.5, 1.0, 1.5, 2.0]),
+ new Float32Array([2.0, 3.0, 4.0, 5.0])
+];
+const edgeWeights = new Float32Array([0.3, 0.7]);
+
+const output = layer.forward(nodeEmbedding, neighbors, edgeWeights);
+console.log('Output dimension:', layer.outputDim);
+```
+
+### Tensor Compression
+
+```typescript
+const compressor = new JsTensorCompress();
+
+// Compress based on access frequency
+const embedding = new Float32Array(128).fill(0.5);
+const compressed = compressor.compress(embedding, 0.5); // 50% access frequency
+
+// Decompress
+const decompressed = compressor.decompress(compressed);
+
+// Or specify compression level explicitly
+const compressedPQ8 = compressor.compressWithLevel(embedding, "pq8");
+
+// Get compression ratio
+const ratio = compressor.getCompressionRatio(0.5); // Returns ~2.0 for half precision
+```
+
+### Compression Levels
+
+Access frequency determines compression:
+- `f > 0.8`: **Full precision** (no compression) - hot data
+- `f > 0.4`: **Half precision** (2x compression) - warm data
+- `f > 0.1`: **8-bit PQ** (4x compression) - cool data
+- `f > 0.01`: **4-bit PQ** (8x compression) - cold data
+- `f <= 0.01`: **Binary** (32x compression) - archive data
+
+### Differentiable Search
+
+```typescript
+const query = new Float32Array([1.0, 0.0, 0.0]);
+const candidates = [
+ new Float32Array([1.0, 0.0, 0.0]), // Perfect match
+ new Float32Array([0.9, 0.1, 0.0]), // Close match
+ new Float32Array([0.0, 1.0, 0.0]) // Orthogonal
+];
+
+const config = new SearchConfig(2, 1.0); // k=2, temperature=1.0
+const result = differentiableSearch(query, candidates, config);
+
+console.log('Top indices:', result.indices);
+console.log('Weights:', result.weights);
+```
+
+## API Reference
+
+### `JsRuvectorLayer`
+
+```typescript
+class JsRuvectorLayer {
+ constructor(
+ inputDim: number,
+ hiddenDim: number,
+ heads: number,
+ dropout: number
+ );
+
+ forward(
+ nodeEmbedding: Float32Array,
+ neighborEmbeddings: Float32Array[],
+ edgeWeights: Float32Array
+ ): Float32Array;
+
+ readonly outputDim: number;
+}
+```
+
+### `JsTensorCompress`
+
+```typescript
+class JsTensorCompress {
+ constructor();
+
+ compress(embedding: Float32Array, accessFreq: number): object;
+ compressWithLevel(embedding: Float32Array, level: string): object;
+ decompress(compressed: object): Float32Array;
+ getCompressionRatio(accessFreq: number): number;
+}
+```
+
+Compression levels: `"none"`, `"half"`, `"pq8"`, `"pq4"`, `"binary"`
+
+### `differentiableSearch`
+
+```typescript
+function differentiableSearch(
+ query: Float32Array,
+ candidateEmbeddings: Float32Array[],
+ config: SearchConfig
+): { indices: number[], weights: number[] };
+```
+
+### `SearchConfig`
+
+```typescript
+class SearchConfig {
+ constructor(k: number, temperature: number);
+ k: number; // Number of results
+ temperature: number; // Softmax temperature (lower = sharper)
+}
+```
+
+### `cosineSimilarity`
+
+```typescript
+function cosineSimilarity(a: Float32Array, b: Float32Array): number;
+```
+
+## Building from Source
+
+```bash
+# Install wasm-pack
+curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
+
+# Build for Node.js
+wasm-pack build --target nodejs
+
+# Build for browser
+wasm-pack build --target web
+
+# Build for bundler (webpack, etc.)
+wasm-pack build --target bundler
+```
+
+## Performance
+
+- GNN layers use efficient attention mechanisms
+- Compression reduces memory usage by 2-32x
+- All operations are optimized for WASM
+- No garbage collection during forward passes
+
+## License
+
+MIT
diff --git a/crates/ruvector-gnn-wasm/src/lib.rs b/crates/ruvector-gnn-wasm/src/lib.rs
new file mode 100644
index 000000000..56a5ead46
--- /dev/null
+++ b/crates/ruvector-gnn-wasm/src/lib.rs
@@ -0,0 +1,397 @@
+//! WebAssembly bindings for RuVector GNN
+//!
+//! This module provides high-performance browser bindings for Graph Neural Network
+//! operations on HNSW topology, including:
+//! - GNN layer forward passes
+//! - Tensor compression with adaptive level selection
+//! - Differentiable search with soft attention
+//! - Hierarchical forward propagation
+
+use wasm_bindgen::prelude::*;
+use serde::{Deserialize, Serialize};
+use ruvector_gnn::{
+ RuvectorLayer, TensorCompress, CompressedTensor, CompressionLevel,
+ differentiable_search as core_differentiable_search,
+ hierarchical_forward as core_hierarchical_forward,
+};
+
+/// Initialize panic hook for better error messages
+#[wasm_bindgen(start)]
+pub fn init() {
+ #[cfg(feature = "console_error_panic_hook")]
+ console_error_panic_hook::set_once();
+}
+
+// ============================================================================
+// Type Definitions for WASM
+// ============================================================================
+
+/// Query configuration for differentiable search
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[wasm_bindgen]
+pub struct SearchConfig {
+ /// Number of top results to return
+ pub k: usize,
+ /// Temperature for softmax (lower = sharper, higher = smoother)
+ pub temperature: f32,
+}
+
+#[wasm_bindgen]
+impl SearchConfig {
+ /// Create a new search configuration
+ #[wasm_bindgen(constructor)]
+ pub fn new(k: usize, temperature: f32) -> Self {
+ Self { k, temperature }
+ }
+}
+
+/// Search results with indices and weights (internal)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct SearchResultInternal {
+ /// Indices of top-k candidates
+ indices: Vec,
+ /// Soft weights for each result
+ weights: Vec,
+}
+
+// ============================================================================
+// JsRuvectorLayer - GNN Layer Wrapper
+// ============================================================================
+
+/// Graph Neural Network layer for HNSW topology
+#[wasm_bindgen]
+pub struct JsRuvectorLayer {
+ inner: RuvectorLayer,
+ hidden_dim: usize,
+}
+
+#[wasm_bindgen]
+impl JsRuvectorLayer {
+ /// Create a new GNN layer
+ ///
+ /// # Arguments
+ /// * `input_dim` - Dimension of input node embeddings
+ /// * `hidden_dim` - Dimension of hidden representations
+ /// * `heads` - Number of attention heads
+ /// * `dropout` - Dropout rate (0.0 to 1.0)
+ #[wasm_bindgen(constructor)]
+ pub fn new(input_dim: usize, hidden_dim: usize, heads: usize, dropout: f32) -> Result {
+ if dropout < 0.0 || dropout > 1.0 {
+ return Err(JsValue::from_str("Dropout must be between 0.0 and 1.0"));
+ }
+
+ Ok(JsRuvectorLayer {
+ inner: RuvectorLayer::new(input_dim, hidden_dim, heads, dropout),
+ hidden_dim,
+ })
+ }
+
+ /// Forward pass through the GNN layer
+ ///
+ /// # Arguments
+ /// * `node_embedding` - Current node's embedding (Float32Array)
+ /// * `neighbor_embeddings` - Embeddings of neighbor nodes (array of Float32Arrays)
+ /// * `edge_weights` - Weights of edges to neighbors (Float32Array)
+ ///
+ /// # Returns
+ /// Updated node embedding (Float32Array)
+ #[wasm_bindgen]
+ pub fn forward(
+ &self,
+ node_embedding: Vec,
+ neighbor_embeddings: JsValue,
+ edge_weights: Vec,
+ ) -> Result, JsValue> {
+ // Convert neighbor embeddings from JS value
+ let neighbors: Vec> = serde_wasm_bindgen::from_value(neighbor_embeddings)
+ .map_err(|e| JsValue::from_str(&format!("Failed to parse neighbor embeddings: {}", e)))?;
+
+ // Validate inputs
+ if neighbors.len() != edge_weights.len() {
+ return Err(JsValue::from_str(&format!(
+ "Number of neighbors ({}) must match number of edge weights ({})",
+ neighbors.len(),
+ edge_weights.len()
+ )));
+ }
+
+ // Call core forward
+ let result = self.inner.forward(&node_embedding, &neighbors, &edge_weights);
+
+ Ok(result)
+ }
+
+ /// Get the output dimension of this layer
+ #[wasm_bindgen(getter, js_name = outputDim)]
+ pub fn output_dim(&self) -> usize {
+ self.hidden_dim
+ }
+}
+
+// ============================================================================
+// JsTensorCompress - Tensor Compression Wrapper
+// ============================================================================
+
+/// Tensor compressor with adaptive level selection
+#[wasm_bindgen]
+pub struct JsTensorCompress {
+ inner: TensorCompress,
+}
+
+#[wasm_bindgen]
+impl JsTensorCompress {
+ /// Create a new tensor compressor
+ #[wasm_bindgen(constructor)]
+ pub fn new() -> Self {
+ Self {
+ inner: TensorCompress::new(),
+ }
+ }
+
+ /// Compress an embedding based on access frequency
+ ///
+ /// # Arguments
+ /// * `embedding` - The input embedding vector (Float32Array)
+ /// * `access_freq` - Access frequency in range [0.0, 1.0]
+ /// - f > 0.8: Full precision (hot data)
+ /// - f > 0.4: Half precision (warm data)
+ /// - f > 0.1: 8-bit PQ (cool data)
+ /// - f > 0.01: 4-bit PQ (cold data)
+ /// - f <= 0.01: Binary (archive)
+ ///
+ /// # Returns
+ /// Compressed tensor as JsValue
+ #[wasm_bindgen]
+ pub fn compress(&self, embedding: Vec, access_freq: f32) -> Result {
+ let compressed = self.inner
+ .compress(&embedding, access_freq)
+ .map_err(|e| JsValue::from_str(&format!("Compression failed: {}", e)))?;
+
+ // Serialize using serde_wasm_bindgen
+ serde_wasm_bindgen::to_value(&compressed)
+ .map_err(|e| JsValue::from_str(&format!("Serialization failed: {}", e)))
+ }
+
+ /// Compress with explicit compression level
+ ///
+ /// # Arguments
+ /// * `embedding` - The input embedding vector
+ /// * `level` - Compression level ("none", "half", "pq8", "pq4", "binary")
+ ///
+ /// # Returns
+ /// Compressed tensor as JsValue
+ #[wasm_bindgen(js_name = compressWithLevel)]
+ pub fn compress_with_level(&self, embedding: Vec, level: &str) -> Result {
+ let compression_level = match level {
+ "none" => CompressionLevel::None,
+ "half" => CompressionLevel::Half { scale: 1.0 },
+ "pq8" => CompressionLevel::PQ8 {
+ subvectors: 8,
+ centroids: 16,
+ },
+ "pq4" => CompressionLevel::PQ4 {
+ subvectors: 8,
+ outlier_threshold: 3.0,
+ },
+ "binary" => CompressionLevel::Binary { threshold: 0.0 },
+ _ => return Err(JsValue::from_str(&format!("Unknown compression level: {}", level))),
+ };
+
+ let compressed = self.inner
+ .compress_with_level(&embedding, &compression_level)
+ .map_err(|e| JsValue::from_str(&format!("Compression failed: {}", e)))?;
+
+ // Serialize using serde_wasm_bindgen
+ serde_wasm_bindgen::to_value(&compressed)
+ .map_err(|e| JsValue::from_str(&format!("Serialization failed: {}", e)))
+ }
+
+ /// Decompress a compressed tensor
+ ///
+ /// # Arguments
+ /// * `compressed` - Serialized compressed tensor (JsValue)
+ ///
+ /// # Returns
+ /// Decompressed embedding vector (Float32Array)
+ #[wasm_bindgen]
+ pub fn decompress(&self, compressed: JsValue) -> Result, JsValue> {
+ let compressed_tensor: CompressedTensor = serde_wasm_bindgen::from_value(compressed)
+ .map_err(|e| JsValue::from_str(&format!("Deserialization failed: {}", e)))?;
+
+ let decompressed = self.inner
+ .decompress(&compressed_tensor)
+ .map_err(|e| JsValue::from_str(&format!("Decompression failed: {}", e)))?;
+
+ Ok(decompressed)
+ }
+
+ /// Get compression ratio estimate for a given access frequency
+ ///
+ /// # Arguments
+ /// * `access_freq` - Access frequency in range [0.0, 1.0]
+ ///
+ /// # Returns
+ /// Estimated compression ratio (original_size / compressed_size)
+ #[wasm_bindgen(js_name = getCompressionRatio)]
+ pub fn get_compression_ratio(&self, access_freq: f32) -> f32 {
+ if access_freq > 0.8 {
+ 1.0 // No compression
+ } else if access_freq > 0.4 {
+ 2.0 // Half precision
+ } else if access_freq > 0.1 {
+ 4.0 // 8-bit PQ
+ } else if access_freq > 0.01 {
+ 8.0 // 4-bit PQ
+ } else {
+ 32.0 // Binary
+ }
+ }
+}
+
+// ============================================================================
+// Standalone Functions
+// ============================================================================
+
+/// Differentiable search using soft attention mechanism
+///
+/// # Arguments
+/// * `query` - The query vector (Float32Array)
+/// * `candidate_embeddings` - List of candidate embedding vectors (array of Float32Arrays)
+/// * `config` - Search configuration (k and temperature)
+///
+/// # Returns
+/// Object with indices and weights for top-k candidates
+#[wasm_bindgen(js_name = differentiableSearch)]
+pub fn differentiable_search(
+ query: Vec,
+ candidate_embeddings: JsValue,
+ config: &SearchConfig,
+) -> Result {
+ // Convert candidate embeddings from JS value
+ let candidates: Vec> = serde_wasm_bindgen::from_value(candidate_embeddings)
+ .map_err(|e| JsValue::from_str(&format!("Failed to parse candidate embeddings: {}", e)))?;
+
+ // Call core search function
+ let (indices, weights) = core_differentiable_search(
+ &query,
+ &candidates,
+ config.k,
+ config.temperature,
+ );
+
+ let result = SearchResultInternal { indices, weights };
+ serde_wasm_bindgen::to_value(&result)
+ .map_err(|e| JsValue::from_str(&format!("Failed to serialize result: {}", e)))
+}
+
+/// Hierarchical forward pass through multiple GNN layers
+///
+/// # Arguments
+/// * `query` - The query vector (Float32Array)
+/// * `layer_embeddings` - Embeddings organized by layer (array of arrays of Float32Arrays)
+/// * `gnn_layers` - Array of GNN layers to process through
+///
+/// # Returns
+/// Final embedding after hierarchical processing (Float32Array)
+#[wasm_bindgen(js_name = hierarchicalForward)]
+pub fn hierarchical_forward(
+ query: Vec,
+ layer_embeddings: JsValue,
+ gnn_layers: Vec,
+) -> Result, JsValue> {
+ // Convert layer embeddings from JS value
+ let embeddings: Vec>> = serde_wasm_bindgen::from_value(layer_embeddings)
+ .map_err(|e| JsValue::from_str(&format!("Failed to parse layer embeddings: {}", e)))?;
+
+ // Extract inner layers
+ let core_layers: Vec = gnn_layers.iter().map(|l| l.inner.clone()).collect();
+
+ // Call core function
+ let result = core_hierarchical_forward(&query, &embeddings, &core_layers);
+
+ Ok(result)
+}
+
+// ============================================================================
+// Utility Functions
+// ============================================================================
+
+/// Get version information
+#[wasm_bindgen]
+pub fn version() -> String {
+ env!("CARGO_PKG_VERSION").to_string()
+}
+
+/// Compute cosine similarity between two vectors
+///
+/// # Arguments
+/// * `a` - First vector (Float32Array)
+/// * `b` - Second vector (Float32Array)
+///
+/// # Returns
+/// Cosine similarity score [-1.0, 1.0]
+#[wasm_bindgen(js_name = cosineSimilarity)]
+pub fn cosine_similarity(a: Vec, b: Vec) -> Result {
+ if a.len() != b.len() {
+ return Err(JsValue::from_str(&format!(
+ "Vector dimensions must match: {} vs {}",
+ a.len(),
+ b.len()
+ )));
+ }
+
+ let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
+ let norm_a: f32 = a.iter().map(|x| x * x).sum::().sqrt();
+ let norm_b: f32 = b.iter().map(|x| x * x).sum::().sqrt();
+
+ if norm_a == 0.0 || norm_b == 0.0 {
+ Ok(0.0)
+ } else {
+ Ok(dot_product / (norm_a * norm_b))
+ }
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use wasm_bindgen_test::*;
+
+ wasm_bindgen_test_configure!(run_in_browser);
+
+ #[wasm_bindgen_test]
+ fn test_version() {
+ assert!(!version().is_empty());
+ }
+
+ #[wasm_bindgen_test]
+ fn test_ruvector_layer_creation() {
+ let layer = JsRuvectorLayer::new(4, 8, 2, 0.1);
+ assert!(layer.is_ok());
+ }
+
+ #[wasm_bindgen_test]
+ fn test_tensor_compress_creation() {
+ let compressor = JsTensorCompress::new();
+ assert_eq!(compressor.get_compression_ratio(1.0), 1.0);
+ assert_eq!(compressor.get_compression_ratio(0.5), 2.0);
+ }
+
+ #[wasm_bindgen_test]
+ fn test_cosine_similarity() {
+ let a = vec![1.0, 0.0, 0.0];
+ let b = vec![1.0, 0.0, 0.0];
+ let sim = cosine_similarity(a, b).unwrap();
+ assert!((sim - 1.0).abs() < 1e-6);
+ }
+
+ #[wasm_bindgen_test]
+ fn test_search_config() {
+ let config = SearchConfig::new(5, 1.0);
+ assert_eq!(config.k, 5);
+ assert_eq!(config.temperature, 1.0);
+ }
+}
diff --git a/crates/ruvector-gnn/Cargo.toml b/crates/ruvector-gnn/Cargo.toml
new file mode 100644
index 000000000..8a8f280fe
--- /dev/null
+++ b/crates/ruvector-gnn/Cargo.toml
@@ -0,0 +1,58 @@
+[package]
+name = "ruvector-gnn"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+authors.workspace = true
+repository.workspace = true
+description = "Graph Neural Network layer for Ruvector on HNSW topology"
+
+[dependencies]
+# Core
+ruvector-core = { path = "../ruvector-core", default-features = false }
+
+# Math and numerics
+ndarray = { workspace = true, features = ["serde"] }
+rand = { workspace = true }
+rand_distr = { workspace = true }
+
+# Serialization
+serde = { workspace = true }
+serde_json = { workspace = true }
+
+# Error handling
+thiserror = { workspace = true }
+anyhow = { workspace = true }
+
+# Performance
+rayon = { workspace = true }
+parking_lot = { workspace = true }
+dashmap = { workspace = true }
+
+# Memory mapping (non-WASM only)
+memmap2 = { workspace = true, optional = true }
+page_size = { version = "0.6", optional = true }
+
+# Platform-specific dependencies
+[target.'cfg(target_os = "linux")'.dependencies]
+libc = "0.2"
+
+# Optional dependencies
+napi = { workspace = true, optional = true }
+napi-derive = { workspace = true, optional = true }
+
+[features]
+default = ["simd", "mmap"]
+simd = []
+wasm = []
+napi = ["dep:napi", "dep:napi-derive"]
+mmap = ["dep:memmap2", "dep:page_size"]
+
+[dev-dependencies]
+criterion = { workspace = true }
+proptest = { workspace = true }
+tempfile = "3.10"
+
+[lib]
+crate-type = ["rlib"]
diff --git a/crates/ruvector-gnn/src/compress.rs b/crates/ruvector-gnn/src/compress.rs
new file mode 100644
index 000000000..20153709b
--- /dev/null
+++ b/crates/ruvector-gnn/src/compress.rs
@@ -0,0 +1,685 @@
+//! Tensor compression with adaptive level selection
+//!
+//! This module provides multi-level tensor compression based on access frequency:
+//! - Hot data (f > 0.8): Full precision
+//! - Warm data (f > 0.4): Half precision
+//! - Cool data (f > 0.1): 8-bit product quantization
+//! - Cold data (f > 0.01): 4-bit product quantization
+//! - Archive (f <= 0.01): Binary quantization
+
+use crate::error::{GnnError, Result};
+use serde::{Deserialize, Serialize};
+
+/// Compression level with associated parameters
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub enum CompressionLevel {
+ /// Full precision - no compression
+ None,
+
+ /// Half precision with scale factor
+ Half { scale: f32 },
+
+ /// Product quantization with 8-bit codes
+ PQ8 { subvectors: u8, centroids: u8 },
+
+ /// Product quantization with 4-bit codes and outlier handling
+ PQ4 {
+ subvectors: u8,
+ outlier_threshold: f32,
+ },
+
+ /// Binary quantization with threshold
+ Binary { threshold: f32 },
+}
+
+/// Compressed tensor data
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum CompressedTensor {
+ /// Uncompressed full precision data
+ Full { data: Vec },
+
+ /// Half precision data
+ Half {
+ data: Vec,
+ scale: f32,
+ dim: usize,
+ },
+
+ /// 8-bit product quantization
+ PQ8 {
+ codes: Vec,
+ codebooks: Vec>,
+ subvector_dim: usize,
+ dim: usize,
+ },
+
+ /// 4-bit product quantization with outliers
+ PQ4 {
+ codes: Vec, // Packed 4-bit codes
+ codebooks: Vec>,
+ outliers: Vec<(usize, f32)>, // (index, value) pairs
+ subvector_dim: usize,
+ dim: usize,
+ },
+
+ /// Binary quantization
+ Binary {
+ bits: Vec,
+ threshold: f32,
+ dim: usize,
+ },
+}
+
+/// Tensor compressor with adaptive level selection
+#[derive(Debug, Clone)]
+pub struct TensorCompress {
+ /// Default compression parameters
+ default_level: CompressionLevel,
+}
+
+impl Default for TensorCompress {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl TensorCompress {
+ /// Create a new tensor compressor with default settings
+ pub fn new() -> Self {
+ Self {
+ default_level: CompressionLevel::None,
+ }
+ }
+
+ /// Compress an embedding based on access frequency
+ ///
+ /// # Arguments
+ /// * `embedding` - The input embedding vector
+ /// * `access_freq` - Access frequency in range [0.0, 1.0]
+ ///
+ /// # Returns
+ /// Compressed tensor using adaptive compression level
+ pub fn compress(&self, embedding: &[f32], access_freq: f32) -> Result {
+ if embedding.is_empty() {
+ return Err(GnnError::InvalidInput(
+ "Empty embedding vector".to_string(),
+ ));
+ }
+
+ let level = self.select_level(access_freq);
+ self.compress_with_level(embedding, &level)
+ }
+
+ /// Compress with explicit compression level
+ pub fn compress_with_level(
+ &self,
+ embedding: &[f32],
+ level: &CompressionLevel,
+ ) -> Result {
+ match level {
+ CompressionLevel::None => self.compress_none(embedding),
+ CompressionLevel::Half { scale } => self.compress_half(embedding, *scale),
+ CompressionLevel::PQ8 {
+ subvectors,
+ centroids,
+ } => self.compress_pq8(embedding, *subvectors, *centroids),
+ CompressionLevel::PQ4 {
+ subvectors,
+ outlier_threshold,
+ } => self.compress_pq4(embedding, *subvectors, *outlier_threshold),
+ CompressionLevel::Binary { threshold } => self.compress_binary(embedding, *threshold),
+ }
+ }
+
+ /// Decompress a compressed tensor
+ pub fn decompress(&self, compressed: &CompressedTensor) -> Result> {
+ match compressed {
+ CompressedTensor::Full { data } => Ok(data.clone()),
+ CompressedTensor::Half { data, scale, dim } => {
+ self.decompress_half(data, *scale, *dim)
+ }
+ CompressedTensor::PQ8 {
+ codes,
+ codebooks,
+ subvector_dim,
+ dim,
+ } => self.decompress_pq8(codes, codebooks, *subvector_dim, *dim),
+ CompressedTensor::PQ4 {
+ codes,
+ codebooks,
+ outliers,
+ subvector_dim,
+ dim,
+ } => self.decompress_pq4(codes, codebooks, outliers, *subvector_dim, *dim),
+ CompressedTensor::Binary {
+ bits,
+ threshold,
+ dim,
+ } => self.decompress_binary(bits, *threshold, *dim),
+ }
+ }
+
+ /// Select compression level based on access frequency
+ ///
+ /// Thresholds:
+ /// - f > 0.8: None (hot data)
+ /// - f > 0.4: Half (warm data)
+ /// - f > 0.1: PQ8 (cool data)
+ /// - f > 0.01: PQ4 (cold data)
+ /// - f <= 0.01: Binary (archive)
+ fn select_level(&self, access_freq: f32) -> CompressionLevel {
+ if access_freq > 0.8 {
+ CompressionLevel::None
+ } else if access_freq > 0.4 {
+ CompressionLevel::Half { scale: 1.0 }
+ } else if access_freq > 0.1 {
+ CompressionLevel::PQ8 {
+ subvectors: 8,
+ centroids: 16,
+ }
+ } else if access_freq > 0.01 {
+ CompressionLevel::PQ4 {
+ subvectors: 8,
+ outlier_threshold: 3.0,
+ }
+ } else {
+ CompressionLevel::Binary { threshold: 0.0 }
+ }
+ }
+
+ // === Compression implementations ===
+
+ fn compress_none(&self, embedding: &[f32]) -> Result {
+ Ok(CompressedTensor::Full {
+ data: embedding.to_vec(),
+ })
+ }
+
+ fn compress_half(&self, embedding: &[f32], scale: f32) -> Result {
+ // Simple half precision: scale and convert to 16-bit
+ let data: Vec = embedding
+ .iter()
+ .map(|&x| {
+ let scaled = x * scale;
+ let clamped = scaled.clamp(-65504.0, 65504.0);
+ // Convert to half precision representation
+ f32_to_f16_bits(clamped)
+ })
+ .collect();
+
+ Ok(CompressedTensor::Half {
+ data,
+ scale,
+ dim: embedding.len(),
+ })
+ }
+
+ fn compress_pq8(&self, embedding: &[f32], subvectors: u8, centroids: u8) -> Result {
+ let dim = embedding.len();
+ let subvectors = subvectors as usize;
+
+ if dim % subvectors != 0 {
+ return Err(GnnError::InvalidInput(format!(
+ "Dimension {} not divisible by subvectors {}",
+ dim, subvectors
+ )));
+ }
+
+ let subvector_dim = dim / subvectors;
+ let mut codes = Vec::with_capacity(subvectors);
+ let mut codebooks = Vec::with_capacity(subvectors);
+
+ // For each subvector, create a codebook and quantize
+ for i in 0..subvectors {
+ let start = i * subvector_dim;
+ let end = start + subvector_dim;
+ let subvector = &embedding[start..end];
+
+ // Simple k-means clustering (k=centroids)
+ let (codebook, code) = self.quantize_subvector(subvector, centroids as usize);
+ codes.push(code);
+ codebooks.push(codebook);
+ }
+
+ Ok(CompressedTensor::PQ8 {
+ codes,
+ codebooks,
+ subvector_dim,
+ dim,
+ })
+ }
+
+ fn compress_pq4(
+ &self,
+ embedding: &[f32],
+ subvectors: u8,
+ outlier_threshold: f32,
+ ) -> Result {
+ let dim = embedding.len();
+ let subvectors = subvectors as usize;
+
+ if dim % subvectors != 0 {
+ return Err(GnnError::InvalidInput(format!(
+ "Dimension {} not divisible by subvectors {}",
+ dim, subvectors
+ )));
+ }
+
+ let subvector_dim = dim / subvectors;
+ let mut codes = Vec::with_capacity(subvectors);
+ let mut codebooks = Vec::with_capacity(subvectors);
+ let mut outliers = Vec::new();
+
+ // Detect outliers based on magnitude
+ let mean = embedding.iter().sum::() / dim as f32;
+ let std_dev =
+ (embedding.iter().map(|&x| (x - mean).powi(2)).sum::() / dim as f32).sqrt();
+
+ // For each subvector
+ for i in 0..subvectors {
+ let start = i * subvector_dim;
+ let end = start + subvector_dim;
+ let subvector = &embedding[start..end];
+
+ // Extract outliers
+ let mut cleaned_subvector = subvector.to_vec();
+ for (j, &val) in subvector.iter().enumerate() {
+ if (val - mean).abs() > outlier_threshold * std_dev {
+ outliers.push((start + j, val));
+ cleaned_subvector[j] = mean; // Replace with mean
+ }
+ }
+
+ // Quantize to 4-bit (16 centroids)
+ let (codebook, code) = self.quantize_subvector(&cleaned_subvector, 16);
+ codes.push(code);
+ codebooks.push(codebook);
+ }
+
+ Ok(CompressedTensor::PQ4 {
+ codes,
+ codebooks,
+ outliers,
+ subvector_dim,
+ dim,
+ })
+ }
+
+ fn compress_binary(&self, embedding: &[f32], threshold: f32) -> Result {
+ let dim = embedding.len();
+ let num_bytes = (dim + 7) / 8;
+ let mut bits = vec![0u8; num_bytes];
+
+ for (i, &val) in embedding.iter().enumerate() {
+ if val > threshold {
+ let byte_idx = i / 8;
+ let bit_idx = i % 8;
+ bits[byte_idx] |= 1 << bit_idx;
+ }
+ }
+
+ Ok(CompressedTensor::Binary {
+ bits,
+ threshold,
+ dim,
+ })
+ }
+
+ // === Decompression implementations ===
+
+ fn decompress_half(&self, data: &[u16], scale: f32, dim: usize) -> Result> {
+ if data.len() != dim {
+ return Err(GnnError::InvalidInput(format!(
+ "Dimension mismatch: expected {}, got {}",
+ dim,
+ data.len()
+ )));
+ }
+
+ Ok(data
+ .iter()
+ .map(|&bits| f16_bits_to_f32(bits) / scale)
+ .collect())
+ }
+
+ fn decompress_pq8(
+ &self,
+ codes: &[u8],
+ codebooks: &[Vec],
+ subvector_dim: usize,
+ dim: usize,
+ ) -> Result> {
+ let subvectors = codes.len();
+ let expected_dim = subvectors * subvector_dim;
+
+ if expected_dim != dim {
+ return Err(GnnError::InvalidInput(format!(
+ "Dimension mismatch: expected {}, got {}",
+ dim, expected_dim
+ )));
+ }
+
+ let mut result = Vec::with_capacity(dim);
+
+ for (code, codebook) in codes.iter().zip(codebooks.iter()) {
+ let centroid_idx = *code as usize;
+ if centroid_idx >= codebook.len() / subvector_dim {
+ return Err(GnnError::InvalidInput(format!(
+ "Invalid centroid index: {}",
+ centroid_idx
+ )));
+ }
+
+ let start = centroid_idx * subvector_dim;
+ let end = start + subvector_dim;
+ result.extend_from_slice(&codebook[start..end]);
+ }
+
+ Ok(result)
+ }
+
+ fn decompress_pq4(
+ &self,
+ codes: &[u8],
+ codebooks: &[Vec],
+ outliers: &[(usize, f32)],
+ subvector_dim: usize,
+ dim: usize,
+ ) -> Result> {
+ // First decompress using PQ8 logic
+ let mut result = self.decompress_pq8(codes, codebooks, subvector_dim, dim)?;
+
+ // Restore outliers
+ for &(idx, val) in outliers {
+ if idx < result.len() {
+ result[idx] = val;
+ }
+ }
+
+ Ok(result)
+ }
+
+ fn decompress_binary(&self, bits: &[u8], _threshold: f32, dim: usize) -> Result> {
+ let expected_bytes = (dim + 7) / 8;
+ if bits.len() != expected_bytes {
+ return Err(GnnError::InvalidInput(format!(
+ "Dimension mismatch: expected {} bytes, got {}",
+ expected_bytes,
+ bits.len()
+ )));
+ }
+
+ let mut result = Vec::with_capacity(dim);
+
+ for i in 0..dim {
+ let byte_idx = i / 8;
+ let bit_idx = i % 8;
+ let is_set = (bits[byte_idx] & (1 << bit_idx)) != 0;
+ result.push(if is_set { 1.0 } else { -1.0 });
+ }
+
+ Ok(result)
+ }
+
+ // === Helper methods ===
+
+ /// Simple quantization using k-means-like approach
+ fn quantize_subvector(&self, subvector: &[f32], k: usize) -> (Vec, u8) {
+ let dim = subvector.len();
+
+ // Initialize centroids using simple range-based approach
+ let min_val = subvector
+ .iter()
+ .cloned()
+ .fold(f32::INFINITY, f32::min);
+ let max_val = subvector
+ .iter()
+ .cloned()
+ .fold(f32::NEG_INFINITY, f32::max);
+ let range = max_val - min_val;
+
+ if range < 1e-6 {
+ // All values are essentially the same
+ let codebook = vec![min_val; dim * k];
+ return (codebook, 0);
+ }
+
+ // Create k centroids evenly spaced across the range
+ let centroids: Vec> = (0..k)
+ .map(|i| {
+ let offset = min_val + (i as f32 / k as f32) * range;
+ vec![offset; dim]
+ })
+ .collect();
+
+ // Find nearest centroid for this subvector
+ let code = self.nearest_centroid(subvector, ¢roids);
+
+ // Flatten codebook
+ let codebook: Vec = centroids.into_iter().flatten().collect();
+
+ (codebook, code as u8)
+ }
+
+ fn nearest_centroid(&self, subvector: &[f32], centroids: &[Vec]) -> usize {
+ centroids
+ .iter()
+ .enumerate()
+ .map(|(i, centroid)| {
+ let dist: f32 = subvector
+ .iter()
+ .zip(centroid.iter())
+ .map(|(a, b)| (a - b).powi(2))
+ .sum();
+ (i, dist)
+ })
+ .min_by(|(_, a), (_, b)| {
+ a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)
+ })
+ .map(|(i, _)| i)
+ .unwrap_or(0)
+ }
+}
+
+// === Half precision conversion helpers ===
+
+/// Convert f32 to f16 bits (simplified implementation)
+fn f32_to_f16_bits(value: f32) -> u16 {
+ // Simple conversion: scale to 16-bit range
+ // This is a simplified version, not IEEE 754 half precision
+ let scaled = (value * 1000.0).clamp(-32768.0, 32767.0);
+ ((scaled as i32) + 32768) as u16
+}
+
+/// Convert f16 bits to f32 (simplified implementation)
+fn f16_bits_to_f32(bits: u16) -> f32 {
+ // Reverse of f32_to_f16_bits
+ let value = bits as i32 - 32768;
+ value as f32 / 1000.0
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_compress_none() {
+ let compressor = TensorCompress::new();
+ let embedding = vec![1.0, 2.0, 3.0, 4.0];
+
+ let compressed = compressor.compress(&embedding, 1.0).unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ assert_eq!(embedding, decompressed);
+ }
+
+ #[test]
+ fn test_compress_half() {
+ let compressor = TensorCompress::new();
+ let embedding = vec![1.0, 2.0, 3.0, 4.0];
+
+ let compressed = compressor.compress(&embedding, 0.5).unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ // Half precision should be close but not exact
+ for (a, b) in embedding.iter().zip(decompressed.iter()) {
+ assert!((a - b).abs() < 0.01, "Expected {}, got {}", a, b);
+ }
+ }
+
+ #[test]
+ fn test_compress_binary() {
+ let compressor = TensorCompress::new();
+ let embedding = vec![1.0, -1.0, 0.5, -0.5];
+
+ let compressed = compressor.compress(&embedding, 0.005).unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ // Binary should be +1 or -1
+ assert_eq!(decompressed.len(), embedding.len());
+ for val in &decompressed {
+ assert!(*val == 1.0 || *val == -1.0);
+ }
+ }
+
+ #[test]
+ fn test_select_level() {
+ let compressor = TensorCompress::new();
+
+ // Hot data
+ assert!(matches!(
+ compressor.select_level(0.9),
+ CompressionLevel::None
+ ));
+
+ // Warm data
+ assert!(matches!(
+ compressor.select_level(0.5),
+ CompressionLevel::Half { .. }
+ ));
+
+ // Cool data
+ assert!(matches!(
+ compressor.select_level(0.2),
+ CompressionLevel::PQ8 { .. }
+ ));
+
+ // Cold data
+ assert!(matches!(
+ compressor.select_level(0.05),
+ CompressionLevel::PQ4 { .. }
+ ));
+
+ // Archive
+ assert!(matches!(
+ compressor.select_level(0.001),
+ CompressionLevel::Binary { .. }
+ ));
+ }
+
+ #[test]
+ fn test_empty_embedding() {
+ let compressor = TensorCompress::new();
+ let result = compressor.compress(&[], 0.5);
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_pq8_compression() {
+ let compressor = TensorCompress::new();
+ let embedding: Vec = (0..64).map(|i| i as f32 * 0.1).collect();
+
+ let compressed = compressor.compress_pq8(&embedding, 8, 16).unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ assert_eq!(decompressed.len(), embedding.len());
+ }
+
+ #[test]
+ fn test_round_trip_all_levels() {
+ let compressor = TensorCompress::new();
+ let embedding: Vec = (0..128).map(|i| (i as f32 - 64.0) * 0.01).collect();
+
+ let access_frequencies = vec![0.9, 0.5, 0.2, 0.05, 0.001];
+
+ for freq in access_frequencies {
+ let compressed = compressor.compress(&embedding, freq).unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+ assert_eq!(decompressed.len(), embedding.len());
+ }
+ }
+
+ #[test]
+ fn test_half_precision_roundtrip() {
+ let compressor = TensorCompress::new();
+ // Use values within the supported range (-32.768 to 32.767)
+ let values = vec![-30.0, -1.0, 0.0, 1.0, 30.0];
+
+ for val in values {
+ let embedding = vec![val; 4];
+ let compressed = compressor
+ .compress_with_level(&embedding, &CompressionLevel::Half { scale: 1.0 })
+ .unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ for (a, b) in embedding.iter().zip(decompressed.iter()) {
+ let diff = (a - b).abs();
+ assert!(
+ diff < 0.1,
+ "Value {} decompressed to {}, diff: {}",
+ a,
+ b,
+ diff
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn test_binary_threshold() {
+ let compressor = TensorCompress::new();
+ let embedding = vec![0.5, -0.5, 1.5, -1.5];
+
+ let compressed = compressor
+ .compress_with_level(&embedding, &CompressionLevel::Binary { threshold: 0.0 })
+ .unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ // Values > 0 should be 1.0, values <= 0 should be -1.0
+ assert_eq!(decompressed, vec![1.0, -1.0, 1.0, -1.0]);
+ }
+
+ #[test]
+ fn test_pq4_with_outliers() {
+ let compressor = TensorCompress::new();
+ // Create embedding with some outliers
+ let mut embedding: Vec = (0..64).map(|i| i as f32 * 0.01).collect();
+ embedding[10] = 100.0; // Outlier
+ embedding[30] = -100.0; // Outlier
+
+ let compressed = compressor
+ .compress_with_level(
+ &embedding,
+ &CompressionLevel::PQ4 {
+ subvectors: 8,
+ outlier_threshold: 2.0,
+ },
+ )
+ .unwrap();
+ let decompressed = compressor.decompress(&compressed).unwrap();
+
+ assert_eq!(decompressed.len(), embedding.len());
+ // Outliers should be preserved
+ assert_eq!(decompressed[10], 100.0);
+ assert_eq!(decompressed[30], -100.0);
+ }
+
+ #[test]
+ fn test_dimension_validation() {
+ let compressor = TensorCompress::new();
+ let embedding = vec![1.0; 10]; // Not divisible by 8
+
+ let result = compressor.compress_pq8(&embedding, 8, 16);
+ assert!(result.is_err());
+ }
+}
diff --git a/crates/ruvector-gnn/src/error.rs b/crates/ruvector-gnn/src/error.rs
new file mode 100644
index 000000000..e25dcf08f
--- /dev/null
+++ b/crates/ruvector-gnn/src/error.rs
@@ -0,0 +1,111 @@
+//! Error types for the GNN module.
+
+use thiserror::Error;
+
+/// Result type alias for GNN operations.
+pub type Result = std::result::Result;
+
+/// Errors that can occur during GNN operations.
+#[derive(Error, Debug)]
+pub enum GnnError {
+ /// Tensor dimension mismatch
+ #[error("Tensor dimension mismatch: expected {expected}, got {actual}")]
+ DimensionMismatch {
+ /// Expected dimension
+ expected: String,
+ /// Actual dimension
+ actual: String,
+ },
+
+ /// Invalid tensor shape
+ #[error("Invalid tensor shape: {0}")]
+ InvalidShape(String),
+
+ /// Layer configuration error
+ #[error("Layer configuration error: {0}")]
+ LayerConfig(String),
+
+ /// Training error
+ #[error("Training error: {0}")]
+ Training(String),
+
+ /// Compression error
+ #[error("Compression error: {0}")]
+ Compression(String),
+
+ /// Search error
+ #[error("Search error: {0}")]
+ Search(String),
+
+ /// Invalid input
+ #[error("Invalid input: {0}")]
+ InvalidInput(String),
+
+ /// Memory mapping error
+ #[cfg(not(target_arch = "wasm32"))]
+ #[error("Memory mapping error: {0}")]
+ Mmap(String),
+
+ /// I/O error
+ #[error("I/O error: {0}")]
+ Io(#[from] std::io::Error),
+
+ /// Core library error
+ #[error("Core error: {0}")]
+ Core(#[from] ruvector_core::error::RuvectorError),
+
+ /// Generic error
+ #[error("{0}")]
+ Other(String),
+}
+
+impl GnnError {
+ /// Create a dimension mismatch error
+ pub fn dimension_mismatch(expected: impl Into, actual: impl Into) -> Self {
+ Self::DimensionMismatch {
+ expected: expected.into(),
+ actual: actual.into(),
+ }
+ }
+
+ /// Create an invalid shape error
+ pub fn invalid_shape(msg: impl Into) -> Self {
+ Self::InvalidShape(msg.into())
+ }
+
+ /// Create a layer config error
+ pub fn layer_config(msg: impl Into) -> Self {
+ Self::LayerConfig(msg.into())
+ }
+
+ /// Create a training error
+ pub fn training(msg: impl Into) -> Self {
+ Self::Training(msg.into())
+ }
+
+ /// Create a compression error
+ pub fn compression(msg: impl Into) -> Self {
+ Self::Compression(msg.into())
+ }
+
+ /// Create a search error
+ pub fn search(msg: impl Into) -> Self {
+ Self::Search(msg.into())
+ }
+
+ /// Create a memory mapping error
+ #[cfg(not(target_arch = "wasm32"))]
+ pub fn mmap(msg: impl Into) -> Self {
+ Self::Mmap(msg.into())
+ }
+
+ /// Create an invalid input error
+ pub fn invalid_input(msg: impl Into) -> Self {
+ Self::InvalidInput(msg.into())
+ }
+
+ /// Create a generic error
+ pub fn other(msg: impl Into) -> Self {
+ Self::Other(msg.into())
+ }
+}
diff --git a/crates/ruvector-gnn/src/layer.rs b/crates/ruvector-gnn/src/layer.rs
new file mode 100644
index 000000000..568a2ad8b
--- /dev/null
+++ b/crates/ruvector-gnn/src/layer.rs
@@ -0,0 +1,521 @@
+//! GNN Layer Implementation for HNSW Topology
+//!
+//! This module implements graph neural network layers that operate on HNSW graph structure,
+//! including attention mechanisms, normalization, and gated recurrent updates.
+
+use ndarray::{Array1, Array2, ArrayView1};
+use rand::Rng;
+use rand_distr::{Distribution, Normal};
+use serde::{Deserialize, Serialize};
+
+/// Linear transformation layer (weight matrix multiplication)
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Linear {
+ weights: Array2,
+ bias: Array1,
+}
+
+impl Linear {
+ /// Create a new linear layer with Xavier/Glorot initialization
+ pub fn new(input_dim: usize, output_dim: usize) -> Self {
+ let mut rng = rand::thread_rng();
+
+ // Xavier initialization: scale = sqrt(2.0 / (input_dim + output_dim))
+ let scale = (2.0 / (input_dim + output_dim) as f32).sqrt();
+ let normal = Normal::new(0.0, scale as f64).unwrap();
+
+ let weights = Array2::from_shape_fn((output_dim, input_dim), |_| {
+ normal.sample(&mut rng) as f32
+ });
+
+ let bias = Array1::zeros(output_dim);
+
+ Self { weights, bias }
+ }
+
+ /// Forward pass: y = Wx + b
+ pub fn forward(&self, input: &[f32]) -> Vec {
+ let x = ArrayView1::from(input);
+ let output = self.weights.dot(&x) + &self.bias;
+ output.to_vec()
+ }
+
+ /// Get output dimension
+ pub fn output_dim(&self) -> usize {
+ self.weights.shape()[0]
+ }
+}
+
+/// Layer normalization
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LayerNorm {
+ gamma: Array1,
+ beta: Array1,
+ eps: f32,
+}
+
+impl LayerNorm {
+ /// Create a new layer normalization layer
+ pub fn new(dim: usize, eps: f32) -> Self {
+ Self {
+ gamma: Array1::ones(dim),
+ beta: Array1::zeros(dim),
+ eps,
+ }
+ }
+
+ /// Forward pass: normalize and scale
+ pub fn forward(&self, input: &[f32]) -> Vec {
+ let x = ArrayView1::from(input);
+
+ // Compute mean and variance
+ let mean = x.mean().unwrap_or(0.0);
+ let variance = x.iter()
+ .map(|&v| (v - mean).powi(2))
+ .sum::() / x.len() as f32;
+
+ // Normalize
+ let normalized = x.mapv(|v| (v - mean) / (variance + self.eps).sqrt());
+
+ // Scale and shift
+ let output = &self.gamma * &normalized + &self.beta;
+ output.to_vec()
+ }
+}
+
+/// Multi-head attention mechanism
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct MultiHeadAttention {
+ num_heads: usize,
+ head_dim: usize,
+ q_linear: Linear,
+ k_linear: Linear,
+ v_linear: Linear,
+ out_linear: Linear,
+}
+
+impl MultiHeadAttention {
+ /// Create a new multi-head attention layer
+ pub fn new(embed_dim: usize, num_heads: usize) -> Self {
+ assert!(
+ embed_dim % num_heads == 0,
+ "Embedding dimension must be divisible by number of heads"
+ );
+
+ let head_dim = embed_dim / num_heads;
+
+ Self {
+ num_heads,
+ head_dim,
+ q_linear: Linear::new(embed_dim, embed_dim),
+ k_linear: Linear::new(embed_dim, embed_dim),
+ v_linear: Linear::new(embed_dim, embed_dim),
+ out_linear: Linear::new(embed_dim, embed_dim),
+ }
+ }
+
+ /// Forward pass: compute multi-head attention
+ ///
+ /// # Arguments
+ /// * `query` - Query vector
+ /// * `keys` - Key vectors from neighbors
+ /// * `values` - Value vectors from neighbors
+ ///
+ /// # Returns
+ /// Attention-weighted output vector
+ pub fn forward(&self, query: &[f32], keys: &[Vec], values: &[Vec]) -> Vec {
+ if keys.is_empty() || values.is_empty() {
+ return query.to_vec();
+ }
+
+ // Project query, keys, and values
+ let q = self.q_linear.forward(query);
+ let k: Vec> = keys.iter().map(|k| self.k_linear.forward(k)).collect();
+ let v: Vec> = values.iter().map(|v| self.v_linear.forward(v)).collect();
+
+ // Reshape for multi-head attention
+ let q_heads = self.split_heads(&q);
+ let k_heads: Vec>> = k.iter()
+ .map(|k_vec| self.split_heads(k_vec))
+ .collect();
+ let v_heads: Vec>> = v.iter()
+ .map(|v_vec| self.split_heads(v_vec))
+ .collect();
+
+ // Compute attention for each head
+ let mut head_outputs = Vec::new();
+ for h in 0..self.num_heads {
+ let q_h = &q_heads[h];
+ let k_h: Vec<&Vec> = k_heads.iter().map(|heads| &heads[h]).collect();
+ let v_h: Vec<&Vec> = v_heads.iter().map(|heads| &heads[h]).collect();
+
+ let head_output = self.scaled_dot_product_attention(q_h, &k_h, &v_h);
+ head_outputs.push(head_output);
+ }
+
+ // Concatenate heads
+ let concat: Vec = head_outputs.into_iter().flatten().collect();
+
+ // Final linear projection
+ self.out_linear.forward(&concat)
+ }
+
+ /// Split vector into multiple heads
+ fn split_heads(&self, x: &[f32]) -> Vec> {
+ let mut heads = Vec::new();
+ for h in 0..self.num_heads {
+ let start = h * self.head_dim;
+ let end = start + self.head_dim;
+ heads.push(x[start..end].to_vec());
+ }
+ heads
+ }
+
+ /// Scaled dot-product attention
+ fn scaled_dot_product_attention(
+ &self,
+ query: &[f32],
+ keys: &[&Vec],
+ values: &[&Vec],
+ ) -> Vec {
+ if keys.is_empty() {
+ return query.to_vec();
+ }
+
+ let scale = (self.head_dim as f32).sqrt();
+
+ // Compute attention scores
+ let scores: Vec = keys.iter()
+ .map(|k| {
+ let dot: f32 = query.iter().zip(k.iter()).map(|(q, k)| q * k).sum();
+ dot / scale
+ })
+ .collect();
+
+ // Softmax
+ let max_score = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+ let exp_scores: Vec = scores.iter().map(|&s| (s - max_score).exp()).collect();
+ let sum_exp: f32 = exp_scores.iter().sum();
+ let attention_weights: Vec = exp_scores.iter().map(|&e| e / sum_exp).collect();
+
+ // Weighted sum of values
+ let mut output = vec![0.0; self.head_dim];
+ for (weight, value) in attention_weights.iter().zip(values.iter()) {
+ for (out, &val) in output.iter_mut().zip(value.iter()) {
+ *out += weight * val;
+ }
+ }
+
+ output
+ }
+}
+
+/// Gated Recurrent Unit (GRU) cell for state updates
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct GRUCell {
+ // Update gate
+ w_z: Linear,
+ u_z: Linear,
+
+ // Reset gate
+ w_r: Linear,
+ u_r: Linear,
+
+ // Candidate hidden state
+ w_h: Linear,
+ u_h: Linear,
+}
+
+impl GRUCell {
+ /// Create a new GRU cell
+ pub fn new(input_dim: usize, hidden_dim: usize) -> Self {
+ Self {
+ // Update gate
+ w_z: Linear::new(input_dim, hidden_dim),
+ u_z: Linear::new(hidden_dim, hidden_dim),
+
+ // Reset gate
+ w_r: Linear::new(input_dim, hidden_dim),
+ u_r: Linear::new(hidden_dim, hidden_dim),
+
+ // Candidate hidden state
+ w_h: Linear::new(input_dim, hidden_dim),
+ u_h: Linear::new(hidden_dim, hidden_dim),
+ }
+ }
+
+ /// Forward pass: update hidden state
+ ///
+ /// # Arguments
+ /// * `input` - Current input
+ /// * `hidden` - Previous hidden state
+ ///
+ /// # Returns
+ /// Updated hidden state
+ pub fn forward(&self, input: &[f32], hidden: &[f32]) -> Vec {
+ // Update gate: z_t = sigmoid(W_z * x_t + U_z * h_{t-1})
+ let z = self.sigmoid_vec(&self.add_vecs(
+ &self.w_z.forward(input),
+ &self.u_z.forward(hidden),
+ ));
+
+ // Reset gate: r_t = sigmoid(W_r * x_t + U_r * h_{t-1})
+ let r = self.sigmoid_vec(&self.add_vecs(
+ &self.w_r.forward(input),
+ &self.u_r.forward(hidden),
+ ));
+
+ // Candidate hidden state: h_tilde = tanh(W_h * x_t + U_h * (r_t β h_{t-1}))
+ let r_hidden = self.mul_vecs(&r, hidden);
+ let h_tilde = self.tanh_vec(&self.add_vecs(
+ &self.w_h.forward(input),
+ &self.u_h.forward(&r_hidden),
+ ));
+
+ // Final hidden state: h_t = (1 - z_t) β h_{t-1} + z_t β h_tilde
+ let one_minus_z: Vec = z.iter().map(|&zval| 1.0 - zval).collect();
+ let term1 = self.mul_vecs(&one_minus_z, hidden);
+ let term2 = self.mul_vecs(&z, &h_tilde);
+
+ self.add_vecs(&term1, &term2)
+ }
+
+ /// Sigmoid activation
+ fn sigmoid(&self, x: f32) -> f32 {
+ 1.0 / (1.0 + (-x).exp())
+ }
+
+ /// Sigmoid for vectors
+ fn sigmoid_vec(&self, v: &[f32]) -> Vec {
+ v.iter().map(|&x| self.sigmoid(x)).collect()
+ }
+
+ /// Tanh activation
+ fn tanh(&self, x: f32) -> f32 {
+ x.tanh()
+ }
+
+ /// Tanh for vectors
+ fn tanh_vec(&self, v: &[f32]) -> Vec {
+ v.iter().map(|&x| self.tanh(x)).collect()
+ }
+
+ /// Element-wise addition
+ fn add_vecs(&self, a: &[f32], b: &[f32]) -> Vec {
+ a.iter().zip(b.iter()).map(|(x, y)| x + y).collect()
+ }
+
+ /// Element-wise multiplication
+ fn mul_vecs(&self, a: &[f32], b: &[f32]) -> Vec {
+ a.iter().zip(b.iter()).map(|(x, y)| x * y).collect()
+ }
+}
+
+/// Main GNN layer operating on HNSW topology
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RuvectorLayer {
+ /// Message weight matrix
+ w_msg: Linear,
+
+ /// Aggregation weight matrix
+ w_agg: Linear,
+
+ /// GRU update cell
+ w_update: GRUCell,
+
+ /// Multi-head attention
+ attention: MultiHeadAttention,
+
+ /// Layer normalization
+ norm: LayerNorm,
+
+ /// Dropout rate
+ dropout: f32,
+}
+
+impl RuvectorLayer {
+ /// Create a new Ruvector GNN layer
+ ///
+ /// # Arguments
+ /// * `input_dim` - Dimension of input node embeddings
+ /// * `hidden_dim` - Dimension of hidden representations
+ /// * `heads` - Number of attention heads
+ /// * `dropout` - Dropout rate (0.0 to 1.0)
+ pub fn new(input_dim: usize, hidden_dim: usize, heads: usize, dropout: f32) -> Self {
+ assert!(
+ dropout >= 0.0 && dropout <= 1.0,
+ "Dropout must be between 0.0 and 1.0"
+ );
+
+ Self {
+ w_msg: Linear::new(input_dim, hidden_dim),
+ w_agg: Linear::new(hidden_dim, hidden_dim),
+ w_update: GRUCell::new(hidden_dim, hidden_dim),
+ attention: MultiHeadAttention::new(hidden_dim, heads),
+ norm: LayerNorm::new(hidden_dim, 1e-5),
+ dropout,
+ }
+ }
+
+ /// Forward pass through the GNN layer
+ ///
+ /// # Arguments
+ /// * `node_embedding` - Current node's embedding
+ /// * `neighbor_embeddings` - Embeddings of neighbor nodes
+ /// * `edge_weights` - Weights of edges to neighbors (e.g., distances)
+ ///
+ /// # Returns
+ /// Updated node embedding
+ pub fn forward(
+ &self,
+ node_embedding: &[f32],
+ neighbor_embeddings: &[Vec],
+ edge_weights: &[f32],
+ ) -> Vec {
+ if neighbor_embeddings.is_empty() {
+ // No neighbors: return normalized projection
+ let projected = self.w_msg.forward(node_embedding);
+ return self.norm.forward(&projected);
+ }
+
+ // Step 1: Message passing - transform node and neighbor embeddings
+ let node_msg = self.w_msg.forward(node_embedding);
+ let neighbor_msgs: Vec> = neighbor_embeddings
+ .iter()
+ .map(|n| self.w_msg.forward(n))
+ .collect();
+
+ // Step 2: Attention-based aggregation
+ let attention_output = self.attention.forward(
+ &node_msg,
+ &neighbor_msgs,
+ &neighbor_msgs,
+ );
+
+ // Step 3: Weighted aggregation using edge weights
+ let weighted_msgs = self.aggregate_messages(&neighbor_msgs, edge_weights);
+
+ // Step 4: Combine attention and weighted aggregation
+ let combined = self.add_vecs(&attention_output, &weighted_msgs);
+ let aggregated = self.w_agg.forward(&combined);
+
+ // Step 5: GRU update
+ let updated = self.w_update.forward(&aggregated, &node_msg);
+
+ // Step 6: Apply dropout (simplified - always apply scaling)
+ let dropped = self.apply_dropout(&updated);
+
+ // Step 7: Layer normalization
+ self.norm.forward(&dropped)
+ }
+
+ /// Aggregate neighbor messages with edge weights
+ fn aggregate_messages(&self, messages: &[Vec], weights: &[f32]) -> Vec {
+ if messages.is_empty() || weights.is_empty() {
+ return vec![0.0; self.w_msg.output_dim()];
+ }
+
+ // Normalize weights to sum to 1
+ let weight_sum: f32 = weights.iter().sum();
+ let normalized_weights: Vec = if weight_sum > 0.0 {
+ weights.iter().map(|&w| w / weight_sum).collect()
+ } else {
+ vec![1.0 / weights.len() as f32; weights.len()]
+ };
+
+ // Weighted sum
+ let dim = messages[0].len();
+ let mut aggregated = vec![0.0; dim];
+
+ for (msg, &weight) in messages.iter().zip(normalized_weights.iter()) {
+ for (agg, &m) in aggregated.iter_mut().zip(msg.iter()) {
+ *agg += weight * m;
+ }
+ }
+
+ aggregated
+ }
+
+ /// Apply dropout (simplified version - just scales by (1-dropout))
+ fn apply_dropout(&self, input: &[f32]) -> Vec {
+ let scale = 1.0 - self.dropout;
+ input.iter().map(|&x| x * scale).collect()
+ }
+
+ /// Element-wise vector addition
+ fn add_vecs(&self, a: &[f32], b: &[f32]) -> Vec {
+ a.iter().zip(b.iter()).map(|(x, y)| x + y).collect()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_linear_layer() {
+ let linear = Linear::new(4, 2);
+ let input = vec![1.0, 2.0, 3.0, 4.0];
+ let output = linear.forward(&input);
+ assert_eq!(output.len(), 2);
+ }
+
+ #[test]
+ fn test_layer_norm() {
+ let norm = LayerNorm::new(4, 1e-5);
+ let input = vec![1.0, 2.0, 3.0, 4.0];
+ let output = norm.forward(&input);
+
+ // Check that output has zero mean (approximately)
+ let mean: f32 = output.iter().sum::() / output.len() as f32;
+ assert!((mean).abs() < 1e-5);
+ }
+
+ #[test]
+ fn test_multihead_attention() {
+ let attention = MultiHeadAttention::new(8, 2);
+ let query = vec![0.5; 8];
+ let keys = vec![vec![0.3; 8], vec![0.7; 8]];
+ let values = vec![vec![0.2; 8], vec![0.8; 8]];
+
+ let output = attention.forward(&query, &keys, &values);
+ assert_eq!(output.len(), 8);
+ }
+
+ #[test]
+ fn test_gru_cell() {
+ let gru = GRUCell::new(4, 8);
+ let input = vec![1.0; 4];
+ let hidden = vec![0.5; 8];
+
+ let new_hidden = gru.forward(&input, &hidden);
+ assert_eq!(new_hidden.len(), 8);
+ }
+
+ #[test]
+ fn test_ruvector_layer() {
+ let layer = RuvectorLayer::new(4, 8, 2, 0.1);
+
+ let node = vec![1.0, 2.0, 3.0, 4.0];
+ let neighbors = vec![
+ vec![0.5, 1.0, 1.5, 2.0],
+ vec![2.0, 3.0, 4.0, 5.0],
+ ];
+ let weights = vec![0.3, 0.7];
+
+ let output = layer.forward(&node, &neighbors, &weights);
+ assert_eq!(output.len(), 8);
+ }
+
+ #[test]
+ fn test_ruvector_layer_no_neighbors() {
+ let layer = RuvectorLayer::new(4, 8, 2, 0.1);
+
+ let node = vec![1.0, 2.0, 3.0, 4.0];
+ let neighbors: Vec> = vec![];
+ let weights: Vec = vec![];
+
+ let output = layer.forward(&node, &neighbors, &weights);
+ assert_eq!(output.len(), 8);
+ }
+}
diff --git a/crates/ruvector-gnn/src/lib.rs b/crates/ruvector-gnn/src/lib.rs
new file mode 100644
index 000000000..ecf7eb182
--- /dev/null
+++ b/crates/ruvector-gnn/src/lib.rs
@@ -0,0 +1,40 @@
+//! # RuVector GNN
+//!
+//! Graph Neural Network capabilities for RuVector, providing tensor operations,
+//! GNN layers, compression, and differentiable search.
+
+#![warn(missing_docs)]
+#![deny(unsafe_op_in_unsafe_fn)]
+
+pub mod error;
+pub mod layer;
+pub mod tensor;
+pub mod compress;
+pub mod search;
+pub mod training;
+pub mod query;
+
+#[cfg(all(not(target_arch = "wasm32"), feature = "mmap"))]
+pub mod mmap;
+
+// Re-export commonly used types
+pub use error::{GnnError, Result};
+pub use layer::RuvectorLayer;
+pub use compress::{CompressedTensor, CompressionLevel, TensorCompress};
+pub use search::{cosine_similarity, differentiable_search, hierarchical_forward};
+pub use training::{TrainConfig, OnlineConfig, info_nce_loss, local_contrastive_loss, sgd_step};
+pub use query::{QueryMode, RuvectorQuery, QueryResult, SubGraph};
+
+#[cfg(all(not(target_arch = "wasm32"), feature = "mmap"))]
+pub use mmap::{AtomicBitmap, MmapManager, MmapGradientAccumulator};
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_basic() {
+ // Basic smoke test to ensure the crate compiles
+ assert!(true);
+ }
+}
diff --git a/crates/ruvector-gnn/src/mmap.rs b/crates/ruvector-gnn/src/mmap.rs
new file mode 100644
index 000000000..5755bb1c6
--- /dev/null
+++ b/crates/ruvector-gnn/src/mmap.rs
@@ -0,0 +1,863 @@
+//! Memory-mapped embedding management for large-scale GNN training.
+//!
+//! This module provides efficient memory-mapped access to embeddings and gradients
+//! that don't fit in RAM. It includes:
+//! - `MmapManager`: Memory-mapped embedding storage with dirty tracking
+//! - `MmapGradientAccumulator`: Lock-free gradient accumulation
+//! - `AtomicBitmap`: Thread-safe bitmap for access/dirty tracking
+//!
+//! Only available on non-WASM targets.
+
+#![cfg(all(not(target_arch = "wasm32"), feature = "mmap"))]
+
+use crate::error::{GnnError, Result};
+use std::fs::{File, OpenOptions};
+use std::io::{self, Write};
+use std::path::Path;
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use parking_lot::RwLock;
+use memmap2::{MmapMut, MmapOptions};
+
+/// Thread-safe bitmap using atomic operations.
+///
+/// Used for tracking which embeddings have been accessed or modified.
+/// Each bit represents one embedding node.
+#[derive(Debug)]
+pub struct AtomicBitmap {
+ /// Array of 64-bit atomic integers, each storing 64 bits
+ bits: Vec,
+ /// Total number of bits (nodes)
+ size: usize,
+}
+
+impl AtomicBitmap {
+ /// Create a new atomic bitmap with the specified capacity.
+ ///
+ /// # Arguments
+ /// * `size` - Number of bits to allocate
+ pub fn new(size: usize) -> Self {
+ let num_words = (size + 63) / 64;
+ let bits = (0..num_words)
+ .map(|_| AtomicU64::new(0))
+ .collect();
+
+ Self { bits, size }
+ }
+
+ /// Set a bit to 1 (mark as accessed/dirty).
+ ///
+ /// # Arguments
+ /// * `index` - Bit index to set
+ pub fn set(&self, index: usize) {
+ if index >= self.size {
+ return;
+ }
+ let word_idx = index / 64;
+ let bit_idx = index % 64;
+ self.bits[word_idx].fetch_or(1u64 << bit_idx, Ordering::Release);
+ }
+
+ /// Clear a bit to 0 (mark as clean/not accessed).
+ ///
+ /// # Arguments
+ /// * `index` - Bit index to clear
+ pub fn clear(&self, index: usize) {
+ if index >= self.size {
+ return;
+ }
+ let word_idx = index / 64;
+ let bit_idx = index % 64;
+ self.bits[word_idx].fetch_and(!(1u64 << bit_idx), Ordering::Release);
+ }
+
+ /// Check if a bit is set.
+ ///
+ /// # Arguments
+ /// * `index` - Bit index to check
+ ///
+ /// # Returns
+ /// `true` if the bit is set, `false` otherwise
+ pub fn get(&self, index: usize) -> bool {
+ if index >= self.size {
+ return false;
+ }
+ let word_idx = index / 64;
+ let bit_idx = index % 64;
+ let word = self.bits[word_idx].load(Ordering::Acquire);
+ (word & (1u64 << bit_idx)) != 0
+ }
+
+ /// Clear all bits in the bitmap.
+ pub fn clear_all(&self) {
+ for word in &self.bits {
+ word.store(0, Ordering::Release);
+ }
+ }
+
+ /// Get all set bit indices (for finding dirty pages).
+ ///
+ /// # Returns
+ /// Vector of indices where bits are set
+ pub fn get_set_indices(&self) -> Vec {
+ let mut indices = Vec::new();
+ for (word_idx, word) in self.bits.iter().enumerate() {
+ let mut w = word.load(Ordering::Acquire);
+ while w != 0 {
+ let bit_idx = w.trailing_zeros() as usize;
+ indices.push(word_idx * 64 + bit_idx);
+ w &= w - 1; // Clear lowest set bit
+ }
+ }
+ indices
+ }
+}
+
+/// Memory-mapped embedding manager with dirty tracking and prefetching.
+///
+/// Manages large embedding matrices that may not fit in RAM using memory-mapped files.
+/// Tracks which embeddings have been accessed and modified for efficient I/O.
+#[derive(Debug)]
+pub struct MmapManager {
+ /// The memory-mapped file
+ file: File,
+ /// Mutable memory mapping
+ mmap: MmapMut,
+ /// Operating system page size
+ page_size: usize,
+ /// Embedding dimension
+ d_embed: usize,
+ /// Bitmap tracking which embeddings have been accessed
+ access_bitmap: AtomicBitmap,
+ /// Bitmap tracking which embeddings have been modified
+ dirty_bitmap: AtomicBitmap,
+ /// Pin count for each page (prevents eviction)
+ pin_count: Vec,
+ /// Maximum number of nodes
+ max_nodes: usize,
+}
+
+impl MmapManager {
+ /// Create a new memory-mapped embedding manager.
+ ///
+ /// # Arguments
+ /// * `path` - Path to the memory-mapped file
+ /// * `d_embed` - Embedding dimension
+ /// * `max_nodes` - Maximum number of nodes to support
+ ///
+ /// # Returns
+ /// A new `MmapManager` instance
+ pub fn new(path: &Path, d_embed: usize, max_nodes: usize) -> Result {
+ // Calculate required file size
+ let embedding_size = d_embed * std::mem::size_of::();
+ let file_size = max_nodes * embedding_size;
+
+ // Create or open the file
+ let file = OpenOptions::new()
+ .read(true)
+ .write(true)
+ .create(true)
+ .open(path)
+ .map_err(|e| GnnError::mmap(format!("Failed to open mmap file: {}", e)))?;
+
+ // Set file size
+ file.set_len(file_size as u64)
+ .map_err(|e| GnnError::mmap(format!("Failed to set file size: {}", e)))?;
+
+ // Create memory mapping
+ let mmap = unsafe {
+ MmapOptions::new()
+ .len(file_size)
+ .map_mut(&file)
+ .map_err(|e| GnnError::mmap(format!("Failed to create mmap: {}", e)))?
+ };
+
+ // Get system page size
+ let page_size = page_size::get();
+ let num_pages = (file_size + page_size - 1) / page_size;
+
+ Ok(Self {
+ file,
+ mmap,
+ page_size,
+ d_embed,
+ access_bitmap: AtomicBitmap::new(max_nodes),
+ dirty_bitmap: AtomicBitmap::new(max_nodes),
+ pin_count: (0..num_pages).map(|_| AtomicU32::new(0)).collect(),
+ max_nodes,
+ })
+ }
+
+ /// Calculate the byte offset for a given node's embedding.
+ ///
+ /// # Arguments
+ /// * `node_id` - Node identifier
+ ///
+ /// # Returns
+ /// Byte offset in the memory-mapped file
+ #[inline]
+ pub fn embedding_offset(&self, node_id: u64) -> usize {
+ (node_id as usize) * self.d_embed * std::mem::size_of::()
+ }
+
+ /// Get a read-only reference to a node's embedding.
+ ///
+ /// # Arguments
+ /// * `node_id` - Node identifier
+ ///
+ /// # Returns
+ /// Slice containing the embedding vector
+ ///
+ /// # Panics
+ /// Panics if node_id is out of bounds
+ pub fn get_embedding(&self, node_id: u64) -> &[f32] {
+ let offset = self.embedding_offset(node_id);
+ let end = offset + self.d_embed * std::mem::size_of::();
+
+ // Mark as accessed
+ self.access_bitmap.set(node_id as usize);
+
+ // Safety: We control the offset and know the data is properly aligned
+ unsafe {
+ let ptr = self.mmap.as_ptr().add(offset) as *const f32;
+ std::slice::from_raw_parts(ptr, self.d_embed)
+ }
+ }
+
+ /// Set a node's embedding data.
+ ///
+ /// # Arguments
+ /// * `node_id` - Node identifier
+ /// * `data` - Embedding vector to write
+ ///
+ /// # Panics
+ /// Panics if node_id is out of bounds or data length doesn't match d_embed
+ pub fn set_embedding(&mut self, node_id: u64, data: &[f32]) {
+ assert_eq!(
+ data.len(),
+ self.d_embed,
+ "Embedding data length must match d_embed"
+ );
+
+ let offset = self.embedding_offset(node_id);
+
+ // Mark as accessed and dirty
+ self.access_bitmap.set(node_id as usize);
+ self.dirty_bitmap.set(node_id as usize);
+
+ // Safety: We control the offset and know the data is properly aligned
+ unsafe {
+ let ptr = self.mmap.as_mut_ptr().add(offset) as *mut f32;
+ std::ptr::copy_nonoverlapping(data.as_ptr(), ptr, self.d_embed);
+ }
+ }
+
+ /// Flush all dirty pages to disk.
+ ///
+ /// # Returns
+ /// `Ok(())` on success, error otherwise
+ pub fn flush_dirty(&self) -> io::Result<()> {
+ let dirty_nodes = self.dirty_bitmap.get_set_indices();
+
+ if dirty_nodes.is_empty() {
+ return Ok(());
+ }
+
+ // Flush the entire mmap for simplicity
+ // In a production system, you might want to flush only dirty pages
+ self.mmap.flush()?;
+
+ // Clear dirty bitmap after successful flush
+ for &node_id in &dirty_nodes {
+ self.dirty_bitmap.clear(node_id);
+ }
+
+ Ok(())
+ }
+
+ /// Prefetch embeddings into memory for better cache locality.
+ ///
+ /// # Arguments
+ /// * `node_ids` - List of node IDs to prefetch
+ pub fn prefetch(&self, node_ids: &[u64]) {
+ #[cfg(target_os = "linux")]
+ {
+ use std::os::unix::io::AsRawFd;
+
+ for &node_id in node_ids {
+ let offset = self.embedding_offset(node_id);
+ let page_offset = (offset / self.page_size) * self.page_size;
+ let length = self.d_embed * std::mem::size_of::();
+
+ unsafe {
+ // Use madvise to hint the kernel to prefetch
+ libc::madvise(
+ self.mmap.as_ptr().add(page_offset) as *mut libc::c_void,
+ length,
+ libc::MADV_WILLNEED,
+ );
+ }
+ }
+ }
+
+ // On non-Linux platforms, just access the data to bring it into cache
+ #[cfg(not(target_os = "linux"))]
+ {
+ for &node_id in node_ids {
+ let _ = self.get_embedding(node_id);
+ }
+ }
+ }
+
+ /// Get the embedding dimension.
+ pub fn d_embed(&self) -> usize {
+ self.d_embed
+ }
+
+ /// Get the maximum number of nodes.
+ pub fn max_nodes(&self) -> usize {
+ self.max_nodes
+ }
+}
+
+/// Memory-mapped gradient accumulator with fine-grained locking.
+///
+/// Allows multiple threads to accumulate gradients concurrently with minimal contention.
+/// Uses reader-writer locks at a configurable granularity.
+pub struct MmapGradientAccumulator {
+ /// Memory-mapped gradient storage (using UnsafeCell for interior mutability)
+ grad_mmap: std::cell::UnsafeCell,
+ /// Number of nodes per lock (lock granularity)
+ lock_granularity: usize,
+ /// Reader-writer locks for gradient regions
+ locks: Vec>,
+ /// Number of nodes
+ n_nodes: usize,
+ /// Embedding dimension
+ d_embed: usize,
+ /// Gradient file
+ _file: File,
+}
+
+impl MmapGradientAccumulator {
+ /// Create a new memory-mapped gradient accumulator.
+ ///
+ /// # Arguments
+ /// * `path` - Path to the gradient file
+ /// * `d_embed` - Embedding dimension
+ /// * `max_nodes` - Maximum number of nodes
+ ///
+ /// # Returns
+ /// A new `MmapGradientAccumulator` instance
+ pub fn new(path: &Path, d_embed: usize, max_nodes: usize) -> Result {
+ // Calculate required file size
+ let grad_size = d_embed * std::mem::size_of::();
+ let file_size = max_nodes * grad_size;
+
+ // Create or open the file
+ let file = OpenOptions::new()
+ .read(true)
+ .write(true)
+ .create(true)
+ .open(path)
+ .map_err(|e| GnnError::mmap(format!("Failed to open gradient file: {}", e)))?;
+
+ // Set file size
+ file.set_len(file_size as u64)
+ .map_err(|e| GnnError::mmap(format!("Failed to set gradient file size: {}", e)))?;
+
+ // Create memory mapping
+ let grad_mmap = unsafe {
+ MmapOptions::new()
+ .len(file_size)
+ .map_mut(&file)
+ .map_err(|e| GnnError::mmap(format!("Failed to create gradient mmap: {}", e)))?
+ };
+
+ // Zero out the gradients
+ for byte in grad_mmap.iter() {
+ // This forces the pages to be allocated and zeroed
+ let _ = byte;
+ }
+
+ // Use a lock granularity of 64 nodes per lock for good parallelism
+ let lock_granularity = 64;
+ let num_locks = (max_nodes + lock_granularity - 1) / lock_granularity;
+ let locks = (0..num_locks).map(|_| RwLock::new(())).collect();
+
+ Ok(Self {
+ grad_mmap: std::cell::UnsafeCell::new(grad_mmap),
+ lock_granularity,
+ locks,
+ n_nodes: max_nodes,
+ d_embed,
+ _file: file,
+ })
+ }
+
+ /// Calculate the byte offset for a node's gradient.
+ ///
+ /// # Arguments
+ /// * `node_id` - Node identifier
+ ///
+ /// # Returns
+ /// Byte offset in the gradient file
+ #[inline]
+ pub fn grad_offset(&self, node_id: u64) -> usize {
+ (node_id as usize) * self.d_embed * std::mem::size_of::()
+ }
+
+ /// Accumulate gradients for a specific node.
+ ///
+ /// # Arguments
+ /// * `node_id` - Node identifier
+ /// * `grad` - Gradient vector to accumulate
+ ///
+ /// # Panics
+ /// Panics if grad length doesn't match d_embed
+ pub fn accumulate(&self, node_id: u64, grad: &[f32]) {
+ assert_eq!(
+ grad.len(),
+ self.d_embed,
+ "Gradient length must match d_embed"
+ );
+
+ let lock_idx = (node_id as usize) / self.lock_granularity;
+ let _lock = self.locks[lock_idx].write();
+
+ let offset = self.grad_offset(node_id);
+
+ // Safety: We hold the write lock for this region, ensuring exclusive access
+ unsafe {
+ let mmap = &mut *self.grad_mmap.get();
+ let ptr = mmap.as_mut_ptr().add(offset) as *mut f32;
+ let grad_slice = std::slice::from_raw_parts_mut(ptr, self.d_embed);
+
+ // Accumulate gradients
+ for (g, &new_g) in grad_slice.iter_mut().zip(grad.iter()) {
+ *g += new_g;
+ }
+ }
+ }
+
+ /// Apply accumulated gradients to embeddings and zero out gradients.
+ ///
+ /// # Arguments
+ /// * `learning_rate` - Learning rate for gradient descent
+ /// * `embeddings` - Embedding manager to update
+ pub fn apply(&mut self, learning_rate: f32, embeddings: &mut MmapManager) {
+ assert_eq!(
+ self.d_embed,
+ embeddings.d_embed,
+ "Gradient and embedding dimensions must match"
+ );
+
+ // Process all nodes
+ for node_id in 0..self.n_nodes.min(embeddings.max_nodes) {
+ let grad = self.get_grad(node_id as u64);
+ let embedding = embeddings.get_embedding(node_id as u64);
+
+ // Apply gradient descent: embedding -= learning_rate * grad
+ let mut updated = vec![0.0f32; self.d_embed];
+ for i in 0..self.d_embed {
+ updated[i] = embedding[i] - learning_rate * grad[i];
+ }
+
+ embeddings.set_embedding(node_id as u64, &updated);
+ }
+
+ // Zero out gradients after applying
+ self.zero_grad();
+ }
+
+ /// Zero out all accumulated gradients.
+ pub fn zero_grad(&mut self) {
+ // Zero the entire gradient buffer
+ unsafe {
+ let mmap = &mut *self.grad_mmap.get();
+ for byte in mmap.iter_mut() {
+ *byte = 0;
+ }
+ }
+ }
+
+ /// Get a read-only reference to a node's accumulated gradient.
+ ///
+ /// # Arguments
+ /// * `node_id` - Node identifier
+ ///
+ /// # Returns
+ /// Slice containing the gradient vector
+ pub fn get_grad(&self, node_id: u64) -> &[f32] {
+ let lock_idx = (node_id as usize) / self.lock_granularity;
+ let _lock = self.locks[lock_idx].read();
+
+ let offset = self.grad_offset(node_id);
+
+ // Safety: We hold the read lock for this region
+ unsafe {
+ let mmap = &*self.grad_mmap.get();
+ let ptr = mmap.as_ptr().add(offset) as *const f32;
+ std::slice::from_raw_parts(ptr, self.d_embed)
+ }
+ }
+
+ /// Get the embedding dimension.
+ pub fn d_embed(&self) -> usize {
+ self.d_embed
+ }
+
+ /// Get the number of nodes.
+ pub fn n_nodes(&self) -> usize {
+ self.n_nodes
+ }
+}
+
+// Implement Drop to ensure proper cleanup
+impl Drop for MmapManager {
+ fn drop(&mut self) {
+ // Try to flush dirty pages before dropping
+ let _ = self.flush_dirty();
+ }
+}
+
+impl Drop for MmapGradientAccumulator {
+ fn drop(&mut self) {
+ // Flush gradient data
+ unsafe {
+ let mmap = &mut *self.grad_mmap.get();
+ let _ = mmap.flush();
+ }
+ }
+}
+
+// Safety: MmapGradientAccumulator is safe to send between threads
+// because access is protected by RwLocks
+unsafe impl Send for MmapGradientAccumulator {}
+unsafe impl Sync for MmapGradientAccumulator {}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use std::fs;
+ use tempfile::TempDir;
+
+ #[test]
+ fn test_atomic_bitmap_basic() {
+ let bitmap = AtomicBitmap::new(128);
+
+ assert!(!bitmap.get(0));
+ assert!(!bitmap.get(127));
+
+ bitmap.set(0);
+ bitmap.set(127);
+ bitmap.set(64);
+
+ assert!(bitmap.get(0));
+ assert!(bitmap.get(127));
+ assert!(bitmap.get(64));
+ assert!(!bitmap.get(1));
+
+ bitmap.clear(0);
+ assert!(!bitmap.get(0));
+ assert!(bitmap.get(127));
+ }
+
+ #[test]
+ fn test_atomic_bitmap_get_set_indices() {
+ let bitmap = AtomicBitmap::new(256);
+
+ bitmap.set(0);
+ bitmap.set(63);
+ bitmap.set(64);
+ bitmap.set(128);
+ bitmap.set(255);
+
+ let mut indices = bitmap.get_set_indices();
+ indices.sort();
+
+ assert_eq!(indices, vec![0, 63, 64, 128, 255]);
+ }
+
+ #[test]
+ fn test_atomic_bitmap_clear_all() {
+ let bitmap = AtomicBitmap::new(128);
+
+ bitmap.set(0);
+ bitmap.set(64);
+ bitmap.set(127);
+
+ assert!(bitmap.get(0));
+
+ bitmap.clear_all();
+
+ assert!(!bitmap.get(0));
+ assert!(!bitmap.get(64));
+ assert!(!bitmap.get(127));
+ }
+
+ #[test]
+ fn test_mmap_manager_creation() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let manager = MmapManager::new(&path, 128, 1000).unwrap();
+
+ assert_eq!(manager.d_embed(), 128);
+ assert_eq!(manager.max_nodes(), 1000);
+ assert!(path.exists());
+ }
+
+ #[test]
+ fn test_mmap_manager_set_get_embedding() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let mut manager = MmapManager::new(&path, 64, 100).unwrap();
+
+ let embedding = vec![1.0f32; 64];
+ manager.set_embedding(0, &embedding);
+
+ let retrieved = manager.get_embedding(0);
+ assert_eq!(retrieved.len(), 64);
+ assert_eq!(retrieved[0], 1.0);
+ assert_eq!(retrieved[63], 1.0);
+ }
+
+ #[test]
+ fn test_mmap_manager_multiple_embeddings() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let mut manager = MmapManager::new(&path, 32, 100).unwrap();
+
+ for i in 0..10 {
+ let embedding: Vec = (0..32).map(|j| (i * 32 + j) as f32).collect();
+ manager.set_embedding(i, &embedding);
+ }
+
+ // Verify each embedding
+ for i in 0..10 {
+ let retrieved = manager.get_embedding(i);
+ assert_eq!(retrieved.len(), 32);
+ assert_eq!(retrieved[0], (i * 32) as f32);
+ assert_eq!(retrieved[31], (i * 32 + 31) as f32);
+ }
+ }
+
+ #[test]
+ fn test_mmap_manager_dirty_tracking() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let mut manager = MmapManager::new(&path, 64, 100).unwrap();
+
+ let embedding = vec![2.0f32; 64];
+ manager.set_embedding(5, &embedding);
+
+ // Should be marked as dirty
+ assert!(manager.dirty_bitmap.get(5));
+
+ // Flush and check it's clean
+ manager.flush_dirty().unwrap();
+ assert!(!manager.dirty_bitmap.get(5));
+ }
+
+ #[test]
+ fn test_mmap_manager_persistence() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ {
+ let mut manager = MmapManager::new(&path, 64, 100).unwrap();
+ let embedding = vec![3.14f32; 64];
+ manager.set_embedding(10, &embedding);
+ manager.flush_dirty().unwrap();
+ }
+
+ // Reopen and verify data persisted
+ {
+ let manager = MmapManager::new(&path, 64, 100).unwrap();
+ let retrieved = manager.get_embedding(10);
+ assert_eq!(retrieved[0], 3.14);
+ assert_eq!(retrieved[63], 3.14);
+ }
+ }
+
+ #[test]
+ fn test_gradient_accumulator_creation() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("gradients.bin");
+
+ let accumulator = MmapGradientAccumulator::new(&path, 128, 1000).unwrap();
+
+ assert_eq!(accumulator.d_embed(), 128);
+ assert_eq!(accumulator.n_nodes(), 1000);
+ assert!(path.exists());
+ }
+
+ #[test]
+ fn test_gradient_accumulator_accumulate() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("gradients.bin");
+
+ let accumulator = MmapGradientAccumulator::new(&path, 64, 100).unwrap();
+
+ let grad1 = vec![1.0f32; 64];
+ let grad2 = vec![2.0f32; 64];
+
+ accumulator.accumulate(0, &grad1);
+ accumulator.accumulate(0, &grad2);
+
+ let accumulated = accumulator.get_grad(0);
+ assert_eq!(accumulated[0], 3.0);
+ assert_eq!(accumulated[63], 3.0);
+ }
+
+ #[test]
+ fn test_gradient_accumulator_zero_grad() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("gradients.bin");
+
+ let mut accumulator = MmapGradientAccumulator::new(&path, 64, 100).unwrap();
+
+ let grad = vec![1.5f32; 64];
+ accumulator.accumulate(0, &grad);
+
+ let accumulated = accumulator.get_grad(0);
+ assert_eq!(accumulated[0], 1.5);
+
+ accumulator.zero_grad();
+
+ let zeroed = accumulator.get_grad(0);
+ assert_eq!(zeroed[0], 0.0);
+ assert_eq!(zeroed[63], 0.0);
+ }
+
+ #[test]
+ fn test_gradient_accumulator_apply() {
+ let temp_dir = TempDir::new().unwrap();
+ let embed_path = temp_dir.path().join("embeddings.bin");
+ let grad_path = temp_dir.path().join("gradients.bin");
+
+ let mut embeddings = MmapManager::new(&embed_path, 32, 100).unwrap();
+ let mut accumulator = MmapGradientAccumulator::new(&grad_path, 32, 100).unwrap();
+
+ // Set initial embedding
+ let initial = vec![10.0f32; 32];
+ embeddings.set_embedding(0, &initial);
+
+ // Accumulate gradient
+ let grad = vec![1.0f32; 32];
+ accumulator.accumulate(0, &grad);
+
+ // Apply with learning rate 0.1
+ accumulator.apply(0.1, &mut embeddings);
+
+ // Check updated embedding: 10.0 - 0.1 * 1.0 = 9.9
+ let updated = embeddings.get_embedding(0);
+ assert!((updated[0] - 9.9).abs() < 1e-6);
+
+ // Check gradients were zeroed
+ let zeroed_grad = accumulator.get_grad(0);
+ assert_eq!(zeroed_grad[0], 0.0);
+ }
+
+ #[test]
+ fn test_gradient_accumulator_concurrent_accumulation() {
+ use std::thread;
+
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("gradients.bin");
+
+ let accumulator = std::sync::Arc::new(
+ MmapGradientAccumulator::new(&path, 64, 100).unwrap()
+ );
+
+ let mut handles = vec![];
+
+ // Spawn 10 threads, each accumulating 1.0 to node 0
+ for _ in 0..10 {
+ let acc = accumulator.clone();
+ let handle = thread::spawn(move || {
+ let grad = vec![1.0f32; 64];
+ acc.accumulate(0, &grad);
+ });
+ handles.push(handle);
+ }
+
+ for handle in handles {
+ handle.join().unwrap();
+ }
+
+ // Should have accumulated 10.0
+ let result = accumulator.get_grad(0);
+ assert_eq!(result[0], 10.0);
+ }
+
+ #[test]
+ fn test_embedding_offset_calculation() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let manager = MmapManager::new(&path, 64, 100).unwrap();
+
+ assert_eq!(manager.embedding_offset(0), 0);
+ assert_eq!(manager.embedding_offset(1), 64 * 4); // 64 floats * 4 bytes
+ assert_eq!(manager.embedding_offset(10), 64 * 4 * 10);
+ }
+
+ #[test]
+ fn test_grad_offset_calculation() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("gradients.bin");
+
+ let accumulator = MmapGradientAccumulator::new(&path, 128, 100).unwrap();
+
+ assert_eq!(accumulator.grad_offset(0), 0);
+ assert_eq!(accumulator.grad_offset(1), 128 * 4); // 128 floats * 4 bytes
+ assert_eq!(accumulator.grad_offset(5), 128 * 4 * 5);
+ }
+
+ #[test]
+ #[should_panic(expected = "Embedding data length must match d_embed")]
+ fn test_set_embedding_wrong_size() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let mut manager = MmapManager::new(&path, 64, 100).unwrap();
+ let wrong_size = vec![1.0f32; 32]; // Should be 64
+ manager.set_embedding(0, &wrong_size);
+ }
+
+ #[test]
+ #[should_panic(expected = "Gradient length must match d_embed")]
+ fn test_accumulate_wrong_size() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("gradients.bin");
+
+ let accumulator = MmapGradientAccumulator::new(&path, 64, 100).unwrap();
+ let wrong_size = vec![1.0f32; 32]; // Should be 64
+ accumulator.accumulate(0, &wrong_size);
+ }
+
+ #[test]
+ fn test_prefetch() {
+ let temp_dir = TempDir::new().unwrap();
+ let path = temp_dir.path().join("embeddings.bin");
+
+ let mut manager = MmapManager::new(&path, 64, 100).unwrap();
+
+ // Set some embeddings
+ for i in 0..10 {
+ let embedding = vec![i as f32; 64];
+ manager.set_embedding(i, &embedding);
+ }
+
+ // Prefetch should not crash
+ manager.prefetch(&[0, 1, 2, 3, 4]);
+
+ // Access should still work
+ let retrieved = manager.get_embedding(2);
+ assert_eq!(retrieved[0], 2.0);
+ }
+}
diff --git a/crates/ruvector-gnn/src/mmap_fixed.rs b/crates/ruvector-gnn/src/mmap_fixed.rs
new file mode 100644
index 000000000..8a5b2014c
--- /dev/null
+++ b/crates/ruvector-gnn/src/mmap_fixed.rs
@@ -0,0 +1,82 @@
+//! Memory-mapped embedding management for large-scale GNN training.
+//!
+//! This module provides efficient memory-mapped access to embeddings and gradients
+//! that don't fit in RAM. It includes:
+//! - `MmapManager`: Memory-mapped embedding storage with dirty tracking
+//! - `MmapGradientAccumulator`: Lock-free gradient accumulation
+//! - `AtomicBitmap`: Thread-safe bitmap for access/dirty tracking
+//!
+//! Only available on non-WASM targets.
+
+#![cfg(all(not(target_arch = "wasm32"), feature = "mmap"))]
+
+use crate::error::{GnnError, Result};
+use std::cell::UnsafeCell;
+use std::fs::{File, OpenOptions};
+use std::io;
+use std::path::Path;
+use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
+use parking_lot::RwLock;
+use memmap2::{MmapMut, MmapOptions};
+
+/// Thread-safe bitmap using atomic operations.
+#[derive(Debug)]
+pub struct AtomicBitmap {
+ bits: Vec,
+ size: usize,
+}
+
+impl AtomicBitmap {
+ pub fn new(size: usize) -> Self {
+ let num_words = (size + 63) / 64;
+ let bits = (0..num_words).map(|_| AtomicU64::new(0)).collect();
+ Self { bits, size }
+ }
+
+ pub fn set(&self, index: usize) {
+ if index >= self.size {
+ return;
+ }
+ let word_idx = index / 64;
+ let bit_idx = index % 64;
+ self.bits[word_idx].fetch_or(1u64 << bit_idx, Ordering::Release);
+ }
+
+ pub fn clear(&self, index: usize) {
+ if index >= self.size {
+ return;
+ }
+ let word_idx = index / 64;
+ let bit_idx = index % 64;
+ self.bits[word_idx].fetch_and(!(1u64 << bit_idx), Ordering::Release);
+ }
+
+ pub fn get(&self, index: usize) -> bool {
+ if index >= self.size {
+ return false;
+ }
+ let word_idx = index / 64;
+ let bit_idx = index % 64;
+ let word = self.bits[word_idx].load(Ordering::Acquire);
+ (word & (1u64 << bit_idx)) != 0
+ }
+
+ pub fn clear_all(&self) {
+ for word in &self.bits {
+ word.store(0, Ordering::Release);
+ }
+ }
+
+ pub fn get_set_indices(&self) -> Vec {
+ let mut indices = Vec::new();
+ for (word_idx, word) in self.bits.iter().enumerate() {
+ let mut w = word.load(Ordering::Acquire);
+ while w != 0 {
+ let bit_idx = w.trailing_zeros() as usize;
+ indices.push(word_idx * 64 + bit_idx);
+ w &= w - 1;
+ }
+ }
+ indices
+ }
+}
diff --git a/crates/ruvector-gnn/src/query.rs b/crates/ruvector-gnn/src/query.rs
new file mode 100644
index 000000000..103be9789
--- /dev/null
+++ b/crates/ruvector-gnn/src/query.rs
@@ -0,0 +1,670 @@
+//! Query API for RuVector GNN
+//!
+//! Provides high-level query interfaces for vector search, neural search,
+//! and subgraph extraction.
+
+use serde::{Serialize, Deserialize};
+
+/// Query mode for different search strategies
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+pub enum QueryMode {
+ /// Pure HNSW vector search
+ VectorSearch,
+ /// GNN-enhanced neural search
+ NeuralSearch,
+ /// Extract k-hop subgraph around results
+ SubgraphExtraction,
+ /// Differentiable search with soft attention
+ DifferentiableSearch,
+}
+
+/// Query configuration for RuVector searches
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RuvectorQuery {
+ /// Query vector for similarity search
+ pub vector: Option>,
+ /// Text query (requires embedding model)
+ pub text: Option,
+ /// Node ID for subgraph extraction
+ pub node_id: Option,
+ /// Search mode
+ pub mode: QueryMode,
+ /// Number of results to return
+ pub k: usize,
+ /// HNSW search parameter (exploration factor)
+ pub ef: usize,
+ /// GNN depth for neural search
+ pub gnn_depth: usize,
+ /// Temperature for differentiable search (higher = softer)
+ pub temperature: f32,
+ /// Whether to return attention weights
+ pub return_attention: bool,
+}
+
+impl Default for RuvectorQuery {
+ fn default() -> Self {
+ Self {
+ vector: None,
+ text: None,
+ node_id: None,
+ mode: QueryMode::VectorSearch,
+ k: 10,
+ ef: 50,
+ gnn_depth: 2,
+ temperature: 1.0,
+ return_attention: false,
+ }
+ }
+}
+
+impl RuvectorQuery {
+ /// Create a basic vector search query
+ ///
+ /// # Arguments
+ /// * `vector` - Query vector
+ /// * `k` - Number of results to return
+ ///
+ /// # Example
+ /// ```
+ /// use ruvector_gnn::query::RuvectorQuery;
+ ///
+ /// let query = RuvectorQuery::vector_search(vec![0.1, 0.2, 0.3], 10);
+ /// assert_eq!(query.k, 10);
+ /// ```
+ pub fn vector_search(vector: Vec, k: usize) -> Self {
+ Self {
+ vector: Some(vector),
+ mode: QueryMode::VectorSearch,
+ k,
+ ..Default::default()
+ }
+ }
+
+ /// Create a GNN-enhanced neural search query
+ ///
+ /// # Arguments
+ /// * `vector` - Query vector
+ /// * `k` - Number of results to return
+ /// * `gnn_depth` - Number of GNN layers to apply
+ ///
+ /// # Example
+ /// ```
+ /// use ruvector_gnn::query::RuvectorQuery;
+ ///
+ /// let query = RuvectorQuery::neural_search(vec![0.1, 0.2, 0.3], 10, 3);
+ /// assert_eq!(query.gnn_depth, 3);
+ /// ```
+ pub fn neural_search(vector: Vec, k: usize, gnn_depth: usize) -> Self {
+ Self {
+ vector: Some(vector),
+ mode: QueryMode::NeuralSearch,
+ k,
+ gnn_depth,
+ ..Default::default()
+ }
+ }
+
+ /// Create a subgraph extraction query
+ ///
+ /// # Arguments
+ /// * `vector` - Query vector
+ /// * `k` - Number of nodes in subgraph
+ ///
+ /// # Example
+ /// ```
+ /// use ruvector_gnn::query::RuvectorQuery;
+ ///
+ /// let query = RuvectorQuery::subgraph_search(vec![0.1, 0.2, 0.3], 20);
+ /// assert_eq!(query.k, 20);
+ /// ```
+ pub fn subgraph_search(vector: Vec, k: usize) -> Self {
+ Self {
+ vector: Some(vector),
+ mode: QueryMode::SubgraphExtraction,
+ k,
+ ..Default::default()
+ }
+ }
+
+ /// Create a differentiable search query with temperature
+ ///
+ /// # Arguments
+ /// * `vector` - Query vector
+ /// * `k` - Number of results
+ /// * `temperature` - Softmax temperature (higher = softer distribution)
+ pub fn differentiable_search(vector: Vec, k: usize, temperature: f32) -> Self {
+ Self {
+ vector: Some(vector),
+ mode: QueryMode::DifferentiableSearch,
+ k,
+ temperature,
+ return_attention: true,
+ ..Default::default()
+ }
+ }
+
+ /// Set text query (requires embedding model)
+ pub fn with_text(mut self, text: String) -> Self {
+ self.text = Some(text);
+ self
+ }
+
+ /// Set node ID for centered queries
+ pub fn with_node(mut self, node_id: u64) -> Self {
+ self.node_id = Some(node_id);
+ self
+ }
+
+ /// Set EF parameter for HNSW search
+ pub fn with_ef(mut self, ef: usize) -> Self {
+ self.ef = ef;
+ self
+ }
+
+ /// Enable attention weight return
+ pub fn with_attention(mut self) -> Self {
+ self.return_attention = true;
+ self
+ }
+}
+
+/// Subgraph representation
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct SubGraph {
+ /// Node IDs in the subgraph
+ pub nodes: Vec,
+ /// Edges as (from, to, weight) tuples
+ pub edges: Vec<(u64, u64, f32)>,
+}
+
+impl SubGraph {
+ /// Create a new empty subgraph
+ pub fn new() -> Self {
+ Self {
+ nodes: Vec::new(),
+ edges: Vec::new(),
+ }
+ }
+
+ /// Create subgraph with nodes and edges
+ pub fn with_edges(nodes: Vec, edges: Vec<(u64, u64, f32)>) -> Self {
+ Self { nodes, edges }
+ }
+
+ /// Get number of nodes
+ pub fn node_count(&self) -> usize {
+ self.nodes.len()
+ }
+
+ /// Get number of edges
+ pub fn edge_count(&self) -> usize {
+ self.edges.len()
+ }
+
+ /// Check if subgraph contains a node
+ pub fn contains_node(&self, node_id: u64) -> bool {
+ self.nodes.contains(&node_id)
+ }
+
+ /// Get average edge weight
+ pub fn average_edge_weight(&self) -> f32 {
+ if self.edges.is_empty() {
+ return 0.0;
+ }
+ let sum: f32 = self.edges.iter().map(|(_, _, w)| w).sum();
+ sum / self.edges.len() as f32
+ }
+}
+
+impl Default for SubGraph {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+/// Query result with nodes, scores, and optional metadata
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct QueryResult {
+ /// Matched node IDs
+ pub nodes: Vec,
+ /// Similarity scores (higher = more similar)
+ pub scores: Vec,
+ /// Optional node embeddings after GNN processing
+ pub embeddings: Option>>,
+ /// Optional attention weights from differentiable search
+ pub attention_weights: Option>>,
+ /// Optional subgraph extraction
+ pub subgraph: Option,
+ /// Query latency in milliseconds
+ pub latency_ms: u64,
+}
+
+impl QueryResult {
+ /// Create a new empty query result
+ pub fn new() -> Self {
+ Self {
+ nodes: Vec::new(),
+ scores: Vec::new(),
+ embeddings: None,
+ attention_weights: None,
+ subgraph: None,
+ latency_ms: 0,
+ }
+ }
+
+ /// Create query result with nodes and scores
+ ///
+ /// # Arguments
+ /// * `nodes` - Node IDs
+ /// * `scores` - Similarity scores
+ ///
+ /// # Example
+ /// ```
+ /// use ruvector_gnn::query::QueryResult;
+ ///
+ /// let result = QueryResult::with_nodes(vec![1, 2, 3], vec![0.9, 0.8, 0.7]);
+ /// assert_eq!(result.nodes.len(), 3);
+ /// ```
+ pub fn with_nodes(nodes: Vec, scores: Vec) -> Self {
+ Self {
+ nodes,
+ scores,
+ embeddings: None,
+ attention_weights: None,
+ subgraph: None,
+ latency_ms: 0,
+ }
+ }
+
+ /// Add embeddings to the result
+ pub fn with_embeddings(mut self, embeddings: Vec>) -> Self {
+ self.embeddings = Some(embeddings);
+ self
+ }
+
+ /// Add attention weights to the result
+ pub fn with_attention(mut self, attention: Vec>) -> Self {
+ self.attention_weights = Some(attention);
+ self
+ }
+
+ /// Add subgraph to the result
+ pub fn with_subgraph(mut self, subgraph: SubGraph) -> Self {
+ self.subgraph = Some(subgraph);
+ self
+ }
+
+ /// Set query latency
+ pub fn with_latency(mut self, latency_ms: u64) -> Self {
+ self.latency_ms = latency_ms;
+ self
+ }
+
+ /// Get number of results
+ pub fn len(&self) -> usize {
+ self.nodes.len()
+ }
+
+ /// Check if result is empty
+ pub fn is_empty(&self) -> bool {
+ self.nodes.is_empty()
+ }
+
+ /// Get top-k results
+ pub fn top_k(&self, k: usize) -> Self {
+ let k = k.min(self.nodes.len());
+ Self {
+ nodes: self.nodes[..k].to_vec(),
+ scores: self.scores[..k].to_vec(),
+ embeddings: self.embeddings.as_ref().map(|e| e[..k].to_vec()),
+ attention_weights: self.attention_weights.as_ref().map(|a| a[..k].to_vec()),
+ subgraph: self.subgraph.clone(),
+ latency_ms: self.latency_ms,
+ }
+ }
+
+ /// Get the best result (highest score)
+ pub fn best(&self) -> Option<(u64, f32)> {
+ if self.nodes.is_empty() {
+ None
+ } else {
+ Some((self.nodes[0], self.scores[0]))
+ }
+ }
+
+ /// Filter results by minimum score
+ pub fn filter_by_score(mut self, min_score: f32) -> Self {
+ let mut filtered_nodes = Vec::new();
+ let mut filtered_scores = Vec::new();
+ let mut filtered_embeddings = Vec::new();
+ let mut filtered_attention = Vec::new();
+
+ for i in 0..self.nodes.len() {
+ if self.scores[i] >= min_score {
+ filtered_nodes.push(self.nodes[i]);
+ filtered_scores.push(self.scores[i]);
+
+ if let Some(ref emb) = self.embeddings {
+ filtered_embeddings.push(emb[i].clone());
+ }
+
+ if let Some(ref att) = self.attention_weights {
+ filtered_attention.push(att[i].clone());
+ }
+ }
+ }
+
+ self.nodes = filtered_nodes;
+ self.scores = filtered_scores;
+
+ if !filtered_embeddings.is_empty() {
+ self.embeddings = Some(filtered_embeddings);
+ }
+
+ if !filtered_attention.is_empty() {
+ self.attention_weights = Some(filtered_attention);
+ }
+
+ self
+ }
+}
+
+impl Default for QueryResult {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_query_mode_serialization() {
+ let mode = QueryMode::NeuralSearch;
+ let json = serde_json::to_string(&mode).unwrap();
+ let deserialized: QueryMode = serde_json::from_str(&json).unwrap();
+ assert_eq!(mode, deserialized);
+ }
+
+ #[test]
+ fn test_ruvector_query_default() {
+ let query = RuvectorQuery::default();
+ assert_eq!(query.k, 10);
+ assert_eq!(query.ef, 50);
+ assert_eq!(query.gnn_depth, 2);
+ assert_eq!(query.temperature, 1.0);
+ assert_eq!(query.mode, QueryMode::VectorSearch);
+ assert!(!query.return_attention);
+ }
+
+ #[test]
+ fn test_vector_search_query() {
+ let vector = vec![0.1, 0.2, 0.3, 0.4];
+ let query = RuvectorQuery::vector_search(vector.clone(), 5);
+
+ assert_eq!(query.vector, Some(vector));
+ assert_eq!(query.k, 5);
+ assert_eq!(query.mode, QueryMode::VectorSearch);
+ }
+
+ #[test]
+ fn test_neural_search_query() {
+ let vector = vec![0.1, 0.2, 0.3];
+ let query = RuvectorQuery::neural_search(vector.clone(), 10, 3);
+
+ assert_eq!(query.vector, Some(vector));
+ assert_eq!(query.k, 10);
+ assert_eq!(query.gnn_depth, 3);
+ assert_eq!(query.mode, QueryMode::NeuralSearch);
+ }
+
+ #[test]
+ fn test_subgraph_search_query() {
+ let vector = vec![0.5, 0.5];
+ let query = RuvectorQuery::subgraph_search(vector.clone(), 20);
+
+ assert_eq!(query.vector, Some(vector));
+ assert_eq!(query.k, 20);
+ assert_eq!(query.mode, QueryMode::SubgraphExtraction);
+ }
+
+ #[test]
+ fn test_differentiable_search_query() {
+ let vector = vec![0.3, 0.4, 0.5];
+ let query = RuvectorQuery::differentiable_search(vector.clone(), 15, 0.5);
+
+ assert_eq!(query.vector, Some(vector));
+ assert_eq!(query.k, 15);
+ assert_eq!(query.temperature, 0.5);
+ assert_eq!(query.mode, QueryMode::DifferentiableSearch);
+ assert!(query.return_attention);
+ }
+
+ #[test]
+ fn test_query_builder_pattern() {
+ let query = RuvectorQuery::vector_search(vec![0.1, 0.2], 5)
+ .with_text("hello world".to_string())
+ .with_node(42)
+ .with_ef(100)
+ .with_attention();
+
+ assert_eq!(query.text, Some("hello world".to_string()));
+ assert_eq!(query.node_id, Some(42));
+ assert_eq!(query.ef, 100);
+ assert!(query.return_attention);
+ }
+
+ #[test]
+ fn test_subgraph_new() {
+ let subgraph = SubGraph::new();
+ assert_eq!(subgraph.node_count(), 0);
+ assert_eq!(subgraph.edge_count(), 0);
+ }
+
+ #[test]
+ fn test_subgraph_with_edges() {
+ let nodes = vec![1, 2, 3];
+ let edges = vec![(1, 2, 0.8), (2, 3, 0.6), (1, 3, 0.5)];
+ let subgraph = SubGraph::with_edges(nodes.clone(), edges.clone());
+
+ assert_eq!(subgraph.nodes, nodes);
+ assert_eq!(subgraph.edges, edges);
+ assert_eq!(subgraph.node_count(), 3);
+ assert_eq!(subgraph.edge_count(), 3);
+ }
+
+ #[test]
+ fn test_subgraph_contains_node() {
+ let nodes = vec![1, 2, 3];
+ let subgraph = SubGraph::with_edges(nodes, vec![]);
+
+ assert!(subgraph.contains_node(1));
+ assert!(subgraph.contains_node(2));
+ assert!(subgraph.contains_node(3));
+ assert!(!subgraph.contains_node(4));
+ }
+
+ #[test]
+ fn test_subgraph_average_edge_weight() {
+ let edges = vec![(1, 2, 0.8), (2, 3, 0.6), (1, 3, 0.4)];
+ let subgraph = SubGraph::with_edges(vec![1, 2, 3], edges);
+
+ let avg = subgraph.average_edge_weight();
+ assert!((avg - 0.6).abs() < 0.001);
+ }
+
+ #[test]
+ fn test_subgraph_empty_average() {
+ let subgraph = SubGraph::new();
+ assert_eq!(subgraph.average_edge_weight(), 0.0);
+ }
+
+ #[test]
+ fn test_query_result_new() {
+ let result = QueryResult::new();
+ assert!(result.is_empty());
+ assert_eq!(result.len(), 0);
+ assert_eq!(result.latency_ms, 0);
+ }
+
+ #[test]
+ fn test_query_result_with_nodes() {
+ let nodes = vec![1, 2, 3];
+ let scores = vec![0.9, 0.8, 0.7];
+ let result = QueryResult::with_nodes(nodes.clone(), scores.clone());
+
+ assert_eq!(result.nodes, nodes);
+ assert_eq!(result.scores, scores);
+ assert_eq!(result.len(), 3);
+ assert!(!result.is_empty());
+ }
+
+ #[test]
+ fn test_query_result_builder_pattern() {
+ let embeddings = vec![vec![0.1, 0.2], vec![0.3, 0.4]];
+ let attention = vec![vec![0.5, 0.5], vec![0.6, 0.4]];
+ let subgraph = SubGraph::with_edges(vec![1, 2], vec![(1, 2, 0.8)]);
+
+ let result = QueryResult::with_nodes(vec![1, 2], vec![0.9, 0.8])
+ .with_embeddings(embeddings.clone())
+ .with_attention(attention.clone())
+ .with_subgraph(subgraph.clone())
+ .with_latency(100);
+
+ assert_eq!(result.embeddings, Some(embeddings));
+ assert_eq!(result.attention_weights, Some(attention));
+ assert_eq!(result.subgraph, Some(subgraph));
+ assert_eq!(result.latency_ms, 100);
+ }
+
+ #[test]
+ fn test_query_result_top_k() {
+ let nodes = vec![1, 2, 3, 4, 5];
+ let scores = vec![0.9, 0.8, 0.7, 0.6, 0.5];
+ let result = QueryResult::with_nodes(nodes, scores);
+
+ let top_3 = result.top_k(3);
+ assert_eq!(top_3.len(), 3);
+ assert_eq!(top_3.nodes, vec![1, 2, 3]);
+ assert_eq!(top_3.scores, vec![0.9, 0.8, 0.7]);
+ }
+
+ #[test]
+ fn test_query_result_top_k_overflow() {
+ let result = QueryResult::with_nodes(vec![1, 2], vec![0.9, 0.8]);
+ let top_10 = result.top_k(10);
+ assert_eq!(top_10.len(), 2); // Should only return available results
+ }
+
+ #[test]
+ fn test_query_result_best() {
+ let result = QueryResult::with_nodes(vec![1, 2, 3], vec![0.9, 0.8, 0.7]);
+ let best = result.best();
+ assert_eq!(best, Some((1, 0.9)));
+ }
+
+ #[test]
+ fn test_query_result_best_empty() {
+ let result = QueryResult::new();
+ assert_eq!(result.best(), None);
+ }
+
+ #[test]
+ fn test_query_result_filter_by_score() {
+ let nodes = vec![1, 2, 3, 4, 5];
+ let scores = vec![0.9, 0.8, 0.7, 0.6, 0.5];
+ let result = QueryResult::with_nodes(nodes, scores);
+
+ let filtered = result.filter_by_score(0.7);
+ assert_eq!(filtered.len(), 3);
+ assert_eq!(filtered.nodes, vec![1, 2, 3]);
+ assert_eq!(filtered.scores, vec![0.9, 0.8, 0.7]);
+ }
+
+ #[test]
+ fn test_query_result_filter_with_embeddings() {
+ let nodes = vec![1, 2, 3];
+ let scores = vec![0.9, 0.6, 0.8];
+ let embeddings = vec![vec![0.1], vec![0.2], vec![0.3]];
+
+ let result = QueryResult::with_nodes(nodes, scores)
+ .with_embeddings(embeddings);
+
+ let filtered = result.filter_by_score(0.7);
+ assert_eq!(filtered.len(), 2);
+ assert_eq!(filtered.nodes, vec![1, 3]);
+ assert_eq!(filtered.embeddings, Some(vec![vec![0.1], vec![0.3]]));
+ }
+
+ #[test]
+ fn test_query_result_filter_with_attention() {
+ let nodes = vec![1, 2, 3];
+ let scores = vec![0.9, 0.5, 0.8];
+ let attention = vec![vec![0.5, 0.5], vec![0.6, 0.4], vec![0.7, 0.3]];
+
+ let result = QueryResult::with_nodes(nodes, scores)
+ .with_attention(attention);
+
+ let filtered = result.filter_by_score(0.75);
+ assert_eq!(filtered.len(), 2);
+ assert_eq!(filtered.nodes, vec![1, 3]);
+ assert_eq!(filtered.attention_weights, Some(vec![vec![0.5, 0.5], vec![0.7, 0.3]]));
+ }
+
+ #[test]
+ fn test_query_serialization() {
+ let query = RuvectorQuery::neural_search(vec![0.1, 0.2], 5, 2);
+ let json = serde_json::to_string(&query).unwrap();
+ let deserialized: RuvectorQuery = serde_json::from_str(&json).unwrap();
+
+ assert_eq!(deserialized.k, query.k);
+ assert_eq!(deserialized.gnn_depth, query.gnn_depth);
+ assert_eq!(deserialized.mode, query.mode);
+ }
+
+ #[test]
+ fn test_result_serialization() {
+ let result = QueryResult::with_nodes(vec![1, 2], vec![0.9, 0.8])
+ .with_latency(50);
+
+ let json = serde_json::to_string(&result).unwrap();
+ let deserialized: QueryResult = serde_json::from_str(&json).unwrap();
+
+ assert_eq!(deserialized.nodes, result.nodes);
+ assert_eq!(deserialized.scores, result.scores);
+ assert_eq!(deserialized.latency_ms, result.latency_ms);
+ }
+
+ #[test]
+ fn test_subgraph_serialization() {
+ let subgraph = SubGraph::with_edges(
+ vec![1, 2, 3],
+ vec![(1, 2, 0.8), (2, 3, 0.6)]
+ );
+
+ let json = serde_json::to_string(&subgraph).unwrap();
+ let deserialized: SubGraph = serde_json::from_str(&json).unwrap();
+
+ assert_eq!(deserialized.nodes, subgraph.nodes);
+ assert_eq!(deserialized.edges, subgraph.edges);
+ }
+
+ #[test]
+ fn test_edge_case_empty_filter() {
+ let result = QueryResult::with_nodes(vec![1, 2], vec![0.5, 0.4]);
+ let filtered = result.filter_by_score(0.9);
+
+ assert!(filtered.is_empty());
+ assert_eq!(filtered.len(), 0);
+ }
+
+ #[test]
+ fn test_query_mode_variants() {
+ // Test all query mode variants
+ assert_eq!(QueryMode::VectorSearch, QueryMode::VectorSearch);
+ assert_ne!(QueryMode::VectorSearch, QueryMode::NeuralSearch);
+ assert_ne!(QueryMode::NeuralSearch, QueryMode::SubgraphExtraction);
+ assert_ne!(QueryMode::SubgraphExtraction, QueryMode::DifferentiableSearch);
+ }
+}
diff --git a/crates/ruvector-gnn/src/search.rs b/crates/ruvector-gnn/src/search.rs
new file mode 100644
index 000000000..44e815180
--- /dev/null
+++ b/crates/ruvector-gnn/src/search.rs
@@ -0,0 +1,244 @@
+use crate::layer::RuvectorLayer;
+
+/// Compute cosine similarity between two vectors
+pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
+ assert_eq!(a.len(), b.len(), "Vectors must have the same length");
+
+ let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
+ let norm_a: f32 = a.iter().map(|x| x * x).sum::().sqrt();
+ let norm_b: f32 = b.iter().map(|x| x * x).sum::().sqrt();
+
+ if norm_a == 0.0 || norm_b == 0.0 {
+ 0.0
+ } else {
+ dot_product / (norm_a * norm_b)
+ }
+}
+
+/// Apply softmax with temperature scaling
+fn softmax(values: &[f32], temperature: f32) -> Vec {
+ if values.is_empty() {
+ return Vec::new();
+ }
+
+ // Scale by temperature and subtract max for numerical stability
+ let max_val = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+ let exp_values: Vec = values
+ .iter()
+ .map(|&x| ((x - max_val) / temperature).exp())
+ .collect();
+
+ let sum: f32 = exp_values.iter().sum();
+
+ if sum == 0.0 {
+ vec![1.0 / values.len() as f32; values.len()]
+ } else {
+ exp_values.iter().map(|&x| x / sum).collect()
+ }
+}
+
+/// Differentiable search using soft attention mechanism
+///
+/// # Arguments
+/// * `query` - The query vector
+/// * `candidate_embeddings` - List of candidate embedding vectors
+/// * `k` - Number of top results to return
+/// * `temperature` - Temperature for softmax (lower = sharper, higher = smoother)
+///
+/// # Returns
+/// * Tuple of (indices, soft_weights) for top-k candidates
+pub fn differentiable_search(
+ query: &[f32],
+ candidate_embeddings: &[Vec],
+ k: usize,
+ temperature: f32,
+) -> (Vec, Vec) {
+ if candidate_embeddings.is_empty() {
+ return (Vec::new(), Vec::new());
+ }
+
+ let k = k.min(candidate_embeddings.len());
+
+ // 1. Compute similarities using cosine similarity
+ let similarities: Vec = candidate_embeddings
+ .iter()
+ .map(|embedding| cosine_similarity(query, embedding))
+ .collect();
+
+ // 2. Apply softmax with temperature to get soft weights
+ let soft_weights = softmax(&similarities, temperature);
+
+ // 3. Get top-k indices by sorting similarities
+ let mut indexed_weights: Vec<(usize, f32)> = soft_weights
+ .iter()
+ .enumerate()
+ .map(|(i, &w)| (i, w))
+ .collect();
+
+ // Sort by weight descending
+ indexed_weights.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
+
+ // Take top-k
+ let top_k: Vec<(usize, f32)> = indexed_weights.into_iter().take(k).collect();
+
+ let indices: Vec = top_k.iter().map(|&(i, _)| i).collect();
+ let weights: Vec = top_k.iter().map(|&(_, w)| w).collect();
+
+ (indices, weights)
+}
+
+/// Hierarchical forward pass through GNN layers
+///
+/// # Arguments
+/// * `query` - The query vector
+/// * `layer_embeddings` - Embeddings organized by layer (outer vec = layers, inner vec = nodes per layer)
+/// * `gnn_layers` - The GNN layers to process through
+///
+/// # Returns
+/// * Final embedding after hierarchical processing
+pub fn hierarchical_forward(
+ query: &[f32],
+ layer_embeddings: &[Vec>],
+ gnn_layers: &[RuvectorLayer],
+) -> Vec {
+ if layer_embeddings.is_empty() || gnn_layers.is_empty() {
+ return query.to_vec();
+ }
+
+ let mut current_embedding = query.to_vec();
+
+ // Process through each layer from top to bottom
+ for (layer_idx, (embeddings, gnn_layer)) in
+ layer_embeddings.iter().zip(gnn_layers.iter()).enumerate()
+ {
+ if embeddings.is_empty() {
+ continue;
+ }
+
+ // Find most relevant nodes at this layer using differentiable search
+ let (top_indices, weights) = differentiable_search(
+ ¤t_embedding,
+ embeddings,
+ 5.min(embeddings.len()), // Top-5 or all if less
+ 1.0, // Default temperature
+ );
+
+ // Aggregate embeddings from top nodes using soft weights
+ let mut aggregated = vec![0.0; current_embedding.len()];
+ for (&idx, &weight) in top_indices.iter().zip(weights.iter()) {
+ for (i, &val) in embeddings[idx].iter().enumerate() {
+ if i < aggregated.len() {
+ aggregated[i] += weight * val;
+ }
+ }
+ }
+
+ // Combine with current embedding
+ let combined: Vec = current_embedding
+ .iter()
+ .zip(&aggregated)
+ .map(|(curr, agg)| (curr + agg) / 2.0)
+ .collect();
+
+ // Apply GNN layer transformation
+ // Extract neighbor embeddings and compute edge weights
+ let neighbor_embs: Vec> = top_indices
+ .iter()
+ .map(|&idx| embeddings[idx].clone())
+ .collect();
+
+ let edge_weights_vec: Vec = weights.clone();
+
+ current_embedding = gnn_layer.forward(&combined, &neighbor_embs, &edge_weights_vec);
+ }
+
+ current_embedding
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_cosine_similarity() {
+ let a = vec![1.0, 0.0, 0.0];
+ let b = vec![1.0, 0.0, 0.0];
+ assert!((cosine_similarity(&a, &b) - 1.0).abs() < 1e-6);
+
+ let c = vec![1.0, 0.0, 0.0];
+ let d = vec![0.0, 1.0, 0.0];
+ assert!((cosine_similarity(&c, &d) - 0.0).abs() < 1e-6);
+ }
+
+ #[test]
+ fn test_softmax() {
+ let values = vec![1.0, 2.0, 3.0];
+ let result = softmax(&values, 1.0);
+
+ // Sum should be 1.0
+ let sum: f32 = result.iter().sum();
+ assert!((sum - 1.0).abs() < 1e-6);
+
+ // Higher values should have higher probabilities
+ assert!(result[2] > result[1]);
+ assert!(result[1] > result[0]);
+ }
+
+ #[test]
+ fn test_softmax_with_temperature() {
+ let values = vec![1.0, 2.0, 3.0];
+
+ // Lower temperature = sharper distribution
+ let sharp = softmax(&values, 0.1);
+ let smooth = softmax(&values, 10.0);
+
+ // Sharp should have more weight on max
+ assert!(sharp[2] > smooth[2]);
+ }
+
+ #[test]
+ fn test_differentiable_search() {
+ let query = vec![1.0, 0.0, 0.0];
+ let candidates = vec![
+ vec![1.0, 0.0, 0.0], // Perfect match
+ vec![0.9, 0.1, 0.0], // Close match
+ vec![0.0, 1.0, 0.0], // Orthogonal
+ ];
+
+ let (indices, weights) = differentiable_search(&query, &candidates, 2, 1.0);
+
+ assert_eq!(indices.len(), 2);
+ assert_eq!(weights.len(), 2);
+
+ // First result should be the perfect match
+ assert_eq!(indices[0], 0);
+
+ // Weights should sum to less than or equal to 1.0 (since we took top-k)
+ let sum: f32 = weights.iter().sum();
+ assert!(sum <= 1.0 + 1e-6);
+ }
+
+ #[test]
+ fn test_hierarchical_forward() {
+ // Use consistent dimensions throughout
+ let query = vec![1.0, 0.0];
+
+ // Layer embeddings should match the output dimensions of each layer
+ let layer_embeddings = vec![
+ // First layer: embeddings are 2-dimensional (match query)
+ vec![
+ vec![1.0, 0.0],
+ vec![0.0, 1.0],
+ ],
+ ];
+
+ // Single GNN layer that maintains dimension
+ let gnn_layers = vec![
+ RuvectorLayer::new(2, 2, 1, 0.0), // input_dim, hidden_dim, heads, dropout
+ ];
+
+ let result = hierarchical_forward(&query, &layer_embeddings, &gnn_layers);
+
+ assert_eq!(result.len(), 2); // Should match hidden_dim of last layer
+ }
+}
diff --git a/crates/ruvector-gnn/src/tensor.rs b/crates/ruvector-gnn/src/tensor.rs
new file mode 100644
index 000000000..32b466747
--- /dev/null
+++ b/crates/ruvector-gnn/src/tensor.rs
@@ -0,0 +1,779 @@
+//! Tensor operations for GNN computations.
+//!
+//! Provides efficient tensor operations including:
+//! - Matrix multiplication
+//! - Element-wise operations
+//! - Activation functions
+//! - Weight initialization
+//! - Normalization
+
+use crate::error::{GnnError, Result};
+use rand::Rng;
+use rand_distr::{Distribution, Normal, Uniform};
+
+/// Basic tensor operations for GNN computations
+#[derive(Debug, Clone, PartialEq)]
+pub struct Tensor {
+ /// Flattened tensor data
+ pub data: Vec,
+ /// Shape of the tensor (dimensions)
+ pub shape: Vec,
+}
+
+impl Tensor {
+ /// Create a new tensor from data and shape
+ ///
+ /// # Arguments
+ /// * `data` - Flattened tensor data
+ /// * `shape` - Dimensions of the tensor
+ ///
+ /// # Returns
+ /// A new `Tensor` instance
+ ///
+ /// # Errors
+ /// Returns `GnnError::InvalidShape` if data length doesn't match shape
+ pub fn new(data: Vec, shape: Vec) -> Result {
+ let expected_len: usize = shape.iter().product();
+ if data.len() != expected_len {
+ return Err(GnnError::invalid_shape(format!(
+ "Data length {} doesn't match shape {:?} (expected {})",
+ data.len(),
+ shape,
+ expected_len
+ )));
+ }
+ Ok(Self { data, shape })
+ }
+
+ /// Create a zero-filled tensor with the given shape
+ ///
+ /// # Arguments
+ /// * `shape` - Dimensions of the tensor
+ ///
+ /// # Returns
+ /// A new zero-filled `Tensor`
+ ///
+ /// # Errors
+ /// Returns `GnnError::InvalidShape` if shape is empty or contains zero
+ pub fn zeros(shape: &[usize]) -> Result {
+ if shape.is_empty() || shape.iter().any(|&d| d == 0) {
+ return Err(GnnError::invalid_shape(format!(
+ "Invalid shape: {:?}",
+ shape
+ )));
+ }
+ let size: usize = shape.iter().product();
+ Ok(Self {
+ data: vec![0.0; size],
+ shape: shape.to_vec(),
+ })
+ }
+
+ /// Create a 1D tensor from a vector
+ ///
+ /// # Arguments
+ /// * `data` - Vector data
+ ///
+ /// # Returns
+ /// A new 1D `Tensor`
+ pub fn from_vec(data: Vec) -> Self {
+ let len = data.len();
+ Self {
+ data,
+ shape: vec![len],
+ }
+ }
+
+ /// Compute dot product with another tensor (both must be 1D)
+ ///
+ /// # Arguments
+ /// * `other` - Another tensor to compute dot product with
+ ///
+ /// # Returns
+ /// The dot product as a scalar
+ ///
+ /// # Errors
+ /// Returns `GnnError::DimensionMismatch` if tensors are not 1D or have different lengths
+ pub fn dot(&self, other: &Tensor) -> Result {
+ if self.shape.len() != 1 || other.shape.len() != 1 {
+ return Err(GnnError::dimension_mismatch(
+ "1D tensors",
+ format!("{}D and {}D", self.shape.len(), other.shape.len()),
+ ));
+ }
+ if self.shape[0] != other.shape[0] {
+ return Err(GnnError::dimension_mismatch(
+ format!("length {}", self.shape[0]),
+ format!("length {}", other.shape[0]),
+ ));
+ }
+
+ let result = self
+ .data
+ .iter()
+ .zip(other.data.iter())
+ .map(|(a, b)| a * b)
+ .sum();
+ Ok(result)
+ }
+
+ /// Matrix multiplication
+ ///
+ /// # Arguments
+ /// * `other` - Another tensor to multiply with
+ ///
+ /// # Returns
+ /// The result of matrix multiplication
+ ///
+ /// # Errors
+ /// Returns `GnnError::DimensionMismatch` if dimensions are incompatible
+ pub fn matmul(&self, other: &Tensor) -> Result {
+ // Support 1D x 1D (dot product), 2D x 1D, 2D x 2D
+ match (self.shape.len(), other.shape.len()) {
+ (1, 1) => {
+ let dot = self.dot(other)?;
+ Ok(Tensor::from_vec(vec![dot]))
+ }
+ (2, 1) => {
+ // Matrix-vector multiplication
+ let m = self.shape[0];
+ let n = self.shape[1];
+ if n != other.shape[0] {
+ return Err(GnnError::dimension_mismatch(
+ format!("{}x{}", m, n),
+ format!("vector of length {}", other.shape[0]),
+ ));
+ }
+
+ let mut result = vec![0.0; m];
+ for i in 0..m {
+ for j in 0..n {
+ result[i] += self.data[i * n + j] * other.data[j];
+ }
+ }
+ Ok(Tensor::from_vec(result))
+ }
+ (2, 2) => {
+ // Matrix-matrix multiplication
+ let m = self.shape[0];
+ let n = self.shape[1];
+ let p = other.shape[1];
+
+ if n != other.shape[0] {
+ return Err(GnnError::dimension_mismatch(
+ format!("{}x{}", m, n),
+ format!("{}x{}", other.shape[0], p),
+ ));
+ }
+
+ let mut result = vec![0.0; m * p];
+ for i in 0..m {
+ for j in 0..p {
+ for k in 0..n {
+ result[i * p + j] += self.data[i * n + k] * other.data[k * p + j];
+ }
+ }
+ }
+ Tensor::new(result, vec![m, p])
+ }
+ _ => Err(GnnError::dimension_mismatch(
+ "1D or 2D tensors",
+ format!("{}D and {}D", self.shape.len(), other.shape.len()),
+ )),
+ }
+ }
+
+ /// Element-wise addition
+ ///
+ /// # Arguments
+ /// * `other` - Another tensor to add
+ ///
+ /// # Returns
+ /// The sum of the two tensors
+ ///
+ /// # Errors
+ /// Returns `GnnError::DimensionMismatch` if shapes don't match
+ pub fn add(&self, other: &Tensor) -> Result {
+ if self.shape != other.shape {
+ return Err(GnnError::dimension_mismatch(
+ format!("{:?}", self.shape),
+ format!("{:?}", other.shape),
+ ));
+ }
+
+ let result: Vec = self
+ .data
+ .iter()
+ .zip(other.data.iter())
+ .map(|(a, b)| a + b)
+ .collect();
+
+ Tensor::new(result, self.shape.clone())
+ }
+
+ /// Scalar multiplication
+ ///
+ /// # Arguments
+ /// * `scalar` - Scalar value to multiply by
+ ///
+ /// # Returns
+ /// A new tensor with all elements scaled
+ pub fn scale(&self, scalar: f32) -> Tensor {
+ let result: Vec = self.data.iter().map(|&x| x * scalar).collect();
+ Tensor {
+ data: result,
+ shape: self.shape.clone(),
+ }
+ }
+
+ /// ReLU activation function (max(0, x))
+ ///
+ /// # Returns
+ /// A new tensor with ReLU applied element-wise
+ pub fn relu(&self) -> Tensor {
+ let result: Vec = self.data.iter().map(|&x| x.max(0.0)).collect();
+ Tensor {
+ data: result,
+ shape: self.shape.clone(),
+ }
+ }
+
+ /// Sigmoid activation function (1 / (1 + e^(-x)))
+ ///
+ /// # Returns
+ /// A new tensor with sigmoid applied element-wise
+ pub fn sigmoid(&self) -> Tensor {
+ let result: Vec = self
+ .data
+ .iter()
+ .map(|&x| 1.0 / (1.0 + (-x).exp()))
+ .collect();
+ Tensor {
+ data: result,
+ shape: self.shape.clone(),
+ }
+ }
+
+ /// Tanh activation function
+ ///
+ /// # Returns
+ /// A new tensor with tanh applied element-wise
+ pub fn tanh(&self) -> Tensor {
+ let result: Vec = self.data.iter().map(|&x| x.tanh()).collect();
+ Tensor {
+ data: result,
+ shape: self.shape.clone(),
+ }
+ }
+
+ /// Compute L2 norm (Euclidean norm)
+ ///
+ /// # Returns
+ /// The L2 norm of the tensor
+ pub fn l2_norm(&self) -> f32 {
+ self.data.iter().map(|&x| x * x).sum::().sqrt()
+ }
+
+ /// Normalize the tensor to unit L2 norm
+ ///
+ /// # Returns
+ /// A normalized tensor
+ ///
+ /// # Errors
+ /// Returns `GnnError::InvalidInput` if norm is zero
+ pub fn normalize(&self) -> Result {
+ let norm = self.l2_norm();
+ if norm == 0.0 {
+ return Err(GnnError::invalid_input(
+ "Cannot normalize zero vector".to_string(),
+ ));
+ }
+ Ok(self.scale(1.0 / norm))
+ }
+
+ /// Get a slice view of the tensor data
+ ///
+ /// # Returns
+ /// A slice reference to the underlying data
+ pub fn as_slice(&self) -> &[f32] {
+ &self.data
+ }
+
+ /// Consume the tensor and return the underlying vector
+ ///
+ /// # Returns
+ /// The vector containing the tensor data
+ pub fn into_vec(self) -> Vec {
+ self.data
+ }
+
+ /// Get the number of elements in the tensor
+ pub fn len(&self) -> usize {
+ self.data.len()
+ }
+
+ /// Check if the tensor is empty
+ pub fn is_empty(&self) -> bool {
+ self.data.is_empty()
+ }
+}
+
+/// Xavier/Glorot initialization for neural network weights
+///
+/// Samples from uniform distribution U(-a, a) where a = sqrt(6 / (fan_in + fan_out))
+///
+/// # Arguments
+/// * `fan_in` - Number of input units
+/// * `fan_out` - Number of output units
+///
+/// # Returns
+/// A vector of initialized weights
+///
+/// # Panics
+/// Panics if fan_in or fan_out is 0
+pub fn xavier_init(fan_in: usize, fan_out: usize) -> Vec {
+ assert!(fan_in > 0 && fan_out > 0, "fan_in and fan_out must be positive");
+
+ let limit = (6.0 / (fan_in + fan_out) as f32).sqrt();
+ let uniform = Uniform::new(-limit, limit);
+ let mut rng = rand::thread_rng();
+
+ (0..fan_in * fan_out)
+ .map(|_| uniform.sample(&mut rng))
+ .collect()
+}
+
+/// He initialization for ReLU networks
+///
+/// Samples from normal distribution N(0, sqrt(2 / fan_in))
+///
+/// # Arguments
+/// * `fan_in` - Number of input units
+///
+/// # Returns
+/// A vector of initialized weights
+///
+/// # Panics
+/// Panics if fan_in is 0
+pub fn he_init(fan_in: usize) -> Vec {
+ assert!(fan_in > 0, "fan_in must be positive");
+
+ let std_dev = (2.0 / fan_in as f32).sqrt();
+ let normal = Normal::new(0.0, std_dev).expect("Invalid normal distribution parameters");
+ let mut rng = rand::thread_rng();
+
+ (0..fan_in)
+ .map(|_| normal.sample(&mut rng))
+ .collect()
+}
+
+/// Element-wise (Hadamard) product
+///
+/// # Arguments
+/// * `a` - First vector
+/// * `b` - Second vector
+///
+/// # Returns
+/// Element-wise product of the two vectors
+///
+/// # Panics
+/// Panics if vectors have different lengths
+pub fn hadamard_product(a: &[f32], b: &[f32]) -> Vec {
+ assert_eq!(a.len(), b.len(), "Vectors must have the same length");
+ a.iter().zip(b.iter()).map(|(x, y)| x * y).collect()
+}
+
+/// Element-wise vector addition
+///
+/// # Arguments
+/// * `a` - First vector
+/// * `b` - Second vector
+///
+/// # Returns
+/// Element-wise sum of the two vectors
+///
+/// # Panics
+/// Panics if vectors have different lengths
+pub fn vector_add(a: &[f32], b: &[f32]) -> Vec {
+ assert_eq!(a.len(), b.len(), "Vectors must have the same length");
+ a.iter().zip(b.iter()).map(|(x, y)| x + y).collect()
+}
+
+/// Scalar multiplication of a vector
+///
+/// # Arguments
+/// * `v` - Input vector
+/// * `scalar` - Scalar multiplier
+///
+/// # Returns
+/// Vector with all elements multiplied by scalar
+pub fn vector_scale(v: &[f32], scalar: f32) -> Vec {
+ v.iter().map(|&x| x * scalar).collect()
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ const EPSILON: f32 = 1e-6;
+
+ fn assert_vec_approx_eq(a: &[f32], b: &[f32], epsilon: f32) {
+ assert_eq!(a.len(), b.len(), "Vectors have different lengths");
+ for (i, (&x, &y)) in a.iter().zip(b.iter()).enumerate() {
+ assert!(
+ (x - y).abs() < epsilon,
+ "Values at index {} differ: {} vs {} (diff: {})",
+ i,
+ x,
+ y,
+ (x - y).abs()
+ );
+ }
+ }
+
+ #[test]
+ fn test_tensor_new() {
+ let data = vec![1.0, 2.0, 3.0, 4.0];
+ let tensor = Tensor::new(data.clone(), vec![2, 2]).unwrap();
+ assert_eq!(tensor.data, data);
+ assert_eq!(tensor.shape, vec![2, 2]);
+ }
+
+ #[test]
+ fn test_tensor_new_invalid_shape() {
+ let data = vec![1.0, 2.0, 3.0];
+ let result = Tensor::new(data, vec![2, 2]);
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_tensor_zeros() {
+ let tensor = Tensor::zeros(&[3, 2]).unwrap();
+ assert_eq!(tensor.data, vec![0.0; 6]);
+ assert_eq!(tensor.shape, vec![3, 2]);
+ }
+
+ #[test]
+ fn test_tensor_zeros_invalid_shape() {
+ let result = Tensor::zeros(&[0, 2]);
+ assert!(result.is_err());
+
+ let result = Tensor::zeros(&[]);
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_tensor_from_vec() {
+ let data = vec![1.0, 2.0, 3.0];
+ let tensor = Tensor::from_vec(data.clone());
+ assert_eq!(tensor.data, data);
+ assert_eq!(tensor.shape, vec![3]);
+ }
+
+ #[test]
+ fn test_dot_product() {
+ let a = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let b = Tensor::from_vec(vec![4.0, 5.0, 6.0]);
+ let result = a.dot(&b).unwrap();
+ assert!((result - 32.0).abs() < EPSILON); // 1*4 + 2*5 + 3*6 = 32
+ }
+
+ #[test]
+ fn test_dot_product_dimension_mismatch() {
+ let a = Tensor::from_vec(vec![1.0, 2.0]);
+ let b = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let result = a.dot(&b);
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_matmul_1d() {
+ let a = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let b = Tensor::from_vec(vec![4.0, 5.0, 6.0]);
+ let result = a.matmul(&b).unwrap();
+ assert_eq!(result.shape, vec![1]);
+ assert!((result.data[0] - 32.0).abs() < EPSILON);
+ }
+
+ #[test]
+ fn test_matmul_2d_1d() {
+ // Matrix-vector multiplication
+ let mat = Tensor::new(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![2, 3]).unwrap();
+ let vec = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let result = mat.matmul(&vec).unwrap();
+
+ assert_eq!(result.shape, vec![2]);
+ // [1,2,3] * [1,2,3]' = 14
+ // [4,5,6] * [1,2,3]' = 32
+ assert_vec_approx_eq(&result.data, &[14.0, 32.0], EPSILON);
+ }
+
+ #[test]
+ fn test_matmul_2d_2d() {
+ // Matrix-matrix multiplication
+ let a = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
+ let b = Tensor::new(vec![5.0, 6.0, 7.0, 8.0], vec![2, 2]).unwrap();
+ let result = a.matmul(&b).unwrap();
+
+ assert_eq!(result.shape, vec![2, 2]);
+ // [[1,2], [3,4]] * [[5,6], [7,8]] = [[19,22], [43,50]]
+ assert_vec_approx_eq(&result.data, &[19.0, 22.0, 43.0, 50.0], EPSILON);
+ }
+
+ #[test]
+ fn test_matmul_dimension_mismatch() {
+ let a = Tensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]).unwrap();
+ let b = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let result = a.matmul(&b);
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_add() {
+ let a = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let b = Tensor::from_vec(vec![4.0, 5.0, 6.0]);
+ let result = a.add(&b).unwrap();
+ assert_eq!(result.data, vec![5.0, 7.0, 9.0]);
+ }
+
+ #[test]
+ fn test_add_dimension_mismatch() {
+ let a = Tensor::from_vec(vec![1.0, 2.0]);
+ let b = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let result = a.add(&b);
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_scale() {
+ let tensor = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let result = tensor.scale(2.0);
+ assert_eq!(result.data, vec![2.0, 4.0, 6.0]);
+ }
+
+ #[test]
+ fn test_relu() {
+ let tensor = Tensor::from_vec(vec![-1.0, 0.0, 1.0, 2.0]);
+ let result = tensor.relu();
+ assert_eq!(result.data, vec![0.0, 0.0, 1.0, 2.0]);
+ }
+
+ #[test]
+ fn test_sigmoid() {
+ let tensor = Tensor::from_vec(vec![0.0, 1.0, -1.0]);
+ let result = tensor.sigmoid();
+
+ assert!((result.data[0] - 0.5).abs() < EPSILON);
+ assert!((result.data[1] - 0.7310586).abs() < EPSILON);
+ assert!((result.data[2] - 0.26894143).abs() < EPSILON);
+ }
+
+ #[test]
+ fn test_tanh() {
+ let tensor = Tensor::from_vec(vec![0.0, 1.0, -1.0]);
+ let result = tensor.tanh();
+
+ assert!((result.data[0] - 0.0).abs() < EPSILON);
+ assert!((result.data[1] - 0.7615942).abs() < EPSILON);
+ assert!((result.data[2] - (-0.7615942)).abs() < EPSILON);
+ }
+
+ #[test]
+ fn test_l2_norm() {
+ let tensor = Tensor::from_vec(vec![3.0, 4.0]);
+ let norm = tensor.l2_norm();
+ assert!((norm - 5.0).abs() < EPSILON);
+ }
+
+ #[test]
+ fn test_normalize() {
+ let tensor = Tensor::from_vec(vec![3.0, 4.0]);
+ let result = tensor.normalize().unwrap();
+ assert_vec_approx_eq(&result.data, &[0.6, 0.8], EPSILON);
+ assert!((result.l2_norm() - 1.0).abs() < EPSILON);
+ }
+
+ #[test]
+ fn test_normalize_zero_vector() {
+ let tensor = Tensor::from_vec(vec![0.0, 0.0]);
+ let result = tensor.normalize();
+ assert!(result.is_err());
+ }
+
+ #[test]
+ fn test_as_slice() {
+ let data = vec![1.0, 2.0, 3.0];
+ let tensor = Tensor::from_vec(data.clone());
+ assert_eq!(tensor.as_slice(), &data[..]);
+ }
+
+ #[test]
+ fn test_into_vec() {
+ let data = vec![1.0, 2.0, 3.0];
+ let tensor = Tensor::from_vec(data.clone());
+ assert_eq!(tensor.into_vec(), data);
+ }
+
+ #[test]
+ fn test_len() {
+ let tensor = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ assert_eq!(tensor.len(), 3);
+ }
+
+ #[test]
+ fn test_is_empty() {
+ let tensor = Tensor::from_vec(vec![]);
+ assert!(tensor.is_empty());
+
+ let tensor = Tensor::from_vec(vec![1.0]);
+ assert!(!tensor.is_empty());
+ }
+
+ #[test]
+ fn test_xavier_init() {
+ let weights = xavier_init(100, 50);
+ assert_eq!(weights.len(), 5000);
+
+ // Check that values are in expected range
+ let limit = (6.0 / 150.0_f32).sqrt();
+ for &w in &weights {
+ assert!(w >= -limit && w <= limit);
+ }
+
+ // Check distribution properties
+ let mean: f32 = weights.iter().sum::() / weights.len() as f32;
+ assert!(mean.abs() < 0.1); // Mean should be close to 0
+ }
+
+ #[test]
+ #[should_panic(expected = "fan_in and fan_out must be positive")]
+ fn test_xavier_init_zero_fan() {
+ xavier_init(0, 10);
+ }
+
+ #[test]
+ fn test_he_init() {
+ let weights = he_init(100);
+ assert_eq!(weights.len(), 100);
+
+ // Check distribution properties
+ let mean: f32 = weights.iter().sum::() / weights.len() as f32;
+ assert!(mean.abs() < 0.2); // Mean should be close to 0
+ }
+
+ #[test]
+ #[should_panic(expected = "fan_in must be positive")]
+ fn test_he_init_zero_fan() {
+ he_init(0);
+ }
+
+ #[test]
+ fn test_hadamard_product() {
+ let a = vec![1.0, 2.0, 3.0];
+ let b = vec![4.0, 5.0, 6.0];
+ let result = hadamard_product(&a, &b);
+ assert_eq!(result, vec![4.0, 10.0, 18.0]);
+ }
+
+ #[test]
+ #[should_panic(expected = "Vectors must have the same length")]
+ fn test_hadamard_product_length_mismatch() {
+ let a = vec![1.0, 2.0];
+ let b = vec![1.0, 2.0, 3.0];
+ hadamard_product(&a, &b);
+ }
+
+ #[test]
+ fn test_vector_add() {
+ let a = vec![1.0, 2.0, 3.0];
+ let b = vec![4.0, 5.0, 6.0];
+ let result = vector_add(&a, &b);
+ assert_eq!(result, vec![5.0, 7.0, 9.0]);
+ }
+
+ #[test]
+ #[should_panic(expected = "Vectors must have the same length")]
+ fn test_vector_add_length_mismatch() {
+ let a = vec![1.0, 2.0];
+ let b = vec![1.0, 2.0, 3.0];
+ vector_add(&a, &b);
+ }
+
+ #[test]
+ fn test_vector_scale() {
+ let v = vec![1.0, 2.0, 3.0];
+ let result = vector_scale(&v, 2.5);
+ assert_vec_approx_eq(&result, &[2.5, 5.0, 7.5], EPSILON);
+ }
+
+ #[test]
+ fn test_complex_operations() {
+ // Test chaining operations
+ let a = Tensor::from_vec(vec![1.0, 2.0, 3.0]);
+ let b = Tensor::from_vec(vec![0.5, 1.0, 1.5]);
+
+ let sum = a.add(&b).unwrap();
+ let scaled = sum.scale(2.0);
+ let activated = scaled.relu();
+ let normalized = activated.normalize().unwrap();
+
+ assert!((normalized.l2_norm() - 1.0).abs() < EPSILON);
+ }
+
+ #[test]
+ fn test_edge_case_single_element() {
+ let tensor = Tensor::from_vec(vec![5.0]);
+ assert_eq!(tensor.len(), 1);
+ assert_eq!(tensor.l2_norm(), 5.0);
+
+ let normalized = tensor.normalize().unwrap();
+ assert_vec_approx_eq(&normalized.data, &[1.0], EPSILON);
+ }
+
+ #[test]
+ fn test_edge_case_negative_values() {
+ let tensor = Tensor::from_vec(vec![-3.0, -4.0]);
+ assert!((tensor.l2_norm() - 5.0).abs() < EPSILON);
+
+ let relu_result = tensor.relu();
+ assert_eq!(relu_result.data, vec![0.0, 0.0]);
+ }
+
+ #[test]
+ fn test_large_matrix_multiplication() {
+ // 10x10 matrix multiplication
+ let size = 10;
+ let a_data: Vec = (0..size * size).map(|i| i as f32).collect();
+ let b_data: Vec = (0..size * size).map(|i| (i % 2) as f32).collect();
+
+ let a = Tensor::new(a_data, vec![size, size]).unwrap();
+ let b = Tensor::new(b_data, vec![size, size]).unwrap();
+
+ let result = a.matmul(&b).unwrap();
+ assert_eq!(result.shape, vec![size, size]);
+ assert_eq!(result.len(), size * size);
+ }
+
+ #[test]
+ fn test_activation_functions_range() {
+ let tensor = Tensor::from_vec(vec![-10.0, -1.0, 0.0, 1.0, 10.0]);
+
+ // Sigmoid should be in (0, 1)
+ let sigmoid = tensor.sigmoid();
+ for &val in &sigmoid.data {
+ assert!(val > 0.0 && val < 1.0);
+ }
+
+ // Tanh should be in [-1, 1]
+ let tanh = tensor.tanh();
+ for &val in &tanh.data {
+ assert!(val >= -1.0 && val <= 1.0);
+ }
+
+ // ReLU should be non-negative
+ let relu = tensor.relu();
+ for &val in &relu.data {
+ assert!(val >= 0.0);
+ }
+ }
+}
diff --git a/crates/ruvector-gnn/src/training.rs b/crates/ruvector-gnn/src/training.rs
new file mode 100644
index 000000000..9a99389fe
--- /dev/null
+++ b/crates/ruvector-gnn/src/training.rs
@@ -0,0 +1,565 @@
+//! Training utilities for GNN models.
+//!
+//! Provides training loop utilities, optimizers, and loss functions.
+
+use crate::error::{GnnError, Result};
+use crate::search::cosine_similarity;
+use ndarray::Array2;
+
+/// Optimizer types
+#[derive(Debug, Clone)]
+pub enum OptimizerType {
+ /// Stochastic Gradient Descent
+ Sgd { learning_rate: f32 },
+ /// Adam optimizer
+ Adam {
+ /// Learning rate
+ learning_rate: f32,
+ /// Beta1 parameter
+ beta1: f32,
+ /// Beta2 parameter
+ beta2: f32,
+ },
+}
+
+/// TODO: Implement optimizer
+pub struct Optimizer {
+ optimizer_type: OptimizerType,
+}
+
+impl Optimizer {
+ /// Create a new optimizer
+ pub fn new(optimizer_type: OptimizerType) -> Self {
+ Self { optimizer_type }
+ }
+
+ /// TODO: Perform optimization step
+ pub fn step(&mut self, params: &mut Array2, grads: &Array2) -> Result<()> {
+ unimplemented!("TODO: Implement optimizer step")
+ }
+}
+
+/// Loss function types
+#[derive(Debug, Clone, Copy)]
+pub enum LossType {
+ /// Mean Squared Error
+ Mse,
+ /// Cross Entropy
+ CrossEntropy,
+ /// Binary Cross Entropy
+ BinaryCrossEntropy,
+}
+
+/// TODO: Implement loss functions
+pub struct Loss;
+
+impl Loss {
+ /// TODO: Compute loss
+ pub fn compute(
+ loss_type: LossType,
+ predictions: &Array2,
+ targets: &Array2,
+ ) -> Result {
+ unimplemented!("TODO: Implement loss computation")
+ }
+
+ /// TODO: Compute loss gradient
+ pub fn gradient(
+ loss_type: LossType,
+ predictions: &Array2,
+ targets: &Array2,
+ ) -> Result> {
+ unimplemented!("TODO: Implement loss gradient")
+ }
+}
+
+/// TODO: Implement training configuration
+#[derive(Debug, Clone)]
+pub struct TrainingConfig {
+ /// Number of epochs
+ pub epochs: usize,
+ /// Batch size
+ pub batch_size: usize,
+ /// Learning rate
+ pub learning_rate: f32,
+ /// Loss type
+ pub loss_type: LossType,
+ /// Optimizer type
+ pub optimizer_type: OptimizerType,
+}
+
+impl Default for TrainingConfig {
+ fn default() -> Self {
+ Self {
+ epochs: 100,
+ batch_size: 32,
+ learning_rate: 0.001,
+ loss_type: LossType::Mse,
+ optimizer_type: OptimizerType::Adam {
+ learning_rate: 0.001,
+ beta1: 0.9,
+ beta2: 0.999,
+ },
+ }
+ }
+}
+
+/// Configuration for contrastive learning training
+#[derive(Debug, Clone)]
+pub struct TrainConfig {
+ /// Batch size for training
+ pub batch_size: usize,
+ /// Number of negative samples per positive
+ pub n_negatives: usize,
+ /// Temperature parameter for contrastive loss
+ pub temperature: f32,
+ /// Learning rate for optimization
+ pub learning_rate: f32,
+ /// Number of updates before flushing to storage
+ pub flush_threshold: usize,
+}
+
+impl Default for TrainConfig {
+ fn default() -> Self {
+ Self {
+ batch_size: 256,
+ n_negatives: 64,
+ temperature: 0.07,
+ learning_rate: 0.001,
+ flush_threshold: 1000,
+ }
+ }
+}
+
+/// Configuration for online learning
+#[derive(Debug, Clone)]
+pub struct OnlineConfig {
+ /// Number of local optimization steps
+ pub local_steps: usize,
+ /// Whether to propagate updates to neighbors
+ pub propagate_updates: bool,
+}
+
+impl Default for OnlineConfig {
+ fn default() -> Self {
+ Self {
+ local_steps: 5,
+ propagate_updates: true,
+ }
+ }
+}
+
+/// Compute InfoNCE contrastive loss
+///
+/// InfoNCE (Information Noise-Contrastive Estimation) loss is used for contrastive learning.
+/// It maximizes agreement between anchor and positive samples while minimizing agreement
+/// with negative samples.
+///
+/// # Arguments
+/// * `anchor` - The anchor embedding vector
+/// * `positives` - Positive example embeddings (similar to anchor)
+/// * `negatives` - Negative example embeddings (dissimilar to anchor)
+/// * `temperature` - Temperature scaling parameter (lower = sharper distinctions)
+///
+/// # Returns
+/// * The computed loss value (lower is better)
+///
+/// # Example
+/// ```
+/// use ruvector_gnn::training::info_nce_loss;
+///
+/// let anchor = vec![1.0, 0.0, 0.0];
+/// let positive = vec![0.9, 0.1, 0.0];
+/// let negative1 = vec![0.0, 1.0, 0.0];
+/// let negative2 = vec![0.0, 0.0, 1.0];
+///
+/// let loss = info_nce_loss(
+/// &anchor,
+/// &[&positive],
+/// &[&negative1, &negative2],
+/// 0.07
+/// );
+/// assert!(loss > 0.0);
+/// ```
+pub fn info_nce_loss(
+ anchor: &[f32],
+ positives: &[&[f32]],
+ negatives: &[&[f32]],
+ temperature: f32,
+) -> f32 {
+ if positives.is_empty() {
+ return 0.0;
+ }
+
+ // Compute similarities with positives (scaled by temperature)
+ let pos_sims: Vec = positives
+ .iter()
+ .map(|pos| cosine_similarity(anchor, pos) / temperature)
+ .collect();
+
+ // Compute similarities with negatives (scaled by temperature)
+ let neg_sims: Vec = negatives
+ .iter()
+ .map(|neg| cosine_similarity(anchor, neg) / temperature)
+ .collect();
+
+ // For each positive, compute the InfoNCE loss using log-sum-exp trick for numerical stability
+ let mut total_loss = 0.0;
+ for &pos_sim in &pos_sims {
+ // Use log-sum-exp trick to avoid overflow
+ // log(exp(pos_sim) / (exp(pos_sim) + sum(exp(neg_sim))))
+ // = pos_sim - log(exp(pos_sim) + sum(exp(neg_sim)))
+ // = pos_sim - log_sum_exp([pos_sim, neg_sims...])
+
+ // Collect all logits for log-sum-exp
+ let mut all_logits = vec![pos_sim];
+ all_logits.extend(&neg_sims);
+
+ // Compute log-sum-exp with numerical stability
+ let max_logit = all_logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+ let log_sum_exp = max_logit + all_logits
+ .iter()
+ .map(|&x| (x - max_logit).exp())
+ .sum::()
+ .ln();
+
+ // Loss = -log(exp(pos_sim) / sum_exp) = -(pos_sim - log_sum_exp)
+ total_loss -= pos_sim - log_sum_exp;
+ }
+
+ // Average over positives
+ total_loss / positives.len() as f32
+}
+
+/// Compute local contrastive loss for graph structures
+///
+/// This loss encourages node embeddings to be similar to their neighbors
+/// and dissimilar to non-neighbors in the graph.
+///
+/// # Arguments
+/// * `node_embedding` - The embedding of the target node
+/// * `neighbor_embeddings` - Embeddings of neighbor nodes
+/// * `non_neighbor_embeddings` - Embeddings of non-neighbor nodes
+/// * `temperature` - Temperature scaling parameter
+///
+/// # Returns
+/// * The computed loss value (lower is better)
+///
+/// # Example
+/// ```
+/// use ruvector_gnn::training::local_contrastive_loss;
+///
+/// let node = vec![1.0, 0.0, 0.0];
+/// let neighbor = vec![0.9, 0.1, 0.0];
+/// let non_neighbor1 = vec![0.0, 1.0, 0.0];
+/// let non_neighbor2 = vec![0.0, 0.0, 1.0];
+///
+/// let loss = local_contrastive_loss(
+/// &node,
+/// &[neighbor],
+/// &[non_neighbor1, non_neighbor2],
+/// 0.07
+/// );
+/// assert!(loss > 0.0);
+/// ```
+pub fn local_contrastive_loss(
+ node_embedding: &[f32],
+ neighbor_embeddings: &[Vec],
+ non_neighbor_embeddings: &[Vec],
+ temperature: f32,
+) -> f32 {
+ if neighbor_embeddings.is_empty() {
+ return 0.0;
+ }
+
+ // Convert to slices for info_nce_loss
+ let positives: Vec<&[f32]> = neighbor_embeddings.iter().map(|v| v.as_slice()).collect();
+ let negatives: Vec<&[f32]> = non_neighbor_embeddings
+ .iter()
+ .map(|v| v.as_slice())
+ .collect();
+
+ info_nce_loss(node_embedding, &positives, &negatives, temperature)
+}
+
+/// Perform a single SGD (Stochastic Gradient Descent) optimization step
+///
+/// Updates the embedding in-place by subtracting the scaled gradient.
+///
+/// # Arguments
+/// * `embedding` - The embedding to update (modified in-place)
+/// * `grad` - The gradient vector
+/// * `learning_rate` - The learning rate (step size)
+///
+/// # Example
+/// ```
+/// use ruvector_gnn::training::sgd_step;
+///
+/// let mut embedding = vec![1.0, 2.0, 3.0];
+/// let gradient = vec![0.1, -0.2, 0.3];
+/// let learning_rate = 0.01;
+///
+/// sgd_step(&mut embedding, &gradient, learning_rate);
+///
+/// // Embedding is now updated: embedding[i] -= learning_rate * grad[i]
+/// assert!((embedding[0] - 0.999).abs() < 1e-6);
+/// assert!((embedding[1] - 2.002).abs() < 1e-6);
+/// assert!((embedding[2] - 2.997).abs() < 1e-6);
+/// ```
+pub fn sgd_step(embedding: &mut [f32], grad: &[f32], learning_rate: f32) {
+ assert_eq!(
+ embedding.len(),
+ grad.len(),
+ "Embedding and gradient must have the same length"
+ );
+
+ for (emb, &g) in embedding.iter_mut().zip(grad.iter()) {
+ *emb -= learning_rate * g;
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn test_train_config_default() {
+ let config = TrainConfig::default();
+ assert_eq!(config.batch_size, 256);
+ assert_eq!(config.n_negatives, 64);
+ assert_eq!(config.temperature, 0.07);
+ assert_eq!(config.learning_rate, 0.001);
+ assert_eq!(config.flush_threshold, 1000);
+ }
+
+ #[test]
+ fn test_online_config_default() {
+ let config = OnlineConfig::default();
+ assert_eq!(config.local_steps, 5);
+ assert!(config.propagate_updates);
+ }
+
+ #[test]
+ fn test_info_nce_loss_basic() {
+ // Anchor and positive are similar
+ let anchor = vec![1.0, 0.0, 0.0];
+ let positive = vec![0.9, 0.1, 0.0];
+
+ // Negatives are orthogonal
+ let negative1 = vec![0.0, 1.0, 0.0];
+ let negative2 = vec![0.0, 0.0, 1.0];
+
+ let loss = info_nce_loss(
+ &anchor,
+ &[&positive],
+ &[&negative1, &negative2],
+ 0.07,
+ );
+
+ // Loss should be positive
+ assert!(loss > 0.0);
+
+ // Loss should be reasonable (not infinite or NaN)
+ assert!(loss.is_finite());
+ }
+
+ #[test]
+ fn test_info_nce_loss_perfect_match() {
+ // Anchor and positive are identical
+ let anchor = vec![1.0, 0.0, 0.0];
+ let positive = vec![1.0, 0.0, 0.0];
+
+ // Negatives are very different
+ let negative1 = vec![0.0, 1.0, 0.0];
+ let negative2 = vec![0.0, 0.0, 1.0];
+
+ let loss = info_nce_loss(
+ &anchor,
+ &[&positive],
+ &[&negative1, &negative2],
+ 0.07,
+ );
+
+ // Loss should be lower for perfect match
+ assert!(loss < 1.0);
+ assert!(loss.is_finite());
+ }
+
+ #[test]
+ fn test_info_nce_loss_no_positives() {
+ let anchor = vec![1.0, 0.0, 0.0];
+ let negative1 = vec![0.0, 1.0, 0.0];
+
+ let loss = info_nce_loss(&anchor, &[], &[&negative1], 0.07);
+
+ // Should return 0.0 when no positives
+ assert_eq!(loss, 0.0);
+ }
+
+ #[test]
+ fn test_info_nce_loss_temperature_effect() {
+ let anchor = vec![1.0, 0.0, 0.0];
+ let positive = vec![0.9, 0.1, 0.0];
+ let negative = vec![0.0, 1.0, 0.0];
+
+ // Test with reasonable temperature values
+ // Very low temperatures can cause numerical issues, so we use 0.07 (standard) and 1.0
+ let loss_low_temp = info_nce_loss(&anchor, &[&positive], &[&negative], 0.07);
+ let loss_high_temp = info_nce_loss(&anchor, &[&positive], &[&negative], 1.0);
+
+ // Both should be positive and finite
+ assert!(loss_low_temp > 0.0 && loss_low_temp.is_finite(),
+ "Low temp loss should be positive and finite, got: {}", loss_low_temp);
+ assert!(loss_high_temp > 0.0 && loss_high_temp.is_finite(),
+ "High temp loss should be positive and finite, got: {}", loss_high_temp);
+
+ // With standard temperature, the loss should be reasonable
+ assert!(loss_low_temp < 10.0, "Loss should not be too large");
+ assert!(loss_high_temp < 10.0, "Loss should not be too large");
+ }
+
+ #[test]
+ fn test_local_contrastive_loss_basic() {
+ let node = vec![1.0, 0.0, 0.0];
+ let neighbor = vec![0.9, 0.1, 0.0];
+ let non_neighbor1 = vec![0.0, 1.0, 0.0];
+ let non_neighbor2 = vec![0.0, 0.0, 1.0];
+
+ let loss = local_contrastive_loss(
+ &node,
+ &[neighbor],
+ &[non_neighbor1, non_neighbor2],
+ 0.07,
+ );
+
+ // Loss should be positive and finite
+ assert!(loss > 0.0);
+ assert!(loss.is_finite());
+ }
+
+ #[test]
+ fn test_local_contrastive_loss_multiple_neighbors() {
+ let node = vec![1.0, 0.0, 0.0];
+ let neighbor1 = vec![0.9, 0.1, 0.0];
+ let neighbor2 = vec![0.95, 0.05, 0.0];
+ let non_neighbor = vec![0.0, 1.0, 0.0];
+
+ let loss = local_contrastive_loss(
+ &node,
+ &[neighbor1, neighbor2],
+ &[non_neighbor],
+ 0.07,
+ );
+
+ assert!(loss > 0.0);
+ assert!(loss.is_finite());
+ }
+
+ #[test]
+ fn test_local_contrastive_loss_no_neighbors() {
+ let node = vec![1.0, 0.0, 0.0];
+ let non_neighbor = vec![0.0, 1.0, 0.0];
+
+ let loss = local_contrastive_loss(&node, &[], &[non_neighbor], 0.07);
+
+ // Should return 0.0 when no neighbors
+ assert_eq!(loss, 0.0);
+ }
+
+ #[test]
+ fn test_sgd_step_basic() {
+ let mut embedding = vec![1.0, 2.0, 3.0];
+ let gradient = vec![0.1, -0.2, 0.3];
+ let learning_rate = 0.01;
+
+ sgd_step(&mut embedding, &gradient, learning_rate);
+
+ // Expected: embedding[i] -= learning_rate * grad[i]
+ assert!((embedding[0] - 0.999).abs() < 1e-6); // 1.0 - 0.01 * 0.1
+ assert!((embedding[1] - 2.002).abs() < 1e-6); // 2.0 - 0.01 * (-0.2)
+ assert!((embedding[2] - 2.997).abs() < 1e-6); // 3.0 - 0.01 * 0.3
+ }
+
+ #[test]
+ fn test_sgd_step_zero_gradient() {
+ let mut embedding = vec![1.0, 2.0, 3.0];
+ let original = embedding.clone();
+ let gradient = vec![0.0, 0.0, 0.0];
+ let learning_rate = 0.01;
+
+ sgd_step(&mut embedding, &gradient, learning_rate);
+
+ // Embedding should not change with zero gradient
+ assert_eq!(embedding, original);
+ }
+
+ #[test]
+ fn test_sgd_step_zero_learning_rate() {
+ let mut embedding = vec![1.0, 2.0, 3.0];
+ let original = embedding.clone();
+ let gradient = vec![0.1, 0.2, 0.3];
+ let learning_rate = 0.0;
+
+ sgd_step(&mut embedding, &gradient, learning_rate);
+
+ // Embedding should not change with zero learning rate
+ assert_eq!(embedding, original);
+ }
+
+ #[test]
+ fn test_sgd_step_large_learning_rate() {
+ let mut embedding = vec![10.0, 20.0, 30.0];
+ let gradient = vec![1.0, 2.0, 3.0];
+ let learning_rate = 5.0;
+
+ sgd_step(&mut embedding, &gradient, learning_rate);
+
+ // Expected: embedding[i] -= learning_rate * grad[i]
+ assert!((embedding[0] - 5.0).abs() < 1e-5); // 10.0 - 5.0 * 1.0
+ assert!((embedding[1] - 10.0).abs() < 1e-5); // 20.0 - 5.0 * 2.0
+ assert!((embedding[2] - 15.0).abs() < 1e-5); // 30.0 - 5.0 * 3.0
+ }
+
+ #[test]
+ #[should_panic(expected = "Embedding and gradient must have the same length")]
+ fn test_sgd_step_mismatched_lengths() {
+ let mut embedding = vec![1.0, 2.0, 3.0];
+ let gradient = vec![0.1, 0.2]; // Wrong length
+
+ sgd_step(&mut embedding, &gradient, 0.01);
+ }
+
+ #[test]
+ fn test_info_nce_loss_multiple_positives() {
+ let anchor = vec![1.0, 0.0, 0.0];
+ let positive1 = vec![0.9, 0.1, 0.0];
+ let positive2 = vec![0.95, 0.05, 0.0];
+ let negative = vec![0.0, 1.0, 0.0];
+
+ let loss = info_nce_loss(
+ &anchor,
+ &[&positive1, &positive2],
+ &[&negative],
+ 0.07,
+ );
+
+ // Loss should be positive and finite
+ assert!(loss > 0.0);
+ assert!(loss.is_finite());
+ }
+
+ #[test]
+ fn test_contrastive_loss_gradient_property() {
+ // Test that loss decreases when positive becomes more similar
+ let anchor = vec![1.0, 0.0, 0.0];
+ let positive_far = vec![0.5, 0.5, 0.0];
+ let positive_close = vec![0.9, 0.1, 0.0];
+ let negative = vec![0.0, 1.0, 0.0];
+
+ let loss_far = info_nce_loss(&anchor, &[&positive_far], &[&negative], 0.07);
+ let loss_close = info_nce_loss(&anchor, &[&positive_close], &[&negative], 0.07);
+
+ // Loss should be lower when positive is closer to anchor
+ assert!(loss_close < loss_far);
+ }
+}
diff --git a/crates/ruvector-graph-wasm/Cargo.toml b/crates/ruvector-graph-wasm/Cargo.toml
index 5851c578e..4bd355cc0 100644
--- a/crates/ruvector-graph-wasm/Cargo.toml
+++ b/crates/ruvector-graph-wasm/Cargo.toml
@@ -13,6 +13,7 @@ crate-type = ["cdylib", "rlib"]
[dependencies]
ruvector-core = { version = "0.1.1", path = "../ruvector-core", default-features = false }
+ruvector-graph = { path = "../ruvector-graph", default-features = false, features = ["wasm"] }
parking_lot = { workspace = true }
getrandom = { workspace = true }
diff --git a/crates/ruvector-graph/Cargo.toml b/crates/ruvector-graph/Cargo.toml
index dbc0c2fb9..7326b87ac 100644
--- a/crates/ruvector-graph/Cargo.toml
+++ b/crates/ruvector-graph/Cargo.toml
@@ -10,18 +10,18 @@ description = "Distributed Neo4j-compatible hypergraph database with SIMD optimi
[dependencies]
# RuVector dependencies
-ruvector-core = { path = "../ruvector-core", features = ["hnsw", "simd"] }
+ruvector-core = { path = "../ruvector-core", default-features = false, features = ["simd"] }
ruvector-raft = { path = "../ruvector-raft", optional = true }
ruvector-cluster = { path = "../ruvector-cluster", optional = true }
ruvector-replication = { path = "../ruvector-replication", optional = true }
-# Storage and indexing
-redb = { workspace = true }
-memmap2 = { workspace = true }
-hnsw_rs = { workspace = true }
+# Storage and indexing (optional for WASM)
+redb = { workspace = true, optional = true }
+memmap2 = { workspace = true, optional = true }
+hnsw_rs = { workspace = true, optional = true }
# SIMD and performance
-simsimd = { workspace = true }
+simsimd = { workspace = true, optional = true }
rayon = { workspace = true }
crossbeam = { workspace = true }
num_cpus = "1.16"
@@ -32,9 +32,9 @@ bincode = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
-# Async runtime
-tokio = { workspace = true, features = ["rt-multi-thread", "sync", "macros", "time", "net"] }
-futures = { workspace = true }
+# Async runtime (optional for WASM)
+tokio = { workspace = true, features = ["rt-multi-thread", "sync", "macros", "time", "net"], optional = true }
+futures = { workspace = true, optional = true }
# Error handling and logging
thiserror = { workspace = true }
@@ -69,11 +69,11 @@ lalrpop-util = { version = "0.21", optional = true }
# Cache
lru = "0.12"
-moka = { version = "0.12", features = ["future"] }
+moka = { version = "0.12", features = ["future"], optional = true }
-# Compression (for storage optimization)
-zstd = "0.13"
-lz4 = "1.24"
+# Compression (for storage optimization, optional for WASM)
+zstd = { version = "0.13", optional = true }
+lz4 = { version = "1.24", optional = true }
# Networking (for federation)
tonic = { version = "0.12", features = ["transport"], optional = true }
@@ -103,13 +103,28 @@ csv = "1.3"
pest_generator = "2.7"
[features]
-default = ["simd"]
+default = ["full"]
+
+# Full feature set (non-WASM)
+full = ["simd", "storage", "async-runtime", "compression", "hnsw_rs", "ruvector-core/hnsw"]
# SIMD optimizations
-simd = ["ruvector-core/simd"]
+simd = ["ruvector-core/simd", "simsimd"]
+
+# Storage backends
+storage = ["redb", "memmap2"]
+
+# Async runtime support
+async-runtime = ["tokio", "futures", "moka"]
+
+# Compression support
+compression = ["zstd", "lz4"]
+
+# WASM-compatible minimal build (parser + core graph operations)
+wasm = []
# Distributed deployment with RAFT
-distributed = ["ruvector-raft", "ruvector-cluster", "ruvector-replication", "blake3", "xxhash-rust"]
+distributed = ["ruvector-raft", "ruvector-cluster", "ruvector-replication", "blake3", "xxhash-rust", "full"]
# Cross-cluster federation
federation = ["tonic", "prost", "tower", "hyper", "distributed"]
diff --git a/crates/ruvector-graph/src/graph.rs b/crates/ruvector-graph/src/graph.rs
index fb5741b85..3fb063345 100644
--- a/crates/ruvector-graph/src/graph.rs
+++ b/crates/ruvector-graph/src/graph.rs
@@ -6,8 +6,10 @@ use crate::hyperedge::{Hyperedge, HyperedgeId};
use crate::index::{AdjacencyIndex, EdgeTypeIndex, HyperedgeNodeIndex, LabelIndex, PropertyIndex};
use crate::node::Node;
use crate::types::{EdgeId, NodeId, PropertyValue};
+#[cfg(feature = "storage")]
use crate::storage::GraphStorage;
use dashmap::DashMap;
+#[cfg(feature = "storage")]
use std::path::Path;
use std::sync::Arc;
@@ -30,6 +32,7 @@ pub struct GraphDB {
/// Hyperedge node index
hyperedge_node_index: HyperedgeNodeIndex,
/// Optional persistent storage
+ #[cfg(feature = "storage")]
storage: Option,
}
@@ -45,11 +48,13 @@ impl GraphDB {
edge_type_index: EdgeTypeIndex::new(),
adjacency_index: AdjacencyIndex::new(),
hyperedge_node_index: HyperedgeNodeIndex::new(),
+ #[cfg(feature = "storage")]
storage: None,
}
}
/// Create a new graph database with persistent storage
+ #[cfg(feature = "storage")]
pub fn with_storage>(path: P) -> anyhow::Result {
let storage = GraphStorage::new(path)?;
@@ -63,6 +68,7 @@ impl GraphDB {
}
/// Load all data from storage into memory
+ #[cfg(feature = "storage")]
fn load_from_storage(&mut self) -> anyhow::Result<()> {
if let Some(storage) = &self.storage {
// Load nodes
@@ -108,6 +114,7 @@ impl GraphDB {
self.nodes.insert(id.clone(), node.clone());
// Persist to storage if available
+ #[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_node(&node)?;
}
@@ -128,6 +135,7 @@ impl GraphDB {
self.property_index.remove_node(&node);
// Delete from storage if available
+ #[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.delete_node(id.as_ref())?;
}
@@ -177,6 +185,7 @@ impl GraphDB {
self.edges.insert(id.clone(), edge.clone());
// Persist to storage if available
+ #[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_edge(&edge)?;
}
@@ -197,6 +206,7 @@ impl GraphDB {
self.adjacency_index.remove_edge(&edge);
// Delete from storage if available
+ #[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.delete_edge(id.as_ref())?;
}
@@ -256,6 +266,7 @@ impl GraphDB {
self.hyperedges.insert(id.clone(), hyperedge.clone());
// Persist to storage if available
+ #[cfg(feature = "storage")]
if let Some(storage) = &self.storage {
storage.insert_hyperedge(&hyperedge)?;
}
diff --git a/crates/ruvector-graph/src/hybrid/vector_index.rs b/crates/ruvector-graph/src/hybrid/vector_index.rs
index dcd07d749..cd907743b 100644
--- a/crates/ruvector-graph/src/hybrid/vector_index.rs
+++ b/crates/ruvector-graph/src/hybrid/vector_index.rs
@@ -1,12 +1,16 @@
//! Vector indexing for graph elements
//!
-//! Integrates RuVector's HNSW index with graph nodes, edges, and hyperedges.
+//! Integrates RuVector's index (HNSW or Flat) with graph nodes, edges, and hyperedges.
use crate::error::{GraphError, Result};
use crate::types::{NodeId, EdgeId, PropertyValue, Properties};
+#[cfg(feature = "hnsw_rs")]
use ruvector_core::index::hnsw::HnswIndex;
+use ruvector_core::index::flat::FlatIndex;
use ruvector_core::index::VectorIndex;
-use ruvector_core::types::{DistanceMetric, HnswConfig, SearchResult};
+use ruvector_core::types::{DistanceMetric, SearchResult};
+#[cfg(feature = "hnsw_rs")]
+use ruvector_core::types::HnswConfig;
use dashmap::DashMap;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
@@ -30,7 +34,8 @@ pub struct EmbeddingConfig {
pub dimensions: usize,
/// Distance metric for similarity
pub metric: DistanceMetric,
- /// HNSW index configuration
+ /// HNSW index configuration (only used when hnsw_rs feature is enabled)
+ #[cfg(feature = "hnsw_rs")]
pub hnsw_config: HnswConfig,
/// Property name where embeddings are stored
pub embedding_property: String,
@@ -41,20 +46,27 @@ impl Default for EmbeddingConfig {
Self {
dimensions: 384, // Common for small models like MiniLM
metric: DistanceMetric::Cosine,
+ #[cfg(feature = "hnsw_rs")]
hnsw_config: HnswConfig::default(),
embedding_property: "embedding".to_string(),
}
}
}
+// Index type alias based on feature flags
+#[cfg(feature = "hnsw_rs")]
+type IndexImpl = HnswIndex;
+#[cfg(not(feature = "hnsw_rs"))]
+type IndexImpl = FlatIndex;
+
/// Hybrid index combining graph structure with vector search
pub struct HybridIndex {
/// Node embeddings index
- node_index: Arc>>,
+ node_index: Arc>>,
/// Edge embeddings index
- edge_index: Arc>>,
+ edge_index: Arc>>,
/// Hyperedge embeddings index
- hyperedge_index: Arc>>,
+ hyperedge_index: Arc>>,
/// Mapping from node IDs to internal vector IDs
node_id_map: Arc