additional index APIs

c_api/IndexIVF_c_ex.cpp

@@ -0,0 +1,22 @@
+/**
+ * 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.h"
+#include <faiss/IndexIVF.h>
+#include "macros_impl.h"
+
+using faiss::IndexIVF;
+
+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
+}

c_api/IndexIVF_c_ex.h

@@ -0,0 +1,30 @@
+/**
+ * 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 -*-
+
+#ifndef FAISS_INDEX_IVF_EX_C_H
+#define FAISS_INDEX_IVF_EX_C_H
+
+#include "Clustering_c.h"
+#include "Index_c.h"
+#include "faiss_c.h"
+#include "IndexIVF_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+int faiss_IndexIVF_set_direct_map(
+        FaissIndexIVF* index,
+        int direct_map_type);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

c_api/Index_c_ex.cpp

@@ -0,0 +1,33 @@
+/**
+ * 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 "Index_c.h"
+#include <faiss/Index.h>
+#include "macros_impl.h"
+
+extern "C" {
+
+int faiss_Index_reconstruct_batch(const FaissIndex* index, idx_t n, const idx_t* keys, float* recons) {
+    try {
+        reinterpret_cast<const faiss::Index*>(index)->reconstruct_batch(n, keys, recons);
+    }
+    CATCH_AND_HANDLE
+}
+
+
+int faiss_Index_merge_from(FaissIndex* index, FaissIndex* other, const idx_t add_id) {
+    try {
+        reinterpret_cast<faiss::Index*>(index)->merge_from(
+                *reinterpret_cast<faiss::Index*>(other), add_id);
+    }
+    CATCH_AND_HANDLE
+}
+
+}

c_api/Index_c_ex.h

@@ -0,0 +1,32 @@
+/**
+ * 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 2004-present Facebook. All Rights Reserved
+// -*- c -*-
+
+#ifndef FAISS_INDEX_EX_C_H
+#define FAISS_INDEX_EX_C_H
+
+#include <stddef.h>
+#include "faiss_c.h"
+#include "Index_c.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+int faiss_Index_reconstruct_batch(const FaissIndex* index, idx_t n,
+        const idx_t* keys, float* recons);
+
+
+int faiss_Index_merge_from(FaissIndex* index, FaissIndex* other, idx_t add_id);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

c_api/index_io_c.cpp

@@ -11,6 +11,7 @@
 
 #include "index_io_c.h"
 #include <faiss/index_io.h>
+#include <faiss/impl/io.h>
 #include "macros_impl.h"
 
 using faiss::Index;

c_api/index_io_c_ex.cpp

@@ -22,24 +22,19 @@ int faiss_write_index_buf(const FaissIndex* idx, int* size, unsigned char** buf)
         faiss::VectorIOWriter writer;
 
         faiss::write_index(reinterpret_cast<const Index*>(idx), &writer);
-        unsigned char* tempBuf = (unsigned char*)malloc((writer.data.size() + 1) * sizeof(uint8_t));
+        unsigned char* tempBuf = (unsigned char*)malloc((writer.data.size()) * sizeof(uint8_t));
         std::copy(writer.data.begin(), writer.data.end(), tempBuf);
-        tempBuf[writer.data.size()] = 0;
-
-        // return the serialized index content to the passed buf
-        // furthermore, return its size (perhaps useful for memory management)
         *buf = tempBuf;
         *size = writer.data.size();
+        writer.data.clear();
     }
     CATCH_AND_HANDLE
 }
 
 int faiss_read_index_buf(const unsigned char* buf, int size, int io_flags, FaissIndex** p_out) {
     try {
         faiss::VectorIOReader reader;
-        reader.data.resize(size);
-        memcpy(reader.data.data(), buf, size);
-
+        reader.data.assign(buf, buf + size);
         auto index = faiss::read_index(&reader, io_flags);
         *p_out = reinterpret_cast<FaissIndex*>(index);
     }

faiss/impl/index_read.cpp

@@ -194,20 +194,18 @@ InvertedLists* read_InvertedLists(IOReader* f, int io_flags) {
                 " WARN! inverted lists not stored with IVF object\n");
         return nullptr;
     } else if (h == fourcc("ilar") && !(io_flags & IO_FLAG_SKIP_IVF_DATA)) {
-        auto ails = new ArrayInvertedLists(0, 0);
-        READ1(ails->nlist);
-        READ1(ails->code_size);
-        ails->ids.resize(ails->nlist);
-        ails->codes.resize(ails->nlist);
+        size_t nlist, code_size;
+        READ1(nlist);
+        READ1(code_size);
+        auto ails = new ArrayInvertedLists(nlist, code_size);
         std::vector<size_t> sizes(ails->nlist);
         read_ArrayInvertedLists_sizes(f, sizes);
+
         for (size_t i = 0; i < ails->nlist; i++) {
-            ails->ids[i].resize(sizes[i]);
-            ails->codes[i].resize(sizes[i] * ails->code_size);
-        }
-        for (size_t i = 0; i < ails->nlist; i++) {
-            size_t n = ails->ids[i].size();
+            size_t n = sizes[i];
             if (n > 0) {
+                ails->ids[i].resize(n);
+                ails->codes[i].resize(n * ails->code_size);
                 READANDCHECK(ails->codes[i].data(), n * ails->code_size);
                 READANDCHECK(ails->ids[i].data(), n);
             }