Skip to content

Commit 98fab2b

Browse files
mgravellmgravell
committed
better support for vector data
1 parent 752c65f commit 98fab2b

File tree

5 files changed

+71
-47
lines changed

5 files changed

+71
-47
lines changed

‎src/NRedisStack/PublicAPI/PublicAPI.Unshipped.txt‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,28 +38,27 @@
3838
[NRS001]NRedisStack.Search.HybridSearchQuery.SortBy(string! field) -> NRedisStack.Search.HybridSearchQuery!
3939
[NRS001]NRedisStack.Search.HybridSearchQuery.Timeout(bool timeout = true) -> NRedisStack.Search.HybridSearchQuery!
4040
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearch(NRedisStack.Search.HybridSearchQuery.VectorSearchConfig config) -> NRedisStack.Search.HybridSearchQuery!
41-
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearch(string! fieldName, NRedisStack.Search.VectorData vectorData) -> NRedisStack.Search.HybridSearchQuery!
41+
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearch(string! fieldName, NRedisStack.Search.VectorData! vectorData) -> NRedisStack.Search.HybridSearchQuery!
4242
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
4343
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.FieldName.get -> string!
4444
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.Filter.get -> string?
4545
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.Method.get -> NRedisStack.Search.VectorSearchMethod?
4646
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.ScoreAlias.get -> string?
47-
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.VectorData.get -> NRedisStack.Search.VectorData
47+
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.VectorData.get -> NRedisStack.Search.VectorData?
4848
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.VectorSearchConfig() -> void
49-
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.VectorSearchConfig(string! fieldName, NRedisStack.Search.VectorData vectorData, NRedisStack.Search.VectorSearchMethod? method = null, string? filter = null, string? scoreAlias = null) -> void
49+
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.VectorSearchConfig(string! fieldName, NRedisStack.Search.VectorData! vectorData, NRedisStack.Search.VectorSearchMethod? method = null, string? filter = null, string? scoreAlias = null) -> void
5050
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.WithFieldName(string! fieldName) -> NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
5151
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.WithFilter(string? filter) -> NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
5252
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.WithMethod(NRedisStack.Search.VectorSearchMethod? method) -> NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
5353
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.WithScoreAlias(string? scoreAlias) -> NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
54-
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.WithVectorData(NRedisStack.Search.VectorData vectorData) -> NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
54+
[NRS001]NRedisStack.Search.HybridSearchQuery.VectorSearchConfig.WithVectorData(NRedisStack.Search.VectorData! vectorData) -> NRedisStack.Search.HybridSearchQuery.VectorSearchConfig
5555
[NRS001]NRedisStack.Search.HybridSearchResult
5656
[NRS001]NRedisStack.Search.HybridSearchResult.ExecutionTime.get -> System.TimeSpan
5757
[NRS001]NRedisStack.Search.HybridSearchResult.Results.get -> System.Collections.Generic.IReadOnlyDictionary<string!, object!>![]!
5858
[NRS001]NRedisStack.Search.HybridSearchResult.TotalResults.get -> long
5959
[NRS001]NRedisStack.Search.HybridSearchResult.Warnings.get -> string![]!
6060
[NRS001]NRedisStack.Search.Scorer
6161
[NRS001]NRedisStack.Search.VectorData
62-
[NRS001]NRedisStack.Search.VectorData.VectorData() -> void
6362
[NRS001]NRedisStack.Search.VectorSearchMethod
6463
[NRS001]NRedisStack.SearchCommands.HybridSearch(string! indexName, NRedisStack.Search.HybridSearchQuery! query) -> NRedisStack.Search.HybridSearchResult!
6564
[NRS001]NRedisStack.SearchCommandsAsync.HybridSearchAsync(string! indexName, NRedisStack.Search.HybridSearchQuery! query) -> System.Threading.Tasks.Task<NRedisStack.Search.HybridSearchResult!>!
@@ -81,7 +80,9 @@
8180
[NRS001]static NRedisStack.Search.Scorer.Hamming.get -> NRedisStack.Search.Scorer!
8281
[NRS001]static NRedisStack.Search.Scorer.TfIdf.get -> NRedisStack.Search.Scorer!
8382
[NRS001]static NRedisStack.Search.Scorer.TfIdfDocNorm.get -> NRedisStack.Search.Scorer!
84-
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData(byte[]! data) -> NRedisStack.Search.VectorData
85-
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData(System.ReadOnlyMemory<byte> data) -> NRedisStack.Search.VectorData
83+
[NRS001]static NRedisStack.Search.VectorData.Create(System.ReadOnlyMemory<float> vector) -> NRedisStack.Search.VectorData!
84+
[NRS001]static NRedisStack.Search.VectorData.FromBase64(string! base64) -> NRedisStack.Search.VectorData!
85+
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData!(float[]! data) -> NRedisStack.Search.VectorData!
86+
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData!(System.ReadOnlyMemory<float> vector) -> NRedisStack.Search.VectorData!
8687
[NRS001]static NRedisStack.Search.VectorSearchMethod.NearestNeighbour(int count = 10, int? maxTopCandidates = null, string? distanceAlias = null) -> NRedisStack.Search.VectorSearchMethod!
8788
[NRS001]static NRedisStack.Search.VectorSearchMethod.Range(double radius, double? epsilon = null, string? distanceAlias = null) -> NRedisStack.Search.VectorSearchMethod!

