[pytx] Implement a new cleaner PDQ index solution #1698

haianhng31 · 2024-11-14T02:53:55Z

Summary

Resolve issue #1613 .
This PR introduces a new PDQIndex2 class for managing and querying a PDQ hash-based index using FAISS. Key changes include:

_PDQHashIndex: A wrapper around the FAISS index to handle adding hashes and performing searches using the FAISS library.
PDQIndex2: A class for managing a PDQ index, which includes:
- The ability to add PDQ hashes and associate them with entries.
- A query method for searching the index for matching entries based on a query hash.
- A serialize and deserialize method for persisting and loading the index from binary streams using pickle.

Test Plan

I have included several test cases for this, currently all in one file.

For the bug in issue #1318, I have created 2 test cases relating to it:

Test how the index handles when it only has one entry. -> Return correctly 1 hash
Test exact matches in sample set -> Doesn't return all hashes

Also ran test_pdq_index with 1,000 banked & 10,000 queries in the unit test. The test passed but it took 4 minutes to run.
And if i tried with 10,000 banked & 100,000 queries, the test took forever.

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

Dcallies

Looking good so far!

Toplevel comments:

I am not convinced yet that this indexing solution is correct yet. To help convince me, you can write a unittest that generates two large sets of random hashes (PDQSignal.get_random_signal()) b and q (e.g. 1,000 and 10,000), uses the brute force method to check exactly how many of hashes in b should match the hashes in q, then loads the b hashes into your indexing solution, and then queries all of the q hashes and makes sure you exactly get the expected results.
Many of your existing tests are tautological or should check to make sure the matching hashes are actually the ones that should be matching. You can likely replace many of your existing tests with variations on 2 above - I put some ideas for you in the IndexRecord class, but there are others that might work.
I recommend adding a method that can return a random hash exactly X bits different from a given hash, and then using that to help in your tests.

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

Dcallies · 2024-11-16T22:40:39Z

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

+    def add(self, pdq_strings: t.Sequence[str]) -> None:
+        """
+        Add PDQ hashes to the FAISS index.
+        Args:
+            pdq_strings (Sequence[str]): PDQ hash strings to add
+        """
+        vectors = self._convert_pdq_strings_to_ndarray(pdq_strings)
+        self.faiss_index.add(vectors)


nit: It's somewhat confusing that this is add which takes multiple here, but the other is add which takes singular.

ignorable: I think you can just move this functionality to PDQIndex2

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

Dcallies · 2024-11-16T23:21:31Z

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

+    # Create a near-match by flipping a few bits
+    base_hash = pdq_hashes[0]
+    near_hash = hex(int(base_hash, 16) ^ 0xF)[2:].zfill(64)  # Flip 4 bits


blocking: Did you forget to query near_hash?

ignorable: In another similar project, I wrote a helper on the signal type itself to provide a random hash that was X distance away from a given hash by flipping random bits.

I query the near_hash on the line below

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

Dcallies · 2024-11-16T23:23:35Z

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

+    pdq_hashes = [
+        "f8f8f0cee0f4a84f06370a22038f63f0b36e2ed596621e1d33e6b39c4e9c9b22",
+        "f" * 64,
+        "0" * 64,
+        "a" * 64,
+    ]
+    index = PDQIndex2[str](
+        entries=[
+            ("f8f8f0cee0f4a84f06370a22038f63f0b36e2ed596621e1d33e6b39c4e9c9b22", 0)
+        ]
+    )
+    return index, pdq_hashes


blocking: Why do you return hashes which aren't in the index?

so that I can use the other hashes to test for unmatch when needed.
For example, in the test: test_one_entry_sample_index

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

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

Dcallies

Nice work! We're nearly there, only a few more blocking comments!

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

Dcallies · 2024-11-20T14:41:47Z

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

+    index = PDQIndex2()
+    for i, base_hash in enumerate(base_hashes):
+        index.add(base_hash, i)


nit: PDQIndex2.build()

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

Dcallies · 2024-11-20T14:43:59Z

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

+from threatexchange.signal_type.pdq.pdq_utils import (
+    BITS_IN_PDQ,
+    convert_pdq_strings_to_ndarray,
+)


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

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

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

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

Dcallies

So close!

