def query(self, x):
n, d = x.shape
assert d == self.rsearch.dimension
sims = None
if self.rsearch.var_type == rs.FLOAT32:
sims = np.empty((n, self.rsearch.topk), dtype=np.float32)
elif self.rsearch.var_type == rs.INT8:
sims = np.empty((n, self.rsearch.topk), dtype=np.int32)
uids = np.empty((n, self.rsearch.topk), dtype=np.int32)
self.rsearch.query(rs.swig_ptr(x), n, rs.swig_ptr(sims),
rs.swig_ptr(uids))
return sims, uids
def query_by_uids(self, uids):
n = uids.shape[0]
x = rs.AreaTimeVector()
nn = rs.int2size_t(n)
x.resize(n)
self.simple_index_areatime.query_by_uids(rs.swig_ptr(uids), n,
x.data())
return x
def query_by_uids(self, uids):
n, d = uids.shape
assert d == self.rsearch.dimension
res = None
if self.rsearch.var_type == rs.FLOAT32:
res = np.empty((n, d), dtype=np.float32)
elif self.rsearch.var_type == rs.INT8:
res = np.empty((n, d), dtype=np.int8)
self.rsearch.query(self, rs.swig_ptr(res), n)
return res
def query_with_uids(self, x, uids):
s = uids.shape[0]
idx = rs.get_int_p_v()
res = rs.get_int_v()
self.simple_index_areatime.query_with_uids(x.data(), x.size(),
rs.swig_ptr(uids), s, idx,
res)
n = rs.get_int(res)
a = intpp2array(idx, n)
return a
def change_by_uids(self, x, uids):
n, d = x.shape
assert d == self.rsearch.dimension
assert uids.shape == (n, ), 'not same nb of vectors as ids'
return self.rsearch.add_with_uids(rs.swig_ptr(x), rs.swig_ptr(uids), n)
def add(self, x):
n, d = x.shape
assert d == self.rsearch.dimension
return self.rsearch.add(rs.swig_ptr(x), n)
def intpp2array(pp, n):
p = rs.get_int_p(pp)
a = np.empty(n, dtype=np.int32)
rs.memcpy(rs.swig_ptr(a), p, a.nbytes)
return a
def remove_by_uids(self, uids):
n = uids.shape[0]
return self.rsearch.remove_by_uids(rs.swig_ptr(uids), n)
def remove_by_uids(self, uids):
n = uids.shape[0]
return self.simple_index_areatime.remove_by_uids(rs.swig_ptr(uids), n)
def change_by_uids(self, x, uids):
n = uids.shape[0]
assert x.size() == n
return self.simple_index_areatime.change_by_uids(
x.data(), rs.swig_ptr(uids), n)
