Skip to content

Commit 4b4225e

Browse files
mgravellmgravell
committed
VLINKS impl
1 parent f20db67 commit 4b4225e

5 files changed

Lines changed: 228 additions & 41 deletions

File tree

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#nullable enable
2+
[SER001]override StackExchange.Redis.VectorSetLink.ToString() -> string!
23
[SER001]StackExchange.Redis.IDatabase.VectorSetAdd(StackExchange.Redis.RedisKey key, StackExchange.Redis.RedisValue element, System.ReadOnlyMemory<float> values, int? reducedDimensions = null, StackExchange.Redis.VectorQuantizationType quantizationType = StackExchange.Redis.VectorQuantizationType.Int8, int? buildExplorationFactor = null, int? maxConnections = null, bool useCheckAndSet = false, string? attributesJson = null, StackExchange.Redis.CommandFlags flags = StackExchange.Redis.CommandFlags.None) -> bool
34
[SER001]StackExchange.Redis.IDatabase.VectorSetContains(StackExchange.Redis.RedisKey key, StackExchange.Redis.RedisValue member, StackExchange.Redis.CommandFlags flags = StackExchange.Redis.CommandFlags.None) -> bool
45
[SER001]StackExchange.Redis.IDatabase.VectorSetDimension(StackExchange.Redis.RedisKey key, StackExchange.Redis.CommandFlags flags = StackExchange.Redis.CommandFlags.None) -> int

