/* Redis implementation for vector sets. The data structure itself * is implemented in hnsw.c. * * Copyright (c) 2009-Present, Redis Ltd. * All rights reserved. * * Licensed under your choice of (a) the Redis Source Available License 2.0 * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the * GNU Affero General Public License v3 (AGPLv3). * Originally authored by: Salvatore Sanfilippo. * * ======================== Understand threading model ========================= * This code implements threaded operarations for two of the commands: * * 1. VSIM, by default. * 2. VADD, if the CAS option is specified. * * Note that even if the second operation, VADD, is a write operation, only * the neighbors collection for the new node is performed in a thread: then, * the actual insert is performed in the reply callback VADD_CASReply(), * which is executed in the main thread. * * Threaded operations need us to protect various operations with mutexes, * even if a certain degree of protection is already provided by the HNSW * library. Here are a few very important things about this implementation * and the way locking is performed. * * 1. All the write operations are performed in the main Redis thread: * this also include VADD_CASReply() callback, that is called by Redis * internals only in the context of the main thread. However the HNSW * library allows background threads in hnsw_search() (VSIM) to modify * nodes metadata to speedup search (to understand if a node was already * visited), but this only happens after acquiring a specific lock * for a given "read slot". * * 2. We use a global lock for each Vector Set object, called "in_use". This * lock is a read-write lock, and is acquired in read mode by all the * threads that perform reads in the background. It is only acquired in * write mode by vectorSetWaitAllBackgroundClients(): the function acquires * the lock and immediately releases it, with the effect of waiting all the * background threads still running from ending their execution. * * Note that no thread can be spawned, since we only call * vectorSetWaitAllBackgroundClients() from the main Redis thread, that * is also the only thread spawning other threads. * * vectorSetWaitAllBackgroundClients() is used in two ways: * A) When we need to delete a vector set because of (DEL) or other * operations destroying the object, we need to wait that all the * background threads working with this object finished their work. * B) When we modify the HNSW nodes bypassing the normal locking * provided by the HNSW library. This only happens when we update * an existing node attribute so far, in VSETATTR and when we call * VADD to update a node with the SETATTR option. * * 3. Often during read operations performed by Redis commands in the * main thread (VCARD, VEMB, VRANDMEMBER, ...) we don't acquire any * lock at all. The commands run in the main Redis thread, we can only * have, at the same time, background reads against the same data * structure. Note that VSIM_thread() and VADD_thread() still modify the * read slot metadata, that is node->visited_epoch[slot], but as long as * our read commands running in the main thread don't need to use * hnsw_search() or other HNSW functions using the visited epochs slots * we are safe. * * 4. There is a race from the moment we create a thread, passing the * vector set object, to the moment the thread can actually lock the * result win the in_use_lock mutex: as the thread starts, in the meanwhile * a DEL/expire could trigger and remove the object. For this reason * we use an atomic counter that protects our object for this small * time in vectorSetWaitAllBackgroundClients(). This prevents removal * of objects that are about to be taken by threads. * * Note that other competing solutions could be used to fix the problem * but have their set of issues, however they are worth documenting here * and evaluating in the future: * * A. Using a conditional variable we could "wait" for the thread to * acquire the lock. However this means waiting before returning * to the event loop, and would make the command execution slower. * B. We could use again an atomic variable, like we did, but this time * as a refcount for the object, with a vsetAcquire() vsetRelease(). * In this case, the command could retain the object in the main thread * before starting the thread, and the thread, after the work is done, * could release it. This way sometimes the object would be freed by * the thread, and it's while now can be safe to do the kind of resource * deallocation that vectorSetReleaseObject() does, given that the * Redis Modules API is not always thread safe this solution may not * be future-proof. However there is to evaluate it better in the * future. * C. We could use the "B" solution but instead of freeing the object * in the thread, in this specific case we could just put it into a * list and defer it for later freeing (for instance in the reply * callback), so that the object is always freed in the main thread. * This would require a list of objects to free. * * However the current solution only disadvantage is the potential busy * loop, but this busy loop in practical terms will almost never do * much: to trigger it, a number of circumnstances must happen: deleting * Vector Set keys while using them, hitting the small window needed to * start the thread and read-lock the mutex. */ #define _DEFAULT_SOURCE #define _USE_MATH_DEFINES #define _POSIX_C_SOURCE 200809L #include "../../src/redismodule.h" #include #include #include #include #include #include #include #include #include #include "hnsw.h" #include "vset_config.h" // We inline directly the expression implementation here so that building // the module is trivial. #include "expr.c" static RedisModuleType *VectorSetType; static uint64_t VectorSetTypeNextId = 0; // Default EF value if not specified during creation. #define VSET_DEFAULT_C_EF 200 // Default EF value if not specified during search. #define VSET_DEFAULT_SEARCH_EF 100 // Default num elements returned by VSIM. #define VSET_DEFAULT_COUNT 10 /* ========================== Internal data structure ====================== */ /* Our abstract data type needs a dual representation similar to Redis * sorted set: the proximity graph, and also a element -> graph-node map * that will allow us to perform deletions and other operations that have * as input the element itself. */ struct vsetObject { HNSW *hnsw; // Proximity graph. RedisModuleDict *dict; // Element -> node mapping. float *proj_matrix; // Random projection matrix, NULL if no projection uint32_t proj_input_size; // Input dimension after projection. // Output dimension is implicit in // hnsw->vector_dim. pthread_rwlock_t in_use_lock; // Lock needed to destroy the object safely. uint64_t id; // Unique ID used by threaded VADD to know the // object is still the same. uint64_t numattribs; // Number of nodes associated with an attribute. atomic_int thread_creation_pending; // Number of threads that are currently // pending to lock the object. }; /* Each node has two associated values: the associated string (the item * in the set) and potentially a JSON string, that is, the attributes, used * for hybrid search with the VSIM FILTER option. */ struct vsetNodeVal { RedisModuleString *item; RedisModuleString *attrib; }; /* Count the number of set bits in an integer (population count/Hamming weight). * This is a portable implementation that doesn't rely on compiler * extensions. */ static inline uint32_t bit_count(uint32_t n) { uint32_t count = 0; while (n) { count += n & 1; n >>= 1; } return count; } /* Create a Hadamard-based projection matrix for dimensionality reduction. * Uses {-1, +1} entries with a pattern based on bit operations. * The pattern is matrix[i][j] = (i & j) % 2 == 0 ? 1 : -1 * Matrix is scaled by 1/sqrt(input_dim) for normalization. * Returns NULL on allocation failure. * * Note that compared to other approaches (random gaussian weights), what * we have here is deterministic, it means that our replicas will have * the same set of weights. Also this approach seems to work much better * in practice, and the distances between elements are better guaranteed. * * Note that we still save the projection matrix in the RDB file, because * in the future we may change the weights generation, and we want everything * to be backward compatible. */ float *createProjectionMatrix(uint32_t input_dim, uint32_t output_dim) { float *matrix = RedisModule_Alloc(sizeof(float) * input_dim * output_dim); /* Scale factor to normalize the projection. */ const float scale = 1.0f / sqrt(input_dim); /* Fill the matrix using Hadamard pattern. */ for (uint32_t i = 0; i < output_dim; i++) { for (uint32_t j = 0; j < input_dim; j++) { /* Calculate position in the flattened matrix. */ uint32_t pos = i * input_dim + j; /* Hadamard pattern: use bit operations to determine sign * If the count of 1-bits in the bitwise AND of i and j is even, * the value is 1, otherwise -1. */ int value = (bit_count(i & j) % 2 == 0) ? 1 : -1; /* Store the scaled value. */ matrix[pos] = value * scale; } } return matrix; } /* Apply random projection to input vector. Returns new allocated vector. */ float *applyProjection(const float *input, const float *proj_matrix, uint32_t input_dim, uint32_t output_dim) { float *output = RedisModule_Alloc(sizeof(float) * output_dim); for (uint32_t i = 0; i < output_dim; i++) { const float *row = &proj_matrix[i * input_dim]; float sum = 0.0f; for (uint32_t j = 0; j < input_dim; j++) { sum += row[j] * input[j]; } output[i] = sum; } return output; } /* Create the vector as HNSW+Dictionary combined data structure. */ struct vsetObject *createVectorSetObject(unsigned int dim, uint32_t quant_type, uint32_t hnsw_M) { struct vsetObject *o; o = RedisModule_Alloc(sizeof(*o)); o->id = VectorSetTypeNextId++; o->hnsw = hnsw_new(dim,quant_type,hnsw_M); if (!o->hnsw) { // May fail because of mutex creation. RedisModule_Free(o); return NULL; } o->dict = RedisModule_CreateDict(NULL); o->proj_matrix = NULL; o->proj_input_size = 0; o->numattribs = 0; o->thread_creation_pending = 0; RedisModule_Assert(pthread_rwlock_init(&o->in_use_lock,NULL) == 0); return o; } void vectorSetReleaseNodeValue(void *v) { struct vsetNodeVal *nv = v; RedisModule_FreeString(NULL,nv->item); if (nv->attrib) RedisModule_FreeString(NULL,nv->attrib); RedisModule_Free(nv); } /* Free the vector set object. */ void vectorSetReleaseObject(struct vsetObject *o) { if (!o) return; if (o->hnsw) hnsw_free(o->hnsw,vectorSetReleaseNodeValue); if (o->dict) RedisModule_FreeDict(NULL,o->dict); if (o->proj_matrix) RedisModule_Free(o->proj_matrix); pthread_rwlock_destroy(&o->in_use_lock); RedisModule_Free(o); } /* Wait for all the threads performing operations on this * index to terminate their work (locking for write will * wait for all the other threads). * * if 'for_del' is set to 1, we also wait for all the pending threads * that still didn't acquire the lock to finish their work. This * is useful only if we are going to call this function to delete * the object, and not if we want to just to modify it. */ void vectorSetWaitAllBackgroundClients(struct vsetObject *vset, int for_del) { if (for_del) { // If we are going to destroy the object, after this call, let's // wait for threads that are being created and still didn't had // a chance to acquire the lock. while (vset->thread_creation_pending > 0); } RedisModule_Assert(pthread_rwlock_wrlock(&vset->in_use_lock) == 0); pthread_rwlock_unlock(&vset->in_use_lock); } /* Return a string representing the quantization type name of a vector set. */ const char *vectorSetGetQuantName(struct vsetObject *o) { switch(o->hnsw->quant_type) { case HNSW_QUANT_NONE: return "f32"; case HNSW_QUANT_Q8: return "int8"; case HNSW_QUANT_BIN: return "bin"; default: return "unknown"; } } /* Insert the specified element into the Vector Set. * If update is '1', the existing node will be updated. * * Returns 1 if the element was added, or 0 if the element was already there * and was just updated. */ int vectorSetInsert(struct vsetObject *o, float *vec, int8_t *qvec, float qrange, RedisModuleString *val, RedisModuleString *attrib, int update, int ef) { hnswNode *node = RedisModule_DictGet(o->dict,val,NULL); if (node != NULL) { if (update) { /* Wait for clients in the background: background VSIM * operations touch the nodes attributes we are going * to touch. */ vectorSetWaitAllBackgroundClients(o,0); struct vsetNodeVal *nv = node->value; /* Pass NULL as value-free function. We want to reuse * the old value. */ hnsw_delete_node(o->hnsw, node, NULL); node = hnsw_insert(o->hnsw,vec,qvec,qrange,0,nv,ef); RedisModule_Assert(node != NULL); RedisModule_DictReplace(o->dict,val,node); /* If attrib != NULL, the user wants that in case of an update we * update the attribute as well (otherwise it remains as it was). * Note that the order of operations is conceinved so that it * works in case the old attrib and the new attrib pointer is the * same. */ if (attrib) { // Empty attribute string means: unset the attribute during // the update. size_t attrlen; RedisModule_StringPtrLen(attrib,&attrlen); if (attrlen != 0) { RedisModule_RetainString(NULL,attrib); o->numattribs++; } else { attrib = NULL; } if (nv->attrib) { o->numattribs--; RedisModule_FreeString(NULL,nv->attrib); } nv->attrib = attrib; } } return 0; } struct vsetNodeVal *nv = RedisModule_Alloc(sizeof(*nv)); nv->item = val; nv->attrib = attrib; node = hnsw_insert(o->hnsw,vec,qvec,qrange,0,nv,ef); if (node == NULL) { // XXX Technically in Redis-land we don't have out of memory, as we // crash on OOM. However the HNSW library may fail for error in the // locking libc call. Probably impossible in practical terms. RedisModule_Free(nv); return 0; } if (attrib != NULL) o->numattribs++; RedisModule_DictSet(o->dict,val,node); RedisModule_RetainString(NULL,val); if (attrib) RedisModule_RetainString(NULL,attrib); return 1; } /* Parse vector from FP32 blob or VALUES format, with optional REDUCE. * Format: [REDUCE dim] FP32|VALUES ... * Returns allocated vector and sets dimension in *dim. * If reduce_dim is not NULL, sets it to the requested reduction dimension. * Returns NULL on parsing error. * * The function sets as a reference *consumed_args, so that the caller * knows how many arguments we consumed in order to parse the input * vector. Remaining arguments are often command options. */ float *parseVector(RedisModuleString **argv, int argc, int start_idx, size_t *dim, uint32_t *reduce_dim, int *consumed_args) { int consumed = 0; // Arguments consumed /* Check for REDUCE option first. */ if (reduce_dim) *reduce_dim = 0; if (reduce_dim && argc > start_idx + 2 && !strcasecmp(RedisModule_StringPtrLen(argv[start_idx],NULL),"REDUCE")) { long long rdim; if (RedisModule_StringToLongLong(argv[start_idx+1],&rdim) != REDISMODULE_OK || rdim <= 0) { return NULL; } if (reduce_dim) *reduce_dim = rdim; start_idx += 2; // Skip REDUCE and its argument. consumed += 2; } /* Now parse the vector format as before. */ float *vec = NULL; const char *vec_format = RedisModule_StringPtrLen(argv[start_idx],NULL); if (!strcasecmp(vec_format,"FP32")) { if (argc < start_idx + 2) return NULL; // Need FP32 + vector + value. size_t vec_raw_len; const char *blob = RedisModule_StringPtrLen(argv[start_idx+1],&vec_raw_len); // Must be 4 bytes per component. if (vec_raw_len % 4 || vec_raw_len < 4) return NULL; *dim = vec_raw_len/4; vec = RedisModule_Alloc(vec_raw_len); if (!vec) return NULL; memcpy(vec,blob,vec_raw_len); consumed += 2; } else if (!strcasecmp(vec_format,"VALUES")) { if (argc < start_idx + 2) return NULL; // Need at least the dimension. long long vdim; // Vector dimension passed by the user. if (RedisModule_StringToLongLong(argv[start_idx+1],&vdim) != REDISMODULE_OK || vdim < 1) return NULL; // Check that all the arguments are available. if (argc < start_idx + 2 + vdim) return NULL; *dim = vdim; vec = RedisModule_Alloc(sizeof(float) * vdim); if (!vec) return NULL; for (int j = 0; j < vdim; j++) { double val; if (RedisModule_StringToDouble(argv[start_idx+2+j],&val) != REDISMODULE_OK) { RedisModule_Free(vec); return NULL; } vec[j] = val; } consumed += vdim + 2; } else { return NULL; // Unknown format. } if (consumed_args) *consumed_args = consumed; return vec; } /* ========================== Commands implementation ======================= */ /* VADD thread handling the "CAS" version of the command, that is * performed blocking the client, accumulating here, in the thread, the * set of potential candidates, and later inserting the element in the * key (if it still exists, and if it is still the *same* vector set) * in the Reply callback. */ void *VADD_thread(void *arg) { pthread_detach(pthread_self()); void **targ = (void**)arg; RedisModuleBlockedClient *bc = targ[0]; struct vsetObject *vset = targ[1]; float *vec = targ[3]; int ef = (uint64_t)targ[6]; /* Lock the object and signal that we are no longer pending * the lock acquisition. */ RedisModule_Assert(pthread_rwlock_rdlock(&vset->in_use_lock) == 0); vset->thread_creation_pending--; /* Look for candidates... */ InsertContext *ic = hnsw_prepare_insert(vset->hnsw, vec, NULL, 0, 0, ef); targ[5] = ic; // Pass the context to the reply callback. /* Unblock the client so that our read reply will be invoked. */ pthread_rwlock_unlock(&vset->in_use_lock); RedisModule_BlockedClientMeasureTimeEnd(bc); RedisModule_UnblockClient(bc,targ); // Use targ as privdata. return NULL; } /* Reply callback for CAS variant of VADD. * Note: this is called in the main thread, in the background thread * we just do the read operation of gathering the neighbors. */ int VADD_CASReply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { (void)argc; RedisModule_AutoMemory(ctx); /* Use automatic memory management. */ int retval = REDISMODULE_OK; void **targ = (void**)RedisModule_GetBlockedClientPrivateData(ctx); uint64_t vset_id = (unsigned long) targ[2]; float *vec = targ[3]; RedisModuleString *val = targ[4]; InsertContext *ic = targ[5]; int ef = (uint64_t)targ[6]; RedisModuleString *attrib = targ[7]; RedisModule_Free(targ); /* Open the key: there are no guarantees it still exists, or contains * a vector set, or even the SAME vector set. */ RedisModuleKey *key = RedisModule_OpenKey(ctx,argv[1], REDISMODULE_READ|REDISMODULE_WRITE); int type = RedisModule_KeyType(key); struct vsetObject *vset = NULL; if (type != REDISMODULE_KEYTYPE_EMPTY && RedisModule_ModuleTypeGetType(key) == VectorSetType) { vset = RedisModule_ModuleTypeGetValue(key); // Same vector set? if (vset->id != vset_id) vset = NULL; /* Also, if the element was already inserted, we just pretend * the other insert won. We don't even start a threaded VADD * if this was an update, since the deletion of the element itself * in order to perform the update would invalidate the CAS state. */ if (vset && RedisModule_DictGet(vset->dict,val,NULL) != NULL) vset = NULL; } if (vset == NULL) { /* If the object does not match the start of the operation, we * just pretend the VADD was performed BEFORE the key was deleted * or replaced. We return success but don't do anything. */ hnsw_free_insert_context(ic); } else { /* Otherwise try to insert the new element with the neighbors * collected in background. If we fail, do it synchronously again * from scratch. */ // First: allocate the dual-ported value for the node. struct vsetNodeVal *nv = RedisModule_Alloc(sizeof(*nv)); nv->item = val; nv->attrib = attrib; /* Then: insert the node in the HNSW data structure. Note that * 'ic' could be NULL in case hnsw_prepare_insert() failed because of * locking failure (likely impossible in practical terms). */ hnswNode *newnode; if (ic == NULL || (newnode = hnsw_try_commit_insert(vset->hnsw, ic, nv)) == NULL) { /* If we are here, the CAS insert failed. We need to insert * again with full locking for neighbors selection and * actual insertion. This time we can't fail: */ newnode = hnsw_insert(vset->hnsw, vec, NULL, 0, 0, nv, ef); RedisModule_Assert(newnode != NULL); } RedisModule_DictSet(vset->dict,val,newnode); val = NULL; // Don't free it later. attrib = NULL; // Don't free it later. RedisModule_ReplicateVerbatim(ctx); } // Whatever happens is a success... :D RedisModule_ReplyWithBool(ctx,1); if (val) RedisModule_FreeString(ctx,val); // Not added? Free it. if (attrib) RedisModule_FreeString(ctx,attrib); // Not added? Free it. RedisModule_Free(vec); return retval; } /* VADD key [REDUCE dim] FP32|VALUES vector value [CAS] [NOQUANT] [BIN] [Q8] * [M count] */ int VADD_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); /* Use automatic memory management. */ if (argc < 5) return RedisModule_WrongArity(ctx); /* Parse vector with optional REDUCE */ size_t dim = 0; uint32_t reduce_dim = 0; int consumed_args; int cas = 0; // Threaded check-and-set style insert. long long ef = VSET_DEFAULT_C_EF; // HNSW creation time EF for new nodes. long long hnsw_create_M = HNSW_DEFAULT_M; // HNSW creation default M value. float *vec = parseVector(argv, argc, 2, &dim, &reduce_dim, &consumed_args); RedisModuleString *attrib = NULL; // Attributes if passed via ATTRIB. if (!vec) return RedisModule_ReplyWithError(ctx,"ERR invalid vector specification"); /* Missing element string at the end? */ if (argc-2-consumed_args < 1) { RedisModule_Free(vec); return RedisModule_WrongArity(ctx); } /* Parse options after the element string. */ uint32_t quant_type = HNSW_QUANT_Q8; // Default quantization type. for (int j = 2 + consumed_args + 1; j < argc; j++) { const char *opt = RedisModule_StringPtrLen(argv[j], NULL); if (!strcasecmp(opt, "CAS")) { cas = 1; } else if (!strcasecmp(opt, "EF") && j+1 < argc) { if (RedisModule_StringToLongLong(argv[j+1], &ef) != REDISMODULE_OK || ef <= 0 || ef > 1000000) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid EF"); } j++; // skip argument. } else if (!strcasecmp(opt, "M") && j+1 < argc) { if (RedisModule_StringToLongLong(argv[j+1], &hnsw_create_M) != REDISMODULE_OK || hnsw_create_M < HNSW_MIN_M || hnsw_create_M > HNSW_MAX_M) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid M"); } j++; // skip argument. } else if (!strcasecmp(opt, "SETATTR") && j+1 < argc) { attrib = argv[j+1]; j++; // skip argument. } else if (!strcasecmp(opt, "NOQUANT")) { quant_type = HNSW_QUANT_NONE; } else if (!strcasecmp(opt, "BIN")) { quant_type = HNSW_QUANT_BIN; } else if (!strcasecmp(opt, "Q8")) { quant_type = HNSW_QUANT_Q8; } else { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx,"ERR invalid option after element"); } } /* Drop CAS if this is a replica and we are getting the command from the * replication link: we want to add/delete items in the same order as * the master, while with CAS the timing would be different. * * Also for Lua scripts and MULTI/EXEC, we want to run the command * on the main thread. */ if (RedisModule_GetContextFlags(ctx) & (REDISMODULE_CTX_FLAGS_REPLICATED| REDISMODULE_CTX_FLAGS_LUA| REDISMODULE_CTX_FLAGS_MULTI)) { cas = 0; } if (VSGlobalConfig.forceSingleThreadExec) { cas = 0; } /* Open/create key */ RedisModuleKey *key = RedisModule_OpenKey(ctx,argv[1], REDISMODULE_READ|REDISMODULE_WRITE); int type = RedisModule_KeyType(key); if (type != REDISMODULE_KEYTYPE_EMPTY && RedisModule_ModuleTypeGetType(key) != VectorSetType) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx,REDISMODULE_ERRORMSG_WRONGTYPE); } /* Get the correct value argument based on format and REDUCE */ RedisModuleString *val = argv[2 + consumed_args]; /* Create or get existing vector set */ struct vsetObject *vset; if (type == REDISMODULE_KEYTYPE_EMPTY) { cas = 0; /* Do synchronous insert at creation, otherwise the * key would be left empty until the threaded part * does not return. It's also pointless to try try * doing threaded first element insertion. */ vset = createVectorSetObject(reduce_dim ? reduce_dim : dim, quant_type, hnsw_create_M); if (vset == NULL) { // We can't fail for OOM in Redis, but the mutex initialization // at least theoretically COULD fail. Likely this code path // is not reachable in practical terms. RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR unable to create a Vector Set: system resources issue?"); } /* Initialize projection if requested */ if (reduce_dim) { vset->proj_matrix = createProjectionMatrix(dim, reduce_dim); vset->proj_input_size = dim; /* Project the vector */ float *projected = applyProjection(vec, vset->proj_matrix, dim, reduce_dim); RedisModule_Free(vec); vec = projected; } RedisModule_ModuleTypeSetValue(key,VectorSetType,vset); } else { vset = RedisModule_ModuleTypeGetValue(key); if (vset->hnsw->quant_type != quant_type) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR asked quantization mismatch with existing vector set"); } if (vset->hnsw->M != hnsw_create_M) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR asked M value mismatch with existing vector set"); } if ((vset->proj_matrix == NULL && vset->hnsw->vector_dim != dim) || (vset->proj_matrix && vset->hnsw->vector_dim != reduce_dim)) { RedisModule_Free(vec); return RedisModule_ReplyWithErrorFormat(ctx, "ERR Vector dimension mismatch - got %d but set has %d", (int)dim, (int)vset->hnsw->vector_dim); } /* Check REDUCE compatibility */ if (reduce_dim) { if (!vset->proj_matrix) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR cannot add projection to existing set without projection"); } if (reduce_dim != vset->hnsw->vector_dim) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR projection dimension mismatch with existing set"); } } /* Apply projection if needed */ if (vset->proj_matrix) { /* Ensure input dimension matches the projection matrix's expected input dimension */ if (dim != vset->proj_input_size) { RedisModule_Free(vec); return RedisModule_ReplyWithErrorFormat(ctx, "ERR Input dimension mismatch for projection - got %d but projection expects %d", (int)dim, (int)vset->proj_input_size); } float *projected = applyProjection(vec, vset->proj_matrix, vset->proj_input_size, vset->hnsw->vector_dim); RedisModule_Free(vec); vec = projected; dim = vset->hnsw->vector_dim; } } /* For existing keys don't do CAS updates. For how things work now, the * CAS state would be invalidated by the deletion before adding back. */ if (cas && RedisModule_DictGet(vset->dict,val,NULL) != NULL) cas = 0; /* Here depending on the CAS option we directly insert in a blocking * way, or use a thread to do candidate neighbors selection and only * later, in the reply callback, actually add the element. */ if (cas) { RedisModuleBlockedClient *bc = RedisModule_BlockClient(ctx,VADD_CASReply,NULL,NULL,0); pthread_t tid; void **targ = RedisModule_Alloc(sizeof(void*)*8); targ[0] = bc; targ[1] = vset; targ[2] = (void*)(unsigned long)vset->id; targ[3] = vec; targ[4] = val; targ[5] = NULL; // Used later for insertion context. targ[6] = (void*)(unsigned long)ef; targ[7] = attrib; RedisModule_RetainString(ctx,val); if (attrib) RedisModule_RetainString(ctx,attrib); RedisModule_BlockedClientMeasureTimeStart(bc); vset->thread_creation_pending++; if (pthread_create(&tid,NULL,VADD_thread,targ) != 0) { vset->thread_creation_pending--; RedisModule_AbortBlock(bc); RedisModule_Free(targ); RedisModule_FreeString(ctx,val); if (attrib) RedisModule_FreeString(ctx,attrib); // Fall back to synchronous insert, see later in the code. } else { return REDISMODULE_OK; } } /* Insert vector synchronously: we reach this place even * if cas was true but thread creation failed. */ int added = vectorSetInsert(vset,vec,NULL,0,val,attrib,1,ef); RedisModule_Free(vec); RedisModule_ReplyWithBool(ctx,added); if (added) RedisModule_ReplicateVerbatim(ctx); return REDISMODULE_OK; } /* HNSW callback to filter items according to a predicate function * (our FILTER expression in this case). */ int vectorSetFilterCallback(void *value, void *privdata) { exprstate *expr = privdata; struct vsetNodeVal *nv = value; if (nv->attrib == NULL) return 0; // No attributes? No match. size_t json_len; char *json = (char*)RedisModule_StringPtrLen(nv->attrib,&json_len); return exprRun(expr,json,json_len); } /* Common path for the execution of the VSIM command both threaded and * not threaded. Note that 'ctx' may be normal context of a thread safe * context obtained from a blocked client. The locking that is specific * to the vset object is handled by the caller, however the function * 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 withattribs, 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 * way to provide enough recall. * * If the user didn't asked for a specific exploration, we use * VSET_DEFAULT_SEARCH_EF as minimum, or we match count if count * is greater than that. Otherwise the minumim will be the specified * EF argument. */ if (ef == 0) ef = VSET_DEFAULT_SEARCH_EF; if (count > ef) ef = count; int slot = hnsw_acquire_read_slot(vset->hnsw); if (ef > vset->hnsw->node_count) ef = vset->hnsw->node_count; /* Perform search */ hnswNode **neighbors = RedisModule_Alloc(sizeof(hnswNode*)*ef); float *distances = RedisModule_Alloc(sizeof(float)*ef); unsigned int found; if (ground_truth) { found = hnsw_ground_truth_with_filter(vset->hnsw, vec, ef, neighbors, distances, slot, 0, filter_expr ? vectorSetFilterCallback : NULL, filter_expr); } else { 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); } } /* Return results */ int resp3 = RedisModule_GetContextFlags(ctx) & REDISMODULE_CTX_FLAGS_RESP3; int reply_with_map = resp3 && (withscores || withattribs); if (reply_with_map) RedisModule_ReplyWithMap(ctx, REDISMODULE_POSTPONED_LEN); else RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_LEN); long long arraylen = 0; for (unsigned int i = 0; i < found && i < count; i++) { if (distances[i]/2 > epsilon) break; struct vsetNodeVal *nv = neighbors[i]->value; RedisModule_ReplyWithString(ctx, nv->item); arraylen++; /* If the user asked for multiple properties at the same time using * the RESP3 protocol, we wrap the value of the map into an N-items * array. Two for now, since we have just two properties that can be * requested. * * So in the case of RESP2 we will just have the flat reply: * item, score, attribute. For RESP3 instead item -> [score, attribute] */ if (resp3 && withscores && withattribs) RedisModule_ReplyWithArray(ctx,2); if (withscores) { /* The similarity score is provided in a 0-1 range. */ RedisModule_ReplyWithDouble(ctx, 1.0 - distances[i]/2.0); } if (withattribs) { /* Return the attributes as well, if any. */ if (nv->attrib) RedisModule_ReplyWithString(ctx, nv->attrib); else RedisModule_ReplyWithNull(ctx); } } hnsw_release_read_slot(vset->hnsw,slot); if (reply_with_map) { RedisModule_ReplySetMapLength(ctx, arraylen); } else { int items_per_ele = 1+withattribs+withscores; RedisModule_ReplySetArrayLength(ctx, arraylen * items_per_ele); } RedisModule_Free(vec); RedisModule_Free(neighbors); RedisModule_Free(distances); if (filter_expr) exprFree(filter_expr); } /* VSIM thread handling the blocked client request. */ void *VSIM_thread(void *arg) { pthread_detach(pthread_self()); // Extract arguments. void **targ = (void**)arg; RedisModuleBlockedClient *bc = targ[0]; struct vsetObject *vset = targ[1]; float *vec = targ[2]; unsigned long count = (unsigned long)targ[3]; float epsilon = *((float*)targ[4]); unsigned long withscores = (unsigned long)targ[5]; unsigned long withattribs = (unsigned long)targ[6]; unsigned long ef = (unsigned long)targ[7]; exprstate *filter_expr = targ[8]; unsigned long filter_ef = (unsigned long)targ[9]; unsigned long ground_truth = (unsigned long)targ[10]; RedisModule_Free(targ[4]); RedisModule_Free(targ); /* Lock the object and signal that we are no longer pending * the lock acquisition. */ RedisModule_Assert(pthread_rwlock_rdlock(&vset->in_use_lock) == 0); vset->thread_creation_pending--; // Accumulate reply in a thread safe context: no contention. RedisModuleCtx *ctx = RedisModule_GetThreadSafeContext(bc); // Run the query. VSIM_execute(ctx, vset, vec, count, epsilon, withscores, withattribs, ef, filter_expr, filter_ef, ground_truth); pthread_rwlock_unlock(&vset->in_use_lock); // Cleanup. RedisModule_FreeThreadSafeContext(ctx); RedisModule_BlockedClientMeasureTimeEnd(bc); RedisModule_UnblockClient(bc,NULL); return NULL; } /* VSIM key [ELE|FP32|VALUES] [WITHSCORES] [WITHATTRIBS] [COUNT num] [EPSILON eps] [EF exploration-factor] [FILTER expression] [FILTER-EF exploration-factor] */ int VSIM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); /* Basic argument check: need at least key and vector specification * method. */ if (argc < 4) return RedisModule_WrongArity(ctx); /* Defaults */ int withscores = 0; int withattribs = 0; 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? */ int no_thread = 0; /* NOTHREAD option: exec on main thread. */ /* Things computed later. */ long long filter_ef = 0; exprstate *filter_expr = NULL; /* Get key and vector type */ RedisModuleString *key = argv[1]; const char *vectorType = RedisModule_StringPtrLen(argv[2], NULL); /* Get vector set */ RedisModuleKey *keyptr = RedisModule_OpenKey(ctx, key, REDISMODULE_READ); int type = RedisModule_KeyType(keyptr); if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithEmptyArray(ctx); if (RedisModule_ModuleTypeGetType(keyptr) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); struct vsetObject *vset = RedisModule_ModuleTypeGetValue(keyptr); /* Vector parsing stage */ float *vec = NULL; size_t dim = 0; int vector_args = 0; /* Number of args consumed by vector specification */ if (!strcasecmp(vectorType, "ELE")) { /* Get vector from existing element */ RedisModuleString *ele = argv[3]; hnswNode *node = RedisModule_DictGet(vset->dict, ele, NULL); if (!node) { return RedisModule_ReplyWithError(ctx, "ERR element not found in set"); } vec = RedisModule_Alloc(sizeof(float) * vset->hnsw->vector_dim); hnsw_get_node_vector(vset->hnsw,node,vec); dim = vset->hnsw->vector_dim; vector_args = 2; /* ELE + element name */ } else { /* Parse vector. */ int consumed_args; vec = parseVector(argv, argc, 2, &dim, NULL, &consumed_args); if (!vec) { return RedisModule_ReplyWithError(ctx, "ERR invalid vector specification"); } vector_args = consumed_args; /* Apply projection if the set uses it, with the exception * of ELE type, that will already have the right dimension. */ if (vset->proj_matrix && dim != vset->hnsw->vector_dim) { /* Ensure input dimension matches the projection matrix's expected input dimension */ if (dim != vset->proj_input_size) { RedisModule_Free(vec); return RedisModule_ReplyWithErrorFormat(ctx, "ERR Input dimension mismatch for projection - got %d but projection expects %d", (int)dim, (int)vset->proj_input_size); } float *projected = applyProjection(vec, vset->proj_matrix, vset->proj_input_size, vset->hnsw->vector_dim); RedisModule_Free(vec); vec = projected; dim = vset->hnsw->vector_dim; } /* Count consumed arguments */ if (!strcasecmp(vectorType, "FP32")) { vector_args = 2; /* FP32 + vector blob */ } else if (!strcasecmp(vectorType, "VALUES")) { long long vdim; if (RedisModule_StringToLongLong(argv[3], &vdim) != REDISMODULE_OK) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid vector dimension"); } vector_args = 2 + vdim; /* VALUES + dim + values */ } else { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR vector type must be ELE, FP32 or VALUES"); } } /* Check vector dimension matches set */ if (dim != vset->hnsw->vector_dim) { RedisModule_Free(vec); return RedisModule_ReplyWithErrorFormat(ctx, "ERR Vector dimension mismatch - got %d but set has %d", (int)dim, (int)vset->hnsw->vector_dim); } /* Parse optional arguments - start after vector specification */ int j = 2 + vector_args; while (j < argc) { const char *opt = RedisModule_StringPtrLen(argv[j], NULL); if (!strcasecmp(opt, "WITHSCORES")) { withscores = 1; j++; } else if (!strcasecmp(opt, "WITHATTRIBS")) { withattribs = 1; j++; } else if (!strcasecmp(opt, "TRUTH")) { ground_truth = 1; j++; } else if (!strcasecmp(opt, "NOTHREAD")) { no_thread = 1; j++; } else if (!strcasecmp(opt, "COUNT") && j+1 < argc) { if (RedisModule_StringToLongLong(argv[j+1], &count) != REDISMODULE_OK || count <= 0) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid COUNT"); } j += 2; } else if (!strcasecmp(opt, "EPSILON") && j+1 < argc) { if (RedisModule_StringToDouble(argv[j+1], &epsilon) != REDISMODULE_OK || epsilon <= 0) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid EPSILON"); } j += 2; } else if (!strcasecmp(opt, "EF") && j+1 < argc) { if (RedisModule_StringToLongLong(argv[j+1], &ef) != REDISMODULE_OK || ef <= 0 || ef > 1000000) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid EF"); } j += 2; } else if (!strcasecmp(opt, "FILTER-EF") && j+1 < argc) { if (RedisModule_StringToLongLong(argv[j+1], &filter_ef) != REDISMODULE_OK || filter_ef <= 0) { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR invalid FILTER-EF"); } j += 2; } else if (!strcasecmp(opt, "FILTER") && j+1 < argc) { RedisModuleString *exprarg = argv[j+1]; size_t exprlen; char *exprstr = (char*)RedisModule_StringPtrLen(exprarg,&exprlen); int errpos; filter_expr = exprCompile(exprstr,&errpos); if (filter_expr == NULL) { if ((size_t)errpos >= exprlen) errpos = 0; RedisModule_Free(vec); return RedisModule_ReplyWithErrorFormat(ctx, "ERR syntax error in FILTER expression near: %s", exprstr+errpos); } j += 2; } else { RedisModule_Free(vec); return RedisModule_ReplyWithError(ctx, "ERR syntax error in VSIM command"); } } int threaded_request = 1; // Run on a thread, by default. if (filter_ef == 0) filter_ef = count * 100; // Max filter visited nodes. /* Disable threaded for MULTI/EXEC and Lua, or if explicitly * requested by the user via the NOTHREAD option. */ if (no_thread || VSGlobalConfig.forceSingleThreadExec || (RedisModule_GetContextFlags(ctx) & (REDISMODULE_CTX_FLAGS_LUA | REDISMODULE_CTX_FLAGS_MULTI))) { threaded_request = 0; } if (threaded_request) { /* Note: even if we create one thread per request, the underlying * HNSW library has a fixed number of slots for the threads, as it's * defined in HNSW_MAX_THREADS (beware that if you increase it, * every node will use more memory). This means that while this request * is threaded, and will NOT block Redis, it may end waiting for a * free slot if all the HNSW_MAX_THREADS slots are used. */ RedisModuleBlockedClient *bc = RedisModule_BlockClient(ctx,NULL,NULL,NULL,0); pthread_t tid; void **targ = RedisModule_Alloc(sizeof(void*)*11); targ[0] = bc; targ[1] = vset; targ[2] = vec; targ[3] = (void*)count; targ[4] = RedisModule_Alloc(sizeof(float)); *((float*)targ[4]) = epsilon; targ[5] = (void*)(unsigned long)withscores; targ[6] = (void*)(unsigned long)withattribs; targ[7] = (void*)(unsigned long)ef; targ[8] = (void*)filter_expr; targ[9] = (void*)(unsigned long)filter_ef; targ[10] = (void*)(unsigned long)ground_truth; RedisModule_BlockedClientMeasureTimeStart(bc); vset->thread_creation_pending++; if (pthread_create(&tid,NULL,VSIM_thread,targ) != 0) { vset->thread_creation_pending--; RedisModule_AbortBlock(bc); RedisModule_Free(targ[4]); RedisModule_Free(targ); VSIM_execute(ctx, vset, vec, count, epsilon, withscores, withattribs, ef, filter_expr, filter_ef, ground_truth); } } else { VSIM_execute(ctx, vset, vec, count, epsilon, withscores, withattribs, ef, filter_expr, filter_ef, ground_truth); } return REDISMODULE_OK; } /* VDIM : return the dimension of vectors in the vector set. */ int VDIM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc != 2) return RedisModule_WrongArity(ctx); RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); int type = RedisModule_KeyType(key); if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithError(ctx, "ERR key does not exist"); if (RedisModule_ModuleTypeGetType(key) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); return RedisModule_ReplyWithLongLong(ctx, vset->hnsw->vector_dim); } /* VCARD : return cardinality (num of elements) of the vector set. */ int VCARD_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc != 2) return RedisModule_WrongArity(ctx); RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); int type = RedisModule_KeyType(key); if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithLongLong(ctx, 0); if (RedisModule_ModuleTypeGetType(key) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); return RedisModule_ReplyWithLongLong(ctx, vset->hnsw->node_count); } /* VREM key element * Remove an element from a vector set. * Returns 1 if the element was found and removed, 0 if not found. */ int VREM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); /* Use automatic memory management. */ if (argc != 3) return RedisModule_WrongArity(ctx); /* Get key and value */ RedisModuleString *key = argv[1]; RedisModuleString *element = argv[2]; /* Open key */ RedisModuleKey *keyptr = RedisModule_OpenKey(ctx, key, REDISMODULE_READ|REDISMODULE_WRITE); int type = RedisModule_KeyType(keyptr); /* Handle non-existing key or wrong type */ if (type == REDISMODULE_KEYTYPE_EMPTY) { return RedisModule_ReplyWithBool(ctx, 0); } if (RedisModule_ModuleTypeGetType(keyptr) != VectorSetType) { return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); } /* Get vector set from key */ struct vsetObject *vset = RedisModule_ModuleTypeGetValue(keyptr); /* Find the node for this element */ hnswNode *node = RedisModule_DictGet(vset->dict, element, NULL); if (!node) { return RedisModule_ReplyWithBool(ctx, 0); } /* Remove from dictionary */ RedisModule_DictDel(vset->dict, element, NULL); /* Remove from HNSW graph using the high-level API that handles * locking and cleanup. We pass RedisModule_FreeString as the value * free function since the strings were retained at insertion time. */ struct vsetNodeVal *nv = node->value; if (nv->attrib != NULL) vset->numattribs--; RedisModule_Assert(hnsw_delete_node(vset->hnsw, node, vectorSetReleaseNodeValue) == 1); /* Destroy empty vector set. */ if (RedisModule_DictSize(vset->dict) == 0) { RedisModule_DeleteKey(keyptr); } /* Reply and propagate the command */ RedisModule_ReplyWithBool(ctx, 1); RedisModule_ReplicateVerbatim(ctx); return REDISMODULE_OK; } /* VEMB key element * Returns the embedding vector associated with an element, or NIL if not * found. The vector is returned in the same format it was added, but the * return value will have some lack of precision due to quantization and * normalization of vectors. Also, if items were added using REDUCE, the * reduced vector is returned instead. */ int VEMB_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); int raw_output = 0; // RAW option. if (argc < 3) return RedisModule_WrongArity(ctx); /* Parse arguments. */ for (int j = 3; j < argc; j++) { const char *opt = RedisModule_StringPtrLen(argv[j], NULL); if (!strcasecmp(opt,"raw")) { raw_output = 1; } else { return RedisModule_ReplyWithError(ctx,"ERR invalid option"); } } /* Get key and element. */ RedisModuleString *key = argv[1]; RedisModuleString *element = argv[2]; /* Open key. */ RedisModuleKey *keyptr = RedisModule_OpenKey(ctx, key, REDISMODULE_READ); int type = RedisModule_KeyType(keyptr); /* Handle non-existing key and key of wrong type. */ if (type == REDISMODULE_KEYTYPE_EMPTY) { return RedisModule_ReplyWithNull(ctx); } else if (RedisModule_ModuleTypeGetType(keyptr) != VectorSetType) { return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); } /* Lookup the node about the specified element. */ struct vsetObject *vset = RedisModule_ModuleTypeGetValue(keyptr); hnswNode *node = RedisModule_DictGet(vset->dict, element, NULL); if (!node) { return RedisModule_ReplyWithNull(ctx); } if (raw_output) { int output_qrange = vset->hnsw->quant_type == HNSW_QUANT_Q8; RedisModule_ReplyWithArray(ctx, 3+output_qrange); RedisModule_ReplyWithSimpleString(ctx, vectorSetGetQuantName(vset)); RedisModule_ReplyWithStringBuffer(ctx, node->vector, hnsw_quants_bytes(vset->hnsw)); RedisModule_ReplyWithDouble(ctx, node->l2); if (output_qrange) RedisModule_ReplyWithDouble(ctx, node->quants_range); } else { /* Get the vector associated with the node. */ float *vec = RedisModule_Alloc(sizeof(float) * vset->hnsw->vector_dim); hnsw_get_node_vector(vset->hnsw, node, vec); // May dequantize/denorm. /* Return as array of doubles. */ RedisModule_ReplyWithArray(ctx, vset->hnsw->vector_dim); for (uint32_t i = 0; i < vset->hnsw->vector_dim; i++) RedisModule_ReplyWithDouble(ctx, vec[i]); RedisModule_Free(vec); } return REDISMODULE_OK; } /* VSETATTR key element json * Set or remove the JSON attribute associated with an element. * Setting an empty string removes the attribute. * The command returns one if the attribute was actually updated or * zero if there is no key or element. */ int VSETATTR_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc != 4) return RedisModule_WrongArity(ctx); RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ|REDISMODULE_WRITE); int type = RedisModule_KeyType(key); if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithBool(ctx, 0); if (RedisModule_ModuleTypeGetType(key) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); hnswNode *node = RedisModule_DictGet(vset->dict, argv[2], NULL); if (!node) return RedisModule_ReplyWithBool(ctx, 0); struct vsetNodeVal *nv = node->value; RedisModuleString *new_attr = argv[3]; /* Background VSIM operations use the node attributes, so * wait for background operations before messing with them. */ vectorSetWaitAllBackgroundClients(vset,0); /* Set or delete the attribute based on the fact it's an empty * string or not. */ size_t attrlen; RedisModule_StringPtrLen(new_attr, &attrlen); if (attrlen == 0) { // If we had an attribute before, decrease the count and free it. if (nv->attrib) { vset->numattribs--; RedisModule_FreeString(NULL, nv->attrib); nv->attrib = NULL; } } else { // If we didn't have an attribute before, increase the count. // Otherwise free the old one. if (nv->attrib) { RedisModule_FreeString(NULL, nv->attrib); } else { vset->numattribs++; } // Set new attribute. RedisModule_RetainString(NULL, new_attr); nv->attrib = new_attr; } RedisModule_ReplyWithBool(ctx, 1); RedisModule_ReplicateVerbatim(ctx); return REDISMODULE_OK; } /* VGETATTR key element * Get the JSON attribute associated with an element. * Returns NIL if the element has no attribute or doesn't exist. */ int VGETATTR_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc != 3) return RedisModule_WrongArity(ctx); RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); int type = RedisModule_KeyType(key); if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithNull(ctx); if (RedisModule_ModuleTypeGetType(key) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); hnswNode *node = RedisModule_DictGet(vset->dict, argv[2], NULL); if (!node) return RedisModule_ReplyWithNull(ctx); struct vsetNodeVal *nv = node->value; if (!nv->attrib) return RedisModule_ReplyWithNull(ctx); return RedisModule_ReplyWithString(ctx, nv->attrib); } /* ============================== Reflection ================================ */ /* VLINKS key element [WITHSCORES] * Returns the neighbors of an element at each layer in the HNSW graph. * Reply is an array of arrays, where each nested array represents one level * of neighbors, from highest level to level 0. If WITHSCORES is specified, * each neighbor is followed by its distance from the element. */ int VLINKS_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc < 3 || argc > 4) return RedisModule_WrongArity(ctx); RedisModuleString *key = argv[1]; RedisModuleString *element = argv[2]; /* Parse WITHSCORES option. */ int withscores = 0; if (argc == 4) { const char *opt = RedisModule_StringPtrLen(argv[3], NULL); if (strcasecmp(opt, "WITHSCORES") != 0) { return RedisModule_WrongArity(ctx); } withscores = 1; } RedisModuleKey *keyptr = RedisModule_OpenKey(ctx, key, REDISMODULE_READ); int type = RedisModule_KeyType(keyptr); /* Handle non-existing key or wrong type. */ if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithNull(ctx); if (RedisModule_ModuleTypeGetType(keyptr) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); /* Find the node for this element. */ struct vsetObject *vset = RedisModule_ModuleTypeGetValue(keyptr); hnswNode *node = RedisModule_DictGet(vset->dict, element, NULL); if (!node) return RedisModule_ReplyWithNull(ctx); /* Reply with array of arrays, one per level. */ RedisModule_ReplyWithArray(ctx, node->level + 1); /* For each level, from highest to lowest: */ for (int i = node->level; i >= 0; i--) { /* Reply with array of neighbors at this level. */ if (withscores) RedisModule_ReplyWithMap(ctx,node->layers[i].num_links); else RedisModule_ReplyWithArray(ctx,node->layers[i].num_links); /* Add each neighbor's element value to the array. */ for (uint32_t j = 0; j < node->layers[i].num_links; j++) { struct vsetNodeVal *nv = node->layers[i].links[j]->value; RedisModule_ReplyWithString(ctx, nv->item); if (withscores) { float distance = hnsw_distance(vset->hnsw, node, node->layers[i].links[j]); /* Convert distance to similarity score to match * VSIM behavior.*/ float similarity = 1.0 - distance/2.0; RedisModule_ReplyWithDouble(ctx, similarity); } } } return REDISMODULE_OK; } /* VINFO key * Returns information about a vector set, both visible and hidden * features of the HNSW data structure. */ int VINFO_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc != 2) return RedisModule_WrongArity(ctx); RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); int type = RedisModule_KeyType(key); if (type == REDISMODULE_KEYTYPE_EMPTY) return RedisModule_ReplyWithNullArray(ctx); if (RedisModule_ModuleTypeGetType(key) != VectorSetType) return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); /* Reply with hash */ RedisModule_ReplyWithMap(ctx, 9); /* Quantization type */ RedisModule_ReplyWithSimpleString(ctx, "quant-type"); RedisModule_ReplyWithSimpleString(ctx, vectorSetGetQuantName(vset)); /* HNSW M value */ RedisModule_ReplyWithSimpleString(ctx, "hnsw-m"); RedisModule_ReplyWithLongLong(ctx, vset->hnsw->M); /* Vector dimensionality. */ RedisModule_ReplyWithSimpleString(ctx, "vector-dim"); RedisModule_ReplyWithLongLong(ctx, vset->hnsw->vector_dim); /* Original input dimension before projection. * This is zero for vector sets without a random projection matrix. */ RedisModule_ReplyWithSimpleString(ctx, "projection-input-dim"); RedisModule_ReplyWithLongLong(ctx, vset->proj_input_size); /* Number of elements. */ RedisModule_ReplyWithSimpleString(ctx, "size"); RedisModule_ReplyWithLongLong(ctx, vset->hnsw->node_count); /* Max level of HNSW. */ RedisModule_ReplyWithSimpleString(ctx, "max-level"); RedisModule_ReplyWithLongLong(ctx, vset->hnsw->max_level); /* Number of nodes with attributes. */ RedisModule_ReplyWithSimpleString(ctx, "attributes-count"); RedisModule_ReplyWithLongLong(ctx, vset->numattribs); /* Vector set ID. */ RedisModule_ReplyWithSimpleString(ctx, "vset-uid"); RedisModule_ReplyWithLongLong(ctx, vset->id); /* HNSW max node ID. */ RedisModule_ReplyWithSimpleString(ctx, "hnsw-max-node-uid"); RedisModule_ReplyWithLongLong(ctx, vset->hnsw->last_id); return REDISMODULE_OK; } /* VRANDMEMBER key [count] * Return random members from a vector set. * * Without count: returns a single random member. * With positive count: N unique random members (no duplicates). * With negative count: N random members (with possible duplicates). * * If the key doesn't exist, returns NULL if count is not given, or * an empty array if a count was given. */ int VRANDMEMBER_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); /* Use automatic memory management. */ /* Check arguments. */ if (argc != 2 && argc != 3) return RedisModule_WrongArity(ctx); /* Parse optional count argument. */ long long count = 1; /* Default is to return a single element. */ int with_count = (argc == 3); if (with_count) { if (RedisModule_StringToLongLong(argv[2], &count) != REDISMODULE_OK) { return RedisModule_ReplyWithError(ctx, "ERR COUNT value is not an integer"); } /* Count = 0 is a special case, return empty array */ if (count == 0) { return RedisModule_ReplyWithEmptyArray(ctx); } } /* Open key. */ RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); int type = RedisModule_KeyType(key); /* Handle non-existing key. */ if (type == REDISMODULE_KEYTYPE_EMPTY) { if (!with_count) { return RedisModule_ReplyWithNull(ctx); } else { return RedisModule_ReplyWithEmptyArray(ctx); } } /* Check key type. */ if (RedisModule_ModuleTypeGetType(key) != VectorSetType) { return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); } /* Get vector set from key. */ struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); uint64_t set_size = vset->hnsw->node_count; /* No elements in the set? */ if (set_size == 0) { if (!with_count) { return RedisModule_ReplyWithNull(ctx); } else { return RedisModule_ReplyWithEmptyArray(ctx); } } /* Case 1: No count specified: return a single element. */ if (!with_count) { hnswNode *random_node = hnsw_random_node(vset->hnsw, 0); if (random_node) { struct vsetNodeVal *nv = random_node->value; return RedisModule_ReplyWithString(ctx, nv->item); } else { return RedisModule_ReplyWithNull(ctx); } } /* Case 2: COUNT option given, return an array of elements. */ int allow_duplicates = (count < 0); long long abs_count = (count < 0) ? -count : count; /* Cap the count to the set size if we are not allowing duplicates. */ if (!allow_duplicates && abs_count > (long long)set_size) abs_count = set_size; /* Prepare reply. */ RedisModule_ReplyWithArray(ctx, abs_count); if (allow_duplicates) { /* Simple case: With duplicates, just pick random nodes * abs_count times. */ for (long long i = 0; i < abs_count; i++) { hnswNode *random_node = hnsw_random_node(vset->hnsw,0); struct vsetNodeVal *nv = random_node->value; RedisModule_ReplyWithString(ctx, nv->item); } } else { /* Case where count is positive: we need unique elements. * But, if the user asked for many elements, selecting so * many (> 20%) random nodes may be too expansive: we just start * from a random element and follow the next link. * * Otherwisem for the <= 20% case, a dictionary is used to * reject duplicates. */ int use_dict = (abs_count <= set_size * 0.2); if (use_dict) { RedisModuleDict *returned = RedisModule_CreateDict(ctx); long long returned_count = 0; while (returned_count < abs_count) { hnswNode *random_node = hnsw_random_node(vset->hnsw, 0); struct vsetNodeVal *nv = random_node->value; /* Check if we've already returned this element. */ if (RedisModule_DictGet(returned, nv->item, NULL) == NULL) { /* Mark as returned and add to results. */ RedisModule_DictSet(returned, nv->item, (void*)1); RedisModule_ReplyWithString(ctx, nv->item); returned_count++; } } RedisModule_FreeDict(ctx, returned); } else { /* For large samples, get a random starting node and walk * the list. * * IMPORTANT: doing so does not really generate random * elements: it's just a linear scan, but we have no choices. * If we generate too many random elements, more and more would * fail the check of being novel (not yet collected in the set * to return) if the % of elements to emit is too large, we would * spend too much CPU. */ hnswNode *start_node = hnsw_random_node(vset->hnsw, 0); hnswNode *current = start_node; long long returned_count = 0; while (returned_count < abs_count) { if (current == NULL) { /* Restart from head if we hit the end. */ current = vset->hnsw->head; } struct vsetNodeVal *nv = current->value; RedisModule_ReplyWithString(ctx, nv->item); returned_count++; current = current->next; } } } return REDISMODULE_OK; } /* VISMEMBER key element * Check if an element exists in a vector set. * Returns 1 if the element exists, 0 if not. */ int VISMEMBER_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); if (argc != 3) return RedisModule_WrongArity(ctx); RedisModuleString *key = argv[1]; RedisModuleString *element = argv[2]; /* Open key. */ RedisModuleKey *keyptr = RedisModule_OpenKey(ctx, key, REDISMODULE_READ); int type = RedisModule_KeyType(keyptr); /* Handle non-existing key or wrong type. */ if (type == REDISMODULE_KEYTYPE_EMPTY) { /* An element of a non existing key does not exist, like * SISMEMBER & similar. */ return RedisModule_ReplyWithBool(ctx, 0); } if (RedisModule_ModuleTypeGetType(keyptr) != VectorSetType) { return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); } /* Get the object and test membership via the dictionary in constant * time (assuming a member of average size). */ struct vsetObject *vset = RedisModule_ModuleTypeGetValue(keyptr); hnswNode *node = RedisModule_DictGet(vset->dict, element, NULL); return RedisModule_ReplyWithBool(ctx, node != NULL); } /* Structure to represent a range boundary. */ struct vsetRangeOp { int incl; /* 1 if inclusive ([), 0 if exclusive ((). */ int min; /* 1 if this is "-" (minimum). */ int max; /* 1 if this is "+" (maximum). */ unsigned char *ele; /* The actual element, NULL if min/max. */ size_t ele_len; /* Length of the element. */ }; /* Parse a range specification like "[foo" or "(bar" or "-" or "+". * Returns 1 on success, 0 on error. */ int vsetParseRangeOp(RedisModuleString *arg, struct vsetRangeOp *op) { size_t len; const char *str = RedisModule_StringPtrLen(arg, &len); if (len == 0) return 0; /* Initialize the structure. */ op->incl = 0; op->min = 0; op->max = 0; op->ele = NULL; op->ele_len = 0; /* Check for special cases "-" and "+". */ if (len == 1 && str[0] == '-') { op->min = 1; return 1; } if (len == 1 && str[0] == '+') { op->max = 1; return 1; } /* Otherwise, must start with ( or [. */ if (str[0] == '[') { op->incl = 1; } else if (str[0] == '(') { op->incl = 0; } else { return 0; /* Invalid format. */ } /* Extract the string part after the bracket. */ if (len > 1) { op->ele = (unsigned char *)(str + 1); op->ele_len = len - 1; } else { return 0; /* Just a bracket with no string. */ } return 1; } /* Check if the current element is within the range defined by the end operator. * Returns 1 if the element is within range, 0 if it has passed the end. */ int vsetIsElementInRange(const void *ele, size_t ele_len, struct vsetRangeOp *end_op) { /* If end is "+", element is always in range. */ if (end_op->max) return 1; /* Compare current element with end boundary. */ size_t minlen = ele_len < end_op->ele_len ? ele_len : end_op->ele_len; int cmp = memcmp(ele, end_op->ele, minlen); if (cmp == 0) { /* If equal up to minlen, shorter string is smaller. */ if (ele_len < end_op->ele_len) { cmp = -1; } else if (ele_len > end_op->ele_len) { cmp = 1; } } /* Check based on inclusive/exclusive. */ if (end_op->incl) { return cmp <= 0; /* Inclusive: element <= end. */ } else { return cmp < 0; /* Exclusive: element < end. */ } } /* VRANGE key start end [count] * Returns elements in the lexicographical range [start, end] * * Elements must be specified in one of the following forms: * * [myelement * (myelement * + * - * * Elements starting with [ are inclusive, so "myelement" would be * returned if present in the set. Elements starting with ( are exclusive * ranges instead. The special - and + elements mean the minimum and maximum * possible element (inclusive), so "VRANGE key - +" will return everything * (depending on COUNT of course). The special - element can be used only * as starting element, the special + element only as ending element. */ int VRANGE_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { RedisModule_AutoMemory(ctx); /* Check arguments. */ if (argc < 4 || argc > 5) return RedisModule_WrongArity(ctx); /* Parse COUNT if provided. */ long long count = -1; /* Default: return all elements. */ if (argc == 5) { if (RedisModule_StringToLongLong(argv[4], &count) != REDISMODULE_OK) { return RedisModule_ReplyWithError(ctx, "ERR invalid COUNT value"); } } /* Parse range operators. */ struct vsetRangeOp start_op, end_op; if (!vsetParseRangeOp(argv[2], &start_op)) { return RedisModule_ReplyWithError(ctx, "ERR invalid start range format"); } if (!vsetParseRangeOp(argv[3], &end_op)) { return RedisModule_ReplyWithError(ctx, "ERR invalid end range format"); } /* Validate: "-" can only be first arg, "+" can only be second. */ if (start_op.max || end_op.min) { return RedisModule_ReplyWithError(ctx, "ERR '-' can only be used as first argument, '+' only as second"); } /* Open the key. */ RedisModuleKey *key = RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ); int type = RedisModule_KeyType(key); if (type == REDISMODULE_KEYTYPE_EMPTY) { return RedisModule_ReplyWithEmptyArray(ctx); } if (RedisModule_ModuleTypeGetType(key) != VectorSetType) { return RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE); } struct vsetObject *vset = RedisModule_ModuleTypeGetValue(key); /* Start the iterator. */ RedisModuleDictIter *iter; if (start_op.min) { /* Start from the beginning. */ iter = RedisModule_DictIteratorStartC(vset->dict, "^", NULL, 0); } else { /* Start from the specified element. */ const char *op = start_op.incl ? ">=" : ">"; iter = RedisModule_DictIteratorStartC(vset->dict, op, start_op.ele, start_op.ele_len); } /* Collect results. */ RedisModule_ReplyWithArray(ctx, REDISMODULE_POSTPONED_LEN); long long returned = 0; void *key_data; size_t key_len; while ((key_data = RedisModule_DictNextC(iter, &key_len, NULL)) != NULL) { /* Check if we've collected enough elements. */ if (count >= 0 && returned >= count) break; /* Check if we've passed the end range. */ if (!vsetIsElementInRange(key_data, key_len, &end_op)) break; /* Add this element to the result. */ RedisModule_ReplyWithStringBuffer(ctx, key_data, key_len); returned++; } RedisModule_ReplySetArrayLength(ctx, returned); /* Cleanup. */ RedisModule_DictIteratorStop(iter); return REDISMODULE_OK; } /* ============================== vset type methods ========================= */ #define SAVE_FLAG_HAS_PROJMATRIX (1<<0) #define SAVE_FLAG_HAS_ATTRIBS (1<<1) /* Save object to RDB */ void VectorSetRdbSave(RedisModuleIO *rdb, void *value) { struct vsetObject *vset = value; RedisModule_SaveUnsigned(rdb, vset->hnsw->vector_dim); RedisModule_SaveUnsigned(rdb, vset->hnsw->node_count); uint32_t hnsw_config = (vset->hnsw->quant_type & 0xff) | ((vset->hnsw->M & 0xffff) << 8); RedisModule_SaveUnsigned(rdb, hnsw_config); uint32_t save_flags = 0; if (vset->proj_matrix) save_flags |= SAVE_FLAG_HAS_PROJMATRIX; if (vset->numattribs != 0) save_flags |= SAVE_FLAG_HAS_ATTRIBS; RedisModule_SaveUnsigned(rdb, save_flags); /* Save projection matrix if present */ if (vset->proj_matrix) { uint32_t input_dim = vset->proj_input_size; uint32_t output_dim = vset->hnsw->vector_dim; RedisModule_SaveUnsigned(rdb, input_dim); // Output dim is the same as the first value saved // above, so we don't save it. // Save projection matrix as binary blob size_t matrix_size = sizeof(float) * input_dim * output_dim; RedisModule_SaveStringBuffer(rdb, (const char *)vset->proj_matrix, matrix_size); } hnswNode *node = vset->hnsw->head; while(node) { struct vsetNodeVal *nv = node->value; RedisModule_SaveString(rdb, nv->item); if (vset->numattribs) { if (nv->attrib) RedisModule_SaveString(rdb, nv->attrib); else RedisModule_SaveStringBuffer(rdb, "", 0); } hnswSerNode *sn = hnsw_serialize_node(vset->hnsw,node); RedisModule_SaveStringBuffer(rdb, (const char *)sn->vector, sn->vector_size); RedisModule_SaveUnsigned(rdb, sn->params_count); for (uint32_t j = 0; j < sn->params_count; j++) RedisModule_SaveUnsigned(rdb, sn->params[j]); hnsw_free_serialized_node(sn); node = node->next; } } /* Load object from RDB. Recover from recoverable errors (read errors) * by performing cleanup. */ void *VectorSetRdbLoad(RedisModuleIO *rdb, int encver) { if (encver != 0) return NULL; // Invalid version uint32_t dim = RedisModule_LoadUnsigned(rdb); uint64_t elements = RedisModule_LoadUnsigned(rdb); uint32_t hnsw_config = RedisModule_LoadUnsigned(rdb); if (RedisModule_IsIOError(rdb)) return NULL; uint32_t quant_type = hnsw_config & 0xff; uint32_t hnsw_m = (hnsw_config >> 8) & 0xffff; /* Check that the quantization type is correct. Otherwise * return ASAP signaling the error. */ if (quant_type != HNSW_QUANT_NONE && quant_type != HNSW_QUANT_Q8 && quant_type != HNSW_QUANT_BIN) return NULL; if (hnsw_m == 0) hnsw_m = 16; // Default, useful for RDB files predating // this configuration parameter: it was fixed // to 16. struct vsetObject *vset = createVectorSetObject(dim,quant_type,hnsw_m); RedisModule_Assert(vset != NULL); /* Load projection matrix if present */ uint32_t save_flags = RedisModule_LoadUnsigned(rdb); if (RedisModule_IsIOError(rdb)) goto ioerr; int has_projection = save_flags & SAVE_FLAG_HAS_PROJMATRIX; int has_attribs = save_flags & SAVE_FLAG_HAS_ATTRIBS; if (has_projection) { uint32_t input_dim = RedisModule_LoadUnsigned(rdb); if (RedisModule_IsIOError(rdb)) goto ioerr; uint32_t output_dim = dim; size_t matrix_size = sizeof(float) * input_dim * output_dim; vset->proj_matrix = RedisModule_Alloc(matrix_size); vset->proj_input_size = input_dim; // Load projection matrix as a binary blob char *matrix_blob = RedisModule_LoadStringBuffer(rdb, NULL); if (matrix_blob == NULL) goto ioerr; memcpy(vset->proj_matrix, matrix_blob, matrix_size); RedisModule_Free(matrix_blob); } while(elements--) { // Load associated string element. RedisModuleString *ele = RedisModule_LoadString(rdb); if (RedisModule_IsIOError(rdb)) goto ioerr; RedisModuleString *attrib = NULL; if (has_attribs) { attrib = RedisModule_LoadString(rdb); if (RedisModule_IsIOError(rdb)) { RedisModule_FreeString(NULL,ele); goto ioerr; } size_t attrlen; RedisModule_StringPtrLen(attrib,&attrlen); if (attrlen == 0) { RedisModule_FreeString(NULL,attrib); attrib = NULL; } } size_t vector_len; void *vector = RedisModule_LoadStringBuffer(rdb, &vector_len); if (RedisModule_IsIOError(rdb)) { RedisModule_FreeString(NULL,ele); if (attrib) RedisModule_FreeString(NULL,attrib); goto ioerr; } uint32_t vector_bytes = hnsw_quants_bytes(vset->hnsw); if (vector_len != vector_bytes) { RedisModule_LogIOError(rdb,"warning", "Mismatching vector dimension"); RedisModule_FreeString(NULL,ele); if (attrib) RedisModule_FreeString(NULL,attrib); RedisModule_Free(vector); goto ioerr; } // Load node parameters back. uint32_t params_count = RedisModule_LoadUnsigned(rdb); if (RedisModule_IsIOError(rdb)) { RedisModule_FreeString(NULL,ele); if (attrib) RedisModule_FreeString(NULL,attrib); RedisModule_Free(vector); goto ioerr; } uint64_t *params = RedisModule_Alloc(params_count*sizeof(uint64_t)); for (uint32_t j = 0; j < params_count; j++) { // Ignore loading errors here: handled at the end of the loop. params[j] = RedisModule_LoadUnsigned(rdb); } if (RedisModule_IsIOError(rdb)) { RedisModule_FreeString(NULL,ele); if (attrib) RedisModule_FreeString(NULL,attrib); RedisModule_Free(vector); RedisModule_Free(params); goto ioerr; } struct vsetNodeVal *nv = RedisModule_Alloc(sizeof(*nv)); nv->item = ele; nv->attrib = attrib; hnswNode *node = hnsw_insert_serialized(vset->hnsw, vector, params, params_count, nv); if (node == NULL) { RedisModule_LogIOError(rdb,"warning", "Vector set node index loading error"); vectorSetReleaseNodeValue(nv); RedisModule_Free(vector); RedisModule_Free(params); goto ioerr; } if (nv->attrib) vset->numattribs++; RedisModule_DictSet(vset->dict,ele,node); RedisModule_Free(vector); RedisModule_Free(params); } uint64_t salt[2]; RedisModule_GetRandomBytes((unsigned char*)salt,sizeof(salt)); if (!hnsw_deserialize_index(vset->hnsw, salt[0], salt[1])) goto ioerr; return vset; ioerr: /* We want to recover from I/O errors and free the partially allocated * data structure to support diskless replication. */ vectorSetReleaseObject(vset); return NULL; } /* Calculate memory usage */ size_t VectorSetMemUsage(const void *value) { const struct vsetObject *vset = value; size_t size = sizeof(*vset); /* Account for HNSW index base structure */ size += sizeof(HNSW); /* Account for projection matrix if present */ if (vset->proj_matrix) { /* For the matrix size, we need the input dimension. We can get it * from the first node if the set is not empty. */ uint32_t input_dim = vset->proj_input_size; uint32_t output_dim = vset->hnsw->vector_dim; size += sizeof(float) * input_dim * output_dim; } /* Account for each node's memory usage. */ hnswNode *node = vset->hnsw->head; if (node == NULL) return size; /* Base node structure. */ size += sizeof(*node) * vset->hnsw->node_count; /* Vector storage. */ uint64_t vec_storage = hnsw_quants_bytes(vset->hnsw); size += vec_storage * vset->hnsw->node_count; /* Layers array. We use 1.33 as average nodes layers count. */ uint64_t layers_storage = sizeof(hnswNodeLayer) * vset->hnsw->node_count; layers_storage = layers_storage * 4 / 3; // 1.33 times. size += layers_storage; /* All the nodes have layer 0 links. */ uint64_t level0_links = node->layers[0].max_links; uint64_t other_levels_links = level0_links/2; size += sizeof(hnswNode*) * level0_links * vset->hnsw->node_count; /* Add the 0.33 remaining part, but upper layers have less links. */ size += (sizeof(hnswNode*) * other_levels_links * vset->hnsw->node_count)/3; /* Associated string value and attributres. * Use Redis Module API to get string size, and guess that all the * elements have similar size as the first few. */ size_t items_scanned = 0, items_size = 0; size_t attribs_scanned = 0, attribs_size = 0; int scan_effort = 20; while(scan_effort > 0 && node) { struct vsetNodeVal *nv = node->value; items_size += RedisModule_MallocSizeString(nv->item); items_scanned++; if (nv->attrib) { attribs_size += RedisModule_MallocSizeString(nv->attrib); attribs_scanned++; } scan_effort--; node = node->next; } /* Add the memory usage due to items. */ if (items_scanned) size += items_size / items_scanned * vset->hnsw->node_count; /* Add memory usage due to attributres. */ if (attribs_scanned == 0) { /* We were not lucky enough to find a single attribute in the * first few items? Let's use a fixed arbitrary value. */ attribs_scanned = 1; attribs_size = 64; } size += attribs_size / attribs_scanned * vset->numattribs; /* Account for dictionary overhead - this is an approximation. */ size += RedisModule_DictSize(vset->dict) * (sizeof(void*) * 2); return size; } /* Free the entire data structure */ void VectorSetFree(void *value) { struct vsetObject *vset = value; vectorSetWaitAllBackgroundClients(vset,1); vectorSetReleaseObject(value); } /* Add object digest to the digest context */ void VectorSetDigest(RedisModuleDigest *md, void *value) { struct vsetObject *vset = value; /* Add consistent order-independent hash of all vectors */ hnswNode *node = vset->hnsw->head; /* Hash the vector dimension and number of nodes. */ RedisModule_DigestAddLongLong(md, vset->hnsw->node_count); RedisModule_DigestAddLongLong(md, vset->hnsw->vector_dim); RedisModule_DigestEndSequence(md); while(node) { struct vsetNodeVal *nv = node->value; /* Hash each vector component */ RedisModule_DigestAddStringBuffer(md, node->vector, hnsw_quants_bytes(vset->hnsw)); /* Hash the associated value */ size_t len; const char *str = RedisModule_StringPtrLen(nv->item, &len); RedisModule_DigestAddStringBuffer(md, (char*)str, len); if (nv->attrib) { str = RedisModule_StringPtrLen(nv->attrib, &len); RedisModule_DigestAddStringBuffer(md, (char*)str, len); } node = node->next; RedisModule_DigestEndSequence(md); } } // int VectorSets_InitModuleConfig(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { int VectorSets_InitModuleConfig(RedisModuleCtx *ctx) { if (RegisterModuleConfig(ctx) == REDISMODULE_ERR) { RedisModule_Log(ctx, "warning", "Error registering module configuration"); return REDISMODULE_ERR; } // Load default values if (RedisModule_LoadDefaultConfigs(ctx) == REDISMODULE_ERR) { RedisModule_Log(ctx, "warning", "Error loading default module configuration"); return REDISMODULE_ERR; } else { RedisModule_Log(ctx, "verbose", "Successfully loaded default module configuration"); } if (RedisModule_LoadConfigs(ctx) == REDISMODULE_ERR) { RedisModule_Log(ctx, "warning", "Error loading user module configuration"); return REDISMODULE_ERR; } else { RedisModule_Log(ctx, "verbose", "Successfully loaded user module configuration"); } return REDISMODULE_OK; } /* This function must be present on each Redis module. It is used in order to * register the commands into the Redis server. */ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { REDISMODULE_NOT_USED(argv); REDISMODULE_NOT_USED(argc); if (RedisModule_Init(ctx,"vectorset",1,REDISMODULE_APIVER_1) == REDISMODULE_ERR) return REDISMODULE_ERR; if (VectorSets_InitModuleConfig(ctx) == REDISMODULE_ERR) { return REDISMODULE_ERR; } RedisModule_SetModuleOptions(ctx, REDISMODULE_OPTIONS_HANDLE_IO_ERRORS|REDISMODULE_OPTIONS_HANDLE_REPL_ASYNC_LOAD); RedisModuleTypeMethods tm = { .version = REDISMODULE_TYPE_METHOD_VERSION, .rdb_load = VectorSetRdbLoad, .rdb_save = VectorSetRdbSave, .aof_rewrite = NULL, .mem_usage = VectorSetMemUsage, .free = VectorSetFree, .digest = VectorSetDigest }; VectorSetType = RedisModule_CreateDataType(ctx,"vectorset",0,&tm); if (VectorSetType == NULL) return REDISMODULE_ERR; // Register command VADD if (RedisModule_CreateCommand(ctx,"VADD", VADD_RedisCommand,"write deny-oom",1,1,1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vadd_cmd = RedisModule_GetCommand(ctx, "VADD"); if (vadd_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vadd_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "reduce", .type = REDISMODULE_ARG_TYPE_BLOCK, .token = "REDUCE", .flags = REDISMODULE_CMD_ARG_OPTIONAL, .subargs = (RedisModuleCommandArg[]) { { .name = "dim", .type = REDISMODULE_ARG_TYPE_INTEGER }, { .name = NULL } } }, { .name = "format", .type = REDISMODULE_ARG_TYPE_ONEOF, .subargs = (RedisModuleCommandArg[]) { { .name = "fp32", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "FP32" }, { .name = "values", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "VALUES" }, { .name = NULL } } }, { .name = "vector", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "cas", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "CAS", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "quant_type", .type = REDISMODULE_ARG_TYPE_ONEOF, .flags = REDISMODULE_CMD_ARG_OPTIONAL, .subargs = (RedisModuleCommandArg[]) { { .name = "noquant", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "NOQUANT" }, { .name = "bin", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "BIN" }, { .name = "q8", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "Q8" }, { .name = NULL } } }, { .name = "build-exploration-factor", .type = REDISMODULE_ARG_TYPE_INTEGER, .token = "EF", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "attributes", .type = REDISMODULE_ARG_TYPE_STRING, .token = "SETATTR", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "numlinks", .type = REDISMODULE_ARG_TYPE_INTEGER, .token = "M", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = NULL } }; RedisModuleCommandInfo vadd_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Add one or more elements to a vector set, or update its vector if it already exists", .since = "8.0.0", .arity = -5, .args = vadd_args, }; if (RedisModule_SetCommandInfo(vadd_cmd, &vadd_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VREM if (RedisModule_CreateCommand(ctx,"VREM", VREM_RedisCommand,"write",1,1,1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vrem_cmd = RedisModule_GetCommand(ctx, "VREM"); if (vrem_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vrem_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = NULL } }; RedisModuleCommandInfo vrem_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Remove an element from a vector set", .since = "8.0.0", .arity = 3, .args = vrem_args, }; if (RedisModule_SetCommandInfo(vrem_cmd, &vrem_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VSIM if (RedisModule_CreateCommand(ctx,"VSIM", VSIM_RedisCommand,"readonly",1,1,1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vsim_cmd = RedisModule_GetCommand(ctx, "VSIM"); if (vsim_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vsim_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "format", .type = REDISMODULE_ARG_TYPE_ONEOF, .subargs = (RedisModuleCommandArg[]) { { .name = "ele", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "ELE" }, { .name = "fp32", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "FP32" }, { .name = "values", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "VALUES" }, { .name = NULL } } }, { .name = "vector_or_element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "withscores", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "WITHSCORES", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "withattribs", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "WITHATTRIBS", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "count", .type = REDISMODULE_ARG_TYPE_INTEGER, .token = "COUNT", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "max_distance", .type = REDISMODULE_ARG_TYPE_DOUBLE, .token = "EPSILON", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "search-exploration-factor", .type = REDISMODULE_ARG_TYPE_INTEGER, .token = "EF", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "expression", .type = REDISMODULE_ARG_TYPE_STRING, .token = "FILTER", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "max-filtering-effort", .type = REDISMODULE_ARG_TYPE_INTEGER, .token = "FILTER-EF", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "truth", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "TRUTH", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = "nothread", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "NOTHREAD", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = NULL } }; RedisModuleCommandInfo vsim_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return elements by vector similarity", .since = "8.0.0", .arity = -4, .args = vsim_args, }; if (RedisModule_SetCommandInfo(vsim_cmd, &vsim_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VDIM if (RedisModule_CreateCommand(ctx, "VDIM", VDIM_RedisCommand, "readonly fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vdim_cmd = RedisModule_GetCommand(ctx, "VDIM"); if (vdim_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vdim_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = NULL } }; RedisModuleCommandInfo vdim_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return the dimension of vectors in the vector set", .since = "8.0.0", .arity = 2, .args = vdim_args, }; if (RedisModule_SetCommandInfo(vdim_cmd, &vdim_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VCARD if (RedisModule_CreateCommand(ctx, "VCARD", VCARD_RedisCommand, "readonly fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vcard_cmd = RedisModule_GetCommand(ctx, "VCARD"); if (vcard_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vcard_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = NULL } }; RedisModuleCommandInfo vcard_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return the number of elements in a vector set", .since = "8.0.0", .arity = 2, .args = vcard_args, }; if (RedisModule_SetCommandInfo(vcard_cmd, &vcard_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VEMB if (RedisModule_CreateCommand(ctx, "VEMB", VEMB_RedisCommand, "readonly fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vemb_cmd = RedisModule_GetCommand(ctx, "VEMB"); if (vemb_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vemb_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "raw", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "RAW", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = NULL } }; RedisModuleCommandInfo vemb_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return the vector associated with an element", .since = "8.0.0", .arity = -3, .args = vemb_args, }; if (RedisModule_SetCommandInfo(vemb_cmd, &vemb_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VLINKS if (RedisModule_CreateCommand(ctx, "VLINKS", VLINKS_RedisCommand, "readonly fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vlinks_cmd = RedisModule_GetCommand(ctx, "VLINKS"); if (vlinks_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vlinks_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "withscores", .type = REDISMODULE_ARG_TYPE_PURE_TOKEN, .token = "WITHSCORES", .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = NULL } }; RedisModuleCommandInfo vlinks_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return the neighbors of an element at each layer in the HNSW graph", .since = "8.0.0", .arity = -3, .args = vlinks_args, }; if (RedisModule_SetCommandInfo(vlinks_cmd, &vlinks_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VINFO if (RedisModule_CreateCommand(ctx, "VINFO", VINFO_RedisCommand, "readonly fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vinfo_cmd = RedisModule_GetCommand(ctx, "VINFO"); if (vinfo_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vinfo_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = NULL } }; RedisModuleCommandInfo vinfo_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return information about a vector set", .since = "8.0.0", .arity = 2, .args = vinfo_args, }; if (RedisModule_SetCommandInfo(vinfo_cmd, &vinfo_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VSETATTR if (RedisModule_CreateCommand(ctx, "VSETATTR", VSETATTR_RedisCommand, "write fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vsetattr_cmd = RedisModule_GetCommand(ctx, "VSETATTR"); if (vsetattr_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vsetattr_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "json", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = NULL } }; RedisModuleCommandInfo vsetattr_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Associate or remove the JSON attributes of elements", .since = "8.0.0", .arity = 4, .args = vsetattr_args, }; if (RedisModule_SetCommandInfo(vsetattr_cmd, &vsetattr_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VGETATTR if (RedisModule_CreateCommand(ctx, "VGETATTR", VGETATTR_RedisCommand, "readonly fast", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vgetattr_cmd = RedisModule_GetCommand(ctx, "VGETATTR"); if (vgetattr_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vgetattr_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = NULL } }; RedisModuleCommandInfo vgetattr_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Retrieve the JSON attributes of elements", .since = "8.0.0", .arity = 3, .args = vgetattr_args, }; if (RedisModule_SetCommandInfo(vgetattr_cmd, &vgetattr_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VRANDMEMBER if (RedisModule_CreateCommand(ctx, "VRANDMEMBER", VRANDMEMBER_RedisCommand, "readonly", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vrandmember_cmd = RedisModule_GetCommand(ctx, "VRANDMEMBER"); if (vrandmember_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vrandmember_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "count", .type = REDISMODULE_ARG_TYPE_INTEGER, .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = NULL } }; RedisModuleCommandInfo vrandmember_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return one or multiple random members from a vector set", .since = "8.0.0", .arity = -2, .args = vrandmember_args, }; if (RedisModule_SetCommandInfo(vrandmember_cmd, &vrandmember_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VISMEMBER if (RedisModule_CreateCommand(ctx, "VISMEMBER", VISMEMBER_RedisCommand, "readonly", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vismember_cmd = RedisModule_GetCommand(ctx, "VISMEMBER"); if (vismember_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vismember_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "element", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = NULL } }; RedisModuleCommandInfo vismember_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Check if an element exists in a vector set", .since = "8.2.0", .arity = 3, .args = vismember_args, }; if (RedisModule_SetCommandInfo(vismember_cmd, &vismember_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Register command VRANGE if (RedisModule_CreateCommand(ctx, "VRANGE", VRANGE_RedisCommand, "readonly", 1, 1, 1) == REDISMODULE_ERR) return REDISMODULE_ERR; RedisModuleCommand *vrange_cmd = RedisModule_GetCommand(ctx, "VRANGE"); if (vrange_cmd == NULL) return REDISMODULE_ERR; RedisModuleCommandArg vrange_args[] = { { .name = "key", .type = REDISMODULE_ARG_TYPE_KEY, .key_spec_index = 0 }, { .name = "start", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "end", .type = REDISMODULE_ARG_TYPE_STRING }, { .name = "count", .type = REDISMODULE_ARG_TYPE_INTEGER, .flags = REDISMODULE_CMD_ARG_OPTIONAL }, { .name = NULL } }; RedisModuleCommandInfo vrange_info = { .version = REDISMODULE_COMMAND_INFO_VERSION, .summary = "Return vector set elements in a lex range", .since = "8.4.0", .arity = -4, .args = vrange_args, }; if (RedisModule_SetCommandInfo(vrange_cmd, &vrange_info) == REDISMODULE_ERR) return REDISMODULE_ERR; // Set the allocator for the HNSW library, so that memory tracking // is correct in Redis. hnsw_set_allocator(RedisModule_Free, RedisModule_Alloc, RedisModule_Realloc); return REDISMODULE_OK; } int VectorSets_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { return RedisModule_OnLoad(ctx, argv, argc); }