VSIM TRUTH option for ground truth results.

This commit is contained in:
antirez 2025-03-07 09:58:16 +01:00
parent 2114c65012
commit 0258e85186
4 changed files with 100 additions and 9 deletions

View File

@ -60,7 +60,7 @@ performed in the background, while the command is executed in the main thread.
**VSIM: return elements by vector similarity**
VSIM key [ELE|FP32|VALUES] <vector or element> [WITHSCORES] [COUNT num] [EF search-exploration-factor] [FILTER expression] [FILTER-EF max-filtering-effort]
VSIM key [ELE|FP32|VALUES] <vector or element> [WITHSCORES] [COUNT num] [EF search-exploration-factor] [FILTER expression] [FILTER-EF max-filtering-effort] [TRUTH]
The command returns similar vectors, for simplicity (and verbosity) in the following example, instead of providing a vector using FP32 or VALUES (like in `VADD`), we will ask for elements having a vector similar to a given element already in the sorted set:
@ -88,6 +88,8 @@ It is possible to specify a `COUNT` and also to get the similarity score (from 1
The `EF` argument is the exploration factor: the higher it is, the slower the command becomes, but the better the index is explored to find nodes that are near to our query. Sensible values are from 50 to 1000.
The `TRUTH` option forces the command to perform a linear scan of all the entries inside the set, without using the graph search inside the HNSW, so it returns the best matching elements (the perfect result set) that can be used in order to easily calculate the recall. Of course the linear scan is `O(N)`, so it is much slower than the `log(N)` (considering a small `COUNT`) provided by the HNSW index.
For `FILTER` and `FILTER-EF` options, please check the filtered search section of this documentation.
**VDIM: return the dimension of the vectors inside the vector set**

66
hnsw.c
View File

@ -2552,3 +2552,69 @@ void hnsw_test_graph_recall(HNSW *index, int test_ef, int verbose) {
unreachable_nodes,
total_nodes ? (float)unreachable_nodes * 100 / total_nodes : 0);
}
/* Return exact K-NN items by performing a linear scan of all nodes.
* This function has the same signature as hnsw_search_with_filter() but
* instead of using the graph structure, it scans all nodes to find the
* true nearest neighbors.
*
* Note that neighbors and distances arrays must have space for at least 'k' items.
* norm_query should be set to 1 if the query vector is already normalized.
*
* If the filter_callback is passed, only elements passing the specified filter
* are returned. The slot parameter is ignored but kept for API consistency. */
int hnsw_ground_truth_with_filter
(HNSW *index, const float *query_vector, uint32_t k,
hnswNode **neighbors, float *distances, uint32_t slot,
int query_vector_is_normalized,
int (*filter_callback)(void *value, void *privdata),
void *filter_privdata)
{
/* Note that we don't really use the slot here: it's a linear scan.
* Yet we want the user to acquire the slot as this will hold the
* global lock in read only mode. */
(void) slot;
/* Take our query vector into a temporary node. */
hnswNode query;
if (hnsw_init_tmp_node(index, &query, query_vector_is_normalized, query_vector) == 0) return -1;
/* Accumulate best results into a priority queue. */
pqueue *results = pq_new(k);
if (!results) {
hnsw_free_tmp_node(&query, query_vector);
return -1;
}
/* Scan all nodes linearly. */
hnswNode *current = index->head;
while (current) {
/* Apply filter if needed. */
if (filter_callback &&
!filter_callback(current->value, filter_privdata))
{
current = current->next;
continue;
}
/* Calculate distance to query. */
float dist = hnsw_distance(index, &query, current);
/* Add to results to pqueue. Will be accepted only if better than
* the current worse or pqueue not full. */
pq_push(results, current, dist);
current = current->next;
}
/* Copy results to output arrays. */
uint32_t found = MIN(k, results->count);
for (uint32_t i = 0; i < found; i++) {
neighbors[i] = pq_get_node(results, i);
if (distances) distances[i] = pq_get_distance(results, i);
}
/* Clean up. */
pq_free(results);
hnsw_free_tmp_node(&query, query_vector);
return found;
}

6
hnsw.h
View File

