Implement VSETATTR

Directory.Packages.props

@@ -32,7 +32,7 @@
     <PackageVersion Include="System.Numerics.Tensors" Version="10.0.5" />
     <PackageVersion Include="Microsoft.Extensions.Hosting" Version="10.0.5" />
     <PackageVersion Include="Microsoft.Extensions.Hosting.WindowsServices" Version="10.0.5" />
-    <PackageVersion Include="diskann-garnet" Version="4.0.2" />
+    <PackageVersion Include="diskann-garnet" Version="4.0.4" />
     <PackageVersion Include="Microsoft.VisualStudio.Threading.Analyzers" Version="17.14.15" />
   </ItemGroup>
 </Project>

libs/server/AOF/AofProcessor.cs

@@ -558,6 +558,13 @@ static void StoreRMW<TStringContext>(
                     vectorManager.HandleVectorSetRemoveReplication(activeServerSession.storageSession, preparedParameters.Key, ref stringInput);
                     return;
                 }
+
+                // VSETATTR too
+                if (stringInput.header.cmd == RespCommand.VSETATTR)
+                {
+                    vectorManager.HandleVectorSetSetAttributeReplication(activeServerSession.storageSession, preparedParameters.Key, ref stringInput);
+                    return;
+                }
             }
 
             // RangeIndex commands need actual execution on replay

libs/server/Resp/Vector/DiskANNService.cs

@@ -318,6 +318,17 @@ public bool CheckExternalIdValid(ulong context, nint index, ReadOnlySpan<byte> e
 
             return NativeDiskANNMethods.check_external_id_valid(context, index, (nint)external_id_data, (nuint)external_id_len) == 1;
         }
+
+        public bool SetAttribute(ulong context, nint index, ReadOnlySpan<byte> externalId, ReadOnlySpan<byte> attribute)
+        {
+            var external_id_data = Unsafe.AsPointer(ref MemoryMarshal.GetReference(externalId));
+            var external_id_len = externalId.Length;
+
+            var attribute_data = Unsafe.AsPointer(ref MemoryMarshal.GetReference(attribute));
+            var attribute_len = attribute.Length;
+
+            return NativeDiskANNMethods.set_attribute(context, index, (nint)external_id_data, (nuint)external_id_len, (nint)attribute_data, (nuint)attribute_len) == 1;
+        }
     }
 
     public static partial class NativeDiskANNMethods

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

@@ -140,6 +140,31 @@ internal void ReplicateVectorSetRemove(ReadOnlySpan<byte> key, ReadOnlySpan<byte
             }
         }
 
+        internal void ReplicateVectorSetSetAttribute(ReadOnlySpan<byte> key, ReadOnlySpan<byte> element, ReadOnlySpan<byte> attribute, ref StringInput input, ref StringBasicContext context)
+        {
+            Debug.Assert(input.header.cmd == RespCommand.VSETATTR, "Shouldn't be called with anything but VSETATTR inputs");
+
+            var inputCopy = input;
+            inputCopy.arg1 = VSETATTRAppendLogArg;
+
+            inputCopy.parseState.InitializeWithArguments(PinnedSpanByte.FromPinnedSpan(element), PinnedSpanByte.FromPinnedSpan(attribute));
+
+            ExceptionInjectionHelper.ResetAndWait(ExceptionInjectionType.VectorSet_Pause_Before_Synthetic_Replication_Rmw);
+
+            var res = context.RMW((FixedSpanByteKey)key, ref inputCopy);
+
+            if (res.IsPending)
+            {
+                CompletePending(ref res, ref context);
+            }
+
+            if (!res.IsCompletedSuccessfully)
+            {
+                logger?.LogCritical("Failed to inject replication write for VSETATTR into log, result was {res}", res);
+                throw new GarnetException("Couldn't synthesize Vector Set attribute set operation for replication, data loss will occur");
+            }
+        }
+
         /// <summary>
         /// Vector Set adds are phrased as reads (once the index is created), so they require special handling.
         /// 
@@ -511,6 +536,31 @@ internal void HandleVectorSetRemoveReplication(StorageSession storageSession, Re
             }
         }
 
