Skip to content

Commit ff7fda3

Browse files
author
tiagonapoli
committed
w
1 parent 354cb8d commit ff7fda3

File tree

10 files changed

+191
-4
lines changed

10 files changed

+191
-4
lines changed

‎libs/server/API/GarnetApi.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -539,6 +539,10 @@ public unsafe GarnetStatus VectorSetDimensions(ArgSlice key, out int dimensions)
539539
public unsafe GarnetStatus VectorSetInfo(ArgSlice key, out VectorQuantType quantType, out VectorDistanceMetricType distanceMetricType, out uint vectorDimensions, out uint reducedDimensions, out uint buildExplorationFactor, out uint numberOfLinks, out long size)
540540
=> storageSession.VectorSetInfo(SpanByte.FromPinnedPointer(key.ptr, key.length), out quantType, out distanceMetricType, out vectorDimensions, out reducedDimensions, out buildExplorationFactor, out numberOfLinks, out size);
541541

542+
/// <inheritdoc/>
543+
public unsafe GarnetStatus VectorSetCheckIsMember(ArgSlice key, ArgSlice elementId)
544+
=> storageSession.VectorSetCheckIsMember(SpanByte.FromPinnedPointer(key.ptr, key.length), elementId.ReadOnlySpan);
545+
542546
#endregion
543547
}
544548
}

‎libs/server/API/GarnetWatchApi.cs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -684,6 +684,13 @@ GarnetStatus IGarnetReadApi.VectorSetInfo(ArgSlice key, out VectorQuantType quan
684684
return garnetApi.VectorSetInfo(key, out quantType, out distanceMetricType, out vectorDimensions, out reducedDimensions, out buildExplorationFactor, out numberOfLinks, out size);
685685
}
686686

687+
/// <inheritdoc/>
688+
public GarnetStatus VectorSetCheckIsMember(ArgSlice key, ArgSlice elementId)
689+
{
690+
garnetApi.WATCH(key, StoreType.Main);
691+
return garnetApi.VectorSetCheckIsMember(key, elementId);
692+
}
693+
687694
#endregion
688695
}
689696
}

‎libs/server/API/IGarnetApi.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2067,6 +2067,11 @@ public bool IterateObjectStore<TScanFunctions>(ref TScanFunctions scanFunctions,
20672067
/// Fetch debugging information about the Vector Set.
20682068
/// </summary>
20692069
GarnetStatus VectorSetInfo(ArgSlice key, out VectorQuantType quantType, out VectorDistanceMetricType distanceMetricType, out uint vectorDimensions, out uint reducedDimensions, out uint buildExplorationFactor, out uint numberOfLinks, out long size);
2070+
2071+
/// <summary>
2072+
/// Check if vector set contains the given element id.
2073+
/// </summary>
2074+
GarnetStatus VectorSetCheckIsMember(ArgSlice key, ArgSlice elementId);
20702075

20712076
#endregion
20722077
}

‎libs/server/Resp/RespServerSessionOutput.cs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,36 @@ private void WriteOne()
187187
SendAndReset();
188188
}
189189

190+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
191+
private void WriteBooleanTrue()
192+
{
193+
if (respProtocolVersion >= 3)
194+
{
195+
while (!RespWriteUtils.TryWriteTrue(ref dcurr, dend))
196+
SendAndReset();
197+
}
198+
else
199+
{
200+
while (!RespWriteUtils.TryWriteOne(ref dcurr, dend))
201+
SendAndReset();
202+
}
203+
}
204+
205+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
206+
private void WriteBooleanFalse()
207+
{
208+
if (respProtocolVersion >= 3)
209+
{
210+
while (!RespWriteUtils.TryWriteFalse(ref dcurr, dend))
211+
SendAndReset();
212+
}
213+
else
214+
{
215+
while (!RespWriteUtils.TryWriteZero(ref dcurr, dend))
216+
SendAndReset();
217+
}
218+
}
219+
190220
[MethodImpl(MethodImplOptions.AggressiveInlining)]
191221
private void WriteNull()
192222
{

‎libs/server/Resp/Vector/DiskANNService.cs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,14 @@ public bool CheckInternalIdValid(ulong context, nint index, ReadOnlySpan<byte> i
214214

215215
return NativeDiskANNMethods.check_internal_id_valid(context, index, (nint)internal_id_data, (nuint)internal_id_len) == 1;
216216
}
217+
218+
public bool CheckExternalIdValid(ulong context, nint index, ReadOnlySpan<byte> externalId)
219+
{
220+
var external_id_data = Unsafe.AsPointer(ref MemoryMarshal.GetReference(externalId));
221+
var externa_id_len = externalId.Length;
222+
223+
return NativeDiskANNMethods.check_external_id_valid(context, index, (nint)external_id_data, (nuint)externa_id_len) == 1;
224+
}
217225
}
218226

