Merge pull request #1303 from kartsan03/fix/ann-delete-numpy-ids

Delete the requested rows when ANN ids are a NumPy array

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")

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 = []

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 = []

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):

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.

test/python/testann/testdense/testmilvus.py

@@ -18,6 +18,7 @@ class TestMilvus(DenseTest):
         """
 
         self.runTests("milvus")
+        self.deletenumpy("milvus")
 
     def testMilvusCustom(self):
         """

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()

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):

test/python/testann/testdense/testzvec.py

@@ -18,6 +18,7 @@ class TestZvec(DenseTest):
         """
 
         self.runTests("zvec")
+        self.deletenumpy("zvec")
 
     def testZvecCustom(self):
         """