+        /// <summary>
+        /// Vector Set attribute sets are phrased as reads (once the index is created), so they require special handling.
+        /// 
+        /// Operations that are faked up by <see cref="ReplicateVectorSetSetAttribute"/> running on the Primary get diverted here on a Replica.
+        /// </summary>
+        internal void HandleVectorSetSetAttributeReplication(StorageSession storageSession, ReadOnlySpan<byte> key, ref StringInput input)
+        {
+            Span<byte> indexSpan = stackalloc byte[IndexSizeBytes];
+            var element = input.parseState.GetArgSliceByRef(0);
+            var attribute = input.parseState.GetArgSliceByRef(1);
+
+            var inputCopy = input;
+            inputCopy.arg1 = default;
+
+            using (ReadVectorIndex(storageSession, key, ref inputCopy, indexSpan, out var status))
+            {
+                Debug.Assert(status == GarnetStatus.OK, "Replication should only occur when a setattr is successful, so index must exist");
+
+                if (!TrySetAttribute(indexSpan, element, attribute))
+                {
+                    throw new GarnetException("Failed to set attribute on vector set during AOF sync, this should never happen but will cause data loss if it does");
+                }
+            }
+        }
+
         /// <summary>
         /// Wait until all ops passed to <see cref="HandleVectorSetAddReplication"/> have completed.
         /// </summary>

libs/server/Resp/Vector/VectorManager.cs

@@ -53,6 +53,7 @@ public sealed partial class VectorManager : IDisposable
         internal const long MigrateIndexKeyLogArg = MigrateElementKeyLogArg + 1; // AOF: YES. InitialUpdater: YES (empty dummy key).
         internal const long VADDSetFlagsArg = MigrateIndexKeyLogArg + 1; // AOF: YES. InitialUpdater: NO (record must exist).
         internal const long CreateIndexArg = VADDSetFlagsArg + 1; // New stub record creation. AOF: NO. InitialUpdater: YES.
+        internal const long VSETATTRAppendLogArg = CreateIndexArg + 1; // User VSETATTR update, replayed on replicas. AOF: Yes. InitialUpdater: NO.
 
         /// <summary>
         /// Byte stored on log records to distinguish the INDEX key as a Vector Set