@ -160,5 +160,11 @@ void hnsw_set_allocator(void (*free_ptr)(void*), void *(*malloc_ptr)(size_t),
int hnsw_validate_graph(HNSW *index, uint64_t *connected_nodes, int *reciprocal_links);
void hnsw_test_graph_recall(HNSW *index, int test_ef, int verbose);
float hnsw_distance(HNSW *index, hnswNode *a, hnswNode *b);
int hnsw_ground_truth_with_filter
(HNSW *index, const float *query_vector, uint32_t k,
hnswNode **neighbors, float *distances, uint32_t slot,
int query_vector_is_normalized,
int (*filter_callback)(void *value, void *privdata),
void *filter_privdata);
#endif /* HNSW_H */

33
vset.c
View File

@ -585,7 +585,8 @@ int vectorSetFilterCallback(void *value, void *privdata) {
* handles the HNSW locking explicitly. */
void VSIM_execute(RedisModuleCtx *ctx, struct vsetObject *vset,
float *vec, unsigned long count, float epsilon, unsigned long withscores,
unsigned long ef, exprstate *filter_expr, unsigned long filter_ef)
unsigned long ef, exprstate *filter_expr, unsigned long filter_ef,
int ground_truth)
{
/* In our scan, we can't just collect 'count' elements as
* if count is small we would explore the graph in an insufficient
@ -603,10 +604,20 @@ void VSIM_execute(RedisModuleCtx *ctx, struct vsetObject *vset,
float *distances = RedisModule_Alloc(sizeof(float)*ef);
int slot = hnsw_acquire_read_slot(vset->hnsw);
unsigned int found;
if (filter_expr == NULL) {
found = hnsw_search(vset->hnsw, vec, ef, neighbors, distances, slot, 0);
if (ground_truth) {
found = hnsw_ground_truth_with_filter(vset->hnsw, vec, ef, neighbors,
distances, slot, 0,
filter_expr ? vectorSetFilterCallback : NULL,
filter_expr);
} else {
found = hnsw_search_with_filter(vset->hnsw, vec, ef, neighbors, distances, slot, 0, vectorSetFilterCallback, filter_expr, filter_ef);
if (filter_expr == NULL) {
found = hnsw_search(vset->hnsw, vec, ef, neighbors,
distances, slot, 0);
} else {
found = hnsw_search_with_filter(vset->hnsw, vec, ef, neighbors,
distances, slot, 0, vectorSetFilterCallback,
filter_expr, filter_ef);
}
}
hnsw_release_read_slot(vset->hnsw,slot);
RedisModule_Free(vec);
@ -654,6 +665,7 @@ void *VSIM_thread(void *arg) {
unsigned long ef = (unsigned long)targ[6];
exprstate *filter_expr = targ[7];
unsigned long filter_ef = (unsigned long)targ[8];
unsigned long ground_truth = (unsigned long)targ[9];
RedisModule_Free(targ[4]);
RedisModule_Free(targ);
@ -661,7 +673,7 @@ void *VSIM_thread(void *arg) {
RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(bc);
// Run the query.
VSIM_execute(ctx, vset, vec, count, epsilon, withscores, ef, filter_expr, filter_ef);
VSIM_execute(ctx, vset, vec, count, epsilon, withscores, ef, filter_expr, filter_ef, ground_truth);
// Cleanup.
RedisModule_FreeThreadSafeContext(ctx);
@ -683,6 +695,7 @@ int VSIM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
long long count = VSET_DEFAULT_COUNT; /* New default value */
long long ef = 0; /* Exploration factor (see HNSW paper) */
double epsilon = 2.0; /* Max cosine distance */
long long ground_truth = 0; /* Linear scan instead of HNSW search? */
/* Things computed later. */
long long filter_ef = 0;
@ -772,6 +785,9 @@ int VSIM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
if (!strcasecmp(opt, "WITHSCORES")) {
withscores = 1;
j++;
} else if (!strcasecmp(opt, "TRUTH")) {
ground_truth = 1;
j++;
} else if (!strcasecmp(opt, "COUNT") && j+1 < argc) {
if (RedisModule_StringToLongLong(argv[j+1], &count)
!= REDISMODULE_OK || count <= 0)
@ -852,7 +868,7 @@ int VSIM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
RedisModuleBlockedClient *bc = RedisModule_BlockClient(ctx,NULL,NULL,NULL,0);
pthread_t tid;
void **targ = RedisModule_Alloc(sizeof(void*)*9);
void **targ = RedisModule_Alloc(sizeof(void*)*10);
targ[0] = bc;
targ[1] = vset;
targ[2] = vec;
@ -863,16 +879,17 @@ int VSIM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
targ[6] = (void*)(unsigned long)ef;
targ[7] = (void*)filter_expr;
targ[8] = (void*)(unsigned long)filter_ef;
targ[9] = (void*)(unsigned long)ground_truth;
if (pthread_create(&tid,NULL,VSIM_thread,targ) != 0) {
pthread_rwlock_unlock(&vset->in_use_lock);
RedisModule_AbortBlock(bc);
RedisModule_Free(vec);
RedisModule_Free(targ[4]);
RedisModule_Free(targ);
VSIM_execute(ctx, vset, vec, count, epsilon, withscores, ef, filter_expr, filter_ef);
VSIM_execute(ctx, vset, vec, count, epsilon, withscores, ef, filter_expr, filter_ef, ground_truth);
}
} else {
VSIM_execute(ctx, vset, vec, count, epsilon, withscores, ef, filter_expr, filter_ef);
VSIM_execute(ctx, vset, vec, count, epsilon, withscores, ef, filter_expr, filter_ef, ground_truth);
}
return REDISMODULE_OK;