From 646df205a7ab2472abc958a993a513c8110914e8 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 4 Dec 2024 17:11:00 +0100 Subject: [PATCH 1/3] Skip IVF-PQ PQ packing test for lists with not enough data in them --- cpp/test/neighbors/ann_ivf_pq.cuh | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index fd4e330db2..ec589305ef 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -380,6 +380,14 @@ class ivf_pq_test : public ::testing::TestWithParam { int row_offset = 5; int n_vec = 3; ASSERT_TRUE(row_offset + n_vec < n_rows); + if (row_offset + n_vec > n_rows) { + RAFT_LOG_INFO( + "Skipping IVF-PQ check_packing/pack test for label %u due to insufficient data (%u " + "records)", + label, + uint32_t(n_rows)); + return; + } size_t offset = row_offset * index->pq_dim(); auto codes_to_pack = raft::make_device_matrix_view( codes.data_handle() + offset, n_vec, index->pq_dim()); @@ -393,7 +401,14 @@ class ivf_pq_test : public ::testing::TestWithParam { // Another test with the API that take list_data directly [[maybe_unused]] auto list_data = index->lists()[label]->data.view(); uint32_t n_take = 4; - ASSERT_TRUE(row_offset + n_take < n_rows); + if (row_offset + n_take > n_rows) { + RAFT_LOG_INFO( + "Skipping IVF-PQ check_packing/take test for label %u due to insufficient data (%u " + "records)", + label, + uint32_t(n_rows)); + return; + } auto codes2 = raft::make_device_matrix(handle_, n_take, index->pq_dim()); ivf_pq::helpers::codepacker::unpack( handle_, list_data, index->pq_bits(), row_offset, codes2.view()); From b2955223d706635ba6d16834cf10ef1128526068 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 4 Dec 2024 17:22:52 +0100 Subject: [PATCH 2/3] Fix types in the comparison --- cpp/test/neighbors/ann_ivf_pq.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index ec589305ef..7192e7274f 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -380,7 +380,7 @@ class ivf_pq_test : public ::testing::TestWithParam { int row_offset = 5; int n_vec = 3; ASSERT_TRUE(row_offset + n_vec < n_rows); - if (row_offset + n_vec > n_rows) { + if (static_cast(row_offset + n_vec) > n_rows) { RAFT_LOG_INFO( "Skipping IVF-PQ check_packing/pack test for label %u due to insufficient data (%u " "records)", @@ -401,7 +401,7 @@ class ivf_pq_test : public ::testing::TestWithParam { // Another test with the API that take list_data directly [[maybe_unused]] auto list_data = index->lists()[label]->data.view(); uint32_t n_take = 4; - if (row_offset + n_take > n_rows) { + if (static_cast(row_offset + n_take) > n_rows) { RAFT_LOG_INFO( "Skipping IVF-PQ check_packing/take test for label %u due to insufficient data (%u " "records)", From 63cef32bd2743a83fd4b802f5f2044b3c96e7623 Mon Sep 17 00:00:00 2001 From: achirkin Date: Wed, 4 Dec 2024 18:01:25 +0100 Subject: [PATCH 3/3] Remove the original assert --- cpp/test/neighbors/ann_ivf_pq.cuh | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 7192e7274f..3a92b5e3d6 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -379,7 +379,6 @@ class ivf_pq_test : public ::testing::TestWithParam { // Pack a few vectors back to the list. int row_offset = 5; int n_vec = 3; - ASSERT_TRUE(row_offset + n_vec < n_rows); if (static_cast(row_offset + n_vec) > n_rows) { RAFT_LOG_INFO( "Skipping IVF-PQ check_packing/pack test for label %u due to insufficient data (%u "