diff --git a/src/python/txtai/ann/dense/milvus.py b/src/python/txtai/ann/dense/milvus.py index 14179b23..dd2b2b7d 100644 --- a/src/python/txtai/ann/dense/milvus.py +++ b/src/python/txtai/ann/dense/milvus.py @@ -83,8 +83,9 @@ class Milvus(ANN): self.metadata() def delete(self, ids): - if ids: - self.collection.delete(ids) + # Check length, ids can be a NumPy array + if len(ids): + self.collection.delete([int(x) for x in ids]) def search(self, queries, limit): matches = self.collection.search(queries.tolist(), top_k=limit, metric_type="IP", anns_field="embedding") diff --git a/src/python/txtai/ann/dense/pgvector.py b/src/python/txtai/ann/dense/pgvector.py index ba79f142..d1c08126 100644 --- a/src/python/txtai/ann/dense/pgvector.py +++ b/src/python/txtai/ann/dense/pgvector.py @@ -71,7 +71,8 @@ class PGVector(ANN): self.metadata() def delete(self, ids): - self.database.execute(delete(self.table).where(self.table.c["indexid"].in_(ids))) + # Bind ids as int, NumPy integers don't match the indexid column + self.database.execute(delete(self.table).where(self.table.c["indexid"].in_([int(x) for x in ids]))) def search(self, queries, limit): results = [] diff --git a/src/python/txtai/ann/dense/sqlite.py b/src/python/txtai/ann/dense/sqlite.py index b890b292..1ced45be 100644 --- a/src/python/txtai/ann/dense/sqlite.py +++ b/src/python/txtai/ann/dense/sqlite.py @@ -55,7 +55,8 @@ class SQLite(ANN): self.metadata() def delete(self, ids): - self.database().executemany(self.deletesql(), [(x,) for x in ids]) + # Bind ids as int, sqlite3 binds NumPy integers as blobs that match the wrong rows + self.database().executemany(self.deletesql(), [(int(x),) for x in ids]) def search(self, queries, limit): results = [] diff --git a/src/python/txtai/ann/dense/zvec.py b/src/python/txtai/ann/dense/zvec.py index 06ac7fad..7d793b8e 100644 --- a/src/python/txtai/ann/dense/zvec.py +++ b/src/python/txtai/ann/dense/zvec.py @@ -80,7 +80,8 @@ class Zvec(ANN): self.metadata() def delete(self, ids): - if ids: + # Check length, ids can be a NumPy array + if len(ids): self.backend.delete([str(uid) for uid in ids]) def search(self, queries, limit): diff --git a/test/python/testann/testdense/base.py b/test/python/testann/testdense/base.py index cc581ffb..0a024e4d 100644 --- a/test/python/testann/testdense/base.py +++ b/test/python/testann/testdense/base.py @@ -128,6 +128,25 @@ class DenseTest(unittest.TestCase): return model + def deletenumpy(self, name, params=None): + """ + Test deleting ids passed as a NumPy array. Only the requested rows should be deleted. + + Args: + name: backend name + params: additional config parameters + """ + + model = self.backend(name, params, 10) + model.delete(np.array([3, 4])) + + # Generate query vector + query = np.random.rand(240).astype(np.float32) + self.normalize(query) + + self.assertEqual(model.count(), 8) + self.assertEqual(sorted(uid for uid, _ in model.search(np.array([query]), 10)[0]), [0, 1, 2, 5, 6, 7, 8, 9]) + def save(self, name, params=None): """ Test save/load. diff --git a/test/python/testann/testdense/testmilvus.py b/test/python/testann/testdense/testmilvus.py index 00a6b27f..ea2bdd7e 100644 --- a/test/python/testann/testdense/testmilvus.py +++ b/test/python/testann/testdense/testmilvus.py @@ -18,6 +18,7 @@ class TestMilvus(DenseTest): """ self.runTests("milvus") + self.deletenumpy("milvus") def testMilvusCustom(self): """ diff --git a/test/python/testann/testdense/testpgvector.py b/test/python/testann/testdense/testpgvector.py index 266a027e..b4cbaed5 100644 --- a/test/python/testann/testdense/testpgvector.py +++ b/test/python/testann/testdense/testpgvector.py @@ -70,5 +70,9 @@ class TestPGVector(DenseTest): ann.delete([0]) self.assertEqual(ann.count(), 1) + # Test delete with NumPy ids + ann.delete(np.array([1])) + self.assertEqual(ann.count(), 0) + # Close ANN ann.close() diff --git a/test/python/testann/testdense/testsqlite.py b/test/python/testann/testdense/testsqlite.py index 87c99fd4..756ee869 100644 --- a/test/python/testann/testdense/testsqlite.py +++ b/test/python/testann/testdense/testsqlite.py @@ -26,6 +26,7 @@ class TestSQLite(DenseTest): """ self.runTests("sqlite") + self.deletenumpy("sqlite") @unittest.skipIf(platform.system() == "Darwin", "SQLite extensions not supported on macOS") def testSQLiteBinaryScores(self): diff --git a/test/python/testann/testdense/testzvec.py b/test/python/testann/testdense/testzvec.py index 6332cdb8..43ad3c21 100644 --- a/test/python/testann/testdense/testzvec.py +++ b/test/python/testann/testdense/testzvec.py @@ -18,6 +18,7 @@ class TestZvec(DenseTest): """ self.runTests("zvec") + self.deletenumpy("zvec") def testZvecCustom(self): """