@@ -601,13 +602,22 @@ internal VectorManagerResult TryRemove(ReadOnlySpan<byte> indexValue, ReadOnlySp
         {
             AssertHaveStorageSession();
 
-            ReadIndex(indexValue, out var context, out _, out _, out var quantType, out _, out _, out _, out _, out var indexPtr);
+            ReadIndex(indexValue, out var context, out _, out _, out _, out _, out _, out _, out _, out var indexPtr);
 
             var del = Service.Remove(context, indexPtr, element);
 
             return del ? VectorManagerResult.OK : VectorManagerResult.MissingElement;
         }
 
+        internal bool TrySetAttribute(ReadOnlySpan<byte> indexValue, ReadOnlySpan<byte> element, ReadOnlySpan<byte> attribute)
+        {
+            AssertHaveStorageSession();
+
+            ReadIndex(indexValue, out var context, out _, out _, out _, out _, out _, out _, out _, out var indexPtr);
+
+            return Service.SetAttribute(context, indexPtr, element, attribute);
+        }
+
         /// <summary>
         /// Request deletion of a Vector Set given the VALUE of the index key.
         /// </summary>

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

@@ -108,6 +108,7 @@ void CopyRespToWithInput<TSourceLogRecord>(in TSourceLogRecord srcLogRecord, ref
                 case RespCommand.VINFO:
                 case RespCommand.VREM:
                 case RespCommand.VDIM:
+                case RespCommand.VSETATTR:
                 case RespCommand.GET:
                 case RespCommand.RIGET:
                 case RespCommand.RISET:

libs/server/Storage/Functions/MainStore/RMWMethods.cs

@@ -880,7 +880,17 @@ private readonly IPUResult InPlaceUpdaterWorker(ref LogRecord logRecord, ref Str
                     // However, we do synthesize some (pointless) writes to implement replication
                     // in a similar manner to VADD.
 
-                    Debug.Assert(input.arg1 == VectorManager.VREMAppendLogArg, "VREM in place update should only happen for replication");                    // Ignore everything else
+                    Debug.Assert(input.arg1 == VectorManager.VREMAppendLogArg, "VREM in place update should only happen for replication");                      // Ignore everything else
+                    return IPUResult.Succeeded;
+                case RespCommand.VSETATTR:
+                    // Same rationale as the VADD & VREM cases above.
+                    if (logRecord.RecordType != VectorManager.RecordType)
+                    {
+                        rmwInfo.Action = RMWAction.CancelOperation;
+                        return IPUResult.Failed;
+                    }
+
+                    Debug.Assert(input.arg1 == VectorManager.VSETATTRAppendLogArg, "VSETATTR in place update should only happen for replication");              // Ignore everything else
                     return IPUResult.Succeeded;
                 default:
                     if (cmd > RespCommandExtensions.LastValidCommand)
@@ -1414,6 +1424,15 @@ public readonly bool CopyUpdater<TSourceLogRecord>(in TSourceLogRecord srcLogRec
                     oldValue.CopyTo(dstLogRecord.ValueSpan);
                     break;
 
+                case RespCommand.VSETATTR:
+                    // NeedCopyUpdate cancels when the record is no longer an index, so CopyUpdater is only reached for a genuine index record.
+                    Debug.Assert(srcLogRecord.RecordType == VectorManager.RecordType, "CopyUpdater reached for VSETATTR on a non-index record");
+                    Debug.Assert(input.arg1 == VectorManager.VSETATTRAppendLogArg, "Unexpected CopyUpdater call on VSETATTR key");
+
+                    // Always copy to avoid corruption of the index record - otherwise the allocated destination will contain garbage data
+                    oldValue.CopyTo(dstLogRecord.ValueSpan);
+                    break;
+
                 default:
                     if (input.header.cmd > RespCommandExtensions.LastValidCommand)
                     {

libs/server/Storage/Functions/MainStore/VarLenInputMethods.cs

@@ -361,7 +361,8 @@ public RecordFieldInfo GetRMWModifiedFieldInfo<TSourceLogRecord>(in TSourceLogRe
 
                     case RespCommand.VADD:
                     case RespCommand.VREM:
-                        if (input.arg1 is VectorManager.VADDAppendLogArg or VectorManager.VREMAppendLogArg or VectorManager.RecreateIndexArg or VectorManager.VADDSetFlagsArg)
+                    case RespCommand.VSETATTR:
+                        if (input.arg1 is VectorManager.VADDAppendLogArg or VectorManager.VREMAppendLogArg or VectorManager.RecreateIndexArg or VectorManager.VADDSetFlagsArg or VectorManager.VSETATTRAppendLogArg)
                         {
                             // A copy-update of the index key copies the whole index value to the new record: this is
                             // triggered when a CU is forced on the index record - during replication (VADD/VREM append

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

@@ -164,7 +164,7 @@ sealed partial class StorageSession : IDisposable
         /// Implement Vector Set Add - this may also create a Vector Set if one does not already exist.
         /// </summary>
         [SkipLocalsInit]
-        public unsafe GarnetStatus VectorSetAdd(PinnedSpanByte key, int reduceDims, VectorValueType valueType, PinnedSpanByte values, PinnedSpanByte element, VectorQuantType quantizer, int buildExplorationFactor, PinnedSpanByte attributes, int numLinks, VectorDistanceMetricType distanceMetric, out VectorManagerResult result, out ReadOnlySpan<byte> errorMsg)
+        public GarnetStatus VectorSetAdd(PinnedSpanByte key, int reduceDims, VectorValueType valueType, PinnedSpanByte values, PinnedSpanByte element, VectorQuantType quantizer, int buildExplorationFactor, PinnedSpanByte attributes, int numLinks, VectorDistanceMetricType distanceMetric, out VectorManagerResult result, out ReadOnlySpan<byte> errorMsg)
         {
             var dims =
                 valueType switch
@@ -216,7 +216,7 @@ public unsafe GarnetStatus VectorSetAdd(PinnedSpanByte key, int reduceDims, Vect
         /// Implement Vector Set Remove - returns not found if the element is not present, or the vector set does not exist.
         /// </summary>
         [SkipLocalsInit]
-        public unsafe GarnetStatus VectorSetRemove(PinnedSpanByte key, PinnedSpanByte element)
+        public GarnetStatus VectorSetRemove(PinnedSpanByte key, PinnedSpanByte element)
         {
             parseState.InitializeWithArgument(key);
 
@@ -255,7 +255,7 @@ public unsafe GarnetStatus VectorSetRemove(PinnedSpanByte key, PinnedSpanByte el
         [SkipLocalsInit]
         public GarnetStatus VectorSetSetAttribute(PinnedSpanByte key, PinnedSpanByte element, PinnedSpanByte attribute)
         {
-            parseState.InitializeWithArgument(key);
+            parseState.InitializeWithArguments([key, element, attribute]);
 
             var input = new StringInput(RespCommand.VSETATTR, ref parseState);
             Span<byte> indexSpan = stackalloc byte[VectorManager.IndexSizeBytes];
@@ -266,7 +266,13 @@ public GarnetStatus VectorSetSetAttribute(PinnedSpanByte key, PinnedSpanByte ele
                     return status;
                 }
 
-                // TODO: Implement!
+                if (vectorManager.TrySetAttribute(indexSpan, element, attribute))
+                {
+                    // On successful update, we need to manually replicate the write
+                    vectorManager.ReplicateVectorSetSetAttribute(key, element, attribute, ref input, ref stringBasicContext);
+
+                    return GarnetStatus.OK;
+                }
 
                 return GarnetStatus.NOTFOUND;
             }
@@ -276,7 +282,7 @@ public GarnetStatus VectorSetSetAttribute(PinnedSpanByte key, PinnedSpanByte ele
         /// Perform a similarity search on an existing Vector Set given a vector as a bunch of floats.
         /// </summary>
         [SkipLocalsInit]
-        public unsafe GarnetStatus VectorSetValueSimilarity(PinnedSpanByte key, VectorValueType valueType, PinnedSpanByte values, int count, float delta, int searchExplorationFactor, ReadOnlySpan<byte> filter, int maxFilteringEffort, bool includeAttributes, ref SpanByteAndMemory outputIds, out VectorIdFormat outputIdFormat, out ReadOnlySpan<byte> errorMsg, ref SpanByteAndMemory outputDistances, ref SpanByteAndMemory outputAttributes, out VectorManagerResult result, ref SpanByteAndMemory filterBitmap)
+        public GarnetStatus VectorSetValueSimilarity(PinnedSpanByte key, VectorValueType valueType, PinnedSpanByte values, int count, float delta, int searchExplorationFactor, ReadOnlySpan<byte> filter, int maxFilteringEffort, bool includeAttributes, ref SpanByteAndMemory outputIds, out VectorIdFormat outputIdFormat, out ReadOnlySpan<byte> errorMsg, ref SpanByteAndMemory outputDistances, ref SpanByteAndMemory outputAttributes, out VectorManagerResult result, ref SpanByteAndMemory filterBitmap)
         {
             parseState.InitializeWithArgument(key);
 
@@ -303,7 +309,7 @@ public unsafe GarnetStatus VectorSetValueSimilarity(PinnedSpanByte key, VectorVa
         /// Perform a similarity search on an existing Vector Set given an element that is already in the Vector Set.
         /// </summary>
         [SkipLocalsInit]
-        public unsafe GarnetStatus VectorSetElementSimilarity(PinnedSpanByte key, ReadOnlySpan<byte> element, int count, float delta, int searchExplorationFactor, ReadOnlySpan<byte> filter, int maxFilteringEffort, bool includeAttributes, ref SpanByteAndMemory outputIds, out VectorIdFormat outputIdFormat, ref SpanByteAndMemory outputDistances, ref SpanByteAndMemory outputAttributes, out VectorManagerResult result, ref SpanByteAndMemory filterBitmap)
+        public GarnetStatus VectorSetElementSimilarity(PinnedSpanByte key, ReadOnlySpan<byte> element, int count, float delta, int searchExplorationFactor, ReadOnlySpan<byte> filter, int maxFilteringEffort, bool includeAttributes, ref SpanByteAndMemory outputIds, out VectorIdFormat outputIdFormat, ref SpanByteAndMemory outputDistances, ref SpanByteAndMemory outputAttributes, out VectorManagerResult result, ref SpanByteAndMemory filterBitmap)
         {
             parseState.InitializeWithArgument(key);
 
@@ -379,7 +385,7 @@ public GarnetStatus VectorSetRawEmbedding(PinnedSpanByte key, ReadOnlySpan<byte>
         }
 
         [SkipLocalsInit]
-        internal unsafe GarnetStatus VectorSetDimensions(PinnedSpanByte key, out int dimensions)
+        internal GarnetStatus VectorSetDimensions(PinnedSpanByte key, out int dimensions)
         {
             parseState.InitializeWithArgument(key);
 
@@ -406,7 +412,7 @@ internal unsafe GarnetStatus VectorSetDimensions(PinnedSpanByte key, out int dim
         /// Get debugging information about the VectorSet
         /// </summary>
         [SkipLocalsInit]
-        internal unsafe GarnetStatus VectorSetInfo(PinnedSpanByte key,
+        internal GarnetStatus VectorSetInfo(PinnedSpanByte key,
             out VectorQuantType quantType,
             out VectorDistanceMetricType distanceMetricType,
             out uint vectorDimensions,
@@ -549,7 +555,7 @@ internal GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref
         /// Get the attributes associated with an element in the VectorSet
         /// </summary>
         [SkipLocalsInit]
-        internal unsafe GarnetStatus VectorSetGetAttribute(PinnedSpanByte key, PinnedSpanByte elementId, ref SpanByteAndMemory outputAttributes)
+        internal GarnetStatus VectorSetGetAttribute(PinnedSpanByte key, PinnedSpanByte elementId, ref SpanByteAndMemory outputAttributes)
         {
             parseState.InitializeWithArgument(key);
 

test/cluster/Garnet.test.cluster.vectorsets/VectorSets/ClusterVectorSetTests.cs

@@ -2312,6 +2312,37 @@ public async Task ReplicaReplaysSyntheticVAddAgainstStringKeyAfterRacedDeleteAsy
             ClassicAssert.AreEqual(primaryType, replicaType, "Replica diverged from primary after replaying a synthetic VADD against a String key");
         }
 
+        [Test]
+        public async Task VSETATTRReplicatesAsync()
+        {
+            const int PrimaryIndex = 0;
+            const int SecondaryIndex = 1;
+            const string Key = nameof(VSETATTRReplicatesAsync);
+            const string Element = Key + "_Element";
+
+            _ = await SimpleSetupClusterAsync(DefaultShards, primaryCount: 1, replicaCount: 1, useTLS: false).ConfigureAwait(false);
+
+            var primary = (IPEndPoint)context.endpoints[PrimaryIndex];
+            var secondary = (IPEndPoint)context.endpoints[SecondaryIndex];
+
+            ClassicAssert.AreEqual("master", context.clusterTestUtils.RoleCommand(primary).Value);
+            ClassicAssert.AreEqual("slave", context.clusterTestUtils.RoleCommand(secondary).Value);
+
+            await using var connection = await ConnectionMultiplexer.ConnectAsync(context.clusterTestUtils.GetRedisConfig(context.endpoints)).ConfigureAwait(false);
+            var primaryServer = connection.GetServer(primary);
+            var secondaryServer = connection.GetServer(secondary);
+
+            var addRes = (int)await primaryServer.ExecuteAsync("VADD", [Key, "VALUES", "3", "1", "2", "3", Element]).ConfigureAwait(false);
+            ClassicAssert.AreEqual(1, addRes);
+            var setRes = (int)await primaryServer.ExecuteAsync("VSETATTR", [Key, Element, "{\"foo\":\"bar\"}"]).ConfigureAwait(false);
+            ClassicAssert.AreEqual(1, setRes);
+
+            context.clusterTestUtils.WaitForReplicaAofSync(PrimaryIndex, SecondaryIndex);
+
+            var getRes = (string)await secondaryServer.ExecuteAsync("VGETATTR", [Key, Element]).ConfigureAwait(false);
+            ClassicAssert.AreEqual("{\"foo\":\"bar\"}", getRes);
+        }
+
         private async Task<(List<ShardInfo> Shards, List<ushort> Slots)> SimpleSetupClusterAsync(int shardCount, int primaryCount, int replicaCount, bool onDemandCheckpoint = false, bool useTLS = true)
         {
             context.CreateInstances(shardCount, useTLS: useTLS, enableAOF: true, AofMemorySize: DefaultAOFMemorySize, OnDemandCheckpoint: onDemandCheckpoint, sublogCount: sublogCount, threadPoolMinIOCompletionThreads: 512);