219227
public static partial class NativeDiskANNMethods
@@ -333,5 +341,13 @@ public static partial byte check_internal_id_valid(
333341
nint internal_id,
334342
nuint internal_id_len
335343
);
344+
345+
[LibraryImport(DISKANN_GARNET)]
346+
public static partial byte check_external_id_valid(
347+
ulong context,
348+
nint index,
349+
nint external_id,
350+
nuint external_id_len
351+
);
336352
}
337353
}

‎libs/server/Resp/Vector/RespServerSessionVectors.cs‎

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1018,7 +1018,10 @@ private bool NetworkVGETATTR<TGarnetApi>(ref TGarnetApi storageApi)
10181018
return AbortWithErrorMessage("ERR Vector Set (preview) commands are not enabled");
10191019
}
10201020

1021-
// TODO: implement!
1021+
if (parseState.Count != 2)
1022+
{
1023+
return AbortWithWrongNumberOfArguments("VGETATTR");
1024+
}
10221025

10231026
while (!RespWriteUtils.TryWriteDirect(CmdStrings.RESP_OK, ref dcurr, dend))
10241027
SendAndReset();
@@ -1099,7 +1102,32 @@ private bool NetworkVISMEMBER<TGarnetApi>(ref TGarnetApi storageApi)
10991102
return AbortWithErrorMessage("ERR Vector Set (preview) commands are not enabled");
11001103
}
11011104

1102-
// TODO: implement!
1105+
if (parseState.Count != 2)
1106+
{
1107+
return AbortWithWrongNumberOfArguments("VISMEMBER");
1108+
}
1109+
1110+
var key = parseState.GetArgSliceByRef(0);
1111+
var elementId = parseState.GetArgSliceByRef(1);
1112+
var res = storageApi.VectorSetCheckIsMember(key, elementId);
1113+
if (res == GarnetStatus.OK)
1114+
{
1115+
WriteBooleanTrue();
1116+
return true;
1117+
}
1118+
else if (res == GarnetStatus.NOTFOUND)
1119+
{
1120+
WriteBooleanFalse();
1121+
return true;
1122+
}
1123+
else if (res == GarnetStatus.WRONGTYPE)
1124+
{
1125+
return AbortVectorSetWrongType();
1126+
}
1127+
else if (res == GarnetStatus.BADSTATE)
1128+
{
1129+
return AbortVectorSetPartiallyDeleted(ref key);
1130+
}
11031131

11041132
while (!RespWriteUtils.TryWriteDirect(CmdStrings.RESP_OK, ref dcurr, dend))
11051133
SendAndReset();

‎libs/server/Resp/Vector/VectorManager.cs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -849,6 +849,21 @@ internal bool TryGetEmbedding(ReadOnlySpan<byte> indexValue, ReadOnlySpan<byte>
849849
}
850850
}
851851

852+
/// <summary>
853+
/// Checks whether a specified element exists in a Vector Set.
854+
/// </summary>
855+
/// <param name="indexValue">The raw index value containing Vector Set metadata.</param>
856+
/// <param name="element">The element identifier to check for existence.</param>
857+
/// <returns>
858+
/// <see langword="true"/> if the element exists in the Vector Set; otherwise, <see langword="false"/>.
859+
/// </returns>
860+
internal bool CheckElementExists(ReadOnlySpan<byte> indexValue, ReadOnlySpan<byte> element)
861+
{
862+
AssertHaveStorageSession();
863+
ReadIndex(indexValue, out var context, out var dimensions, out _, out _, out _, out _, out _, out var indexPtr, out _);
864+
return Service.CheckExternalIdValid(context, indexPtr, element);
865+
}
866+
852867
/// <summary>
853868
/// Determine the dimensions of a vector given its <see cref="VectorValueType"/> and its raw data.
854869
/// </summary>

‎libs/server/Storage/Functions/MainStore/PrivateMethods.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ void CopyRespToWithInput(ref RawStringInput input, ref SpanByte value, ref SpanB
125125
case RespCommand.VDIM:
126126
case RespCommand.GET:
127127
case RespCommand.VINFO:
128+
case RespCommand.VISMEMBER:
128129
// Get value without RESP header; exclude expiration
129130
if (value.LengthWithoutMetadata <= dst.Length)
130131
{

‎libs/server/Storage/Session/MainStore/VectorStoreOps.cs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,5 +341,28 @@ internal unsafe GarnetStatus VectorSetInfo(SpanByte key,
341341
return GarnetStatus.OK;
342342
}
343343
}
344+
345+
/// <summary>
346+
/// Check if vector exists in vector set
347+
/// </summary>
348+
[SkipLocalsInit]
349+
internal unsafe GarnetStatus VectorSetCheckIsMember(SpanByte key, ReadOnlySpan<byte> element)
350+
{
351+
parseState.InitializeWithArgument(new(ref key));
352+
353+
var input = new RawStringInput(RespCommand.VISMEMBER, ref parseState);
354+
Span<byte> indexSpan = stackalloc byte[VectorManager.IndexSizeBytes];
355+
using (vectorManager.ReadVectorIndex(this, ref key, ref input, indexSpan, out var status))
356+
{
357+
if (status != GarnetStatus.OK)
358+
{
359+
return status;
360+
}
361+
362+
return vectorManager.CheckElementExists(indexSpan, element) ?
363+
GarnetStatus.OK :
364+
GarnetStatus.NOTFOUND;
365+
}
366+
}
344367
}
345368
}

