[py-tx] Implementation of IVF faiss indices in PDQ

python-threatexchange/threatexchange/signal_type/pdq/pdq_index2.py

@@ -8,7 +8,6 @@
 import faiss
 import numpy as np
 
-
 from threatexchange.signal_type.index import (
     IndexMatchUntyped,
     SignalSimilarityInfoWithIntDistance,
@@ -33,6 +32,29 @@ class PDQIndex2(SignalTypeIndex[IndexT]):
     Purpose of this class: to replace the original index in pytx 2.0
     """
 
+    IVF_THRESHOLD = 1000
+
+    @classmethod
+    def build(
+        cls: t.Type["PDQIndex2"], entries: t.Iterable[t.Tuple[str, IndexT]]
+    ) -> "PDQIndex2":
+        """
+        Build an index from a set of entries.
+        Selects between flat and IVF index based on number of entries.
+        """
+        entries_list = list(entries)
+
+        index = faiss.IndexFlatL2(BITS_IN_PDQ)
+        if len(entries_list) >= cls.IVF_THRESHOLD:
+            nlist = len(entries_list) // 2
+            index = faiss.IndexIVFFlat(index, BITS_IN_PDQ, nlist)
+            vectors = convert_pdq_strings_to_ndarray([h for h, _ in entries_list])
+            index.train(vectors)
+        else:
+            index = faiss.IndexFlatL2(BITS_IN_PDQ)
+
+        return cls(index=index, entries=entries_list)
+
     def __init__(
         self,
         index: t.Optional[faiss.Index] = None,

python-threatexchange/threatexchange/signal_type/tests/test_pdq_index2.py

@@ -178,3 +178,31 @@ def test_one_entry_sample_index():
 
     results = index.query(unmatching_test_hash)
     assert len(results) == 0
+
+
+def test_flat_index_selection():
+    """Test flat index selection for small datasets"""
+    get_random_hashes = _get_hash_generator()
+    base_hashes = get_random_hashes(100)
+    index = PDQIndex2.build([(h, i) for i, h in enumerate(base_hashes)])
+    assert isinstance(index._index.faiss_index, faiss.IndexFlatL2)
+
+
+def test_ivf_index_selection():
+    """Test IVF index selection for large datasets"""
+    get_random_hashes = _get_hash_generator()
+    base_hashes = get_random_hashes(2000)
+    index = PDQIndex2.build([(h, i) for i, h in enumerate(base_hashes)])
+    assert isinstance(index._index.faiss_index, faiss.IndexIVFFlat)
+
+
+def test_build_method():
+    """Test build method creates valid index"""
+    get_random_hashes = _get_hash_generator()
+    base_hashes = get_random_hashes(500)
+    index = PDQIndex2.build([(h, i) for i, h in enumerate(base_hashes)])
+
+    query_hash = base_hashes[0]
+    results = index.query(query_hash)
+    assert len(results) >= 1
+    assert any(r.metadata == 0 for r in results)