mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
121 lines
2.6 KiB
C++
121 lines
2.6 KiB
C++
/**
|
|
* Copyright (c) Facebook, Inc. and its affiliates.
|
|
*
|
|
* This source code is licensed under the MIT license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
// -*- c++ -*-
|
|
|
|
#ifndef FAISS_DIRECT_MAP_H
|
|
#define FAISS_DIRECT_MAP_H
|
|
|
|
#include <faiss/InvertedLists.h>
|
|
#include <unordered_map>
|
|
|
|
|
|
namespace faiss {
|
|
|
|
// When offsets list id + offset are encoded in an uint64
|
|
// we call this LO = list-offset
|
|
|
|
inline uint64_t lo_build (uint64_t list_id, uint64_t offset) {
|
|
return list_id << 32 | offset;
|
|
}
|
|
|
|
inline uint64_t lo_listno (uint64_t lo) {
|
|
return lo >> 32;
|
|
}
|
|
|
|
inline uint64_t lo_offset (uint64_t lo) {
|
|
return lo & 0xffffffff;
|
|
}
|
|
|
|
/**
|
|
* Direct map: a way to map back from ids to inverted lists
|
|
*/
|
|
struct DirectMap {
|
|
typedef Index::idx_t idx_t;
|
|
|
|
enum Type {
|
|
NoMap = 0, // default
|
|
Array = 1, // sequential ids (only for add, no add_with_ids)
|
|
Hashtable = 2 // arbitrary ids
|
|
};
|
|
Type type;
|
|
|
|
/// map for direct access to the elements. Map ids to LO-encoded entries.
|
|
std::vector <idx_t> array;
|
|
std::unordered_map <idx_t, idx_t> hashtable;
|
|
|
|
DirectMap();
|
|
|
|
/// set type and initialize
|
|
void set_type (Type new_type, const InvertedLists *invlists, size_t ntotal);
|
|
|
|
/// get an entry
|
|
idx_t get (idx_t id) const;
|
|
|
|
/// for quick checks
|
|
bool no () const {return type == NoMap; }
|
|
|
|
/**
|
|
* update the direct_map
|
|
*/
|
|
|
|
/// throw if Array and ids is not NULL
|
|
void check_can_add (const idx_t *ids);
|
|
|
|
/// non thread-safe version
|
|
void add_single_id (idx_t id, idx_t list_no, size_t offset);
|
|
|
|
/// remove all entries
|
|
void clear();
|
|
|
|
/**
|
|
* operations on inverted lists that require translation with a DirectMap
|
|
*/
|
|
|
|
/// remove ids from the InvertedLists, possibly using the direct map
|
|
size_t remove_ids(const IDSelector& sel, InvertedLists *invlists);
|
|
|
|
/// update entries, using the direct map
|
|
void update_codes (InvertedLists *invlists,
|
|
int n, const idx_t *ids,
|
|
const idx_t *list_nos,
|
|
const uint8_t *codes);
|
|
|
|
|
|
|
|
};
|
|
|
|
/// Thread-safe way of updating the direct_map
|
|
struct DirectMapAdd {
|
|
|
|
typedef Index::idx_t idx_t;
|
|
|
|
using Type = DirectMap::Type;
|
|
|
|
DirectMap &direct_map;
|
|
DirectMap::Type type;
|
|
size_t ntotal;
|
|
size_t n;
|
|
const idx_t *xids;
|
|
|
|
std::vector<idx_t> all_ofs;
|
|
|
|
DirectMapAdd (DirectMap &direct_map, size_t n, const idx_t *xids);
|
|
|
|
/// add vector i (with id xids[i]) at list_no and offset
|
|
void add (size_t i, idx_t list_no, size_t offset);
|
|
|
|
~DirectMapAdd ();
|
|
};
|
|
|
|
|
|
|
|
}
|
|
|
|
|
|
#endif
|