‎src/StackExchange.Redis/RedisDatabase.cs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5777,13 +5777,13 @@ public bool VectorSetContains(RedisKey key, RedisValue member, CommandFlags flag
57775777
public Lease<RedisValue>? VectorSetGetLinks(RedisKey key, RedisValue member, CommandFlags flags = CommandFlags.None)
57785778
{
57795779
var msg = Message.Create(Database, flags, RedisCommand.VLINKS, key, member);
5780-
return ExecuteSync(msg, ResultProcessor.RedisValueLease);
5780+
return ExecuteSync(msg, ResultProcessor.VectorSetLinks);
57815781
}
57825782

57835783
public Lease<VectorSetLink>? VectorSetGetLinksWithScores(RedisKey key, RedisValue member, CommandFlags flags = CommandFlags.None)
57845784
{
57855785
var msg = Message.Create(Database, flags, RedisCommand.VLINKS, key, member, RedisLiterals.WITHSCORES);
5786-
return ExecuteSync(msg, ResultProcessor.LeaseVectorSetLink);
5786+
return ExecuteSync(msg, ResultProcessor.VectorSetLinksWithScores);
57875787
}
57885788

57895789
public RedisValue VectorSetRandomMember(RedisKey key, CommandFlags flags = CommandFlags.None)
@@ -5899,13 +5899,13 @@ public Task<bool> VectorSetContainsAsync(RedisKey key, RedisValue member, Comman
58995899
public Task<Lease<RedisValue>?> VectorSetGetLinksAsync(RedisKey key, RedisValue member, CommandFlags flags = CommandFlags.None)
59005900
{
59015901
var msg = Message.Create(Database, flags, RedisCommand.VLINKS, key, member);
5902-
return ExecuteAsync(msg, ResultProcessor.RedisValueLease);
5902+
return ExecuteAsync(msg, ResultProcessor.VectorSetLinks);
59035903
}
59045904

59055905
public Task<Lease<VectorSetLink>?> VectorSetGetLinksWithScoresAsync(RedisKey key, RedisValue member, CommandFlags flags = CommandFlags.None)
59065906
{
59075907
var msg = Message.Create(Database, flags, RedisCommand.VLINKS, key, member, RedisLiterals.WITHSCORES);
5908-
return ExecuteAsync(msg, ResultProcessor.LeaseVectorSetLink);
5908+
return ExecuteAsync(msg, ResultProcessor.VectorSetLinksWithScores);
59095909
}
59105910

59115911
public Task<RedisValue> VectorSetRandomMemberAsync(RedisKey key, CommandFlags flags = CommandFlags.None)

‎src/StackExchange.Redis/ResultProcessor.cs‎

Lines changed: 100 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,8 @@ public static readonly ResultProcessor<long>
6363
public static readonly ResultProcessor<int> Int32 = new Int32Processor();
6464
public static readonly ResultProcessor<Lease<float>?> LeaseFloat32 = new LeaseFloat32Processor();
6565

66-
public static readonly ResultProcessor<Lease<VectorSetLink>?> LeaseVectorSetLink = new LeaseVectorSetLinkProcessor();
66+
public static readonly ResultProcessor<Lease<VectorSetLink>?> VectorSetLinksWithScores = new VectorSetLinksWithScoresProcessor();
67+
public static readonly ResultProcessor<Lease<RedisValue>?> VectorSetLinks = new VectorSetLinksProcessor();
6768

6869
public static readonly ResultProcessor<double?>
6970
NullableDouble = new NullableDoubleProcessor();
@@ -104,8 +105,6 @@ public static readonly ResultProcessor<Lease<byte>>
104105

105106
public static readonly ResultProcessor<RedisValue[]>
106107
RedisValueArray = new RedisValueArrayProcessor();
107-
public static readonly ResultProcessor<Lease<RedisValue>?>
108-
RedisValueLease = new RedisValueLeaseProcessor();
109108

110109
public static readonly ResultProcessor<long[]>
111110
Int64Array = new Int64ArrayProcessor();
@@ -1713,7 +1712,7 @@ protected override bool SetResultCore(PhysicalConnection connection, Message mes
17131712

17141713
private abstract class LeaseProcessor<T> : ResultProcessor<Lease<T>?>
17151714
{
1716-
protected sealed override bool SetResultCore(PhysicalConnection connection, Message message, in RawResult result)
1715+
protected override bool SetResultCore(PhysicalConnection connection, Message message, in RawResult result)
17171716
{
17181717
if (result.Resp2TypeArray != ResultType.Array)
17191718
{
@@ -1752,7 +1751,7 @@ protected sealed override bool SetResultCore(PhysicalConnection connection, Mess
17521751

17531752
private abstract class InterleavedLeaseProcessor<T> : ResultProcessor<Lease<T>?>
17541753
{
1755-
protected sealed override bool SetResultCore(PhysicalConnection connection, Message message, in RawResult result)
1754+
protected override bool SetResultCore(PhysicalConnection connection, Message message, in RawResult result)
17561755
{
17571756
if (result.Resp2TypeArray != ResultType.Array)
17581757
{
@@ -1794,32 +1793,114 @@ protected sealed override bool SetResultCore(PhysicalConnection connection, Mess
17941793
protected abstract bool TryParse(in RawResult first, in RawResult second, out T parsed);
17951794
}
17961795

1797-
private sealed class RedisValueLeaseProcessor : LeaseProcessor<RedisValue>
1796+
// takes a nested vector of the form [[A],[B,C],[D]] and exposes it as [A,B,C,D]; this is
1797+
// especially useful for VLINKS
1798+
private abstract class FlattenedLeaseProcessor<T> : ResultProcessor<Lease<T>?>
17981799
{
1799-
protected override bool TryParse(in RawResult raw, out RedisValue parsed)
1800+
protected virtual long GetArrayLength(in RawResult array) => array.GetItems().Length;
1801+
1802+
protected virtual bool TryReadOne(ref Sequence<RawResult>.Enumerator reader, out T value)
18001803
{
1801-
parsed = raw.AsRedisValue();
1802-
return true;
1804+
if (reader.MoveNext() && TryReadOne(reader.Current, out value))
1805+
{
1806+
return true;
1807+
}
1808+
value = default!;
1809+
return false;
1810+
}
1811+
1812+
protected virtual bool TryReadOne(in RawResult result, out T value)
1813+
{
1814+
value = default!;
1815+
return false;
1816+
}
1817+
1818+
protected override bool SetResultCore(PhysicalConnection connection, Message message, in RawResult result)
1819+
{
1820+
if (result.Resp2TypeArray != ResultType.Array)
1821+
{
1822+
return false; // not an array
1823+
}
1824+
if (result.IsNull)
1825+
{
1826+
SetResult(message, Lease<T>.Empty);
1827+
return true;
1828+
}
1829+
var items = result.GetItems();
1830+
long length = 0;
1831+
foreach (ref RawResult item in items)
1832+
{
1833+
if (item.Resp2TypeArray == ResultType.Array && !item.IsNull)
1834+
{
1835+
length += GetArrayLength(item);
1836+
}
1837+
}
1838+
1839+
if (length == 0)
1840+
{
1841+
SetResult(message, Lease<T>.Empty);
1842+
return true;
1843+
}
1844+
var lease = Lease<T>.Create(checked((int)length), clear: false);
1845+
int index = 0;
1846+
var target = lease.Span;
1847+
foreach (ref RawResult item in items)
1848+
{
1849+
if (item.Resp2TypeArray == ResultType.Array && !item.IsNull)
1850+
{
1851+
var iter = item.GetItems().GetEnumerator();
1852+
while (index < target.Length && TryReadOne(ref iter, out T value))
1853+
{
1854+
target[index++] = value;
1855+
}
1856+
}
1857+
}
1858+
1859+
if (index == length)
1860+
{
1861+
SetResult(message, lease);
1862+
return true;
1863+
}
1864+
lease.Dispose(); // failed to fill?
1865+
return false;
18031866
}
18041867
}
18051868

1806-
private sealed class LeaseFloat32Processor : LeaseProcessor<float>
1869+
private sealed class VectorSetLinksWithScoresProcessor : FlattenedLeaseProcessor<VectorSetLink>
18071870
{
1808-
protected override bool TryParse(in RawResult raw, out float parsed)
1871+
protected override long GetArrayLength(in RawResult array) => array.GetItems().Length / 2;
1872+
1873+
protected override bool TryReadOne(ref Sequence<RawResult>.Enumerator reader, out VectorSetLink value)
18091874
{
1810-
var result = raw.TryGetDouble(out double val);
1811-
parsed = (float)val;
1812-
return result;
1875+
if (reader.MoveNext())
1876+
{
1877+
ref readonly RawResult first = ref reader.Current;
1878+
if (reader.MoveNext() && reader.Current.TryGetDouble(out var score))
1879+
{
1880+
value = new VectorSetLink(first.AsRedisValue(), score);
1881+
return true;
1882+
}
1883+
}
1884+
value = default;
1885+
return false;
1886+
}
1887+
}
1888+
1889+
private sealed class VectorSetLinksProcessor : FlattenedLeaseProcessor<RedisValue>
1890+
{
1891+
protected override bool TryReadOne(in RawResult result, out RedisValue value)
1892+
{
1893+
value = result.AsRedisValue();
1894+
return true;
18131895
}
18141896
}
18151897

1816-
private sealed class LeaseVectorSetLinkProcessor : InterleavedLeaseProcessor<VectorSetLink>
1898+
private sealed class LeaseFloat32Processor : LeaseProcessor<float>
18171899
{
1818-
protected override bool TryParse(in RawResult first, in RawResult second, out VectorSetLink parsed)
1900+
protected override bool TryParse(in RawResult raw, out float parsed)
18191901
{
1820-
var member = first.AsRedisValue();
1821-
bool result = second.TryGetDouble(out var score);
1822-
parsed = new VectorSetLink(member, score);
1902+
var result = raw.TryGetDouble(out double val);
1903+
parsed = (float)val;
18231904
return result;
18241905
}
18251906
}

‎src/StackExchange.Redis/VectorSetLink.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,4 +18,7 @@ public readonly struct VectorSetLink(RedisValue member, double score)
1818
/// The similarity score between the queried member and this linked member.
1919
/// </summary>
2020
public double Score { get; } = score;
21+
22+
/// <inheritdoc/>
23+
public override string ToString() => $"{Member}: {Score}";
2124
}

0 commit comments

Comments
 (0)