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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
|
from test import TestCase, fill_redis_with_vectors, generate_random_vector
import random
class HNSWPersistence(TestCase):
def getname(self):
return "HNSW Persistence"
def estimated_runtime(self):
return 30
def _verify_results(self, key, dim, query_vec, reduced_dim=None):
"""Run a query and return results dict"""
k = 10
args = ['VSIM', key]
if reduced_dim:
args.extend(['VALUES', dim])
args.extend([str(x) for x in query_vec])
else:
args.extend(['VALUES', dim])
args.extend([str(x) for x in query_vec])
args.extend(['COUNT', k, 'WITHSCORES'])
results = self.redis.execute_command(*args)
results_dict = {}
for i in range(0, len(results), 2):
key = results[i].decode()
score = float(results[i+1])
results_dict[key] = score
return results_dict
def test(self):
# Setup dimensions
dim = 128
reduced_dim = 32
count = 5000
random.seed(42)
# Create two datasets - one normal and one with dimension reduction
normal_data = fill_redis_with_vectors(self.redis, f"{self.test_key}:normal", count, dim)
projected_data = fill_redis_with_vectors(self.redis, f"{self.test_key}:projected",
count, dim, reduced_dim)
# Generate query vectors we'll use before and after reload
query_vec_normal = generate_random_vector(dim)
query_vec_projected = generate_random_vector(dim)
# Get initial results for both sets
initial_normal = self._verify_results(f"{self.test_key}:normal",
dim, query_vec_normal)
initial_projected = self._verify_results(f"{self.test_key}:projected",
dim, query_vec_projected, reduced_dim)
# Force Redis to save and reload the dataset
self.redis.execute_command('DEBUG', 'RELOAD')
# Verify results after reload
reloaded_normal = self._verify_results(f"{self.test_key}:normal",
dim, query_vec_normal)
reloaded_projected = self._verify_results(f"{self.test_key}:projected",
dim, query_vec_projected, reduced_dim)
# Verify normal vectors results
assert len(initial_normal) == len(reloaded_normal), \
"Normal vectors: Result count mismatch before/after reload"
for key in initial_normal:
assert key in reloaded_normal, f"Normal vectors: Missing item after reload: {key}"
assert abs(initial_normal[key] - reloaded_normal[key]) < 0.0001, \
f"Normal vectors: Score mismatch for {key}: " + \
f"before={initial_normal[key]:.6f}, after={reloaded_normal[key]:.6f}"
# Verify projected vectors results
assert len(initial_projected) == len(reloaded_projected), \
"Projected vectors: Result count mismatch before/after reload"
for key in initial_projected:
assert key in reloaded_projected, \
f"Projected vectors: Missing item after reload: {key}"
assert abs(initial_projected[key] - reloaded_projected[key]) < 0.0001, \
f"Projected vectors: Score mismatch for {key}: " + \
f"before={initial_projected[key]:.6f}, after={reloaded_projected[key]:.6f}"
|
