diff options
Diffstat (limited to 'examples/redis-unstable/modules/vector-sets/tests/dimension_validation.py')
| -rw-r--r-- | examples/redis-unstable/modules/vector-sets/tests/dimension_validation.py | 67 |
1 files changed, 67 insertions, 0 deletions
diff --git a/examples/redis-unstable/modules/vector-sets/tests/dimension_validation.py b/examples/redis-unstable/modules/vector-sets/tests/dimension_validation.py new file mode 100644 index 0000000..f081152 --- /dev/null +++ b/examples/redis-unstable/modules/vector-sets/tests/dimension_validation.py | |||
| @@ -0,0 +1,67 @@ | |||
| 1 | from test import TestCase, generate_random_vector | ||
| 2 | import struct | ||
| 3 | import redis.exceptions | ||
| 4 | |||
| 5 | class DimensionValidation(TestCase): | ||
| 6 | def getname(self): | ||
| 7 | return "[regression] Dimension Validation with Projection" | ||
| 8 | |||
| 9 | def estimated_runtime(self): | ||
| 10 | return 0.5 | ||
| 11 | |||
| 12 | def test(self): | ||
| 13 | # Test scenario 1: Create a set with projection | ||
| 14 | original_dim = 100 | ||
| 15 | reduced_dim = 50 | ||
| 16 | |||
| 17 | # Create the initial vector and set with projection | ||
| 18 | vec1 = generate_random_vector(original_dim) | ||
| 19 | vec1_bytes = struct.pack(f'{original_dim}f', *vec1) | ||
| 20 | |||
| 21 | # Add first vector with projection | ||
| 22 | result = self.redis.execute_command('VADD', self.test_key, | ||
| 23 | 'REDUCE', reduced_dim, | ||
| 24 | 'FP32', vec1_bytes, f'{self.test_key}:item:1') | ||
| 25 | assert result == 1, "First VADD with REDUCE should return 1" | ||
| 26 | |||
| 27 | # Check VINFO returns the correct projection information | ||
| 28 | info = self.redis.execute_command('VINFO', self.test_key) | ||
| 29 | info_map = {k.decode('utf-8'): v for k, v in zip(info[::2], info[1::2])} | ||
| 30 | assert 'vector-dim' in info_map, "VINFO should contain vector-dim" | ||
| 31 | assert info_map['vector-dim'] == reduced_dim, f"Expected reduced dimension {reduced_dim}, got {info['vector-dim']}" | ||
| 32 | assert 'projection-input-dim' in info_map, "VINFO should contain projection-input-dim" | ||
| 33 | assert info_map['projection-input-dim'] == original_dim, f"Expected original dimension {original_dim}, got {info['projection-input-dim']}" | ||
| 34 | |||
| 35 | # Test scenario 2: Try adding a mismatched vector - should fail | ||
| 36 | wrong_dim = 80 | ||
| 37 | wrong_vec = generate_random_vector(wrong_dim) | ||
| 38 | wrong_vec_bytes = struct.pack(f'{wrong_dim}f', *wrong_vec) | ||
| 39 | |||
| 40 | # This should fail with dimension mismatch error | ||
| 41 | try: | ||
| 42 | self.redis.execute_command('VADD', self.test_key, | ||
| 43 | 'REDUCE', reduced_dim, | ||
| 44 | 'FP32', wrong_vec_bytes, f'{self.test_key}:item:2') | ||
| 45 | assert False, "VADD with wrong dimension should fail" | ||
| 46 | except redis.exceptions.ResponseError as e: | ||
| 47 | assert "Input dimension mismatch for projection" in str(e), f"Expected dimension mismatch error, got: {e}" | ||
| 48 | |||
| 49 | # Test scenario 3: Add a correctly-sized vector | ||
| 50 | vec2 = generate_random_vector(original_dim) | ||
| 51 | vec2_bytes = struct.pack(f'{original_dim}f', *vec2) | ||
| 52 | |||
| 53 | # This should succeed | ||
| 54 | result = self.redis.execute_command('VADD', self.test_key, | ||
| 55 | 'REDUCE', reduced_dim, | ||
| 56 | 'FP32', vec2_bytes, f'{self.test_key}:item:3') | ||
| 57 | assert result == 1, "VADD with correct dimensions should succeed" | ||
| 58 | |||
| 59 | # Check VSIM also validates input dimensions | ||
| 60 | wrong_query = generate_random_vector(wrong_dim) | ||
| 61 | try: | ||
| 62 | self.redis.execute_command('VSIM', self.test_key, | ||
| 63 | 'VALUES', wrong_dim, *[str(x) for x in wrong_query], | ||
| 64 | 'COUNT', 10) | ||
| 65 | assert False, "VSIM with wrong dimension should fail" | ||
| 66 | except redis.exceptions.ResponseError as e: | ||
| 67 | assert "Input dimension mismatch for projection" in str(e), f"Expected dimension mismatch error in VSIM, got: {e}" | ||
