Skip to content

Commit

Permalink
Id selector c api (facebookresearch#2755)
Browse files Browse the repository at this point in the history 
Summary:
Implements c_api interfaces for IDSelectorNot, IDSelectorAnd, IDSelectorOr, IDSelectorXOr

Pull Request resolved: facebookresearch#2755

Reviewed By: alexanderguzhva

Differential Revision: D44366433

Pulled By: mdouze

fbshipit-source-id: b28eb32852188888f51b7ba08b37f76e174dcf3c
  • Loading branch information
@kaelen-moda @facebook-github-bot
kaelen-moda authored and facebook-github-bot committed Mar 24, 2023
1 parent 0f07e85 commit aae6ff5
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 0 deletions.
57 changes: 57 additions & 0 deletions c_api/example_c.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,63 @@ int main() {
faiss_IDSelectorRange_free(sel);
}

{ // search xb first 5 but search parameters of id range [20,40] OR
// [45,60]
idx_t* I = malloc(k * nq * sizeof(idx_t));
float* D = malloc(k * nq * sizeof(float));
FaissIDSelectorRange* lhs_sel = NULL;
FAISS_TRY(faiss_IDSelectorRange_new(&lhs_sel, 20, 40));
FaissIDSelectorRange* rhs_sel = NULL;
FAISS_TRY(faiss_IDSelectorRange_new(&rhs_sel, 45, 60));
FaissIDSelectorOr* sel = NULL;
FAISS_TRY(faiss_IDSelectorOr_new(&sel, lhs_sel, rhs_sel));
FaissSearchParameters* params = NULL;
FAISS_TRY(faiss_SearchParameters_new(&params, sel));
FAISS_TRY(
faiss_Index_search_with_params(index, nq, xq, k, params, D, I));
printf("Searching w/ IDSelectorRange [20,40] OR [45,60] \n");
printf("I=\n");
for (int i = 0; i < 5; i++) {
for (int j = 0; j < k; j++)
printf("%5lld (d=%2.3f) ", I[i * k + j], D[i * k + j]);
printf("\n");
}
free(I);
free(D);
faiss_SearchParameters_free(params);
faiss_IDSelectorRange_free(lhs_sel);
faiss_IDSelectorRange_free(rhs_sel);
faiss_IDSelector_free(sel);
}
{ // search xb first 5 but search parameters of id range [20,40] AND
// [15,35] = [20,35]
idx_t* I = malloc(k * nq * sizeof(idx_t));
float* D = malloc(k * nq * sizeof(float));
FaissIDSelectorRange* lhs_sel = NULL;
FAISS_TRY(faiss_IDSelectorRange_new(&lhs_sel, 20, 40));
FaissIDSelectorRange* rhs_sel = NULL;
FAISS_TRY(faiss_IDSelectorRange_new(&rhs_sel, 15, 35));
FaissIDSelectorAnd* sel = NULL;
FAISS_TRY(faiss_IDSelectorAnd_new(&sel, lhs_sel, rhs_sel));
FaissSearchParameters* params = NULL;
FAISS_TRY(faiss_SearchParameters_new(&params, sel));
FAISS_TRY(
faiss_Index_search_with_params(index, nq, xq, k, params, D, I));
printf("Searching w/ IDSelectorRange [20,40] AND [15,35] = [20,35]\n");
printf("I=\n");
for (int i = 0; i < 5; i++) {
for (int j = 0; j < k; j++)
printf("%5lld (d=%2.3f) ", I[i * k + j], D[i * k + j]);
printf("\n");
}
free(I);
free(D);
faiss_SearchParameters_free(params);
faiss_IDSelectorRange_free(lhs_sel);
faiss_IDSelectorRange_free(rhs_sel);
faiss_IDSelector_free(sel);
}

printf("Saving index to disk...\n");
FAISS_TRY(faiss_write_index_fname(index, "example.index"));

Expand Down
50 changes: 50 additions & 0 deletions c_api/impl/AuxIndexStructures_c.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,12 @@
using faiss::BufferList;
using faiss::DistanceComputer;
using faiss::IDSelector;
using faiss::IDSelectorAnd;
using faiss::IDSelectorBatch;
using faiss::IDSelectorNot;
using faiss::IDSelectorOr;
using faiss::IDSelectorRange;
using faiss::IDSelectorXOr;
using faiss::RangeQueryResult;
using faiss::RangeSearchPartialResult;
using faiss::RangeSearchResult;
Expand Down Expand Up @@ -116,6 +120,52 @@ int faiss_IDSelectorBatch_new(
CATCH_AND_HANDLE
}

int faiss_IDSelectorNot_new(
FaissIDSelectorNot** p_sel,
const FaissIDSelector* sel) {
try {
*p_sel = reinterpret_cast<FaissIDSelectorNot*>(
new IDSelectorNot(reinterpret_cast<const IDSelector*>(sel)));
}
CATCH_AND_HANDLE
}

int faiss_IDSelectorAnd_new(
FaissIDSelectorAnd** p_sel,
const FaissIDSelector* lhs_sel,
const FaissIDSelector* rhs_sel) {
try {
*p_sel = reinterpret_cast<FaissIDSelectorAnd*>(new IDSelectorAnd(
reinterpret_cast<const IDSelector*>(lhs_sel),
reinterpret_cast<const IDSelector*>(rhs_sel)));
}
CATCH_AND_HANDLE
}

int faiss_IDSelectorOr_new(
FaissIDSelectorOr** p_sel,
const FaissIDSelector* lhs_sel,
const FaissIDSelector* rhs_sel) {
try {
*p_sel = reinterpret_cast<FaissIDSelectorOr*>(new IDSelectorOr(
reinterpret_cast<const IDSelector*>(lhs_sel),
reinterpret_cast<const IDSelector*>(rhs_sel)));
}
CATCH_AND_HANDLE
}

int faiss_IDSelectorXOr_new(
FaissIDSelectorXOr** p_sel,
const FaissIDSelector* lhs_sel,
const FaissIDSelector* rhs_sel) {
try {
*p_sel = reinterpret_cast<FaissIDSelectorXOr*>(new IDSelectorXOr(
reinterpret_cast<const IDSelector*>(lhs_sel),
reinterpret_cast<const IDSelector*>(rhs_sel)));
}
CATCH_AND_HANDLE
}

// Below are structures used only by Index implementations

DEFINE_DESTRUCTOR(BufferList)
Expand Down
23 changes: 23 additions & 0 deletions c_api/impl/AuxIndexStructures_c.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,29 @@ int faiss_IDSelectorBatch_new(
size_t n,
const idx_t* indices);

FAISS_DECLARE_CLASS(IDSelectorNot)
int faiss_IDSelectorNot_new(
FaissIDSelectorNot** p_sel,
const FaissIDSelector* sel);

FAISS_DECLARE_CLASS(IDSelectorAnd)
int faiss_IDSelectorAnd_new(
FaissIDSelectorAnd** p_sel,
const FaissIDSelector* lhs_sel,
const FaissIDSelector* rhs_sel);

FAISS_DECLARE_CLASS(IDSelectorOr)
int faiss_IDSelectorOr_new(
FaissIDSelectorOr** p_sel,
const FaissIDSelector* lhs_sel,
const FaissIDSelector* rhs_sel);

FAISS_DECLARE_CLASS(IDSelectorXOr)
int faiss_IDSelectorXOr_new(
FaissIDSelectorXOr** p_sel,
const FaissIDSelector* lhs_sel,
const FaissIDSelector* rhs_sel);

// Below are structures used only by Index implementations

/** List of temporary buffers used to store results before they are
Expand Down

0 comments on commit aae6ff5

Please sign in to comment.