HNSW basic neighbor selection

Author: adish29

.vscode/c_cpp_properties.json

@@ -0,0 +1,23 @@
+{
+    "configurations": [
+        {
+            "name": "Win32",
+            "includePath": [
+                "${workspaceFolder}/**",
+                "C:\\users\\adish\\appdata\\local\\packages\\pythonsoftwarefoundation.python.3.9_qbz5n2kfra8p0\\localcache\\local-packages\\python39\\site-packages\\pybind11\\include",
+                "C:\\Users\\Adish\\AppData\\Local\\Programs\\Python\\Python311\\include"
+            ],
+            "defines": [
+                "_DEBUG",
+                "UNICODE",
+                "_UNICODE"
+            ],
+            "windowsSdkVersion": "10.0.18362.0",
+            "compilerPath": "cl.exe",
+            "cStandard": "c17",
+            "cppStandard": "c++17",
+            "intelliSenseMode": "windows-msvc-x64"
+        }
+    ],
+    "version": 4
+}

.vscode/settings.json

@@ -0,0 +1,6 @@
+{
+    "files.associations": {
+        "vector": "cpp",
+        "iostream": "cpp"
+    }
+}

compile.sh

@@ -0,0 +1,3 @@
+g++ -c main.cpp -o main.o
+g++ -c hnsw.cpp -o hnsw.o
+g++ main.o hnsw.o -o my_program

hnsw.cpp

@@ -0,0 +1,91 @@
+#include "hnsw.h"
+
+#include <algorithm>
+#include <iostream>
+#include <queue>
+#include <random>
+#include <set>
+#include <unordered_set>
+#include <vector>
+using namespace std;
+
+vector<int> HNSWGraph::searchLayer(Item& q, int ep, int ef, int lc) {
+	set<pair<double, int>> candidates;
+	set<pair<double, int>> nearestNeighbors;
+	unordered_set<int> isVisited;
+
+	double td = q.dist(items[ep]);
+	candidates.insert(make_pair(td, ep));
+	nearestNeighbors.insert(make_pair(td, ep));
+	isVisited.insert(ep);
+	while (!candidates.empty()) {
+		auto ci = candidates.begin(); candidates.erase(candidates.begin());
+		int nid = ci->second;
+		auto fi = nearestNeighbors.end(); fi--;
+		if (ci->first > fi->first) break;
+		for (int ed: layerEdgeLists[lc][nid]) {
+			if (isVisited.find(ed) != isVisited.end()) continue;
+			fi = nearestNeighbors.end(); fi--;
+			isVisited.insert(ed);
+			td = q.dist(items[ed]);
+			if ((td < fi->first) || nearestNeighbors.size() < ef) {
+				candidates.insert(make_pair(td, ed));
+				nearestNeighbors.insert(make_pair(td, ed));
+				if (nearestNeighbors.size() > ef) nearestNeighbors.erase(fi);
+			}
+		}
+	}
+	vector<int> results;
+	for(auto &p: nearestNeighbors) results.push_back(p.second);
+	return results;
+}
+
+vector<int> HNSWGraph::KNNSearch(Item& q, int K) {
+	int maxLyer = layerEdgeLists.size() - 1;
+	int ep = enterNode;
+	for (int l = maxLyer; l >= 1; l--) ep = searchLayer(q, ep, 1, l)[0];
+	return searchLayer(q, ep, K, 0);
+}
+
+void HNSWGraph::addEdge(int st, int ed, int lc) {
+	if (st == ed) return;
+	layerEdgeLists[lc][st].push_back(ed);
+	layerEdgeLists[lc][ed].push_back(st);
+}
+
+void HNSWGraph::Insert(Item& q) {
+	int nid = items.size();
+	itemNum++; items.push_back(q);
+	// sample layer
+	int maxLyer = layerEdgeLists.size() - 1;
+	int l = 0;
+	uniform_real_distribution<double> distribution(0.0,1.0);
+	while(l < ml && (1.0 / ml <= distribution(generator))) {
+		l++;
+		if (layerEdgeLists.size() <= l) layerEdgeLists.push_back(unordered_map<int, vector<int>>());
+	}
+	if (nid == 0) {
+		enterNode = nid;
+		return;
+	}
+	// search up layer entrance
+	int ep = enterNode;
+	for (int i = maxLyer; i > l; i--) ep = searchLayer(q, ep, 1, i)[0];
+	for (int i = min(l, maxLyer); i >= 0; i--) {
+		int MM = l == 0 ? MMax0 : MMax;
+		vector<int> neighbors = searchLayer(q, ep, efConstruction, i);
+		vector<int> selectedNeighbors = vector<int>(neighbors.begin(), neighbors.begin()+min(int(neighbors.size()), M));
+		for (int n: selectedNeighbors) addEdge(n, nid, i);
+		for (int n: selectedNeighbors) {
+			if (layerEdgeLists[i][n].size() > MM) {
+				vector<pair<double, int>> distPairs;
+				for (int nn: layerEdgeLists[i][n]) distPairs.emplace_back(items[n].dist(items[nn]), nn);
+				sort(distPairs.begin(), distPairs.end());
+				layerEdgeLists[i][n].clear();
+				for (int d = 0; d < min(int(distPairs.size()), MM); d++) layerEdgeLists[i][n].push_back(distPairs[d].second);
+			}
+		}
+		ep = selectedNeighbors[0];
+	}
+	if (l == layerEdgeLists.size() - 1) enterNode = nid;
+}

