Skip to content

Commit

Permalink
Support reset & dispose
Browse files Browse the repository at this point in the history
  • Loading branch information
@asilvas
asilvas committed Sep 21, 2023
1 parent 05488b7 commit 7b60682
Show file tree
Hide file tree
Showing 3 changed files with 159 additions and 12 deletions.
23 changes: 15 additions & 8 deletions lib/index.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,16 +10,16 @@ export interface SearchResult {
export enum MetricType {
METRIC_INNER_PRODUCT = 0, ///< maximum inner product search
METRIC_L2 = 1, ///< squared L2 search
METRIC_L1, ///< L1 (aka cityblock)
METRIC_Linf, ///< infinity distance
METRIC_Lp, ///< L_p distance, p is given by a faiss::Index
METRIC_L1 = 2, ///< L1 (aka cityblock)
METRIC_Linf = 3, ///< infinity distance
METRIC_Lp = 4, ///< L_p distance, p is given by a faiss::Index
/// metric_arg

/// some additional metrics defined in scipy.spatial.distance
METRIC_Canberra = 20,
METRIC_BrayCurtis,
METRIC_JensenShannon,
METRIC_Jaccard, ///< defined as: sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i))
METRIC_BrayCurtis = 21,
METRIC_JensenShannon = 22,
METRIC_Jaccard = 23, ///< defined as: sum_i(min(a_i, b_i)) / sum_i(max(a_i, b_i))
///< where a_i, b_i > 0
}

Expand Down Expand Up @@ -106,8 +106,15 @@ export class Index {
* @param {number[]} ids IDs to read.
* @return {number} number of IDs removed.
*/
removeIds(ids: number[]): number

removeIds(ids: number[]): number;
/**
* Reset the index, resulting in a ntotal of 0.
*/
reset(): void;
/**
* Free all resources associated with the index. Further calls to the index will throw.
*/
dispose(): void;
}

/**
Expand Down
97 changes: 94 additions & 3 deletions src/faiss.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,13 +142,25 @@ class IndexBase : public Napi::ObjectWrap<T>

Napi::Value isTrained(const Napi::CallbackInfo &info)
{
return Napi::Boolean::New(info.Env(), index_->is_trained);
Napi::Env env = info.Env();
if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}

return Napi::Boolean::New(env, index_->is_trained);
}

Napi::Value add(const Napi::CallbackInfo &info)
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 1)
{
Napi::Error::New(env, "Expected 1 argument, but got " + std::to_string(info.Length()) + ".")
Expand Down Expand Up @@ -190,10 +202,45 @@ class IndexBase : public Napi::ObjectWrap<T>
return env.Undefined();
}

Napi::Value reset(const Napi::CallbackInfo &info)
{
Napi::Env env = info.Env();
if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}

index_->reset();

return env.Undefined();
}

Napi::Value dispose(const Napi::CallbackInfo &info)
{
Napi::Env env = info.Env();
if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}

auto idx = index_.release();
delete idx;
index_ = nullptr;

return env.Undefined();
}

Napi::Value train(const Napi::CallbackInfo &info)
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 1)
{
Napi::Error::New(env, "Expected 1 argument, but got " + std::to_string(info.Length()) + ".")
Expand Down Expand Up @@ -239,6 +286,11 @@ class IndexBase : public Napi::ObjectWrap<T>
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 2)
{
Napi::Error::New(env, "Expected 2 arguments, but got " + std::to_string(info.Length()) + ".")
Expand Down Expand Up @@ -314,18 +366,36 @@ class IndexBase : public Napi::ObjectWrap<T>

Napi::Value ntotal(const Napi::CallbackInfo &info)
{
return Napi::Number::New(info.Env(), index_->ntotal);
Napi::Env env = info.Env();
if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}

return Napi::Number::New(env, index_->ntotal);
}

Napi::Value getDimension(const Napi::CallbackInfo &info)
{
return Napi::Number::New(info.Env(), index_->d);
Napi::Env env = info.Env();
if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
return Napi::Number::New(env, index_->d);
}

Napi::Value write(const Napi::CallbackInfo &info)
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 1)
{
Napi::Error::New(env, "Expected 1 argument, but got " + std::to_string(info.Length()) + ".")
Expand All @@ -349,6 +419,11 @@ class IndexBase : public Napi::ObjectWrap<T>
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 1)
{
Napi::Error::New(env, "Expected 1 argument, but got " + std::to_string(info.Length()) + ".")
Expand Down Expand Up @@ -387,6 +462,11 @@ class IndexBase : public Napi::ObjectWrap<T>
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 1)
{
Napi::Error::New(env, "Expected 1 argument, but got " + std::to_string(info.Length()) + ".")
Expand Down Expand Up @@ -425,6 +505,11 @@ class IndexBase : public Napi::ObjectWrap<T>
{
Napi::Env env = info.Env();

if (!index_)
{
Napi::Error::New(env, "Index has been disposed").ThrowAsJavaScriptException();
return env.Undefined();
}
if (info.Length() != 0)
{
Napi::Error::New(env, "Expected 0 arguments, but got " + std::to_string(info.Length()) + ".")
Expand Down Expand Up @@ -461,6 +546,8 @@ class Index : public IndexBase<Index, faiss::IndexFlatL2>
InstanceMethod("getDimension", &Index::getDimension),
InstanceMethod("isTrained", &Index::isTrained),
InstanceMethod("add", &Index::add),
InstanceMethod("reset", &Index::reset),
InstanceMethod("dispose", &Index::dispose),
InstanceMethod("train", &Index::train),
InstanceMethod("search", &Index::search),
InstanceMethod("write", &Index::write),
Expand Down Expand Up @@ -496,6 +583,8 @@ class IndexFlatL2 : public IndexBase<IndexFlatL2, faiss::IndexFlatL2>
InstanceMethod("getDimension", &IndexFlatL2::getDimension),
InstanceMethod("isTrained", &IndexFlatL2::isTrained),
InstanceMethod("add", &IndexFlatL2::add),
InstanceMethod("reset", &IndexFlatL2::reset),
InstanceMethod("dispose", &IndexFlatL2::dispose),
InstanceMethod("train", &IndexFlatL2::train),
InstanceMethod("search", &IndexFlatL2::search),
InstanceMethod("write", &IndexFlatL2::write),
Expand Down Expand Up @@ -530,6 +619,8 @@ class IndexFlatIP : public IndexBase<IndexFlatIP, faiss::IndexFlatIP>
InstanceMethod("getDimension", &IndexFlatIP::getDimension),
InstanceMethod("isTrained", &IndexFlatIP::isTrained),
InstanceMethod("add", &IndexFlatIP::add),
InstanceMethod("reset", &IndexFlatIP::reset),
InstanceMethod("dispose", &IndexFlatIP::dispose),
InstanceMethod("train", &IndexFlatIP::train),
InstanceMethod("search", &IndexFlatIP::search),
InstanceMethod("write", &IndexFlatIP::write),
Expand Down
51 changes: 50 additions & 1 deletion test/Index.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@ describe('Index', () => {
});
});


describe('#toBuffer', () => {
it('new index is same size as old', () => {
const index = Index.fromFactory(2, 'Flat');
Expand All @@ -43,4 +42,54 @@ describe('Index', () => {
expect(index.ntotal()).toBe(newIndex.ntotal());
});
});

describe('#reset', () => {
let index;

beforeEach(() => {
index = Index.fromFactory(2, 'Flat');
index.add([1, 0, 0, 1]);
});

it('reset the index', () => {
expect(index.ntotal()).toBe(2);
index.reset();
expect(index.ntotal()).toBe(0);
});

it('reset the index and add new elements', () => {
expect(index.ntotal()).toBe(2);
index.reset();
expect(index.ntotal()).toBe(0);

index.add([1, 0]);
index.add([1, 2]);
expect(index.ntotal()).toBe(2);
});
});

describe('#dispose', () => {
let index;

beforeEach(() => {
index = Index.fromFactory(2, 'Flat');
index.add([1, 0, 0, 1]);
});

it('disposing an index does not throw', () => {
index.dispose();
});

it('disposing twice will throw', () => {
index.dispose();

expect(() => index.dispose()).toThrow('Index has been disposed');
});

it('invoking a function after dispose will throw', () => {
index.dispose();

expect(() => index.ntotal()).toThrow('Index has been disposed');
});
});
});

0 comments on commit 7b60682

Please sign in to comment.