Native-Like Performance for Nearest Neighbors Search using Hnswlib in Java Applications

Written by: Casper Davies, Germán Hurtado, Henri David, Hussama Ismail, Katarzyna Konewka, Stefan Skoruppa, Vinitha Venugopal and Zhenhua Mai.

Photo by Jacek Dylag on Unsplash

What is a Nearest Neighbor Search?

Let’s imagine you planned an adventure trip for your next vacation. You booked a Bed&Breakfast in Cefalù, Sicily, Italy. The year is 1996, you didn’t know anyone over there; Google Maps, OpenStreetMap and no other kind of digital support existed yet. Worse still, you don’t have a physical map to guide you except that the B&B is next to some vague reference points. How do you find the place? Surely enough, you started by going to those reference points and look around from there.

Your room hunting could be more formally translated into “Given a set of n points and a query point q, find the points closest to the query point” which is one of the most fundamental problems in computational geometry. It has applications in several areas of computer science (e.g., machine learning, pattern recognition, computer vision, and many others) and it is called the nearest neighbor (NN) problem.

Hierarchical Navigable Small World Graph (Hnswlib)

Finding the nearest neighbor (or neighbors) becomes much more complex and costly within large datasets. In order to speed up the search, variations like the approximate k-nearest neighbors (k-nn) are proposed.

Hnswlib is a library written in C++ that implements an efficient and totally graph-based approach for approximate k-nn search. It relies on small world graphs in a multi-layer and hierarchical fashion which is able to offer logarithmic complexity, better scalability, and performance in comparison to vector-only approaches. Moreover, Hnswlib allows incremental index building which is not possible with some other libraries (e.g., Annoy).

In this post, we aren’t going to dive into the theoretical/mathematical details, so for more information please have a look at this research paper.

Using Hnswlib in Java

After the first experiments with Hnswlib, we got quite interested in the idea and definitely wanted to try it out for our application. However, we knew that bringing a Native or Python program to some of our projects would lead us towards extra complexities in the development process. Previously, we used JEP to embed Python and run our scripts. This time we decided to first have a look at what (at the time we are writing this post) was available for Java.

During the exploration, we found many other implementations of Hnswlib, including one completely written in Java. Unfortunately, the performance did not meet our requirements. Furthermore, using a different library, we would lose the ability of reusing an index previously created with the native implementation.

In order to keep compatibility, we could create our own binding to native with JNI (Java Native Interface) but it can be a bit tricky and difficult to maintain. Hence, we quickly looked into the new trends for interoperability in Java (e.g., Panama and GraalVM). Panama’s lack of documentation pushed us away. GraalVM on the other hand appeared to be very well documented, included many code examples and has a very responsive community.

After trying running the LLVM bytecode of Hnswlib on GraalVM, we saw that the performance was still not comparable to native. On the other hand, GraalVM will certainly be a great option for interoperability in the coming years. To make this story short, we also looked at JNA (Java Native Access) which offers an abstraction on top of JNI and that makes the access to native code much easier. With JNA’s help, we managed to create a Java biding to Hnswlib that offers performance close to native and that meets our needs.

Hnswlib with Java Native Access

The Java binding project is called Hnswlib-jna and it is available on GitHub. It was created based on the Python bindings provided by Hnswlib’s authors. Differently from the original Python implementation, the multi-thread support is not included in the bindings itself but it can be easily implemented in Java. Additionally, since it uses native code, a shared library must be provided. More information about how to generate the shared library and about how to link it to your project can be found in the project documentation.

Importing Hnswlib-jna into Your Project

The jar file has been built and distributed using Maven. To have access to the classes, please add the following dependency in your pom.xml:

<dependency>
<groupId>com.stepstone.search.hnswlib.jna</groupId>
<artifactId>hnswlib-jna</artifactId>
<version>1.0.0</version>
</dependency>

Creating and Using a Small World Index

Once the classes (and the shared library) are available to your project, creating a new index is straightforward.

import com.stepstone.search.hnswlib.jna.Index;...Index index = new Index(SpaceName.COSINE, 3);
index.initialize(20);

The Index constructor expects 2 arguments: the first one is the SpaceName (e.g., L2, IP, COSINE) and the second is the dimension (i.e., the length of the arrays defined in the vector space).

Note: This example includes a call to the initialize() method. This method call is a must since it is responsible for initializing the index on the native side. It also allows users to tweak Hnswlib performance.

Adding Items and Querying an Existing Index

Since the index is created and initialized, we are now ready to add items into the vector space and perform queries. Items are added via addItem() method calls. This method expects an array with length equal to the dimension previously set and also a label (integer) which helps Hnswlib identifying the item in the space.

index.addItem(new float[] { 1.3f, 1.2f, 1.5f  }, 1);
index.addItem(new float[] { 1.3f, 1.2f, 1.49f }, 2);
index.addItem(new float[] { 1.3f, 1.2f, 1.48f }, 3);
QueryTuple qt = index.knnQuery(new float[] {1.3f, 1.2f, 1.48f}, 2);

Nearest neighbor searches, on the other hand, are performed via the knnQuery() method. It expects an input array and a k (i.e., the desired number of neighbors). This method returns an instance of QueryTuple which contains two arrays: one with the labels (from the nearest to the farthest) and another one with the coefficients (scores) for the respective labels. After printing out the array’s contents we can see:

[3, 2]
[1.1920929E-7, 5.6028366E-6]

From this example, considering the input array [1.3, 1.2, 1.48], the item with label 3 is the nearest, followed by the item with label 2.

One thing to keep in mind is that SpaceName.COSINE requires normalized vectors. The normalization can be performed on the native side but if you already have the items in a normalized form you should use addNormalizedItem() and knnNormalizedQuery() methods in order avoid redundant normalization.

Other Methods Available in the Index class

  • getLength(): returns the number of elements already inserted in the index;
  • save(path): stores the content of the index into a file;
  • load(path, maxNumberOfElements): loads the content stored in a file path onto the index;
  • clear(): frees the memory allocated for this index in the native context;
  • Index.normalize(array): static function that normalizes an array in Java.

Conclusions

In this article, we discussed a little bit about the nearest neighbors problem, Hierarchical Navigable Small World Graphs, Hnswlib, and our attempts trying to use native calls with the latest Java technologies. As a result and contribution back to the community, we created a Java binding to the native Hnswlib which allows native-like performance and compatibility to Java applications.

The code is open-source shared under Apache License 2.0 and fully available on Github (https://github.com/stepstone-tech/hnswlib-jna).

Read more about the technologies we use or take an inside look at our organisation & processes.
Interested in working at StepStone? Check out our careers page.

--

--