Skip to content

Commit 62a5ee7

Browse files
mgravellmgravell
committed
VectorData; make API usable
1 parent 0c19979 commit 62a5ee7

File tree

4 files changed

+105
-35
lines changed

4 files changed

+105
-35
lines changed

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
NRedisStack.Search.Parameters
22
static NRedisStack.Search.Parameters.From<T>(T obj) -> System.Collections.Generic.IReadOnlyDictionary<string!, object!>!
3+
[NRS001]abstract NRedisStack.Search.VectorData<T>.Dispose() -> void
4+
[NRS001]abstract NRedisStack.Search.VectorData<T>.Span.get -> System.Span<T>
35
[NRS001]const NRedisStack.Search.HybridSearchQuery.Fields.Key = "@__key" -> string!
46
[NRS001]const NRedisStack.Search.HybridSearchQuery.Fields.Score = "@__score" -> string!
57
[NRS001]NRedisStack.ISearchCommands.HybridSearch(string! indexName, NRedisStack.Search.HybridSearchQuery! query, System.Collections.Generic.IReadOnlyDictionary<string!, object!>? parameters = null) -> NRedisStack.Search.HybridSearchResult!
@@ -63,6 +65,7 @@ static NRedisStack.Search.Parameters.From<T>(T obj) -> System.Collections.Generi
6365
[NRS001]NRedisStack.Search.HybridSearchResult.TotalResults.get -> long
6466
[NRS001]NRedisStack.Search.Scorer
6567
[NRS001]NRedisStack.Search.VectorData
68+
[NRS001]NRedisStack.Search.VectorData<T>
6669
[NRS001]NRedisStack.Search.VectorSearchMethod
6770
[NRS001]NRedisStack.SearchCommands.HybridSearch(string! indexName, NRedisStack.Search.HybridSearchQuery! query, System.Collections.Generic.IReadOnlyDictionary<string!, object!>? parameters = null) -> NRedisStack.Search.HybridSearchResult!
6871
[NRS001]NRedisStack.SearchCommandsAsync.HybridSearchAsync(string! indexName, NRedisStack.Search.HybridSearchQuery! query, System.Collections.Generic.IReadOnlyDictionary<string!, object!>? parameters = null) -> System.Threading.Tasks.Task<NRedisStack.Search.HybridSearchResult!>!
@@ -84,10 +87,9 @@ static NRedisStack.Search.Parameters.From<T>(T obj) -> System.Collections.Generi
8487
[NRS001]static NRedisStack.Search.Scorer.Hamming.get -> NRedisStack.Search.Scorer!
8588
[NRS001]static NRedisStack.Search.Scorer.TfIdf.get -> NRedisStack.Search.Scorer!
8689
[NRS001]static NRedisStack.Search.Scorer.TfIdfDocNorm.get -> NRedisStack.Search.Scorer!
87-
[NRS001]static NRedisStack.Search.VectorData.Create(System.ReadOnlyMemory<float> vector) -> NRedisStack.Search.VectorData!
88-
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData!(float[]! data) -> NRedisStack.Search.VectorData!
8990
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData!(string! name) -> NRedisStack.Search.VectorData!
90-
[NRS001]static NRedisStack.Search.VectorData.implicit operator NRedisStack.Search.VectorData!(System.ReadOnlyMemory<float> vector) -> NRedisStack.Search.VectorData!
91+
[NRS001]static NRedisStack.Search.VectorData.Lease<T>(int dimension) -> NRedisStack.Search.VectorData<T>!
92+
[NRS001]static NRedisStack.Search.VectorData.LeaseWithValues<T>(params System.ReadOnlySpan<T> values) -> NRedisStack.Search.VectorData<T>!
9193
[NRS001]static NRedisStack.Search.VectorData.Parameter(string! name) -> NRedisStack.Search.VectorData!
9294
[NRS001]static NRedisStack.Search.VectorData.Raw(System.ReadOnlyMemory<byte> bytes) -> NRedisStack.Search.VectorData!
9395
[NRS001]static NRedisStack.Search.VectorSearchMethod.NearestNeighbour(int count = 10, int? maxCandidates = null) -> NRedisStack.Search.VectorSearchMethod!

‎src/NRedisStack/Search/VectorData.cs‎

Lines changed: 54 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,69 @@
11
using System.Buffers;
22
using System.Diagnostics.CodeAnalysis;
3+
using System.Runtime.CompilerServices;
34
using System.Runtime.InteropServices;
45
using StackExchange.Redis;
56

67
namespace NRedisStack.Search;
78

