Add sharding convenience function for IVF indexes (#4150)

Summary:

Creates a sharding convenience function for IVF indexes.
- The centroids on the quantizer are sharded based on the given sharding function.
- The output is written to files based on the template filename generator param.
- The default sharding function is simply the ith vector mod the total shard count.

This would called by Laser here: https://www.internalfb.com/code/fbsource/[ce1f2e028e79]/fbcode/fblearner/flow/projects/laser/laser_sim_search/knn_trainer.py?lines=295-296. This convenience function will do the file writing, and return the created file names.

There's a few key required changes in FAISS:
1. Allow `std::vector<std::string>` to be used. Updates swigfaiss.swig and array_conversions.py to accommodate. These have to be numpy dtype of `object` instead of the more correct `unicode`, because unicode dtype is fixed length. I couldn't figure out how to create a numpy array with each of the output file names where they have different dtypes. (Say the file names are like file1, file11, file111. The dtype would need to be U5, U6, U7 respectively, as the dtype for unicode contains the length). I tried structured arrays : this does not work either, as numpy makes it into a matrix instead: the `file1 file11 file111` example with explicit setting of U5, U6, U7 turns into `[[file1 file1 file1], [file1 file11 file11], [file1 file11 file111]]`, which we do not want. If someone knows the right syntax, please yell at me
2. Create Python callback for sharding and template filename. Users of this function would inherit from the FilenameTemplateGenerator or ShardingFunction to pass to `shard_ivf_index_centroids`. See the other examples in python_callbacks.cpp. This is required because Python functions cannot be passed through SWIG to C++ (i.e. no std::function or function pointers), so we have to use this approach. This approach allows it to be called from both C++ and Python.

Differential Revision: D68534991

Author: Michael Norris

faiss/CMakeLists.txt

@@ -91,6 +91,7 @@ set(FAISS_SRC
   utils/quantize_lut.cpp
   utils/random.cpp
   utils/sorting.cpp
+  utils/sharding.cpp
   utils/utils.cpp
   utils/distances_fused/avx512.cpp
   utils/distances_fused/distances_fused.cpp
@@ -207,6 +208,7 @@ set(FAISS_HEADERS
   utils/prefetch.h
   utils/quantize_lut.h
   utils/random.h
+  utils/sharding.h
   utils/sorting.h
   utils/simdlib.h
   utils/simdlib_avx2.h

faiss/python/__init__.py

@@ -170,7 +170,7 @@ def replacement_function(*args):
     add_ref_in_constructor(GpuIndexIVFPQ, 1)
     add_ref_in_constructor(GpuIndexIVFScalarQuantizer, 1)
 except NameError as e:
-    logger.info("Failed to load GPU Faiss: %s. Will not load constructor refs for GPU indexes." % e.args[0])
+    logger.info("Failed to load GPU Faiss: %s. Will not load constructor refs for GPU indexes. This is only an error if you're trying to use GPU Faiss." % e.args[0])
 
 add_ref_in_constructor(IndexIVFFlat, 0)
 add_ref_in_constructor(IndexIVFFlatDedup, 0)

faiss/python/array_conversions.py

@@ -97,6 +97,7 @@ def new_meth(cls, *args, **kwargs):
     'UInt16': 'uint16',
     'UInt32': 'uint32',
     'UInt64': 'uint64',
+    'String': 'str',
     **{k: v.lower() for k, v in deprecated_name_map.items()}
 }
 
@@ -107,10 +108,19 @@ def vector_to_array(v):
     if classname.startswith('AlignedTable'):
         return AlignedTable_to_array(v)
     assert classname.endswith('Vector')
-    dtype = np.dtype(vector_name_map[classname[:-6]])
-    a = np.empty(v.size(), dtype=dtype)
-    if v.size() > 0:
-        memcpy(swig_ptr(a), v.data(), a.nbytes)
+    vector_name = vector_name_map[classname[:-6]]
+    # TODO: Remove this hack with 'object' after upgrading to
+    # Numpy 2, which can support variable length strings.
+    if vector_name == 'str':
+        values = []
+        for i in range (0, v.size()):
+            values.append(v.at(i))
+        a = np.array(values, dtype='object')
+    else:
+        dtype = np.dtype(vector_name)
+        a = np.empty(v.size(), dtype=dtype)
+        if v.size() > 0:
+            memcpy(swig_ptr(a), v.data(), a.nbytes)
     return a
 
 

faiss/python/python_callbacks.cpp

@@ -134,3 +134,54 @@ PyCallbackIDSelector::~PyCallbackIDSelector() {
     PyThreadLock gil;
     Py_DECREF(callback);
 }
+
+/***********************************************************
+ * Callbacks for IVF index sharding
+ ***********************************************************/
+
+PyCallbackFilenameTemplateGenerator::PyCallbackFilenameTemplateGenerator(
+        PyObject* callback)
+        : callback(callback) {
+    PyThreadLock gil;
+    Py_INCREF(callback);
+}
+
+std::string PyCallbackFilenameTemplateGenerator::operator()() {
+    PyThreadLock gil;
+    PyObject* template_filename = PyObject_CallFunction(callback, NULL);
+    if (template_filename == nullptr) {
+        FAISS_THROW_MSG("propagate py error");
+    }
+    const char* cstr = PyUnicode_AsUTF8(template_filename);
+    if (cstr == NULL) {
+        // handle error or return empty string
+        return "";
+    }
+    std::string result(cstr);
+    return result;
+}
+
+PyCallbackFilenameTemplateGenerator::~PyCallbackFilenameTemplateGenerator() {
+    PyThreadLock gil;
+    Py_DECREF(callback);
+}
+
+PyCallbackShardingFunction::PyCallbackShardingFunction(PyObject* callback)
+        : callback(callback) {
+    PyThreadLock gil;
+    Py_INCREF(callback);
+}
+
+int64_t PyCallbackShardingFunction::operator()(int64_t i, int64_t shard_count) {
+    PyThreadLock gil;
+    PyObject* shard_id = PyObject_CallFunction(callback, "LL", i, shard_count);
+    if (shard_id == nullptr) {
+        FAISS_THROW_MSG("propagate py error");
+    }
+    return PyLong_AsLongLong(shard_id);
+}
+
+PyCallbackShardingFunction::~PyCallbackShardingFunction() {
+    PyThreadLock gil;
+    Py_DECREF(callback);
+}

faiss/python/python_callbacks.h

@@ -10,6 +10,7 @@
 #include <faiss/impl/IDSelector.h>
 #include <faiss/impl/io.h>
 #include <faiss/invlists/InvertedLists.h>
+#include <faiss/utils/sharding.h>
 #include "Python.h"
 
 //  all callbacks have to acquire the GIL on input
@@ -58,3 +59,27 @@ struct PyCallbackIDSelector : faiss::IDSelector {
 
     ~PyCallbackIDSelector() override;
 };
+
+/***********************************************************
+ * Callbacks for IVF index sharding
+ ***********************************************************/
+
+struct PyCallbackFilenameTemplateGenerator : faiss::FilenameTemplateGenerator {
+    PyObject* callback;
+
+    explicit PyCallbackFilenameTemplateGenerator(PyObject* callback);
+
+    std::string operator()();
+
+    ~PyCallbackFilenameTemplateGenerator();
+};
+
+struct PyCallbackShardingFunction : faiss::ShardingFunction {
+    PyObject* callback;
+
+    explicit PyCallbackShardingFunction(PyObject* callback);
+
+    int64_t operator()(int64_t i, int64_t shard_count);
+
+    ~PyCallbackShardingFunction();
+};

faiss/python/swigfaiss.swig

@@ -145,6 +145,7 @@ typedef uint64_t size_t;
 #include <faiss/IVFlib.h>
 #include <faiss/utils/utils.h>
 
+#include <faiss/utils/sharding.h>
 #include <faiss/utils/sorting.h>
 #include <faiss/utils/distances.h>
 #include <faiss/utils/extra_distances.h>
@@ -270,6 +271,8 @@ namespace std {
 %template(UInt32Vector) std::vector<uint32_t>;
 %template(UInt64Vector) std::vector<uint64_t>;
 
+%template(StringVector) std::vector<std::string>;
+
 %template(Float32VectorVector) std::vector<std::vector<float> >;
 %template(UInt8VectorVector) std::vector<std::vector<uint8_t> >;
 %template(Int32VectorVector) std::vector<std::vector<int32_t> >;
@@ -539,6 +542,17 @@ void gpu_sync_all_devices()
 %include  <faiss/invlists/BlockInvertedLists.h>
 %include  <faiss/invlists/DirectMap.h>
 %include  <faiss/IndexIVF.h>
+
+// Include after IndexIVF.
+%shared_ptr(faiss::FilenameTemplateGenerator);
+%shared_ptr(PyCallbackFilenameTemplateGenerator);
+%shared_ptr(faiss::DefaultFilenameTemplateGenerator);
+%shared_ptr(faiss::ShardingFunction);
+%shared_ptr(PyCallbackShardingFunction);
+%shared_ptr(faiss::DefaultShardingFunction);
+
+%include  <faiss/utils/sharding.h>
+
 // NOTE(hoss): SWIG (wrongly) believes the overloaded const version shadows the
 //   non-const one.
 %warnfilter(509) extract_index_ivf;
@@ -1163,6 +1177,9 @@ PyObject *swig_ptr (PyObject *a)
     if(PyArray_TYPE(ao) == NPY_INT32) {
         return SWIG_NewPointerObj(data, SWIGTYPE_p_int, 0);
     }
+    if(PyArray_TYPE(ao) == NPY_OBJECT) {
+        return SWIG_NewPointerObj(data, SWIGTYPE_p_char, 0);
+    }
     if(PyArray_TYPE(ao) == NPY_BOOL) {
         return SWIG_NewPointerObj(data, SWIGTYPE_p_bool, 0);
     }

faiss/utils/sharding.cpp

@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/clone_index.h>
+#include <faiss/index_io.h>
+#include <faiss/utils/sharding.h>
+#include <cstdio>
+
+namespace faiss {
+
+std::string DefaultFilenameTemplateGenerator::operator()() {
+    return "shard_%d.faissindex";
+}
+
+int64_t DefaultShardingFunction::operator()(int64_t i, int64_t shard_count) {
+    return i % shard_count;
+}
+
+std::vector<std::string> shard_ivf_index_centroids(
+        IndexIVF* index,
+        int64_t shard_count,
+        std::shared_ptr<FilenameTemplateGenerator> filename_template_generator,
+        std::shared_ptr<ShardingFunction> sharding_function) {
+    if (index->quantizer->ntotal == 0) {
+        return std::vector<std::string>();
+    }
+
+    if (filename_template_generator == nullptr) {
+        filename_template_generator =
+                std::make_shared<DefaultFilenameTemplateGenerator>();
+    }
+    if (sharding_function == nullptr) {
+        sharding_function = std::make_shared<DefaultShardingFunction>();
+    }
+
+    IndexIVF* sharded_indexes[shard_count];
+    for (int i = 0; i < shard_count; i++) {
+        sharded_indexes[i] = static_cast<IndexIVF*>(clone_index(index));
+        sharded_indexes[i]->quantizer->reset();
+    }
+
+    // assign centroids to each sharded Index based on sharding_function, and
+    // add them to the quantizer of each sharded index
+    std::vector<float> sharded_centroids[shard_count];
+    for (int i = 0; i < index->quantizer->ntotal; i++) {
+        int shard_id = (*sharding_function)(i, shard_count);
+        float reconstructed[index->quantizer->d];
+        index->quantizer->reconstruct(i, reconstructed);
+        sharded_centroids[shard_id].insert(
+                sharded_centroids[shard_id].end(),
+                &reconstructed[0],
+                &reconstructed[index->quantizer->d]);
+    }
+    for (int i = 0; i < shard_count; i++) {
+        sharded_indexes[i]->quantizer->add(
+                sharded_centroids[i].size() / index->quantizer->d,
+                sharded_centroids[i].data());
+    }
+
+    std::vector<std::string> result;
+    for (int i = 0; i < shard_count; i++) {
+        char fname[256];
+        std::string template_filename = (*filename_template_generator)();
+        snprintf(fname, 256, template_filename.c_str(), i);
+        result.emplace_back(fname);
+        write_index(sharded_indexes[i], fname);
+    }
+
+    for (int i = 0; i < shard_count; i++) {
+        delete sharded_indexes[i];
+    }
+
+    return result;
+}
+
+} // namespace faiss

faiss/utils/sharding.h

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// -*- c++ -*-
+
+/*
+ * Utilities related to sharding indexes.
+ */
+
+#ifndef FAISS_SHARDING_H
+#define FAISS_SHARDING_H
+
+#include <faiss/IndexIVF.h>
+
+namespace faiss {
+
+struct FilenameTemplateGenerator {
+    virtual std::string operator()() = 0;
+    virtual ~FilenameTemplateGenerator() {}
+};
+struct DefaultFilenameTemplateGenerator : FilenameTemplateGenerator {
+    std::string operator()() override;
+};
+
+struct ShardingFunction {
+    virtual int64_t operator()(int64_t i, int64_t shard_count) = 0;
+    virtual ~ShardingFunction() {}
+};
+struct DefaultShardingFunction : ShardingFunction {
+    int64_t operator()(int64_t i, int64_t shard_count) override;
+};
+
+/**
+ * Shards an IVF index centroids by the given sharding function, and writes
+ * the index to the path given by filename_generator. The centroids must already
+ * be added to the index quantizer.
+ *
+ * @param index                       The IVF index containing centroids to
+ *                                    shard.
+ * @param shard_count                 Number of shards.
+ * @param filename_template_generator Function that generates a filename
+ *                                    template of the output indexes.
+ * @param sharding_function           The function to shard by. The default is
+ *                                    ith vector mod shard_count.
+ * @return                            The list of output filenames.
+ */
+std::vector<std::string> shard_ivf_index_centroids(
+        IndexIVF* index,
+        int64_t shard_count = 20,
+        std::shared_ptr<FilenameTemplateGenerator> filename_template_generator =
+                nullptr,
+        std::shared_ptr<ShardingFunction> sharding_function = nullptr);
+} // namespace faiss
+
+#endif

tests/test_sharding.py

@@ -0,0 +1,62 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import faiss
+import unittest
+import numpy as np
+import os
+
+faiss.omp_set_num_threads(4)
+
+
+def custom_filename_template_generator_callback() -> str:
+    randomInt = np.random.randint(0, 1000000000)
+    return str("shard_%d.faissindex" + str(randomInt))
+
+
+def custom_sharding_callback(i, shard_count):
+    del i
+    del shard_count
+    return 3
+
+
+class TestSharding(unittest.TestCase):
+    """
+    Test the sharding.h from Python to ensure integration is working.
+    No correctness checks are done here.
+    """
+    def test_python_sharding(self):
+        d = 32
+        nb = 1000
+        xb = np.random.rand(nb, d).astype('float32')
+        nlist = 100
+        quantizer = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFFlat(quantizer, d, nlist)
+
+        index.quantizer.add(xb)
+
+        res = (
+            faiss.shard_ivf_index_centroids(
+                index,
+                10,
+                faiss.PyCallbackFilenameTemplateGenerator(
+                    custom_filename_template_generator_callback)
+            )
+        )
+        for file in faiss.vector_to_array(res):
+            os.remove(file)
+
+        res = (
+            faiss.shard_ivf_index_centroids(
+                index,
+                10,
+                faiss.PyCallbackFilenameTemplateGenerator(
+                    custom_filename_template_generator_callback),
+                faiss.PyCallbackShardingFunction(
+                    custom_sharding_callback)
+            )
+        )
+        for file in faiss.vector_to_array(res):
+            os.remove(file)