The way you are seeding your hashes is causing duplicates in your query set - see inline blocking comment
As part of your test plan (doesn't have to be part of the test itself), please try generate a very large number of hashes and queries - e.g. 10,000 banked and 100,000 queries - something that takes 5-10 seconds to check. We want to be certain that the index is correct!
Please add a test that serializes and deserializes a non-empty index - you can do something cheeky like:

original = ... # do stuff to prepare index
deserialized = PDQIndex2,deserialize(index.serialize(...))
for index in (original, deserialized):
   # The original test goes here, should get same results on original or deserialized

Dcallies · 2024-11-21T14:38:15Z

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

+def _generate_sample_hashes(size: int, seed: int = 42):
+    random.seed(seed)
+    return [PdqSignal.get_random_signal() for _ in range(size)]
+
+
+SAMPLE_HASHES = _generate_sample_hashes(100)


This works for me!

Dcallies · 2024-11-21T14:38:57Z

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

+from threatexchange.signal_type.pdq.pdq_utils import simple_distance
+
+
+def _generate_sample_hashes(size: int, seed: int = 42):


I'm sorry what do you mean by the image:P ?

42 is the "Answer to the Ultimate Question of Life, The Universe, and Everything" in the book series "The Hitchhiker's Guide to the Galaxy", of which the image is on the cover. I thought you were a fan :P

Read more: https://en.wikipedia.org/wiki/Phrases_from_The_Hitchhiker%27s_Guide_to_the_Galaxy

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

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

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

haianhng31 · 2024-11-21T21:15:54Z

@Dcallies
I have tried with 1,000 banked & 10,000 queries in the unit test. The test passed but it took 4 minutes to run.
And if i tried with 10,000 banked & 100,000 queries, the test took forever.

Dcallies

All blocking comments resolved!

create pdq_index2 and test cases

300df97

haianhng31 requested a review from Dcallies as a code owner November 14, 2024 02:53

facebook-github-bot added the CLA Signed label Nov 14, 2024

haianhng31 commented Nov 14, 2024

View reviewed changes

python-threatexchange/threatexchange/signal_type/pdq/pdq_index2.py Outdated Show resolved Hide resolved

Dcallies changed the title ~~[pytx] Implement a new cleaner PDQ index solution (2)~~ [pytx] Implement a new cleaner PDQ index solution Nov 16, 2024

Dcallies requested changes Nov 16, 2024

View reviewed changes

Dcallies reviewed Nov 18, 2024

View reviewed changes

python-threatexchange/threatexchange/signal_type/pdq/pdq_index2.py Outdated Show resolved Hide resolved

haianhng31 added 3 commits November 18, 2024 18:56

pytx - continue fixing reimplementation pdq index

9fd7744

pytx - add unittest for pdq index2

3cf072b

pytx fix lint

cd44cbd

haianhng31 requested a review from Dcallies November 18, 2024 19:59

pytx - add test case for pdq index2

a5bfb6d

Dcallies requested changes Nov 20, 2024

View reviewed changes

haianhng31 added 2 commits November 20, 2024 20:58

pytx continue edit pdq index2

0a32d65

pytx resolve from nit comments for pdq index2

610af2e

haianhng31 requested a review from Dcallies November 20, 2024 21:20

Dcallies requested changes Nov 21, 2024

View reviewed changes

pytx - edit pdq index2

2ea86a1

haianhng31 requested a review from Dcallies November 21, 2024 23:08

Dcallies approved these changes Nov 25, 2024

View reviewed changes

Dcallies merged commit 2ae7f50 into facebook:main Nov 25, 2024
6 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[pytx] Implement a new cleaner PDQ index solution #1698

[pytx] Implement a new cleaner PDQ index solution #1698

haianhng31 commented Nov 14, 2024 •

edited by Dcallies

Loading

Dcallies left a comment

Dcallies Nov 16, 2024

Dcallies Nov 16, 2024

haianhng31 Nov 18, 2024

Dcallies Nov 16, 2024

haianhng31 Nov 18, 2024

Dcallies left a comment

Dcallies Nov 20, 2024

Dcallies Nov 20, 2024

Dcallies left a comment

Dcallies Nov 21, 2024

Dcallies Nov 21, 2024

haianhng31 Nov 21, 2024 •

edited

Loading

Dcallies Nov 25, 2024

haianhng31 commented Nov 21, 2024

Dcallies left a comment

		from threatexchange.signal_type.pdq.pdq_utils import simple_distance


		def _generate_sample_hashes(size: int, seed: int = 42):

[pytx] Implement a new cleaner PDQ index solution #1698

[pytx] Implement a new cleaner PDQ index solution #1698

Conversation

haianhng31 commented Nov 14, 2024 • edited by Dcallies Loading

Summary

Test Plan

Dcallies left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dcallies left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Dcallies left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

haianhng31 Nov 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

haianhng31 commented Nov 21, 2024

Dcallies left a comment

Choose a reason for hiding this comment

haianhng31 commented Nov 14, 2024 •

edited by Dcallies

Loading

haianhng31 Nov 21, 2024 •

edited

Loading