Skip to content

Commit 2157db1

Browse files
mgravellmgravell
committed
change VSIM API to allow future extensivility:
 - make VectorSetSimilaritySearchRequest abstract - remove Member and Vector - add ByMember and ByVector factory methods - (internal changes to support the above, in the message etc)
1 parent d53f1ab commit 2157db1

File tree

5 files changed

+198
-148
lines changed

5 files changed

+198
-148
lines changed

‎src/StackExchange.Redis/PublicAPI/PublicAPI.Unshipped.txt‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,10 @@
1010
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.FilterExpression.set -> void
1111
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.MaxFilteringEffort.get -> int?
1212
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.MaxFilteringEffort.set -> void
13-
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.Member.get -> StackExchange.Redis.RedisValue
14-
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.Member.set -> void
1513
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.SearchExplorationFactor.get -> int?
1614
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.SearchExplorationFactor.set -> void
1715
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.UseExactSearch.get -> bool
1816
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.UseExactSearch.set -> void
19-
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.Vector.get -> System.ReadOnlyMemory<float>
20-
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.Vector.set -> void
21-
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.VectorSetSimilaritySearchRequest() -> void
2217
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.WithAttributes.get -> bool
2318
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.WithAttributes.set -> void
2419
[SER001]StackExchange.Redis.VectorSetSimilaritySearchRequest.WithScores.get -> bool
@@ -79,3 +74,5 @@
7974
[SER001]StackExchange.Redis.VectorSetSimilaritySearchResult.Score.get -> double
8075
[SER001]StackExchange.Redis.VectorSetSimilaritySearchResult.VectorSetSimilaritySearchResult() -> void
8176
[SER001]StackExchange.Redis.VectorSetSimilaritySearchResult.VectorSetSimilaritySearchResult(StackExchange.Redis.RedisValue member, double score = NaN, string? attributesJson = null) -> void
77+
[SER001]static StackExchange.Redis.VectorSetSimilaritySearchRequest.ByMember(StackExchange.Redis.RedisValue member) -> StackExchange.Redis.VectorSetSimilaritySearchRequest!
78+
[SER001]static StackExchange.Redis.VectorSetSimilaritySearchRequest.ByVector(System.ReadOnlyMemory<float> vector) -> StackExchange.Redis.VectorSetSimilaritySearchRequest!

‎src/StackExchange.Redis/VectorSetSimilaritySearchMessage.cs‎

Lines changed: 79 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -2,20 +2,85 @@
22

33
namespace StackExchange.Redis;
44

