diff --git a/c_api/index_io_c_ex.cpp b/c_api/index_io_c_ex.cpp index 03525c0a37..22fd748e9d 100644 --- a/c_api/index_io_c_ex.cpp +++ b/c_api/index_io_c_ex.cpp @@ -20,7 +20,6 @@ using faiss::IndexBinary; int faiss_write_index_buf(const FaissIndex* idx, int* size, unsigned char** buf) { try { faiss::VectorIOWriter writer; - faiss::write_index(reinterpret_cast(idx), &writer); unsigned char* tempBuf = (unsigned char*)malloc((writer.data.size()) * sizeof(uint8_t)); std::copy(writer.data.begin(), writer.data.end(), tempBuf); @@ -31,10 +30,11 @@ int faiss_write_index_buf(const FaissIndex* idx, int* size, unsigned char** buf) CATCH_AND_HANDLE } -int faiss_read_index_buf(const unsigned char* buf, int size, int io_flags, FaissIndex** p_out) { +int faiss_read_index_buf(const uint8_t* buf, int size, int io_flags, FaissIndex** p_out) { try { - faiss::VectorIOReader reader; - reader.data.assign(buf, buf + size); + faiss::BufIOReader reader; + reader.buf = buf; + reader.buf_size = size; auto index = faiss::read_index(&reader, io_flags); *p_out = reinterpret_cast(index); } diff --git a/faiss/impl/io.cpp b/faiss/impl/io.cpp index 3837525138..9d29b0b807 100644 --- a/faiss/impl/io.cpp +++ b/faiss/impl/io.cpp @@ -55,6 +55,35 @@ size_t VectorIOReader::operator()(void* ptr, size_t size, size_t nitems) { return nitems; } +/*********************************************************************** + * IO Buffer + ***********************************************************************/ + +size_t BufIOReader::operator()(void* ptr, size_t size, size_t nitems) { + // if the read pointer has passed the buffer size, exit out since we've + // read the complete index + if (rp >= buf_size) + return 0; + + // check how many "items" of a particular size are to be read from the buffer + // the size of the item depends on the datatype of field being populated. + size_t nremain = (buf_size - rp) / size; + if (nremain < nitems) // we don't have enough items to be read, in which case + nitems = nremain; // read all the remaining ones + if (size * nitems > 0) { + // finally memcpy the data from buffer to the field of the index being + // populated and increment the read pointer. + memcpy(ptr, &buf[rp], size * nitems); + rp += size * nitems; + } + + return nitems; +} + +BufIOReader::~BufIOReader() { + buf = NULL; +} + /*********************************************************************** * IO File ***********************************************************************/ diff --git a/faiss/impl/io.h b/faiss/impl/io.h index 8d0605a5a6..bcc97f5ab9 100644 --- a/faiss/impl/io.h +++ b/faiss/impl/io.h @@ -61,6 +61,14 @@ struct VectorIOWriter : IOWriter { size_t operator()(const void* ptr, size_t size, size_t nitems) override; }; +struct BufIOReader : IOReader { + const uint8_t* buf; + size_t rp = 0; + size_t buf_size; + size_t operator()(void* ptr, size_t size, size_t nitems) override; + ~BufIOReader() override; +}; + struct FileIOReader : IOReader { FILE* f = nullptr; bool need_close = false;