All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
/data/users/matthijs/github_faiss/faiss/AuxIndexStructures.cpp
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the BSD+Patents license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved
// -*- c++ -*-
#include "AuxIndexStructures.h"
#include <cstring>
namespace faiss {
/***********************************************************************
* RangeSearchResult
***********************************************************************/
RangeSearchResult::RangeSearchResult (idx_t nq, bool alloc_lims): nq (nq) {
if (alloc_lims) {
lims = new size_t [nq + 1];
memset (lims, 0, sizeof(*lims) * (nq + 1));
} else {
lims = nullptr;
}
labels = nullptr;
distances = nullptr;
buffer_size = 1024 * 256;
}
/// called when lims contains the nb of elements result entries
/// for each query
void RangeSearchResult::do_allocation () {
size_t ofs = 0;
for (int i = 0; i < nq; i++) {
size_t n = lims[i];
lims [i] = ofs;
ofs += n;
}
lims [nq] = ofs;
labels = new idx_t [ofs];
distances = new float [ofs];
}
RangeSearchResult::~RangeSearchResult () {
delete [] labels;
delete [] distances;
delete [] lims;
}
/***********************************************************************
* BufferList
***********************************************************************/
BufferList::BufferList (size_t buffer_size):
buffer_size (buffer_size)
{
wp = buffer_size;
}
BufferList::~BufferList ()
{
for (int i = 0; i < buffers.size(); i++) {
delete [] buffers[i].ids;
delete [] buffers[i].dis;
}
}
void BufferList::append_buffer ()
{
Buffer buf = {new idx_t [buffer_size], new float [buffer_size]};
buffers.push_back (buf);
wp = 0;
}
/// copy elemnts ofs:ofs+n-1 seen as linear data in the buffers to
/// tables dest_ids, dest_dis
void BufferList::copy_range (size_t ofs, size_t n,
idx_t * dest_ids, float *dest_dis)
{
size_t bno = ofs / buffer_size;
ofs -= bno * buffer_size;
while (n > 0) {
size_t ncopy = ofs + n < buffer_size ? n : buffer_size - ofs;
Buffer buf = buffers [bno];
memcpy (dest_ids, buf.ids + ofs, ncopy * sizeof(*dest_ids));
memcpy (dest_dis, buf.dis + ofs, ncopy * sizeof(*dest_dis));
dest_ids += ncopy;
dest_dis += ncopy;
ofs = 0;
bno ++;
n -= ncopy;
}
}
/***********************************************************************
* RangeSearchPartialResult
***********************************************************************/
RangeSearchPartialResult::RangeSearchPartialResult (RangeSearchResult * res_in):
BufferList(res_in->buffer_size),
res(res_in)
{}
/// begin a new result
RangeSearchPartialResult::QueryResult &
RangeSearchPartialResult::new_result (idx_t qno)
{
QueryResult qres = {qno, 0, this};
queries.push_back (qres);
return queries.back();
}
void RangeSearchPartialResult::finalize ()
{
set_lims ();
#pragma omp barrier
#pragma omp single
res->do_allocation ();
#pragma omp barrier
set_result ();
}
/// called by range_search before do_allocation
void RangeSearchPartialResult::set_lims ()
{
for (int i = 0; i < queries.size(); i++) {
QueryResult & qres = queries[i];
res->lims[qres.qno] = qres.nres;
}
}
/// called by range_search after do_allocation
void RangeSearchPartialResult::set_result (bool incremental)
{
size_t ofs = 0;
for (int i = 0; i < queries.size(); i++) {
QueryResult & qres = queries[i];
copy_range (ofs, qres.nres,
res->labels + res->lims[qres.qno],
res->distances + res->lims[qres.qno]);
if (incremental) {
res->lims[qres.qno] += qres.nres;
}
ofs += qres.nres;
}
}
/***********************************************************************
* IDSelectorRange
***********************************************************************/
IDSelectorRange::IDSelectorRange (idx_t imin, idx_t imax):
imin (imin), imax (imax)
{
}
bool IDSelectorRange::is_member (idx_t id) const
{
return id >= imin && id < imax;
}
/***********************************************************************
* IDSelectorBatch
***********************************************************************/
IDSelectorBatch::IDSelectorBatch (long n, const idx_t *indices)
{
nbits = 0;
while (n > (1L << nbits)) nbits++;
nbits += 5;
// for n = 1M, nbits = 25 is optimal, see P56659518
mask = (1L << nbits) - 1;
bloom.resize (1UL << (nbits - 3), 0);
for (long i = 0; i < n; i++) {
long id = indices[i];
set.insert(id);
id &= mask;
bloom[id >> 3] |= 1 << (id & 7);
}
}
bool IDSelectorBatch::is_member (idx_t i) const
{
long im = i & mask;
if(!(bloom[im>>3] & (1 << (im & 7)))) {
return 0;
}
return set.count(i);
}
}
faiss::RangeSearchResult::nq
size_t nq
nb of queries
Definition: AuxIndexStructures.h:30
faiss::RangeSearchPartialResult::QueryResult
result structure for a single query
Definition: AuxIndexStructures.h:145
faiss::BufferList::copy_range
void copy_range(size_t ofs, size_t n, idx_t *dest_ids, float *dest_dis)
Definition: AuxIndexStructures.cpp:85
faiss::RangeSearchPartialResult::set_result
void set_result(bool incremental=false)
called by range_search after do_allocation
Definition: AuxIndexStructures.cpp:148
faiss::RangeSearchPartialResult::set_lims
void set_lims()
called by range_search before do_allocation
Definition: AuxIndexStructures.cpp:139
faiss::RangeSearchResult::RangeSearchResult
RangeSearchResult(idx_t nq, bool alloc_lims=true)
lims must be allocated on input to range_search.
Definition: AuxIndexStructures.cpp:23
faiss::BufferList::wp
size_t wp
write pointer in the last buffer.
Definition: AuxIndexStructures.h:109
faiss::RangeSearchResult::distances
float * distances
corresponding distances (not sorted)
Definition: AuxIndexStructures.h:36
faiss::RangeSearchPartialResult::new_result
QueryResult & new_result(idx_t qno)
begin a new result
Definition: AuxIndexStructures.cpp:117
faiss::RangeSearchResult::do_allocation
virtual void do_allocation()
Definition: AuxIndexStructures.cpp:37
faiss::RangeSearchResult
Definition: AuxIndexStructures.h:29
faiss::BufferList::Buffer
Definition: AuxIndexStructures.h:103
faiss::RangeSearchResult::buffer_size
size_t buffer_size
size of the result buffers used
Definition: AuxIndexStructures.h:38
faiss::RangeSearchResult::lims
size_t * lims
size (nq + 1)
Definition: AuxIndexStructures.h:31
faiss::BufferList
Definition: AuxIndexStructures.h:97
faiss::RangeSearchResult::labels
idx_t * labels
result for query i is labels[lims[i]:lims[i+1]]
Definition: AuxIndexStructures.h:35