Support KNN search for FAISS IVF indices #13258

ltamasi · 2024-12-28T01:18:02Z

Summary:
The patch is the read-side counterpart of #13197 . It adds support for K-nearest-neighbor vector similarity searches to FaissIVFIndex. There are two main pieces to this:

KNNIterator is an Iterator implementation that is returned by FaissIVFIndex upon a call to NewIterator. KNNIterator treats its Seek target as a vector embedding and passes it to FAISS along with the number of neighbors requested k as well as the number of probes to use (i.e. the number of inverted lists to check). Applications can then use Next (and Prev) to iterate over the the vectors in the result set. KNNIterator exposes the primary keys associated with the result vectors (see below how this is done), while value and columns are empty. The iterator also supports a property rocksdb.faiss.ivf.index.distance that can be used to retrieve the distance/similarity metric for the current result vector.
IteratorAdapter takes a RocksDB secondary index iterator (see Add a new interface method SecondaryIndex::NewIterator to enable querying the index #13257) and adapts it to the interface required by FAISS (faiss::InvertedListsIterator), enabling FAISS to read the inverted lists stored in RocksDB. Since FAISS only supports numerical vector ids of type faiss::idx_t, IteratorAdapter uses KNNIterator to assign ephemeral (per-query) ids to the inverted list items read during iteration, which are later mapped back to the original primary keys by KNNIterator.

Differential Revision: D67684898

facebook-github-bot · 2024-12-28T01:18:12Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

facebook-github-bot · 2024-12-28T01:28:23Z

This pull request was exported from Phabricator. Differential Revision: D67684898

facebook-github-bot · 2024-12-28T01:32:16Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

facebook-github-bot · 2024-12-30T19:24:20Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

facebook-github-bot · 2024-12-30T19:24:56Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

facebook-github-bot · 2024-12-30T23:10:18Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: The patch is the read-side counterpart of facebook#13197 . It adds support for K-nearest-neighbor vector similarity searches to `FaissIVFIndex`. There are two main pieces to this: 1) `KNNIterator` is an `Iterator` implementation that is returned by `FaissIVFIndex` upon a call to `NewIterator`. `KNNIterator` treats its `Seek` target as a vector embedding and passes it to FAISS along with the number of neighbors requested `k` as well as the number of probes to use (i.e. the number of inverted lists to check). Applications can then use `Next` (and `Prev`) to iterate over the the vectors in the result set. `KNNIterator` exposes the primary keys associated with the result vectors (see below how this is done), while `value` and `columns` are empty. The iterator also supports a property `rocksdb.faiss.ivf.index.distance` that can be used to retrieve the distance/similarity metric for the current result vector. 2) `IteratorAdapter` takes a RocksDB secondary index iterator (see facebook#13257) and adapts it to the interface required by FAISS (`faiss::InvertedListsIterator`), enabling FAISS to read the inverted lists stored in RocksDB. Since FAISS only supports numerical vector ids of type `faiss::idx_t`, `IteratorAdapter` uses `KNNIterator` to assign ephemeral (per-query) ids to the inverted list items read during iteration, which are later mapped back to the original primary keys by `KNNIterator`. Differential Revision: D67684898

facebook-github-bot · 2024-12-31T00:27:03Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: The patch is the read-side counterpart of facebook#13197 . It adds support for K-nearest-neighbor vector similarity searches to `FaissIVFIndex`. There are two main pieces to this: 1) `KNNIterator` is an `Iterator` implementation that is returned by `FaissIVFIndex` upon a call to `NewIterator`. `KNNIterator` treats its `Seek` target as a vector embedding and passes it to FAISS along with the number of neighbors requested `k` as well as the number of probes to use (i.e. the number of inverted lists to check). Applications can then use `Next` (and `Prev`) to iterate over the the vectors in the result set. `KNNIterator` exposes the primary keys associated with the result vectors (see below how this is done), while `value` and `columns` are empty. The iterator also supports a property `rocksdb.faiss.ivf.index.distance` that can be used to retrieve the distance/similarity metric for the current result vector. 2) `IteratorAdapter` takes a RocksDB secondary index iterator (see facebook#13257) and adapts it to the interface required by FAISS (`faiss::InvertedListsIterator`), enabling FAISS to read the inverted lists stored in RocksDB. Since FAISS only supports numerical vector ids of type `faiss::idx_t`, `IteratorAdapter` uses `KNNIterator` to assign ephemeral (per-query) ids to the inverted list items read during iteration, which are later mapped back to the original primary keys by `KNNIterator`. Differential Revision: D67684898

facebook-github-bot · 2025-01-02T20:04:41Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: The patch is the read-side counterpart of facebook#13197 . It adds support for K-nearest-neighbor vector similarity searches to `FaissIVFIndex`. There are two main pieces to this: 1) `KNNIterator` is an `Iterator` implementation that is returned by `FaissIVFIndex` upon a call to `NewIterator`. `KNNIterator` treats its `Seek` target as a vector embedding and passes it to FAISS along with the number of neighbors requested `k` as well as the number of probes to use (i.e. the number of inverted lists to check). Applications can then use `Next` (and `Prev`) to iterate over the the vectors in the result set. `KNNIterator` exposes the primary keys associated with the result vectors (see below how this is done), while `value` and `columns` are empty. The iterator also supports a property `rocksdb.faiss.ivf.index.distance` that can be used to retrieve the distance/similarity metric for the current result vector. 2) `IteratorAdapter` takes a RocksDB secondary index iterator (see facebook#13257) and adapts it to the interface required by FAISS (`faiss::InvertedListsIterator`), enabling FAISS to read the inverted lists stored in RocksDB. Since FAISS only supports numerical vector ids of type `faiss::idx_t`, `IteratorAdapter` uses `KNNIterator` to assign ephemeral (per-query) ids to the inverted list items read during iteration, which are later mapped back to the original primary keys by `KNNIterator`. Differential Revision: D67684898

facebook-github-bot · 2025-01-02T22:29:47Z

This pull request was exported from Phabricator. Differential Revision: D67684898

Summary: The patch is the read-side counterpart of facebook#13197 . It adds support for K-nearest-neighbor vector similarity searches to `FaissIVFIndex`. There are two main pieces to this: 1) `KNNIterator` is an `Iterator` implementation that is returned by `FaissIVFIndex` upon a call to `NewIterator`. `KNNIterator` treats its `Seek` target as a vector embedding and passes it to FAISS along with the number of neighbors requested `k` as well as the number of probes to use (i.e. the number of inverted lists to check). Applications can then use `Next` (and `Prev`) to iterate over the the vectors in the result set. `KNNIterator` exposes the primary keys associated with the result vectors (see below how this is done), while `value` and `columns` are empty. The iterator also supports a property `rocksdb.faiss.ivf.index.distance` that can be used to retrieve the distance/similarity metric for the current result vector. 2) `IteratorAdapter` takes a RocksDB secondary index iterator (see facebook#13257) and adapts it to the interface required by FAISS (`faiss::InvertedListsIterator`), enabling FAISS to read the inverted lists stored in RocksDB. Since FAISS only supports numerical vector ids of type `faiss::idx_t`, `IteratorAdapter` uses `KNNIterator` to assign ephemeral (per-query) ids to the inverted list items read during iteration, which are later mapped back to the original primary keys by `KNNIterator`. Differential Revision: D67684898

facebook-github-bot · 2025-01-02T22:40:12Z

This pull request was exported from Phabricator. Differential Revision: D67684898

jaykorean · 2025-01-03T21:01:58Z

include/rocksdb/options.h

+  // `similarity_search_probes` below.
+  //
+  // Default: none
+  std::optional<size_t> similarity_search_neighbors;


I'm wondering if we want similarity search options can be in its own struct.

SecondaryIndex::NewIterator() may take both ReadOptions and SecondaryIndexOptions (or better name) and these two similarity search options can be part of the new Options struct.

I'm just little more cautious as we are adding more options to ReadOptions which will be used for a very specific use case.

Yeah, I actually thought about that too - let me take a stab at it. btw, it would probably be sufficient for NewIterator to take just the secondary index read options (the "regular" read options would pertain to the underlying iterator).

Thanks! Updated

jaykorean · 2025-01-03T21:55:24Z

utilities/secondary_index/faiss_ivf_index.cc

+    ++pos_;
+  }
+
+  void Prev() override {


faiss::InvertedListsIterator doesn't seem to have prev() API. Just curious where we would use this.

This is KNNIterator, which is an implementation of rocksdb::Iterator, not faiss::InvertedListsIterator. Prev can be used to move back to a previous vector in the result set (see the unit test)

Ah okay. I thought the users would only access the InvertedList objects via FAISS (e.g. InvertedListScanner) and the underlying implementation uses this KNNIterator.

Reading through the unit test again helped. Thanks!

jaykorean · 2025-01-03T22:09:50Z

utilities/secondary_index/faiss_ivf_index.cc

+
+      const Status status = secondary_index_it_->status();
+      if (!status.ok()) {
+        throw std::runtime_error(status.ToString());


I wonder why we throw for errors here (and below). Wouldn't it be sufficient if is_available() returns false after Update() fails?

We have to differentiate between "out of items" and "some kind of error happened", similarly to RocksDB iterators. We wouldn't want to silently ignore I/O errors for example and produce an incorrect result. This is similar to how we expect RocksDB users to check an iterator's status() after Valid() returns false; it is possible that the iterator got invalidated because of an error (and not because we have exhausted all the elements).

For a bit more context, IteratorAdapter's code executes in the context of FAISS, and FAISS uses exceptions to communicate errors. However, we eventually do catch all exceptions coming from FAISS and turn them into error statuses (regardless of whether the exception originated from FAISS itself or from our class).

Summary: The patch is the read-side counterpart of facebook#13197 . It adds support for K-nearest-neighbor vector similarity searches to `FaissIVFIndex`. There are two main pieces to this: 1) `KNNIterator` is an `Iterator` implementation that is returned by `FaissIVFIndex` upon a call to `NewIterator`. `KNNIterator` treats its `Seek` target as a vector embedding and passes it to FAISS along with the number of neighbors requested `k` as well as the number of probes to use (i.e. the number of inverted lists to check). Applications can then use `Next` (and `Prev`) to iterate over the the vectors in the result set. `KNNIterator` exposes the primary keys associated with the result vectors (see below how this is done), while `value` and `columns` are empty. The iterator also supports a property `rocksdb.faiss.ivf.index.distance` that can be used to retrieve the distance/similarity metric for the current result vector. 2) `IteratorAdapter` takes a RocksDB secondary index iterator (see facebook#13257) and adapts it to the interface required by FAISS (`faiss::InvertedListsIterator`), enabling FAISS to read the inverted lists stored in RocksDB. Since FAISS only supports numerical vector ids of type `faiss::idx_t`, `IteratorAdapter` uses `KNNIterator` to assign ephemeral (per-query) ids to the inverted list items read during iteration, which are later mapped back to the original primary keys by `KNNIterator`. Differential Revision: D67684898

facebook-github-bot · 2025-01-04T01:27:48Z

This pull request was exported from Phabricator. Differential Revision: D67684898

jaykorean

LGTM. Thank you!

jaykorean · 2025-01-06T18:46:02Z

utilities/secondary_index/faiss_ivf_index.cc

+    ++pos_;
+  }
+
+  void Prev() override {


Ah okay. I thought the users would only access the InvertedList objects via FAISS (e.g. InvertedListScanner) and the underlying implementation uses this KNNIterator.

Reading through the unit test again helped. Thanks!

jaykorean · 2025-01-06T18:49:55Z

utilities/secondary_index/faiss_ivf_index_test.cc

+  std::unique_ptr<Iterator> underlying_it(db->NewIterator(ReadOptions(), cfh2));
+
+  SecondaryIndexReadOptions read_options;
+  read_options.similarity_search_neighbors = 8;


nit: can be a follow up. One more case for a different number of neighbors would be nice

jaykorean · 2025-01-06T18:55:58Z

utilities/secondary_index/faiss_ivf_index.cc

+    return Slice();
+  }
+
+  Status GetProperty(std::string prop_name, std::string* prop) override {


Not strictly related to this diff, but curious.

I am not sure what was the reason we used the string copy instead of const & for prop_name when this was added 9 years ago.

If the user checks distance for every single embeddings, this can be an unnecessary overhead, right? I'm wondering if we want to introduce a better version of keeping iterator properties in the next major release.

Yeah, this goes way back, so I don't have any context about why the interface is this way but it is something we can fix in a major release (even though the cost of the copy is probably negligible).

facebook-github-bot · 2025-01-06T19:42:49Z

This pull request has been merged in c2de783.

facebook-github-bot added the CLA Signed label Dec 28, 2024

facebook-github-bot added the fb-exported label Dec 28, 2024

ltamasi marked this pull request as draft December 28, 2024 01:18

ltamasi force-pushed the export-D67684898 branch from 88c4cbd to 09e87d1 Compare December 28, 2024 01:28

ltamasi added a commit to ltamasi/rocksdb that referenced this pull request Dec 28, 2024

Support KNN search for FAISS IVF indices (facebook#13258)

09e87d1

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

ltamasi force-pushed the export-D67684898 branch from 09e87d1 to fc2af33 Compare December 28, 2024 01:28

ltamasi added a commit to ltamasi/rocksdb that referenced this pull request Dec 28, 2024

Support KNN search for FAISS IVF indices (facebook#13258)

fc2af33

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

ltamasi added a commit to ltamasi/rocksdb that referenced this pull request Dec 28, 2024

Support KNN search for FAISS IVF indices (facebook#13258)

bd905aa

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

ltamasi force-pushed the export-D67684898 branch 2 times, most recently from bd905aa to 1c9a496 Compare December 30, 2024 19:24

ltamasi added a commit to ltamasi/rocksdb that referenced this pull request Dec 30, 2024

Support KNN search for FAISS IVF indices (facebook#13258)

1c9a496

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

ltamasi added a commit to ltamasi/rocksdb that referenced this pull request Dec 30, 2024

Support KNN search for FAISS IVF indices (facebook#13258)

10ac7a5

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

ltamasi force-pushed the export-D67684898 branch from 1c9a496 to 10ac7a5 Compare December 30, 2024 19:24

ltamasi added a commit to ltamasi/rocksdb that referenced this pull request Dec 30, 2024

Support KNN search for FAISS IVF indices (facebook#13258)

a0d8031

Summary: Pull Request resolved: facebook#13258 Differential Revision: D67684898

ltamasi force-pushed the export-D67684898 branch from 10ac7a5 to a0d8031 Compare December 30, 2024 23:09

ltamasi force-pushed the export-D67684898 branch from a0d8031 to b54c27b Compare December 31, 2024 00:26

ltamasi force-pushed the export-D67684898 branch from b54c27b to 9007290 Compare January 2, 2025 20:04

ltamasi force-pushed the export-D67684898 branch from 9007290 to 574541a Compare January 2, 2025 22:29

ltamasi force-pushed the export-D67684898 branch from 574541a to 660cc4c Compare January 2, 2025 22:39

ltamasi marked this pull request as ready for review January 2, 2025 22:50

ltamasi requested a review from jaykorean January 2, 2025 23:01

jaykorean reviewed Jan 3, 2025

View reviewed changes

ltamasi force-pushed the export-D67684898 branch from 660cc4c to c8b833b Compare January 4, 2025 01:27

jaykorean approved these changes Jan 6, 2025

View reviewed changes

facebook-github-bot closed this in c2de783 Jan 6, 2025

facebook-github-bot added the Merged label Jan 6, 2025

theolivenbaum mentioned this pull request Jan 7, 2025

Documentation for KNN search missing #13275

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support KNN search for FAISS IVF indices #13258

Support KNN search for FAISS IVF indices #13258

ltamasi commented Dec 28, 2024 •

edited

Loading

facebook-github-bot commented Dec 28, 2024

facebook-github-bot commented Dec 28, 2024

facebook-github-bot commented Dec 28, 2024

facebook-github-bot commented Dec 30, 2024

facebook-github-bot commented Dec 30, 2024

facebook-github-bot commented Dec 30, 2024

facebook-github-bot commented Dec 31, 2024

facebook-github-bot commented Jan 2, 2025

facebook-github-bot commented Jan 2, 2025

facebook-github-bot commented Jan 2, 2025

jaykorean Jan 3, 2025

ltamasi Jan 3, 2025

ltamasi Jan 4, 2025

jaykorean Jan 3, 2025

ltamasi Jan 3, 2025

jaykorean Jan 6, 2025

jaykorean Jan 3, 2025

ltamasi Jan 3, 2025 •

edited

Loading

ltamasi Jan 4, 2025

facebook-github-bot commented Jan 4, 2025

jaykorean left a comment

jaykorean Jan 6, 2025

jaykorean Jan 6, 2025

jaykorean Jan 6, 2025

ltamasi Jan 6, 2025

facebook-github-bot commented Jan 6, 2025

Support KNN search for FAISS IVF indices #13258

Support KNN search for FAISS IVF indices #13258

Conversation

ltamasi commented Dec 28, 2024 • edited Loading

facebook-github-bot commented Dec 28, 2024

facebook-github-bot commented Dec 28, 2024

facebook-github-bot commented Dec 28, 2024

facebook-github-bot commented Dec 30, 2024

facebook-github-bot commented Dec 30, 2024

facebook-github-bot commented Dec 30, 2024

facebook-github-bot commented Dec 31, 2024

facebook-github-bot commented Jan 2, 2025

facebook-github-bot commented Jan 2, 2025

facebook-github-bot commented Jan 2, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ltamasi Jan 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

facebook-github-bot commented Jan 4, 2025

jaykorean left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

facebook-github-bot commented Jan 6, 2025

ltamasi commented Dec 28, 2024 •

edited

Loading

ltamasi Jan 3, 2025 •

edited

Loading