add RAW variant for VEMB -> VEMB_RAW

Author: nkaradzhov

packages/client/lib/commands/VEMB_RAW.spec.ts

@@ -0,0 +1,66 @@
+import { strict as assert } from 'node:assert';
+import testUtils, { GLOBAL } from '../test-utils';
+import VEMB_RAW from './VEMB_RAW';
+import { parseArgs } from './generic-transformers';
+
+describe('VEMB_RAW', () => {
+  it('transformArguments', () => {
+    assert.deepEqual(
+      parseArgs(VEMB_RAW, 'key', 'element'),
+      ['VEMB', 'key', 'element', 'RAW']
+    );
+  });
+
+  testUtils.testAll('vEmbRaw', async client => {
+    await client.vAdd('key1', [1.0, 2.0, 3.0], 'element');
+    const result1 = await client.vEmbRaw('key1', 'element');
+    assert.equal(result1.quantization, 'int8');
+    assert.ok(result1.quantizationRange !== undefined);
+
+    await client.vAdd('key2', [1.0, 2.0, 3.0], 'element', { QUANT: 'Q8' });
+    const result2 = await client.vEmbRaw('key2', 'element');
+    assert.equal(result2.quantization, 'int8');
+    assert.ok(result2.quantizationRange !== undefined);
+
+    await client.vAdd('key3', [1.0, 2.0, 3.0], 'element', { QUANT: 'NOQUANT' });
+    const result3 = await client.vEmbRaw('key3', 'element');
+    assert.equal(result3.quantization, 'f32');
+    assert.equal(result3.quantizationRange, undefined);
+
+    await client.vAdd('key4', [1.0, 2.0, 3.0], 'element', { QUANT: 'BIN' });
+    const result4 = await client.vEmbRaw('key4', 'element');
+    assert.equal(result4.quantization, 'bin');
+    assert.equal(result4.quantizationRange, undefined);
+  }, {
+    client: { ...GLOBAL.SERVERS.OPEN, minimumDockerVersion: [8, 0] },
+    cluster: { ...GLOBAL.CLUSTERS.OPEN, minimumDockerVersion: [8, 0] }
+  });
+
+  testUtils.testWithClient('vEmbRaw with RESP3', async client => {
+    await client.vAdd('key1', [1.0, 2.0, 3.0], 'element');
+    const result1 = await client.vEmbRaw('key1', 'element');
+    assert.equal(result1.quantization, 'int8');
+    assert.ok(result1.quantizationRange !== undefined);
+
+    await client.vAdd('key2', [1.0, 2.0, 3.0], 'element', { QUANT: 'Q8' });
+    const result2 = await client.vEmbRaw('key2', 'element');
+    assert.equal(result2.quantization, 'int8');
+    assert.ok(result2.quantizationRange !== undefined);
+
+    await client.vAdd('key3', [1.0, 2.0, 3.0], 'element', { QUANT: 'NOQUANT' });
+    const result3 = await client.vEmbRaw('key3', 'element');
+    assert.equal(result3.quantization, 'f32');
+    assert.equal(result3.quantizationRange, undefined);
+
+    await client.vAdd('key4', [1.0, 2.0, 3.0], 'element', { QUANT: 'BIN' });
+    const result4 = await client.vEmbRaw('key4', 'element');
+    assert.equal(result4.quantization, 'bin');
+    assert.equal(result4.quantizationRange, undefined);
+  }, {
+    ...GLOBAL.SERVERS.OPEN,
+    clientOptions: {
+      RESP: 3
+    },
+    minimumDockerVersion: [8, 0]
+  });
+});

packages/client/lib/commands/VEMB_RAW.ts

@@ -0,0 +1,57 @@
+import { CommandParser } from '../client/parser';
+import {
+  RedisArgument,
+  Command,
+  BlobStringReply,
+  SimpleStringReply,
+  DoubleReply
+} from '../RESP/types';
+import { transformDoubleReply } from './generic-transformers';
+import VEMB from './VEMB';
+
+type RawVembReply = {
+  quantization: SimpleStringReply;
+  raw: BlobStringReply;
+  l2Norm: DoubleReply;
+  quantizationRange?: DoubleReply;
+};
+
+const transformRawVembReply = {
+  2: (reply: any[]): RawVembReply => {
+    return {
+      quantization: reply[0],
+      raw: reply[1],
+      l2Norm: transformDoubleReply[2](reply[2]),
+      ...(reply[3] !== undefined && { quantizationRange: transformDoubleReply[2](reply[3]) })
+    };
+  },
+  3: (reply: any[]): RawVembReply => {
+    return {
+      quantization: reply[0],
+      raw: reply[1],
+      l2Norm: reply[2],
+      quantizationRange: reply[3]
+    };
+  },
+};
+
+export default {
+  IS_READ_ONLY: true,
+  /**
+   * Retrieve the RAW approximate vector associated with a vector set element
+   *
+   * @param parser - The command parser
+   * @param key - The key of the vector set
+   * @param element - The name of the element to retrieve the vector for
+   * @see https://redis.io/commands/vemb/
+   */
+  parseCommand(
+    parser: CommandParser,
+    key: RedisArgument,
+    element: RedisArgument
+  ) {
+    VEMB.parseCommand(parser, key, element);
+    parser.push('RAW');
+  },
+  transformReply: transformRawVembReply
+} as const satisfies Command;

packages/client/lib/commands/index.ts

@@ -348,6 +348,7 @@ import VADD from './VADD';
 import VCARD from './VCARD';
 import VDIM from './VDIM';
 import VEMB from './VEMB';
+import VEMB_RAW from './VEMB_RAW';
 import VGETATTR from './VGETATTR';
 import VINFO from './VINFO';
 import VLINKS from './VLINKS';
@@ -1059,6 +1060,8 @@ export default {
   vDim: VDIM,
   VEMB,
   vEmb: VEMB,
+  VEMB_RAW,
+  vEmbRaw: VEMB_RAW,
   VGETATTR,
   vGetAttr: VGETATTR,
   VINFO,