From 6c1e55d07ca41cd5b679fa0066160900876ee89e Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Wed, 5 Feb 2025 11:21:06 +0100
Subject: [PATCH] VEMB RAW implemented.

---
 README.md | 13 +++++++++++++
 vset.c    | 39 +++++++++++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 34134daa6..1a1a7992d 100644
--- a/README.md
+++ b/README.md
@@ -138,6 +138,19 @@ Because vector sets perform insertion time normalization and optional
 quantization, the returned vector could be approximated. `VEMB` will take
 care to de-quantized and de-normalize the vector before returning it.
 
+It is possible to ask VEMB to return raw data, that is, the interal representation used by the vector: fp32, int8, or a bitmap for binary quantization. This behavior is triggered by the `RAW` option of of VEMB:
+
+    VEMB word_embedding apple RAW
+
+In this case the return value of the command is an array of three or more elements:
+1. The name of the quantization used, that is one of: "fp32", "bin", "q8".
+2. The a string blob containing the raw data, 4 bytes fp32 floats for fp32, a bitmap for binary quants, or int8 bytes array for q8 quants.
+3. A float representing the l2 of the vector before normalization. You need to multiply by this vector if you want to de-normalize the value for any reason.
+
+For q8 quantization, an additional elements is also returned: the quantization
+range, so the integers from -127 to 127 represent (normalized) components
+in the range `-range`, `+range`.
+
 **VLINKS: introspection command that shows neighbors for a node**
 
     VLINKS key element [WITHSCORES]
diff --git a/vset.c b/vset.c
index 6a57f035c..6d9971c46 100644
--- a/vset.c
+++ b/vset.c
@@ -860,8 +860,19 @@ int VREM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
  * 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);
+    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];
@@ -885,16 +896,24 @@ int VEMB_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
         return RedisModule_ReplyWithNull(ctx);
     }
 
-    /* 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.
+    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 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;
 }