Fix tests

Author: Nikos Papailiou

apis/python/test/common.py

@@ -181,3 +181,11 @@ def create_array(path: str, data):
     tiledb.Array.create(path, schema)
     with tiledb.open(path, "w") as A:
         A[:] = data
+
+def accuracy(result, gt):
+    found = 0
+    total = 0
+    for i in range(len(result)):
+        total+=len(result[i])
+        found+=len(np.intersect1d(result[i], gt[i]))
+    return found/total

apis/python/test/test_ingestion.py

@@ -22,7 +22,7 @@ def test_flat_ingestion_u8(tmp_path):
         source_type=source_type,
     )
     result = np.transpose(index.query(np.transpose(query_vectors), k=k))
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
 
 def test_flat_ingestion_f32(tmp_path):
@@ -43,7 +43,7 @@ def test_flat_ingestion_f32(tmp_path):
         source_type=source_type,
     )
     result = np.transpose(index.query(np.transpose(query_vectors), k=k))
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
 
 def test_ivf_flat_ingestion_u8(tmp_path):
@@ -52,13 +52,12 @@ def test_ivf_flat_ingestion_u8(tmp_path):
     k = 10
     size = 100000
     partitions = 100
-    create_random_dataset_u8(nb=size, d=100, nq=2, k=k, path=dataset_dir)
+    create_random_dataset_u8(nb=size, d=100, nq=10, k=k, path=dataset_dir)
     source_type = "U8BIN"
     dtype = np.uint8
 
     query_vectors = get_queries(dataset_dir, dtype=dtype)
     gt_i, gt_d = get_groundtruth(dataset_dir, k)
-
     index = ingest(
         index_type="IVF_FLAT",
         array_uri=array_uri,
@@ -68,15 +67,15 @@ def test_ivf_flat_ingestion_u8(tmp_path):
         input_vectors_per_work_item=int(size / 10),
     )
     result = np.transpose(
-        index.query(np.transpose(query_vectors), k=k, nprobe=partitions)
+        index.query(np.transpose(query_vectors), k=k, nprobe=10)
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
     index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
     result = np.transpose(
         index_ram.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
     result = np.transpose(
         index_ram.query(
             np.transpose(query_vectors),
@@ -85,7 +84,7 @@ def test_ivf_flat_ingestion_u8(tmp_path):
             use_nuv_implementation=True,
         )
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
 
 def test_ivf_flat_ingestion_f32(tmp_path):
@@ -112,13 +111,13 @@ def test_ivf_flat_ingestion_f32(tmp_path):
     result = np.transpose(
         index.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
     index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
     result = np.transpose(
         index_ram.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
     result = np.transpose(
         index_ram.query(
             np.transpose(query_vectors),
@@ -127,7 +126,7 @@ def test_ivf_flat_ingestion_f32(tmp_path):
             use_nuv_implementation=True,
         )
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
 
 def test_ivf_flat_ingestion_fvec(tmp_path):
@@ -157,13 +156,13 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
     result = np.transpose(
         index.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
 
-    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype, memory_budget=int(size / 10))
+    index_ram = IVFFlatIndex(uri=array_uri, dtype=dtype)
     result = np.transpose(
         index_ram.query(np.transpose(query_vectors), k=k, nprobe=partitions)
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98
     result = np.transpose(
         index_ram.query(
             np.transpose(query_vectors),
@@ -172,4 +171,4 @@ def test_ivf_flat_ingestion_fvec(tmp_path):
             use_nuv_implementation=True,
         )
     )
-    assert np.array_equal(np.sort(result, axis=1), np.sort(gt_i, axis=1))
+    assert accuracy(result, gt_i) > 0.98