‎test/Garnet.test/RespVectorSetTests.cs‎

Lines changed: 60 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,8 @@ public void WrongTypeForVectorSetOpsOnNonVectorSetKeys()
120120
exc = ClassicAssert.Throws<RedisServerException>(() => db.Execute("VINFO", ["foo"]));
121121
break;
122122
case RespCommand.VISMEMBER:
123-
// TODO: Implement when VISMEMBER works
124-
continue;
123+
exc = ClassicAssert.Throws<RedisServerException>(() => db.Execute("VISMEMBER", ["foo", "bar"]));
124+
break;
125125
case RespCommand.VLINKS:
126126
// TODO: Implement when VLINKS works
127127
continue;
@@ -1746,6 +1746,7 @@ public void VINFO()
17461746
}
17471747
}
17481748
}
1749+
17491750

17501751
static object[] GenerateVADDOptions(string key, string quantizer, int reduce, int buildExplorationFactor, int numLinks, object[] values, byte[] elementId)
17511752
{
@@ -1791,6 +1792,63 @@ static Dictionary<string, string> BuildDictionaryFromResponse(RedisValue[] respo
17911792
}
17921793
}
17931794

1795+
[Test]
1796+
public void VISMEMBER()
1797+
{
1798+
using var redis = ConnectionMultiplexer.Connect(TestUtils.GetConfig());
1799+
var db = redis.GetDatabase(0);
1800+
1801+
// VISMEMBER on non-existent vector set returns 0
1802+
var res = db.Execute("VISMEMBER", ["nonexistent", new byte[] { 1, 0, 0, 0 }]);
1803+
ClassicAssert.AreEqual(0, (int)res);
1804+
1805+
// Create a vector set with one element
1806+
res = db.Execute("VADD", ["foo", "VALUES", "3", "1.0", "2.0", "3.0", new byte[] { 1, 0, 0, 0 }, "XPREQ8"]);
1807+
ClassicAssert.AreEqual(1, (int)res);
1808+
1809+
// VISMEMBER on element that exists returns 1
1810+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 1, 0, 0, 0 }]);
1811+
ClassicAssert.AreEqual(1, (int)res);
1812+
1813+
// VISMEMBER on element that doesn't exist in the set returns 0
1814+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 2, 0, 0, 0 }]);
1815+
ClassicAssert.AreEqual(0, (int)res);
1816+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 0, 0, 0, 0 }]);
1817+
ClassicAssert.AreEqual(0, (int)res);
1818+
1819+
// Remove the element
1820+
res = db.Execute("VREM", ["foo", new byte[] { 1, 0, 0, 0 }]);
1821+
ClassicAssert.AreEqual(1, (int)res);
1822+
1823+
// VISMEMBER on deleted element returns 0
1824+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 1, 0, 0, 0 }]);
1825+
ClassicAssert.AreEqual(0, (int)res);
1826+
1827+
// Add another member with different ID
1828+
res = db.Execute("VADD", ["foo", "VALUES", "3", "4.0", "5.0", "6.0", new byte[] { 2, 0, 0, 0 }, "XPREQ8"]);
1829+
ClassicAssert.AreEqual(1, (int)res);
1830+
1831+
// VISMEMBER on new element returns 1
1832+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 2, 0, 0, 0 }]);
1833+
ClassicAssert.AreEqual(1, (int)res);
1834+
1835+
// VISMEMBER on old element still returns 0
1836+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 1, 0, 0, 0 }]);
1837+
ClassicAssert.AreEqual(0, (int)res);
1838+
1839+
// Add first element back
1840+
res = db.Execute("VADD", ["foo", "VALUES", "3", "1.0", "2.0", "3.0", new byte[] { 1, 0, 0, 0 }, "XPREQ8"]);
1841+
ClassicAssert.AreEqual(1, (int)res);
1842+
1843+
// VISMEMBER on first element returns 1
1844+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 1, 0, 0, 0 }]);
1845+
ClassicAssert.AreEqual(1, (int)res);
1846+
1847+
// VISMEMBER on second element still returns 1
1848+
res = db.Execute("VISMEMBER", ["foo", new byte[] { 2, 0, 0, 0 }]);
1849+
ClassicAssert.AreEqual(1, (int)res);
1850+
}
1851+
17941852
[Test]
17951853
public void VREM()
17961854
{

0 commit comments

Comments
 (0)