faiss/c_api/IndexIVF_c_ex.cpp at 61dd0bdf546485598eb9cd31d0370ef050f9df75 · deepkaran/faiss · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
/**
 * 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.
 */

// Copyright 2023-present Facebook. All Rights Reserved.
// -*- c++ -*-
#include "IndexIVF_c_ex.h"
#include <faiss/IndexIVF.h>
#include "macros_impl.h"

using faiss::IndexIVF;
using faiss::SearchParameters;
using faiss::SearchParametersIVF;

int faiss_IndexIVF_set_direct_map(FaissIndexIVF* index, int direct_map_type) {
    try {
        reinterpret_cast<IndexIVF*>(index)->set_direct_map_type(
                static_cast<faiss::DirectMap::Type>(direct_map_type));
    }
    CATCH_AND_HANDLE
}

int faiss_SearchParametersIVF_new_with_sel(
        FaissSearchParametersIVF** p_sp,
        FaissIDSelector* sel) {
    try {
        SearchParametersIVF* sp = new SearchParametersIVF;
        sp->sel = reinterpret_cast<faiss::IDSelector*>(sel);
        *p_sp = reinterpret_cast<FaissSearchParametersIVF*>(sp);
    }
    CATCH_AND_HANDLE
}

int faiss_Search_closest_eligible_centroids(
        FaissIndex* index,
        idx_t n,
        const float* query,
        idx_t k,
        float* centroid_distances,
        idx_t* centroid_ids,
        const FaissSearchParameters* params) {
    try {
        faiss::IndexIVF* index_ivf = reinterpret_cast<IndexIVF*>(index);
        assert(index_ivf);

        index_ivf->quantizer->search(n, query, k, centroid_distances, centroid_ids,
            reinterpret_cast<const faiss::SearchParameters*>(params));
    }
    CATCH_AND_HANDLE
}

int faiss_get_lists_for_keys(
        FaissIndexIVF* index,
        idx_t* keys,
        size_t n_keys,
        idx_t* lists) {
    try {
        reinterpret_cast<IndexIVF*>(index)->get_lists_for_keys(keys, n_keys, lists);
    }
    CATCH_AND_HANDLE
}

int faiss_IndexIVF_search_preassigned_with_params(
        const FaissIndexIVF* index,
        idx_t n,
        const float* x,
        idx_t k,
        const idx_t* assign,
        const float* centroid_dis,
        float* distances,
        idx_t* labels,
        int store_pairs,
        const FaissSearchParametersIVF* params) {
    try {
        reinterpret_cast<const IndexIVF*>(index)->search_preassigned(
                n, x, k, assign, centroid_dis, distances, labels, store_pairs,
                reinterpret_cast<const faiss::SearchParametersIVF*>(params));
    }
    CATCH_AND_HANDLE
}

int faiss_IndexIVF_compute_distance_to_codes_for_list(
        FaissIndexIVF* index,
        idx_t list_no,
        const float* x,
        idx_t n,
        const uint8_t* codes,
        float* dists,
        float* dist_table) {
    try {
        reinterpret_cast<IndexIVF*>(index)->compute_distance_to_codes_for_list(
               list_no, x, n, codes, dists, dist_table);
        return 0;
    }
    CATCH_AND_HANDLE
}

int faiss_IndexIVF_compute_distance_table(
        FaissIndexIVF* index,
        const float* x,
        float* dist_table) {
    try {
        reinterpret_cast<IndexIVF*>(index)->compute_distance_table(
               x, dist_table);
        return 0;
    }
    CATCH_AND_HANDLE
}