diff --git a/crates/ruvector-gnn-node/src/lib.rs b/crates/ruvector-gnn-node/src/lib.rs index 9a1b81e87..e73faa05b 100644 --- a/crates/ruvector-gnn-node/src/lib.rs +++ b/crates/ruvector-gnn-node/src/lib.rs @@ -64,39 +64,37 @@ impl RuvectorLayer { /// 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 + /// * `node_embedding` - Current node's embedding (Float32Array) + /// * `neighbor_embeddings` - Embeddings of neighbor nodes (Array of Float32Array) + /// * `edge_weights` - Weights of edges to neighbors (Float32Array) /// /// # Returns - /// Updated node embedding + /// Updated node embedding as Float32Array /// /// # 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 node = 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 weights = new Float32Array([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()) + node_embedding: Float32Array, + neighbor_embeddings: Vec, + edge_weights: Float32Array, + ) -> Result { + let node_slice = node_embedding.as_ref(); + let neighbors_vec: Vec> = neighbor_embeddings + .into_iter() + .map(|arr| arr.to_vec()) .collect(); - let weights_f32: Vec = edge_weights.iter().map(|&x| x as f32).collect(); + let weights_slice = edge_weights.as_ref(); - let result = self.inner.forward(&node_f32, &neighbors_f32, &weights_f32); + let result = self.inner.forward(node_slice, &neighbors_vec, weights_slice); - // Convert back to f64 - Ok(result.iter().map(|&x| x as f64).collect()) + Ok(Float32Array::new(result)) } /// Serialize the layer to JSON @@ -192,7 +190,7 @@ impl TensorCompress { /// Compress an embedding based on access frequency /// /// # Arguments - /// * `embedding` - The input embedding vector + /// * `embedding` - The input embedding vector (Float32Array) /// * `access_freq` - Access frequency in range [0.0, 1.0] /// /// # Returns @@ -200,16 +198,16 @@ impl TensorCompress { /// /// # Example /// ```javascript - /// const embedding = [1.0, 2.0, 3.0, 4.0]; + /// const embedding = new Float32Array([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(); + pub fn compress(&self, embedding: Float32Array, access_freq: f64) -> Result { + let embedding_slice = embedding.as_ref(); let compressed = self .inner - .compress(&embedding_f32, access_freq as f32) + .compress(embedding_slice, access_freq as f32) .map_err(|e| Error::new(Status::GenericFailure, format!("Compression error: {}", e)))?; serde_json::to_string(&compressed).map_err(|e| { @@ -223,7 +221,7 @@ impl TensorCompress { /// Compress with explicit compression level /// /// # Arguments - /// * `embedding` - The input embedding vector + /// * `embedding` - The input embedding vector (Float32Array) /// * `level` - Compression level configuration /// /// # Returns @@ -231,22 +229,22 @@ impl TensorCompress { /// /// # Example /// ```javascript - /// const embedding = [1.0, 2.0, 3.0, 4.0]; + /// const embedding = new Float32Array([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, + embedding: Float32Array, level: CompressionLevelConfig, ) -> Result { - let embedding_f32: Vec = embedding.iter().map(|&x| x as f32).collect(); + let embedding_slice = embedding.as_ref(); let rust_level = level.to_rust()?; let compressed = self .inner - .compress_with_level(&embedding_f32, &rust_level) + .compress_with_level(embedding_slice, &rust_level) .map_err(|e| Error::new(Status::GenericFailure, format!("Compression error: {}", e)))?; serde_json::to_string(&compressed).map_err(|e| { @@ -263,14 +261,14 @@ impl TensorCompress { /// * `compressed_json` - Compressed tensor as JSON string /// /// # Returns - /// Decompressed embedding vector + /// Decompressed embedding vector as Float32Array /// /// # Example /// ```javascript /// const decompressed = compressor.decompress(compressed); /// ``` #[napi] - pub fn decompress(&self, compressed_json: String) -> Result> { + pub fn decompress(&self, compressed_json: String) -> Result { let compressed: RustCompressedTensor = serde_json::from_str(&compressed_json).map_err(|e| { Error::new( @@ -286,7 +284,7 @@ impl TensorCompress { ) })?; - Ok(result.iter().map(|&x| x as f64).collect()) + Ok(Float32Array::new(result)) } } @@ -304,8 +302,8 @@ pub struct SearchResult { /// Differentiable search using soft attention mechanism /// /// # Arguments -/// * `query` - The query vector -/// * `candidate_embeddings` - List of candidate embedding vectors +/// * `query` - The query vector (Float32Array) +/// * `candidate_embeddings` - List of candidate embedding vectors (Array of Float32Array) /// * `k` - Number of top results to return /// * `temperature` - Temperature for softmax (lower = sharper, higher = smoother) /// @@ -314,27 +312,27 @@ pub struct SearchResult { /// /// # 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 query = new Float32Array([1.0, 0.0, 0.0]); +/// const candidates = [new Float32Array([1.0, 0.0, 0.0]), new Float32Array([0.9, 0.1, 0.0]), new Float32Array([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>, + query: Float32Array, + 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()) + let query_slice = query.as_ref(); + let candidates_vec: Vec> = candidate_embeddings + .into_iter() + .map(|arr| arr.to_vec()) .collect(); let (indices, weights) = - rust_differentiable_search(&query_f32, &candidates_f32, k as usize, temperature as f32); + rust_differentiable_search(query_slice, &candidates_vec, k as usize, temperature as f32); Ok(SearchResult { indices: indices.iter().map(|&i| i as u32).collect(), @@ -345,35 +343,35 @@ pub fn differentiable_search( /// Hierarchical forward pass through GNN layers /// /// # Arguments -/// * `query` - The query vector -/// * `layer_embeddings` - Embeddings organized by layer +/// * `query` - The query vector (Float32Array) +/// * `layer_embeddings` - Embeddings organized by layer (Array of Array of Float32Array) /// * `gnn_layers_json` - JSON array of serialized GNN layers /// /// # Returns -/// Final embedding after hierarchical processing +/// Final embedding after hierarchical processing as Float32Array /// /// # Example /// ```javascript -/// const query = [1.0, 0.0]; -/// const layerEmbeddings = [[[1.0, 0.0], [0.0, 1.0]]]; +/// const query = new Float32Array([1.0, 0.0]); +/// const layerEmbeddings = [[new Float32Array([1.0, 0.0]), new Float32Array([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>>, + query: Float32Array, + layer_embeddings: Vec>, gnn_layers_json: Vec, -) -> Result> { - let query_f32: Vec = query.iter().map(|&x| x as f32).collect(); +) -> Result { + let query_slice = query.as_ref(); let embeddings_f32: Vec>> = layer_embeddings - .iter() + .into_iter() .map(|layer| { layer - .iter() - .map(|v| v.iter().map(|&x| x as f32).collect()) + .into_iter() + .map(|arr| arr.to_vec()) .collect() }) .collect(); @@ -390,9 +388,9 @@ pub fn hierarchical_forward( }) .collect::>>()?; - let result = rust_hierarchical_forward(&query_f32, &embeddings_f32, &gnn_layers); + let result = rust_hierarchical_forward(query_slice, &embeddings_f32, &gnn_layers); - Ok(result.iter().map(|&x| x as f64).collect()) + Ok(Float32Array::new(result)) } // ==================== Helper Functions ==================== diff --git a/crates/ruvector-gnn-node/test/basic.test.js b/crates/ruvector-gnn-node/test/basic.test.js index 8cea55bd9..b18a5f267 100644 --- a/crates/ruvector-gnn-node/test/basic.test.js +++ b/crates/ruvector-gnn-node/test/basic.test.js @@ -25,20 +25,20 @@ test('RuvectorLayer creation', () => { 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 node = 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 weights = new Float32Array([0.3, 0.7]); const output = layer.forward(node, neighbors, weights); assert.strictEqual(output.length, 8); - assert.ok(output.every(x => typeof x === 'number')); + assert.ok(output instanceof Float32Array); }); 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 node = new Float32Array([1.0, 2.0, 3.0, 4.0]); const neighbors = []; - const weights = []; + const weights = new Float32Array([]); const output = layer.forward(node, neighbors, weights); assert.strictEqual(output.length, 8); @@ -59,17 +59,17 @@ test('RuvectorLayer deserialization', () => { 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 node = new Float32Array([1.0, 2.0, 3.0, 4.0]); + const neighbors = [new Float32Array([0.5, 1.0, 1.5, 2.0])]; + const weights = new Float32Array([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); - }); + for (let i = 0; i < output1.length; i++) { + assert.ok(Math.abs(output1[i] - output2[i]) < 1e-6); + } }); test('TensorCompress creation', () => { @@ -79,7 +79,7 @@ test('TensorCompress creation', () => { test('TensorCompress adaptive compression', () => { const compressor = new TensorCompress(); - const embedding = [1.0, 2.0, 3.0, 4.0]; + const embedding = new Float32Array([1.0, 2.0, 3.0, 4.0]); const compressed = compressor.compress(embedding, 0.5); assert.strictEqual(typeof compressed, 'string'); @@ -88,20 +88,21 @@ test('TensorCompress adaptive compression', () => { test('TensorCompress round-trip', () => { const compressor = new TensorCompress(); - const embedding = [1.0, 2.0, 3.0, 4.0]; + const embedding = new Float32Array([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); - }); + assert.ok(decompressed instanceof Float32Array); + for (let i = 0; i < decompressed.length; i++) { + assert.ok(Math.abs(decompressed[i] - embedding[i]) < 1e-6); + } }); test('TensorCompress with explicit level', () => { const compressor = new TensorCompress(); - const embedding = Array.from({ length: 64 }, (_, i) => i * 0.1); + const embedding = new Float32Array(Array.from({ length: 64 }, (_, i) => i * 0.1)); const level = { level_type: 'half', @@ -123,11 +124,11 @@ test('getCompressionLevel', () => { }); test('differentiableSearch', () => { - const query = [1.0, 0.0, 0.0]; + const query = new Float32Array([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], + new Float32Array([1.0, 0.0, 0.0]), + new Float32Array([0.9, 0.1, 0.0]), + new Float32Array([0.0, 1.0, 0.0]), ]; const result = differentiableSearch(query, candidates, 2, 1.0); @@ -147,7 +148,7 @@ test('differentiableSearch', () => { }); test('differentiableSearch with empty candidates', () => { - const query = [1.0, 0.0, 0.0]; + const query = new Float32Array([1.0, 0.0, 0.0]); const candidates = []; const result = differentiableSearch(query, candidates, 2, 1.0); @@ -157,9 +158,9 @@ test('differentiableSearch with empty candidates', () => { }); test('hierarchicalForward', () => { - const query = [1.0, 0.0]; + const query = new Float32Array([1.0, 0.0]); const layerEmbeddings = [ - [[1.0, 0.0], [0.0, 1.0]], + [new Float32Array([1.0, 0.0]), new Float32Array([0.0, 1.0])], ]; const layer = new RuvectorLayer(2, 2, 1, 0.0); @@ -167,9 +168,8 @@ test('hierarchicalForward', () => { const result = hierarchicalForward(query, layerEmbeddings, layers); - assert.ok(Array.isArray(result)); + assert.ok(result instanceof Float32Array); assert.strictEqual(result.length, 2); - assert.ok(result.every(x => typeof x === 'number')); }); test('invalid dropout rate throws error', () => { @@ -185,13 +185,13 @@ test('invalid dropout rate throws error', () => { test('compression with empty embedding throws error', () => { const compressor = new TensorCompress(); assert.throws(() => { - compressor.compress([], 0.5); + compressor.compress(new Float32Array([]), 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 embedding = new Float32Array(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