@@ -22,10 +22,11 @@ class TestSelector(unittest.TestCase):
2222 combinations as possible.
2323 """
2424
25- def do_test_id_selector (self , index_key , id_selector_type = "batch" , mt = faiss .METRIC_L2 , k = 10 ):
25+ def do_test_id_selector (self , index_key , id_selector_type = "batch" , mt = faiss .METRIC_L2 , k = 10 , params = None ):
2626 """ Verify that the id selector returns the subset of results that are
2727 members according to the IDSelector.
2828 Supports id_selector_type="batch", "bitmap", "range", "range_sorted", "and", "or", "xor"
29+ params: optional SearchParameters object to override default settings
2930 """
3031 d = 32 # make sure dimension is multiple of 8 for binary
3132 ds = datasets .SyntheticDataset (d , 1000 , 100 , 20 )
@@ -73,6 +74,8 @@ def do_test_id_selector(self, index_key, id_selector_type="batch", mt=faiss.METR
7374 subset = rs .choice (ds .nb , 50 , replace = False ).astype ('int64' )
7475
7576 index .add (xb [subset ])
77+ if "IVF" in index_key and id_selector_type == "range_sorted" :
78+ self .assertTrue (index .check_ids_sorted ())
7679 Dref , Iref0 = index .search (xq , k )
7780 Iref = subset [Iref0 ]
7881 Iref [Iref0 < 0 ] = - 1
@@ -134,11 +137,16 @@ def do_test_id_selector(self, index_key, id_selector_type="batch", mt=faiss.METR
134137 else :
135138 sel = faiss .IDSelectorBatch (subset )
136139
137- params = (
138- faiss .SearchParametersIVF (sel = sel ) if "IVF" in index_key else
139- faiss .SearchParametersPQ (sel = sel ) if "PQ" in index_key else
140- faiss .SearchParameters (sel = sel )
141- )
140+ if params is None :
141+ params = (
142+ faiss .SearchParametersIVF (sel = sel ) if "IVF" in index_key else
143+ faiss .SearchParametersPQ (sel = sel ) if "PQ" in index_key else
144+ faiss .SearchParameters (sel = sel )
145+ )
146+ else :
147+ # Use provided params but ensure selector is set
148+ params .sel = sel
149+
142150 Dnew , Inew = index .search (xq , k , params = params )
143151 np .testing .assert_array_equal (Iref , Inew )
144152 np .testing .assert_almost_equal (Dref , Dnew , decimal = 5 )
@@ -308,6 +316,11 @@ def test_BinaryFlat_id_range(self):
308316 def test_BinaryFlat_id_array (self ):
309317 self .do_test_id_selector ("BinaryFlat" , id_selector_type = "array" )
310318
319+ def test_BinaryFlat_no_heap (self ):
320+ params = faiss .SearchParameters ()
321+ params .use_heap = False
322+ self .do_test_id_selector ("BinaryFlat" , params = params )
323+
311324
312325class TestSearchParams (unittest .TestCase ):
313326
@@ -528,4 +541,4 @@ def test_knn_and_range_PQ(self):
528541 self .do_test_knn_and_range ("IVF32,PQ8x4np" )
529542
530543 def test_knn_and_range_FS (self ):
531- self .do_test_knn_and_range ("IVF32,PQ8x4fs" , range = False )
544+ self .do_test_knn_and_range ("IVF32,PQ8x4fs" , range = False )
0 commit comments