89
[Experimental(Experiments.Server_8_4, UrlFormat = Experiments.UrlFormat)]
9-
public abstract class VectorData
10+
public abstract class VectorData<T> : VectorData, IDisposable where T : unmanaged
1011
{
1112
private protected VectorData()
1213
{
1314
}
1415

16+
public abstract Span<T> Span { get; }
17+
internal sealed class VectorBytesData(int byteLength) : VectorData<T>
18+
{
19+
private byte[]? _oversized = ArrayPool<byte>.Shared.Rent(byteLength);
20+
public override Span<T> Span => MemoryMarshal.Cast<byte, T>(Array.AsSpan(0, byteLength));
21+
internal override object GetSingleArg() => (RedisValue)new ReadOnlyMemory<byte>(Array, 0, byteLength);
22+
23+
private byte[] Array => _oversized ?? ThrowDisposed();
24+
static byte[] ThrowDisposed() => throw new ObjectDisposedException(nameof(VectorData));
25+
public override void Dispose()
26+
{
27+
var tmp = _oversized;
28+
_oversized = null;
29+
if (tmp is not null) ArrayPool<byte>.Shared.Return(tmp);
30+
}
31+
}
32+
33+
public abstract void Dispose();
34+
}
35+
36+
[Experimental(Experiments.Server_8_4, UrlFormat = Experiments.UrlFormat)]
37+
public abstract class VectorData
38+
{
1539
/// <summary>
16-
/// A vector of <see cref="Single"/> entries.
40+
/// Lease a vector that can hold values of <typeparamref name="T"/>.
41+
/// No quantization occurs - the data is transmitted as the raw bytes of the corresponding size.
1742
/// </summary>
18-
public static VectorData Create(ReadOnlyMemory<float> vector) => new VectorDataSingle(vector);
43+
/// <param name="dimension">The number of values to be held.</param>
44+
/// <typeparam name="T">The data type to be represented</typeparam>
45+
public static VectorData<T> Lease<T>(int dimension) where T : unmanaged
46+
{
47+
if (dimension < 0) ThrowDimension();
48+
if (!BitConverter.IsLittleEndian) ThrowBigEndian();
49+
return new VectorData<T>.VectorBytesData(Unsafe.SizeOf<T>() * dimension);
50+
51+
static void ThrowDimension() => throw new ArgumentOutOfRangeException(nameof(dimension));
52+
}
53+
54+
/// <summary>
55+
/// Lease a vector that can hold values of <typeparamref name="T"/>, copying in the supplied values.
56+
/// </summary>
57+
public static VectorData<T> LeaseWithValues<T>(params ReadOnlySpan<T> values) where T : unmanaged
58+
{
59+
var lease = Lease<T>(values.Length);
60+
values.CopyTo(lease.Span);
61+
return lease;
62+
}
63+
64+
private protected VectorData()
65+
{
66+
}
1967

2068
/// <summary>
2169
/// A raw vector payload.
@@ -27,14 +75,16 @@ private protected VectorData()
2775
/// </summary>
2876
public static VectorData Parameter(string name) => new VectorParameter(name);
2977

78+
/*
3079
/// <summary>
3180
/// A vector of <see cref="Single"/> entries.
3281
/// </summary>
3382
public static implicit operator VectorData(float[] data) => new VectorDataSingle(data);
3483
3584
/// <inheritdoc cref="Create"/>
3685
public static implicit operator VectorData(ReadOnlyMemory<float> vector) => new VectorDataSingle(vector);
37-
86+
*/
87+
3888
/// <inheritdoc cref="Parameter"/>
3989
public static implicit operator VectorData(string name) => new VectorParameter(name);
4090

@@ -43,27 +93,6 @@ private protected VectorData()
4393
/// <inheritdoc/>
4494
public override string ToString() => GetType().Name;
4595

46-
private sealed class VectorDataSingle(ReadOnlyMemory<float> vector) : VectorData
47-
{
48-
internal override object GetSingleArg() => ToBase64();
49-
public override string ToString() => ToBase64();
50-
51-
private string ToBase64()
52-
{
53-
if (!BitConverter.IsLittleEndian) ThrowBigEndian(); // we could loop and reverse each, but...how to test?
54-
var bytes = MemoryMarshal.AsBytes(vector.Span);
55-
#if NET || NETSTANDARD2_1_OR_GREATER
56-
return Convert.ToBase64String(bytes);
57-
#else
58-
var oversized = ArrayPool<byte>.Shared.Rent(bytes.Length);
59-
bytes.CopyTo(oversized);
60-
var result = Convert.ToBase64String(oversized, 0, bytes.Length);
61-
ArrayPool<byte>.Shared.Return(oversized);
62-
return result;
63-
#endif
64-
}
65-
}
66-
6796
private sealed class VectorDataRaw(ReadOnlyMemory<byte> bytes) : VectorData
6897
{
6998
internal override object GetSingleArg() => (RedisValue)bytes;

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,9 +92,10 @@ public async Task TestSetup(string endpointId)
9292
{
9393
var api = await CreateIndexAsync(endpointId, populate: false);
9494
Dictionary<string, object> args = new() { ["x"] = "abc" };
95+
using var vector = VectorData.LeaseWithValues<float>(1, 2, 3, 4);
9596
var query = new HybridSearchQuery()
9697
.Search("*")
97-
.VectorSearch("@vector1", new float[] { 1, 2, 3, 4 })
98+
.VectorSearch("@vector1", vector)
9899
.ReturnFields("@text1");
99100
var result = api.FT.HybridSearch(api.Index, query, args);
100101
Assert.Equal(0, result.TotalResults);

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

Lines changed: 44 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
1+
using System.Buffers;
2+
using System.Diagnostics;
13
using NRedisStack.Search;
24
using NRedisStack.Search.Aggregation;
5+
using StackExchange.Redis;
36
using Xunit;
47
using Xunit.Abstractions;
58

@@ -12,7 +15,40 @@ public class HybridSearchUnitTests(ITestOutputHelper log)
1215
private ICollection<object> GetArgs(HybridSearchQuery query, IReadOnlyDictionary<string, object>? parameters = null)
1316
{
1417
Assert.Equal("FT.HYBRID", query.Command);
15-
var args = query.GetArgs(Index, parameters);
18+
var args = query.GetArgs(Index, parameters).ToArray();
19+
byte[] buffer = ArrayPool<byte>.Shared.Rent(128);
20+
for (int i = 0; i < args.Length; i++)
21+
{
22+
var val = args[i];
23+
if (val is RedisValue v)
24+
{
25+
// force non-readable values to an ASCII-printable string
26+
var bytes = v.GetByteCount();
27+
if (bytes > buffer.Length)
28+
{
29+
ArrayPool<byte>.Shared.Return(buffer);
30+
buffer = ArrayPool<byte>.Shared.Rent(bytes);
31+
}
32+
33+
var span = buffer.AsSpan(0, v.CopyTo(buffer));
34+
Debug.Assert(span.Length == bytes, $"expected {bytes}, wrote {span.Length}");
35+
var format = false;
36+
foreach (var b in span)
37+
{
38+
if (b < 32 | b >= 127)
39+
{
40+
format = true;
41+
break;
42+
}
43+
}
44+
45+
if (format) // not how resp-cli works, but: useful enough for unit tests
46+
{
47+
args[i] = Convert.ToBase64String(buffer, 0, bytes);
48+
}
49+
}
50+
}
51+
ArrayPool<byte>.Shared.Return(buffer);
1652
log.WriteLine(query.Command + " " + string.Join(" ", args));
1753
return args;
1854
}
@@ -117,13 +153,14 @@ public void BasicSearch_WithBM25StdTanh()
117153
public void BasicVectorSearch()
118154
{
119155
HybridSearchQuery query = new();
120-
query.VectorSearch("vfield", Array.Empty<float>());
156+
byte[] blob = [];
157+
query.VectorSearch("vfield", VectorData.Raw(blob));
121158

122159
object[] expected = [Index, "VSIM", "vfield", ""];
123160
Assert.Equivalent(expected, GetArgs(query));
124161
}
125162

126-
private static readonly ReadOnlyMemory<float> SomeRandomDataHere = new float[] { 1, 2, 3, 4 };
163+
private static readonly VectorData<float> SomeRandomDataHere = VectorData.LeaseWithValues<float>(1, 2, 3, 4 );
127164

128165
private const string SomeRandomVectorValue = "AACAPwAAAEAAAEBAAACAQA==";
129166

@@ -622,15 +659,15 @@ public void ParameterizedQuery()
622659
IReadOnlyDictionary<string, object> args = new Dictionary<string, object>
623660
{
624661
{ "s", "abc"},
625-
{ "v", VectorData.Create(SomeRandomDataHere) }
662+
{ "v", SomeRandomDataHere }
626663
};
627664
object[] expected = [Index, "SEARCH", "$s", "VSIM", "@field", "$v", "PARAMS", 4, "s", "abc", "v", SomeRandomVectorValue];
628665
var idx = Array.IndexOf(expected, "abc");
629666
Assert.True(idx >= 0);
630667
Assert.Equivalent(expected, GetArgs(query, args));
631668

632669
// issue a second query against the same "query" instance, with different parameter values, this time from an object
633-
args = Parameters.From(new { s = "def", v = VectorData.Create(SomeRandomDataHere) });
670+
args = Parameters.From(new { s = "def", v = SomeRandomDataHere });
634671

635672
expected[idx] = "def"; // update our expectations
636673
expected[idx + 2] = SomeRandomVectorValue;
@@ -646,8 +683,9 @@ public void MakeMeOneWithEverything()
646683
["x"] = 42,
647684
["y"] = "abc"
648685
};
686+
using var vector = VectorData.LeaseWithValues<float>(1, 2, 3);
649687
query.Search(new("foo", Scorer.BM25StdTanh(5), "text_score_alias"))
650-
.VectorSearch(new HybridSearchQuery.VectorSearchConfig("bar", new float[] { 1, 2, 3 },
688+
.VectorSearch(new HybridSearchQuery.VectorSearchConfig("bar", vector,
651689
VectorSearchMethod.NearestNeighbour(10, 100, "vector_distance_alias"))
652690
.WithFilter("@foo:bar").WithScoreAlias("vector_score_alias"))
653691
.Combine(HybridSearchQuery.Combiner.ReciprocalRankFusion(10, 0.5), "my_combined_alias")

0 commit comments

Comments
 (0)