File: vemb.py

package info (click to toggle)
redis 5%3A8.0.2-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid, trixie
  • size: 22,304 kB
  • sloc: ansic: 216,903; tcl: 51,562; sh: 4,625; perl: 4,214; cpp: 3,568; python: 2,954; makefile: 2,055; ruby: 639; javascript: 30; csh: 7
file content (41 lines) | stat: -rw-r--r-- 1,650 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
from test import TestCase
import struct
import math

class VEMB(TestCase):
    def getname(self):
        return "VEMB Command"

    def test(self):
        dim = 4

        # Add same vector in both formats
        vec = [1, 0, 0, 0]
        norm = math.sqrt(sum(x*x for x in vec))
        vec = [x/norm for x in vec]  # Normalize the vector

        # Add using FP32
        vec_bytes = struct.pack(f'{dim}f', *vec)
        self.redis.execute_command('VADD', self.test_key, 'FP32', vec_bytes, f'{self.test_key}:item:1')

        # Add using VALUES
        self.redis.execute_command('VADD', self.test_key, 'VALUES', dim, 
                         *[str(x) for x in vec], f'{self.test_key}:item:2')

        # Get both back with VEMB
        result1 = self.redis.execute_command('VEMB', self.test_key, f'{self.test_key}:item:1')
        result2 = self.redis.execute_command('VEMB', self.test_key, f'{self.test_key}:item:2')

        retrieved_vec1 = [float(x) for x in result1]
        retrieved_vec2 = [float(x) for x in result2]

        # Compare both vectors with original (allow for small quantization errors)
        for i in range(dim):
            assert abs(vec[i] - retrieved_vec1[i]) < 0.01, \
                f"FP32 vector component {i} mismatch: expected {vec[i]}, got {retrieved_vec1[i]}"
            assert abs(vec[i] - retrieved_vec2[i]) < 0.01, \
                f"VALUES vector component {i} mismatch: expected {vec[i]}, got {retrieved_vec2[i]}"

        # Test non-existent item
        result = self.redis.execute_command('VEMB', self.test_key, 'nonexistent')
        assert result is None, "Non-existent item should return nil"