‎src/NRedisStack/Search/HybridSearchQuery.VectorSearchConfig.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ public sealed partial class HybridSearchQuery
99
{
1010
public readonly struct VectorSearchConfig(string fieldName, VectorData vectorData, VectorSearchMethod? method = null, string? filter = null, string? scoreAlias = null)
1111
{
12-
internal bool HasValue => _vectorData.HasValue & _fieldName is not null;
12+
internal bool HasValue => _vectorData is not null & _fieldName is not null;
1313

1414
private readonly string _fieldName = fieldName;
1515
private readonly VectorData _vectorData = vectorData;
@@ -40,7 +40,7 @@ public readonly struct VectorSearchConfig(string fieldName, VectorData vectorDat
4040
/// <summary>
4141
/// The vector data to search for.
4242
/// </summary>
43-
public VectorData VectorData => _vectorData;
43+
public VectorData? VectorData => _vectorData;
4444

4545
/// <summary>
4646
/// Specify the vector search method.
@@ -97,7 +97,7 @@ internal int GetOwnArgsCount()
9797
int count = 0;
9898
if (HasValue)
9999
{
100-
count += 2 + _vectorData.GetOwnArgsCount();
100+
count += 2 + _vectorData.ArgsCount();
101101
if (_method != null) count += _method.GetOwnArgsCount();
102102
if (_filter != null) count += 2;
103103

@@ -112,7 +112,7 @@ internal void AddOwnArgs(List<object> args)
112112
{
113113
args.Add("VSIM");
114114
args.Add(_fieldName);
115-
_vectorData.AddOwnArgs(args);
115+
_vectorData.AddArgs(args);
116116

117117
_method?.AddOwnArgs(args);
118118
if (_filter != null)

‎src/NRedisStack/Search/VectorData.cs‎

Lines changed: 45 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -5,36 +5,58 @@
55
namespace NRedisStack.Search;
66

77
[Experimental(Experiments.Server_8_4, UrlFormat = Experiments.UrlFormat)]
8-
public readonly struct VectorData
8+
public abstract class VectorData
99
{
10-
// intended to allow future flexibility in how we express vectors
11-
private readonly ReadOnlyMemory<byte> _data;
12-
private VectorData(ReadOnlyMemory<byte> data)
10+
private protected VectorData()
1311
{
14-
_data = data;
1512
}
1613

17-
public static implicit operator VectorData(byte[] data) => new(data);
18-
public static implicit operator VectorData(ReadOnlyMemory<byte> data) => new(data);
19-
internal void AddOwnArgs(List<object> args)
14+
/// <summary>
15+
/// A vector of <see cref="Single"/> entries.
16+
/// </summary>
17+
public static VectorData Create(ReadOnlyMemory<float> vector) => new VectorDataSingle(vector);
18+
19+
/// <summary>
20+
/// A pre-formatted base-64 value.
21+
/// </summary>
22+
public static VectorData FromBase64(string base64) => new VectorDataBase64(base64);
23+
24+
/// <summary>
25+
/// A vector of <see cref="Single"/> entries.
26+
/// </summary>
27+
public static implicit operator VectorData(float[] data) => new VectorDataSingle(data);
28+
29+
/// <summary>
30+
/// A vector of <see cref="Single"/> entries.
31+
/// </summary>
32+
public static implicit operator VectorData(ReadOnlyMemory<float> vector) => new VectorDataSingle(vector);
33+
34+
internal virtual void AddArgs(List<object> args) => args.Add(ToString() ?? "");
35+
internal virtual int ArgsCount() => 1;
36+
37+
private sealed class VectorDataSingle(ReadOnlyMemory<float> vector) : VectorData
2038
{
39+
public override string ToString()
40+
{
41+
if (!BitConverter.IsLittleEndian) ThrowBigEndian(); // we could loop and reverse each, but...how to test?
42+
var bytes = MemoryMarshal.AsBytes(vector.Span);
2143
#if NET || NETSTANDARD2_1_OR_GREATER
22-
args.Add(Convert.ToBase64String(_data.Span));
44+
return Convert.ToBase64String(bytes);
2345
#else
24-
if (MemoryMarshal.TryGetArray(_data, out ArraySegment<byte> segment))
25-
{
26-
args.Add(Convert.ToBase64String(segment.Array!, segment.Offset, segment.Count));
27-
}
28-
else
29-
{
30-
var span = _data.Span;
31-
var oversized = ArrayPool<byte>.Shared.Rent(span.Length);
32-
span.CopyTo(oversized);
33-
args.Add(Convert.ToBase64String(oversized, 0, span.Length));
34-
ArrayPool<byte>.Shared.Return(oversized);
35-
}
46+
var oversized = ArrayPool<byte>.Shared.Rent(bytes.Length);
47+
bytes.CopyTo(oversized);
48+
var result = Convert.ToBase64String(oversized, 0, bytes.Length);
49+
ArrayPool<byte>.Shared.Return(oversized);
50+
return result;
3651
#endif
52+
}
3753
}
38-
internal int GetOwnArgsCount() => 1;
39-
internal bool HasValue => _data.Length > 0;
54+
55+
private sealed class VectorDataBase64(string vector) : VectorData
56+
{
57+
public override string ToString() => vector;
58+
}
59+
60+
private protected static void ThrowBigEndian() =>
61+
throw new PlatformNotSupportedException("Big-endian CPUs are not currently supported for this operation");
4062
}

‎tests/NRedisStack.Tests/Search/HybridSearchIntegrationTests.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ public void TestSetup(string endpointId)
5050
Dictionary<string, object> args = new() { ["x"] = "abc" };
5151
var query = new HybridSearchQuery()
5252
.Search("*")
53-
.VectorSearch("@vector1", new byte[] {1,2,3,4})
53+
.VectorSearch("@vector1", new float[] {1,2,3,4})
5454
.Parameters(args)
5555
.ReturnFields("@text1");
5656
var result = api.FT.HybridSearch(api.Index, query);

‎tests/NRedisStack.Tests/Search/HybridSearchUnitTests.cs‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -118,22 +118,23 @@ public void BasicSearch_WithBM25StdTanh()
118118
public void BasicVectorSearch()
119119
{
120120
HybridSearchQuery query = new();
121-
byte[] data = [1, 2, 3];
122-
query.VectorSearch("vfield", data);
121+
query.VectorSearch("vfield", Array.Empty<float>());
123122

124-
object[] expected = [Index, "VSIM", "vfield", "AQID"];
123+
object[] expected = [Index, "VSIM", "vfield", ""];
125124
Assert.Equivalent(expected, GetArgs(query));
126125
}
127126

128-
private static readonly ReadOnlyMemory<byte> SomeRandomDataHere = Encoding.UTF8.GetBytes("some random data here!");
127+
private static readonly ReadOnlyMemory<float> SomeRandomDataHere = new float[] {1, 2, 3, 4};
128+
129+
private const string SomeRandomBase64 = "AACAPwAAAEAAAEBAAACAQA==";
129130

130131
[Fact]
131132
public void BasicNonZeroLengthVectorSearch()
132133
{
133134
HybridSearchQuery query = new();
134135
query.VectorSearch("vfield", SomeRandomDataHere);
135136

136-
object[] expected = [Index, "VSIM", "vfield", "c29tZSByYW5kb20gZGF0YSBoZXJlIQ=="];
137+
object[] expected = [Index, "VSIM", "vfield", SomeRandomBase64];
137138
Assert.Equivalent(expected, GetArgs(query));
138139
}
139140

@@ -152,7 +153,7 @@ public void BasicVectorSearch_WithKNN(bool withScoreAlias, bool withDistanceAlia
152153
query.VectorSearch(searchConfig);
153154

154155
object[] expected =
155-
[Index, "VSIM", "vField", "c29tZSByYW5kb20gZGF0YSBoZXJlIQ==", "KNN", withDistanceAlias ? 4 : 2, "K", 10];
156+
[Index, "VSIM", "vField", SomeRandomBase64, "KNN", withDistanceAlias ? 4 : 2, "K", 10];
156157
if (withDistanceAlias)
157158
{
158159
expected = [..expected, "YIELD_DISTANCE_AS", "my_distance_alias"];
@@ -184,7 +185,7 @@ public void BasicVectorSearch_WithKNN_WithEF(bool withScoreAlias, bool withDista
184185

185186
object[] expected =
186187
[
187-
Index, "VSIM", "vfield", "c29tZSByYW5kb20gZGF0YSBoZXJlIQ==", "KNN", withDistanceAlias ? 6 : 4, "K", 16,
188+
Index, "VSIM", "vfield", SomeRandomBase64, "KNN", withDistanceAlias ? 6 : 4, "K", 16,
188189
"EF_RUNTIME", 100
189190
];
190191
if (withDistanceAlias)
@@ -216,7 +217,7 @@ public void BasicVectorSearch_WithRange(bool withScoreAlias, bool withDistanceAl
216217

217218
object[] expected =
218219
[
219-
Index, "VSIM", "vfield", "c29tZSByYW5kb20gZGF0YSBoZXJlIQ==", "RANGE", withDistanceAlias ? 4 : 2, "RADIUS",
220+
Index, "VSIM", "vfield", SomeRandomBase64, "RANGE", withDistanceAlias ? 4 : 2, "RADIUS",
220221
4.2
221222
];
222223
if (withDistanceAlias)
@@ -249,7 +250,7 @@ public void BasicVectorSearch_WithRange_WithEpsilon(bool withScoreAlias, bool wi
249250

250251
object[] expected =
251252
[
252-
Index, "VSIM", "vfield", "c29tZSByYW5kb20gZGF0YSBoZXJlIQ==", "RANGE", withDistanceAlias ? 6 : 4, "RADIUS",
253+
Index, "VSIM", "vfield", SomeRandomBase64, "RANGE", withDistanceAlias ? 6 : 4, "RADIUS",
253254
4.2, "EPSILON", 0.06
254255
];
255256
if (withDistanceAlias)
@@ -273,7 +274,7 @@ public void BasicVectorSearch_WithFilter_NoPolicy()
273274

274275
object[] expected =
275276
[
276-
Index, "VSIM", "vfield", "c29tZSByYW5kb20gZGF0YSBoZXJlIQ==", "FILTER", "@foo:bar"
277+
Index, "VSIM", "vfield", SomeRandomBase64, "FILTER", "@foo:bar"
277278
];
278279

279280
Assert.Equivalent(expected, GetArgs(query));
@@ -608,7 +609,7 @@ public void MakeMeOneWithEverything()
608609
["y"] = "abc"
609610
};
610611
query.Search(new("foo", Scorer.BM25StdTanh(5), "text_score_alias"))
611-
.VectorSearch(new HybridSearchQuery.VectorSearchConfig("bar", new byte[] { 1, 2, 3 },
612+
.VectorSearch(new HybridSearchQuery.VectorSearchConfig("bar", new float[] { 1, 2, 3 },
612613
VectorSearchMethod.NearestNeighbour(10, 100, "vector_distance_alias"))
613614
.WithFilter("@foo:bar").WithScoreAlias("vector_score_alias"))
614615
.Combine(HybridSearchQuery.Combiner.ReciprocalRankFusion(10, 0.5), "my_combined_alias")
@@ -626,7 +627,7 @@ public void MakeMeOneWithEverything()
626627
[
627628
Index, "SEARCH", "foo", "SCORER", "BM25STD.TANH", "BM25STD_TANH_FACTOR", 5, "YIELD_SCORE_AS",
628629
"text_score_alias", "VSIM", "bar",
629-
"AQID", "KNN", 6, "K", 10, "EF_RUNTIME", 100, "YIELD_DISTANCE_AS", "vector_distance_alias", "FILTER",
630+
"AACAPwAAAEAAAEBA", "KNN", 6, "K", 10, "EF_RUNTIME", 100, "YIELD_DISTANCE_AS", "vector_distance_alias", "FILTER",
630631
"@foo:bar", "YIELD_SCORE_AS", "vector_score_alias", "COMBINE", "RRF", 4, "WINDOW", 10, "CONSTANT", 0.5,
631632
"YIELD_SCORE_AS", "my_combined_alias", "LOAD", 2, "field1", "field2", "GROUPBY", 1, "field1", "REDUCE",
632633
"QUANTILE", 2, "@field3", 0.5, "AS", "reducer_alias", "APPLY", "@field1 + @field2", "AS", "apply_alias",

0 commit comments

Comments
 (0)