hnsw.h

@@ -0,0 +1,54 @@
+#ifndef HNSW_H
+#define HNSW_H
+
+#include <random>
+#include <vector>
+#include <unordered_map>
+#include <iostream>
+using namespace std;
+
+struct Item {
+	Item(vector<double> _values):values(_values) {}
+	vector<double> values;
+	double dist(Item& other) {
+		double result = 0.0;
+		for (int i = 0; i < values.size(); i++){ 
+            result += (values[i] - other.values[i]) * (values[i] - other.values[i]);
+        }
+		return result;
+	}
+};
+
+struct HNSWGraph {
+	HNSWGraph(int _M, int _MMax, int _MMax0, int _efConstruction, int _ml):M(_M),MMax(_MMax),MMax0(_MMax0),efConstruction(_efConstruction),ml(_ml){
+		layerEdgeLists.push_back(unordered_map<int, vector<int>>());
+	}
+	int M;
+	int MMax;
+	int MMax0;
+	int efConstruction;
+	int ml;
+	int itemNum;
+	vector<Item> items;
+	vector<unordered_map<int, vector<int>>> layerEdgeLists;
+	int enterNode;
+
+	default_random_engine generator;
+	void addEdge(int st, int ed, int lc);
+	vector<int> searchLayer(Item& q, int ep, int ef, int lc);
+	void Insert(Item& q);
+	vector<int> KNNSearch(Item& q, int K);
+
+	void printGraph() {
+		for (int l = 0; l < layerEdgeLists.size(); l++) {
+			cout << "Layer:" << l << endl;
+			for (auto it = layerEdgeLists[l].begin(); it != layerEdgeLists[l].end(); ++it) {
+				cout << it->first << ":";
+				for (auto ed: it->second) cout << ed << " ";
+				cout << endl;
+			}
+		}
+	}
+};
+
+#endif

hnsw.o

@@ -0,0 +1,18 @@
+#include <pybind11/stl.h>
+#include <pybind11/pybind11.h>
+#include "hnsw.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(pyhnsw, m) {
+    py::class_<Item>(m, "Item")
+        .def(py::init<std::vector<double>>());
+
+    py::class_<HNSWGraph>(m, "HNSWGraph")
+        .def(py::init<int, int, int, int, int>())
+        .def("addEdge", &HNSWGraph::addEdge)
+        .def("searchLayer", &HNSWGraph::searchLayer)
+        .def("Insert", &HNSWGraph::Insert)
+        .def("KNNSearch", &HNSWGraph::KNNSearch)
+        .def("printGraph", &HNSWGraph::printGraph);
+}

hnsw_py.cpp

@@ -0,0 +1,71 @@
+#include "hnsw.h"
+
+#include <algorithm>
+#include <ctime>
+#include <iostream>
+#include <random>
+#include <vector>
+using namespace std;
+
+void randomTest(int numItems, int dim, int numQueries, int K) {
+	std::default_random_engine generator;
+    std::uniform_real_distribution<double> distribution(0.0, 1.0);
+    std::vector<Item> randomItems;
+    randomItems.reserve(numItems);
+
+    for (int i = 0; i < numItems; i++) {
+        std::vector<double> temp(dim);
+        for (int d = 0; d < dim; d++) {
+            temp[d] = distribution(generator);
+        }
+        randomItems.emplace_back(temp);
+    }
+
+    std::shuffle(randomItems.begin(), randomItems.end(), generator);
+	HNSWGraph myHNSWGraph(10, 30, 30, 10, 2);
+	for (int i = 0; i < numItems; i++) {
+		if (i % 10000 == 0) cout << i << endl;
+		myHNSWGraph.Insert(randomItems[i]);
+	}
+	
+	double total_brute_force_time = 0.0;
+	double total_hnsw_time = 0.0;
+
+	cout << "START QUERY" << endl;
+	int numHits = 0;
+	for (int i = 0; i < numQueries; i++) {
+		vector<double> temp(dim);
+		for (int d = 0; d < dim; d++) {
+			temp[d] = distribution(generator);
+		}
+		Item query(temp);
+
+		// Brute force
+		clock_t begin_time = clock();
+		vector<pair<double, int>> distPairs;
+		for (int j = 0; j < numItems; j++) {
+			if (j == i) continue;
+			distPairs.emplace_back(query.dist(randomItems[j]), j);
+		}
+		sort(distPairs.begin(), distPairs.end());
+		total_brute_force_time += double( clock () - begin_time ) /  CLOCKS_PER_SEC;
+
+		begin_time = clock();
+		vector<int> knns = myHNSWGraph.KNNSearch(query, K);
+		// cout << "Printing vectors";
+		// std::cout << "Contents of knns vector:" << std::endl;
+		// for (size_t i = 0; i < knns.size(); ++i) {
+		// 	std::cout << "knns[" << i << "] = " << knns[i] << std::endl;
+		// }
+		// cout << "\nPrinted Vectors";
+		total_hnsw_time += double( clock () - begin_time ) /  CLOCKS_PER_SEC;
+
+		if (knns[0] == distPairs[0].second) numHits++;
+	}
+	cout << numHits << " " << total_brute_force_time / numQueries  << " " << total_hnsw_time / numQueries << endl;
+}
+
+int main() {
+	randomTest(10000, 4, 100, 5);
+	return 0;
+}