Skip to content

Commit

Permalink
Introducing a new reader to read index using a pointer (#8)
Browse files Browse the repository at this point in the history
* draft implementation bufIOWriter over uint8_t* buf

* code cleanup

* implementation for writing using a ptr instead of vector - BufIOWriter

* revert BufIOWriter implementation

* code comments
  • Loading branch information
Thejas-bhat authored Dec 1, 2023
1 parent 4dd26f8 commit a623ec6
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 4 deletions.
8 changes: 4 additions & 4 deletions c_api/index_io_c_ex.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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<const Index*>(idx), &writer);
unsigned char* tempBuf = (unsigned char*)malloc((writer.data.size()) * sizeof(uint8_t));
std::copy(writer.data.begin(), writer.data.end(), tempBuf);
Expand All @@ -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<FaissIndex*>(index);
}
Expand Down
29 changes: 29 additions & 0 deletions faiss/impl/io.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
***********************************************************************/
Expand Down
8 changes: 8 additions & 0 deletions faiss/impl/io.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down

0 comments on commit a623ec6

Please sign in to comment.