Vector Search with HNSW

This notebook demonstrates the vector indexing and nearest-neighbor search capabilities of NetworkXGraph. It uses hnswlib for approximate nearest-neighbor (ANN) queries.

Key methods covered:

  • enable_vector_index() — create an ANN index on a node property

  • query_vector_index() — find the k nearest nodes to a query vector

  • Supported distance spaces: cosine, l2, ip (inner product)

[ ]:

import importlib.util

package_to_check = 'drm'
spec = importlib.util.find_spec(package_to_check)

if spec is None:
    print(f'⚠️ {package_to_check} no està instal·lat. Iniciant instal·lació...')
    %pip install -q --upgrade drm-tools
    print("✅ Instal·lació completada. L'estat del kernel PODRIA requerir un reinici.")
else:
    print(f'✅ {package_to_check} ja està present al sistema. Saltant instal·lació.')

[ ]:

import numpy as np

try:
    import hnswlib
    print("✅ hnswlib is available")
except ImportError:
    print("⚠️ hnswlib not installed. Run: pip install hnswlib")
    hnswlib = None

from drm import NetworkXGraph, Node

graph = NetworkXGraph()

# Create a small set of items with 4-dimensional feature vectors
# Each item has a 'features' property that will be indexed for vector search
items = [
    {"name": "apple",    "features": [0.9, 0.1, 0.05, 0.0],   "category": "fruit"},
    {"name": "banana",   "features": [0.85, 0.15, 0.0, 0.0],  "category": "fruit"},
    {"name": "carrot",   "features": [0.0, 0.9, 0.05, 0.05],  "category": "vegetable"},
    {"name": "broccoli", "features": [0.05, 0.85, 0.05, 0.05], "category": "vegetable"},
    {"name": "salmon",   "features": [0.1, 0.1, 0.7, 0.1],    "category": "protein"},
    {"name": "chicken",  "features": [0.1, 0.1, 0.75, 0.05],  "category": "protein"},
    {"name": "rice",     "features": [0.2, 0.2, 0.1, 0.5],    "category": "grain"},
]

node_ids = {}
for item in items:
    node = Node(pk={"name": item["name"]}, main_label="Food", features=item["features"], category=item["category"])
    nid = graph.insertNode(node)
    node_ids[item["name"]] = nid

print(f"Inserted {len(items)} food items")
print("Node ids:", node_ids)

Enabling a vector index

Call enable_vector_index(property_name, dimensions, space) to create an HNSW index on a specific node property. The property values must be 1D vectors of the given dimensionality.

Supported space values:

  • "cosine" — cosine similarity (default, recommended for normalized vectors)

  • "l2" — Euclidean distance

  • "ip" — inner product

[ ]:

graph.enable_vector_index(
    property_name="features",
    dimensions=4,
    space="cosine",
)
print("Vector index enabled for property 'features' (4D, cosine distance)")

Querying nearest neighbors

query_vector_index(property_name, vector, top_k) returns a list of (node_id, distance) tuples sorted by ascending distance.

We’ll query for the nearest neighbors of a “mystery” food item.

[ ]:

def lookup_name(node_id):
    """Reverse lookup: find the food name for a node id."""
    for name, nid in node_ids.items():
        if nid == node_id:
            return name
    return f"id={node_id}"

def show_results(query_vec, top_k=3):
    results = graph.query_vector_index(
        property_name="features",
        vector=query_vec,
        top_k=top_k,
    )
    print(f"Query: {query_vec}")
    print(f"Top {top_k} nearest neighbors:")
    for nid, dist in results:
        name = lookup_name(nid)
        print(f"  {name:10s} distance={dist:.4f}")
    return results

# Query 1: Find fruits (similar to apple)
print("=== Query 1: Find fruits ===")
show_results([0.9, 0.05, 0.0, 0.0], top_k=3)

print()

# Query 2: Find vegetables
print("=== Query 2: Find vegetables ===")
show_results([0.0, 0.8, 0.1, 0.1], top_k=3)

print()

# Query 3: Find protein
print("=== Query 3: Find protein ===")
show_results([0.05, 0.05, 0.8, 0.1], top_k=3)

Adding new nodes updates the index automatically

When you insert a new node with the indexed property, it is automatically added to the HNSW index.

[ ]:

# Add a new food item
mango = Node(pk={"name": "mango"}, main_label="Food", features=[0.88, 0.1, 0.02, 0.0], category="fruit")
graph.insertNode(mango)
node_ids["mango"] = graph.checkNode(mango)
print("Added 'mango' to the graph")

# Query again — mango should appear near apple and banana
print()
print("=== Query after adding mango ===")
show_results([0.9, 0.1, 0.0, 0.0], top_k=4)

Multiple vector indexes

You can create separate indexes on different properties.

[ ]:

# Create a second index on a different property (if you had 2D embedding)
# For this demo, we'll just show the API:
# graph.enable_vector_index(
#     property_name="other_features",
#     dimensions=8,
#     space="l2",
# )
print("Multiple vector indexes are supported. See the API for details.")

Cleanup

[ ]:

graph.close()
print("Graph closed.")