b339d089ed
Summary: Pull Request resolved: https://github.com/facebook/rocksdb/pull/13197 The patch adds initial support for backing FAISS's inverted file based indices with data stored in RocksDB. It introduces a `SecondaryIndex` implementation called `FaissIVFIndex` which takes ownership of a `faiss::IndexIVF` object. During indexing, `FaissIVFIndex` treats the original value of the specified primary column as an embedding vector, and passes it to the provided FAISS index object to perform quantization. It replaces the original embedding vector with the result of the coarse quantizer (i.e. the inverted list id), and puts the result of the fine quantizer (if any) into the secondary index value. Note that this patch is only one half of the equation; it provides a way of storing FAISS inverted lists in RocksDB but there is currently no retrieval/search support (this will be a follow-up change). Also, the integration currently works only with our internal Buck build. I plan to add support for `cmake` / `make` based builds similarly to how we handle Folly. Reviewed By: jowlyzhang Differential Revision: D66907065 fbshipit-source-id: 63fdf29895d5feeffc230254a7ddfb0aac050967
60 lines
2.1 KiB
C++
60 lines
2.1 KiB
C++
// Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
//
|
|
// This source code is licensed under both the GPLv2 (found in the
|
|
// COPYING file in the root directory) and Apache 2.0 License
|
|
// (found in the LICENSE.Apache file in the root directory).
|
|
|
|
#pragma once
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
|
|
#include "faiss/IndexIVF.h"
|
|
#include "rocksdb/utilities/secondary_index.h"
|
|
|
|
namespace ROCKSDB_NAMESPACE {
|
|
|
|
// A SecondaryIndex implementation that wraps a FAISS inverted file index.
|
|
class FaissIVFIndex : public SecondaryIndex {
|
|
public:
|
|
explicit FaissIVFIndex(std::unique_ptr<faiss::IndexIVF>&& index,
|
|
std::string primary_column_name);
|
|
~FaissIVFIndex() override;
|
|
|
|
void SetPrimaryColumnFamily(ColumnFamilyHandle* column_family) override;
|
|
void SetSecondaryColumnFamily(ColumnFamilyHandle* column_family) override;
|
|
|
|
ColumnFamilyHandle* GetPrimaryColumnFamily() const override;
|
|
ColumnFamilyHandle* GetSecondaryColumnFamily() const override;
|
|
|
|
Slice GetPrimaryColumnName() const override;
|
|
|
|
Status UpdatePrimaryColumnValue(
|
|
const Slice& primary_key, const Slice& primary_column_value,
|
|
std::optional<std::variant<Slice, std::string>>* updated_column_value)
|
|
const override;
|
|
|
|
Status GetSecondaryKeyPrefix(
|
|
const Slice& primary_key, const Slice& primary_column_value,
|
|
std::variant<Slice, std::string>* secondary_key_prefix) const override;
|
|
|
|
Status GetSecondaryValue(const Slice& primary_key,
|
|
const Slice& primary_column_value,
|
|
const Slice& original_column_value,
|
|
std::optional<std::variant<Slice, std::string>>*
|
|
secondary_value) const override;
|
|
|
|
private:
|
|
class Adapter;
|
|
|
|
static std::string SerializeLabel(faiss::idx_t label);
|
|
static faiss::idx_t DeserializeLabel(Slice label_slice);
|
|
|
|
std::unique_ptr<Adapter> adapter_;
|
|
std::unique_ptr<faiss::IndexIVF> index_;
|
|
std::string primary_column_name_;
|
|
ColumnFamilyHandle* primary_column_family_{};
|
|
ColumnFamilyHandle* secondary_column_family_{};
|
|
};
|
|
|
|
} // namespace ROCKSDB_NAMESPACE
|