5-
internal sealed class VectorSetSimilaritySearchMessage(
5+
internal abstract class VectorSetSimilaritySearchMessage(
66
int db,
77
CommandFlags flags,
88
VectorSetSimilaritySearchMessage.VsimFlags vsimFlags,
99
RedisKey key,
10-
RedisValue member,
11-
ReadOnlyMemory<float> vector,
1210
int count,
1311
double epsilon,
1412
int searchExplorationFactor,
1513
string? filterExpression,
1614
int maxFilteringEffort) : Message(db, flags, RedisCommand.VSIM)
1715
{
18-
public ResultProcessor<Lease<VectorSetSimilaritySearchResult>?> GetResultProcessor() => VectorSetSimilaritySearchProcessor.Instance;
16+
// For "FP32" and "VALUES" scenarios; in the future we might want other vector sizes / encodings - for
17+
// example, there could be some "FP16" or "FP8" transport that requires a ROM-short or ROM-sbyte from
18+
// the calling code. Or, as a convenience, we might want to allow ROM-double input, but transcode that
19+
// to FP32 on the way out.
20+
internal sealed class VectorSetSimilaritySearchBySingleVectorMessage(
21+
int db,
22+
CommandFlags flags,
23+
VsimFlags vsimFlags,
24+
RedisKey key,
25+
ReadOnlyMemory<float> vector,
26+
int count,
27+
double epsilon,
28+
int searchExplorationFactor,
29+
string? filterExpression,
30+
int maxFilteringEffort) : VectorSetSimilaritySearchMessage(db, flags, vsimFlags, key, count, epsilon,
31+
searchExplorationFactor, filterExpression, maxFilteringEffort)
32+
{
33+
internal override int GetSearchTargetArgCount(bool packed) =>
34+
packed ? 2 : 2 + vector.Length; // FP32 {vector} or VALUES {num} {vector}
35+
36+
internal override void WriteSearchTarget(bool packed, PhysicalConnection physical)
37+
{
38+
if (packed)
39+
{
40+
physical.WriteBulkString("FP32"u8);
41+
physical.WriteBulkString(System.Runtime.InteropServices.MemoryMarshal.AsBytes(vector.Span));
42+
}
43+
else
44+
{
45+
physical.WriteBulkString("VALUES"u8);
46+
physical.WriteBulkString(vector.Length);
47+
foreach (var val in vector.Span)
48+
{
49+
physical.WriteBulkString(val);
50+
}
51+
}
52+
}
53+
}
54+
55+
// for "ELE" scenarios
56+
internal sealed class VectorSetSimilaritySearchByMemberMessage(
57+
int db,
58+
CommandFlags flags,
59+
VsimFlags vsimFlags,
60+
RedisKey key,
61+
RedisValue member,
62+
int count,
63+
double epsilon,
64+
int searchExplorationFactor,
65+
string? filterExpression,
66+
int maxFilteringEffort) : VectorSetSimilaritySearchMessage(db, flags, vsimFlags, key, count, epsilon,
67+
searchExplorationFactor, filterExpression, maxFilteringEffort)
68+
{
69+
internal override int GetSearchTargetArgCount(bool packed) => 2; // ELE {member}
70+
71+
internal override void WriteSearchTarget(bool packed, PhysicalConnection physical)
72+
{
73+
physical.WriteBulkString("ELE"u8);
74+
physical.WriteBulkString(member);
75+
}
76+
}
77+
78+
internal abstract int GetSearchTargetArgCount(bool packed);
79+
internal abstract void WriteSearchTarget(bool packed, PhysicalConnection physical);
80+
81+
public ResultProcessor<Lease<VectorSetSimilaritySearchResult>?> GetResultProcessor() =>
82+
VectorSetSimilaritySearchProcessor.Instance;
83+
1984
private sealed class VectorSetSimilaritySearchProcessor : ResultProcessor<Lease<VectorSetSimilaritySearchResult>?>
2085
{
2186
// keep local, since we need to know what flags were being sent
@@ -38,7 +103,8 @@ protected override bool SetResultCore(PhysicalConnection connection, Message mes
38103
// in RESP3 mode (only), when both are requested, we get a sub-array per item; weird, but true
39104
bool internalNesting = withScores && withAttribs && connection.Protocol is RedisProtocol.Resp3;
40105

41-
int rowsPerItem = internalNesting ? 2
106+
int rowsPerItem = internalNesting
107+
? 2
42108
: 1 + ((withScores ? 1 : 0) + (withAttribs ? 1 : 0)); // each value is separate root element
43109

44110
var items = result.GetItems();
@@ -80,11 +146,13 @@ protected override bool SetResultCore(PhysicalConnection connection, Message mes
80146
target[i] = new VectorSetSimilaritySearchResult(member, score, attributesJson);
81147
count++;
82148
}
149+
83150
if (count == target.Length)
84151
{
85152
SetResult(message, lease);
86153
return true;
87154
}
155+
88156
lease.Dispose(); // failed to fill?
89157
}
90158

@@ -111,14 +179,9 @@ internal enum VsimFlags
111179

112180
public override int ArgCount => GetArgCount(VectorSetAddMessage.UseFp32);
113181

114-
private int GetArgCount(bool useFp32)
182+
private int GetArgCount(bool packed)
115183
{
116-
int argCount = 3; // {key} and "ELE {member}", "FP32 {vector}" or "VALUES {num}"
117-
if (member.IsNull && !useFp32)
118-
{
119-
argCount += vector.Length; // {vector} in the VALUES case
120-
}
121-
184+
int argCount = 1 + GetSearchTargetArgCount(packed); // {key} and whatever we need for the vector/element portion
122185
if (HasFlag(VsimFlags.WithScores)) argCount++; // [WITHSCORES]
123186
if (HasFlag(VsimFlags.WithAttributes)) argCount++; // [WITHATTRIBS]
124187
if (HasFlag(VsimFlags.Count)) argCount += 2; // [COUNT {count}]
@@ -133,37 +196,15 @@ private int GetArgCount(bool useFp32)
133196

134197
protected override void WriteImpl(PhysicalConnection physical)
135198
{
136-
var useFp32 = VectorSetAddMessage.UseFp32; // avoid race in debug mode
137-
physical.WriteHeader(Command, GetArgCount(useFp32));
199+
// snapshot to avoid race in debug scenarios
200+
bool packed = VectorSetAddMessage.UseFp32;
201+
physical.WriteHeader(Command, GetArgCount(packed));
138202

139203
// Write key
140204
physical.Write(key);
141205

142206
// Write search target: either "ELE {member}" or vector data
143-
if (!member.IsNull)
144-
{
145-
// Member-based search: "ELE {member}"
146-
physical.WriteBulkString("ELE"u8);
147-
physical.WriteBulkString(member);
148-
}
149-
else
150-
{
151-
// Vector-based search: either "FP32 {vector}" or "VALUES {num} {vector}"
152-
if (useFp32)
153-
{
154-
physical.WriteBulkString("FP32"u8);
155-
physical.WriteBulkString(System.Runtime.InteropServices.MemoryMarshal.AsBytes(vector.Span));
156-
}
157-
else
158-
{
159-
physical.WriteBulkString("VALUES"u8);
160-
physical.WriteBulkString(vector.Length);
161-
foreach (var val in vector.Span)
162-
{
163-
physical.WriteBulkString(val);
164-
}
165-
}
166-
}
207+
WriteSearchTarget(packed, physical);
167208

168209
if (HasFlag(VsimFlags.WithScores))
169210
{

‎src/StackExchange.Redis/VectorSetSimilaritySearchRequest.cs‎

Lines changed: 53 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,17 +9,65 @@ namespace StackExchange.Redis;
99
/// Represents the request for a vector similarity search operation.
1010
/// </summary>
1111
[Experimental(Experiments.VectorSets, UrlFormat = Experiments.UrlFormat)]
12-
public class VectorSetSimilaritySearchRequest
12+
public abstract class VectorSetSimilaritySearchRequest
1313
{
14+
internal VectorSetSimilaritySearchRequest()
15+
{
16+
} // polymorphism left open for future, but needs to be handled internally
17+
18+
private sealed class VectorSetSimilarityByMemberSearchRequest(RedisValue member) : VectorSetSimilaritySearchRequest
19+
{
20+
public RedisValue Member => member;
21+
22+
internal override VectorSetSimilaritySearchMessage ToMessage(RedisKey key, int db, CommandFlags flags)
23+
=> new VectorSetSimilaritySearchMessage.VectorSetSimilaritySearchByMemberMessage(
24+
db,
25+
flags,
26+
_vsimFlags,
27+
key,
28+
Member,
29+
_count,
30+
_epsilon,
31+
_searchExplorationFactor,
32+
_filterExpression,
33+
_maxFilteringEffort);
34+
}
35+
36+
private sealed class VectorSetSimilarityVectorSingleSearchRequest(ReadOnlyMemory<float> vector)
37+
: VectorSetSimilaritySearchRequest
38+
{
39+
public ReadOnlyMemory<float> Vector => vector;
40+
41+
internal override VectorSetSimilaritySearchMessage ToMessage(RedisKey key, int db, CommandFlags flags)
42+
=> new VectorSetSimilaritySearchMessage.VectorSetSimilaritySearchBySingleVectorMessage(
43+
db,
44+
flags,
45+
_vsimFlags,
46+
key,
47+
Vector,
48+
_count,
49+
_epsilon,
50+
_searchExplorationFactor,
51+
_filterExpression,
52+
_maxFilteringEffort);
53+
}
54+
55+
// snapshot the values; I don't trust people not to mutate the object behind my back
56+
internal abstract VectorSetSimilaritySearchMessage ToMessage(RedisKey key, int db, CommandFlags flags);
57+
1458
/// <summary>
15-
/// The query vector.
59+
/// Create a request to search by an existing member in the index.
1660
/// </summary>
17-
public ReadOnlyMemory<float> Vector { get; set; }
61+
/// <param name="member">The member to search for.</param>
62+
public static VectorSetSimilaritySearchRequest ByMember(RedisValue member)
63+
=> new VectorSetSimilarityByMemberSearchRequest(member);
1864

1965
/// <summary>
20-
/// The member to find similar vectors for.
66+
/// Create a request to search by a vector value.
2167
/// </summary>
22-
public RedisValue Member { get; set; }
68+
/// <param name="vector">The vector value to search for.</param>
69+
public static VectorSetSimilaritySearchRequest ByVector(ReadOnlyMemory<float> vector)
70+
=> new VectorSetSimilarityVectorSingleSearchRequest(vector);
2371

2472
private VsimFlags _vsimFlags;
2573

@@ -172,19 +220,4 @@ public bool DisableThreading
172220
get => HasFlag(VsimFlags.DisableThreading);
173221
set => SetFlag(VsimFlags.DisableThreading, value);
174222
}
175-
176-
// snapshot the values; I don't trust people not to mutate the object behind my back
177-
internal VectorSetSimilaritySearchMessage ToMessage(RedisKey key, int db, CommandFlags flags)
178-
=> new(
179-
db,
180-
flags,
181-
_vsimFlags,
182-
key,
183-
Member,
184-
Vector,
185-
_count,
186-
_epsilon,
187-
_searchExplorationFactor,
188-
_filterExpression,
189-
_maxFilteringEffort);
190223
}

‎tests/StackExchange.Redis.Tests/KeyPrefixedVectorSetTests.cs‎

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -171,10 +171,7 @@ public void VectorSetSimilaritySearchByVector()
171171
{
172172
var vector = new[] { 1.0f, 2.0f, 3.0f }.AsMemory();
173173

174-
var query = new VectorSetSimilaritySearchRequest
175-
{
176-
Vector = vector,
177-
};
174+
var query = VectorSetSimilaritySearchRequest.ByVector(vector);
178175
prefixed.VectorSetSimilaritySearch(
179176
"vectorset",
180177
query);
@@ -186,19 +183,16 @@ public void VectorSetSimilaritySearchByVector()
186183
[Fact]
187184
public void VectorSetSimilaritySearchByMember()
188185
{
189-
var query = new VectorSetSimilaritySearchRequest
190-
{
191-
Member = "member1",
192-
Count = 5,
193-
WithScores = true,
194-
WithAttributes = true,
195-
Epsilon = 0.1,
196-
SearchExplorationFactor = 400,
197-
FilterExpression = "category='test'",
198-
MaxFilteringEffort = 1000,
199-
UseExactSearch = true,
200-
DisableThreading = true,
201-
};
186+
var query = VectorSetSimilaritySearchRequest.ByMember("member1");
187+
query.Count = 5;
188+
query.WithScores = true;
189+
query.WithAttributes = true;
190+
query.Epsilon = 0.1;
191+
query.SearchExplorationFactor = 400;
192+
query.FilterExpression = "category='test'";
193+
query.MaxFilteringEffort = 1000;
194+
query.UseExactSearch = true;
195+
query.DisableThreading = true;
202196
prefixed.VectorSetSimilaritySearch(
203197
"vectorset",
204198
query,
@@ -215,10 +209,7 @@ public void VectorSetSimilaritySearchByVector_DefaultParameters()
215209
var vector = new[] { 1.0f, 2.0f }.AsMemory();
216210

217211
// Test that default parameters work correctly
218-
var query = new VectorSetSimilaritySearchRequest
219-
{
220-
Vector = vector,
221-
};
212+
var query = VectorSetSimilaritySearchRequest.ByVector(vector);
222213
prefixed.VectorSetSimilaritySearch("vectorset", query);
223214
mock.Received().VectorSetSimilaritySearch(
224215
"prefix:vectorset",

0 commit comments

Comments
 (0)