Update DiskANN FFI

libs/cluster/Server/Migration/MigrateOperation.cs

@@ -107,7 +107,6 @@ public async Task<bool> TransmitSlotsAsync()
                     input.arg1 = session.NetworkBufferSettings.sendBufferSize - common.NetworkBufferSettings.SendBufferOverheadReserve;
 
                     VectorInput vectorInput = new();
-                    vectorInput.AlignmentExpected = true; // We're moving DiskANN sourced data, so alignment is expected
                     vectorInput.MaxMigrationHeapAllocationSize = session.NetworkBufferSettings.sendBufferSize - common.NetworkBufferSettings.SendBufferOverheadReserve;
 
                     foreach (var (ns, key, hasNs) in sketch.argSliceVector)

libs/server/API/GarnetApi.cs

@@ -386,12 +386,13 @@ public GarnetStatus VectorSetIsMember(PinnedSpanByte key, PinnedSpanByte element
         => storageSession.VectorSetIsMember(key, element);
 
         /// <inheritdoc/>
-        public GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, bool withScores, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults)
-        => storageSession.VectorSetLinks(key, element, withScores, ref idResults, ref distanceResults);
+        public GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults)
+        => storageSession.VectorSetLinks(key, element, ref idResults, ref distanceResults);
 
         /// <inheritdoc/>
-        public GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults)
-        => storageSession.VectorSetRandomMembers(key, count, ref idResults);
+        public GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults, out int actualCount)
+        => storageSession.VectorSetRandomMembers(key, count, ref idResults, out actualCount);
+
         /// <inheritdoc />
         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)
         => storageSession.VectorSetAdd(key, reduceDims, valueType, values, element, quantizer, buildExplorationFactor, attributes, numLinks, distanceMetric, out result, out errorMsg);

libs/server/API/GarnetWatchApi.cs

@@ -638,17 +638,17 @@ public GarnetStatus VectorSetIsMember(PinnedSpanByte key, PinnedSpanByte element
         }
 
         /// <inheritdoc/>
-        public GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, bool withScores, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults)
+        public GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults)
         {
             garnetApi.WATCH(key, StoreType.Main);
-            return garnetApi.VectorSetLinks(key, element, withScores, ref idResults, ref distanceResults);
+            return garnetApi.VectorSetLinks(key, element, ref idResults, ref distanceResults);
         }
 
         /// <inheritdoc/>
-        public GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults)
+        public GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults, out int actualCount)
         {
             garnetApi.WATCH(key, StoreType.Main);
-            return garnetApi.VectorSetRandomMembers(key, count, ref idResults);
+            return garnetApi.VectorSetRandomMembers(key, count, ref idResults, out actualCount);
         }
 
         /// <inheritdoc/>

libs/server/API/IGarnetApi.cs

@@ -2142,9 +2142,9 @@ public bool IterateStore<TScanFunctions>(ref TScanFunctions scanFunctions, ref l
         /// <summary>
         /// For a given element, find all neighbors and (optionally) the distance to those neighbors.
         /// 
-        /// On success, <paramref name="idResults"/> has length prefixed element names, and <paramref name="distanceResults"/> (if <paramref name="withScores"/> is true) has a float for each of those elements.
+        /// On success, <paramref name="idResults"/> has length prefixed element names, and <paramref name="distanceResults"/> has a float for each of those elements.
         /// </summary>
-        GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, bool withScores, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults);
+        GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults);
 
         /// <summary>
         /// Fetch random elements from the given Vector Set.
@@ -2155,7 +2155,7 @@ public bool IterateStore<TScanFunctions>(ref TScanFunctions scanFunctions, ref l
         /// 
         /// On success, <paramref name="idResults"/> has length prefixed element names.
         /// </summary>
-        GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults);
+        GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults, out int actualCount);
 
         /// <summary>
         /// Perform a similarity search given a vector and these parameters.

libs/server/InputHeader.cs

@@ -633,8 +633,6 @@ public struct VectorInput : IStoreInput
         public nint CallbackContext { get; set; }
         public nint Callback { get; set; }
 
-        public bool AlignmentExpected { get; set; }
-
         [MemberNotNullWhen(returnValue: true, member: nameof(MaxMigrationHeapAllocationSize))]
         public bool IsMigrationRead => MaxMigrationHeapAllocationSize != null;
 

libs/server/Resp/Vector/DiskANNService.cs

@@ -40,11 +40,12 @@ public nint CreateIndex(
             uint buildExplorationFactor,
             uint numLinks,
             VectorDistanceMetricType distanceMetric,
-            delegate* unmanaged[Cdecl]<ulong, uint, nint, nuint, nint, nint, void> readCallback,
+            delegate* unmanaged[Cdecl]<ulong, uint, uint, nint, nuint, nint, nint, void> readCallback,
             delegate* unmanaged[Cdecl]<ulong, nint, nuint, nint, nuint, byte> writeCallback,
             delegate* unmanaged[Cdecl]<ulong, nint, nuint, byte> deleteCallback,
             delegate* unmanaged[Cdecl]<ulong, nint, nuint, nuint, nint, nint, byte> readModifyWriteCallback,
-            delegate* unmanaged[Cdecl]<ulong, uint, byte> filterCallback,
+            delegate* unmanaged[Cdecl]<ulong, nint, nuint, byte> filterCallback,
+            delegate* unmanaged[Cdecl]<ulong, nint, nuint, void> logCallback,
             out bool quantizationRequested
         )
         {
@@ -53,7 +54,7 @@ out bool quantizationRequested
 #endif
             unsafe
             {
-                var ret = NativeDiskANNMethods.create_index(context, dimensions, reduceDims, quantType, distanceMetric, buildExplorationFactor, numLinks, (nint)readCallback, (nint)writeCallback, (nint)deleteCallback, (nint)readModifyWriteCallback, (nint)filterCallback, out quantizationRequested);
+                var ret = NativeDiskANNMethods.create_index(context, dimensions, reduceDims, quantType, distanceMetric, buildExplorationFactor, numLinks, (nint)readCallback, (nint)writeCallback, (nint)deleteCallback, (nint)readModifyWriteCallback, (nint)filterCallback, (nint)logCallback, out quantizationRequested);
 
                 Debug.Assert(ret != 0, "create_index failed, returning a null pointer - this shouldn't be possible");
 
@@ -69,14 +70,15 @@ public nint RecreateIndex(
             uint buildExplorationFactor,
             uint numLinks,
             VectorDistanceMetricType distanceMetricType,
-            delegate* unmanaged[Cdecl]<ulong, uint, nint, nuint, nint, nint, void> readCallback,
+            delegate* unmanaged[Cdecl]<ulong, uint, uint, nint, nuint, nint, nint, void> readCallback,
             delegate* unmanaged[Cdecl]<ulong, nint, nuint, nint, nuint, byte> writeCallback,
             delegate* unmanaged[Cdecl]<ulong, nint, nuint, byte> deleteCallback,
             delegate* unmanaged[Cdecl]<ulong, nint, nuint, nuint, nint, nint, byte> readModifyWriteCallback,
-            delegate* unmanaged[Cdecl]<ulong, uint, byte> filterCallback,
+            delegate* unmanaged[Cdecl]<ulong, nint, nuint, byte> filterCallback,
+            delegate* unmanaged[Cdecl]<ulong, nint, nuint, void> logCallback,
             out bool quantizationRequested
         )
-        => CreateIndex(context, dimensions, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetricType, readCallback, writeCallback, deleteCallback, readModifyWriteCallback, filterCallback, out quantizationRequested);
+        => CreateIndex(context, dimensions, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetricType, readCallback, writeCallback, deleteCallback, readModifyWriteCallback, filterCallback, logCallback, out quantizationRequested);
 
         public void DropIndex(ulong context, nint index)
         {
@@ -113,10 +115,8 @@ public bool BuildQuantizationTable(ulong context, nint index)
             return NativeDiskANNMethods.build_quant_table(context, index) == 1;
         }
 
-        public void BackfillQuantizedVectors(ulong context, nint index, int taskIndex, int taskCount)
-        {
-            NativeDiskANNMethods.backfill_quant_vectors(context, index, (nuint)taskIndex, (nuint)taskCount);
-        }
+        public bool BackfillQuantizedVectors(ulong context, nint index, int taskIndex, int taskCount)
+        => NativeDiskANNMethods.backfill_quant_vectors(context, index, (nuint)taskIndex, (nuint)taskCount) == 1;
 
         public bool Remove(ulong context, nint index, ReadOnlySpan<byte> id)
         {
@@ -298,9 +298,111 @@ out nint continuation
             }
         }
 
-        public int ContinueSearch(ulong context, nint index, nint continuation, Span<byte> outputIds, Span<float> outputDistances, out nint newContinuation)
+        public int SearchNeighbors(
+            ulong context,
+            nint index,
+            ReadOnlySpan<byte> id,
+            SpanByteAndMemory outputIds,
+            SpanByteAndMemory outputDistances,
+            out nint continuation
+        )
         {
-            throw new NotImplementedException();
+            var id_data = Unsafe.AsPointer(ref MemoryMarshal.GetReference(id));
+            var id_len = id.Length;
+
+            void* output_ids;
+            void* output_distances;
+
+            GCHandle? outputIdsHandle = null;
+            GCHandle? outputDistancesHandle = null;
+            try
+            {
+                if (!outputIds.IsSpanByte)
+                {
+                    var getRes = MemoryMarshal.TryGetArray<byte>(outputIds.Memory.Memory, out var arrSeg);
+                    Debug.Assert(getRes, "Should always be able to get array to pin");
+
+                    outputIdsHandle = GCHandle.Alloc(arrSeg.Array, GCHandleType.Pinned);
+                    output_ids = Unsafe.AsPointer(ref MemoryMarshal.GetArrayDataReference(arrSeg.Array));
+                }
+                else
+                {
+                    outputIdsHandle = null;
+                    output_ids = Unsafe.AsPointer(ref MemoryMarshal.GetReference(outputIds.Span));
+                }
+
+                var output_ids_len = outputIds.Length;
+
+                if (!outputDistances.IsSpanByte)
+                {
+                    var getRes = MemoryMarshal.TryGetArray<byte>(outputDistances.Memory.Memory, out var arrSeg);
+                    Debug.Assert(getRes, "Should always be able to get array to pin");
+
+                    outputDistancesHandle = GCHandle.Alloc(arrSeg.Array, GCHandleType.Pinned);
+                    output_distances = Unsafe.AsPointer(ref MemoryMarshal.GetArrayDataReference(arrSeg.Array));
+                }
+                else
+                {
+                    outputDistancesHandle = null;
+                    output_distances = Unsafe.AsPointer(ref MemoryMarshal.GetReference(outputDistances.Span));
+                }
+
+                var output_distances_len = outputDistances.Length / sizeof(float);
+
+                continuation = 0;
+                ref var continuationRef = ref continuation;
+                var continuationAddr = (nint)Unsafe.AsPointer(ref continuationRef);
+
+                return NativeDiskANNMethods.search_neighbors(
+                    context,
+                    index,
+                    (nint)id_data,
+                    (nuint)id_len,
+                    (nint)output_ids,
+                    (nuint)output_ids_len,
+                    (nint)output_distances,
+                    (nuint)output_distances_len,
+                    continuationAddr
+                );
+            }
+            finally
+            {
+                outputIdsHandle?.Free();
+                outputDistancesHandle?.Free();
+            }
+        }
+
+        public bool RandomMembers(
+            ulong context,
+            nint index,
+            int count,
+            Span<byte> outputIds
+        )
+        {
+            var output_ids = Unsafe.AsPointer(ref MemoryMarshal.GetReference(outputIds));
+            var output_ids_len = outputIds.Length;
+
+            return NativeDiskANNMethods.random_members(
+                context,
+                index,
+                (uint)count,
+                (nint)output_ids,
+                (nuint)output_ids_len
+            ) == 1;
+        }
+
+        public int ContinueSearch(ulong context, nint index, nint continuation, Span<byte> outputIds, Span<byte> outputDistances, out nint newContinuation)
+        {
+            var output_ids_data = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(outputIds));
+            var output_ids_len = (nuint)outputIds.Length;
+
+            var output_distances_data = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(outputDistances));
+            var output_distances_len = (nuint)outputDistances.Length;
+
+            newContinuation = 0;
+            var newContinuationPtr = (nint)Unsafe.AsPointer(ref newContinuation);
+
+            return NativeDiskANNMethods.continue_search(context, index, continuation, output_ids_data, output_ids_len, output_distances_data, output_distances_len, newContinuationPtr);
         }
 
         public bool CheckInternalIdValid(ulong context, nint index, ReadOnlySpan<byte> internalId)
@@ -356,6 +458,7 @@ public static partial nint create_index(
             nint deleteCallback,
             nint readModifyWriteCallback,
             nint filterCallback,
+            nint logCallback,
             [MarshalAs(UnmanagedType.U1)] out bool quantizationNeeded
         );
 
@@ -431,6 +534,19 @@ public static partial int search_element(
             nint continuation
         );
 
+        [LibraryImport(DISKANN_GARNET)]
+        public static partial int search_neighbors(
+            ulong context,
+            nint index,
+            nint id_data,
+            nuint id_len,
+            nint output_ids,
+            nuint output_ids_len,
+            nint output_distances,
+            nuint output_distances_len,
+            nint continuation
+        );
+
         [LibraryImport(DISKANN_GARNET)]
         public static partial int continue_search(
             ulong context,
@@ -472,11 +588,20 @@ nint index
         );
 
         [LibraryImport(DISKANN_GARNET)]
-        public static partial void backfill_quant_vectors(
+        public static partial byte backfill_quant_vectors(
             ulong context,
             nint index,
             nuint task_index,
             nuint task_count
         );
+
+        [LibraryImport(DISKANN_GARNET)]
+        public static partial byte random_members(
+            ulong context,
+            nint index,
+            uint count,
+            nint output_ids,
+            nuint output_ids_len
+        );
     }
 }

libs/server/Resp/Vector/RespServerSessionVectors.cs

@@ -1706,7 +1706,7 @@ private bool NetworkVLINKS<TGarnetApi>(ref TGarnetApi storageApi)
             var distanceResult = SpanByteAndMemory.FromPinnedSpan(distanceSpace);
             try
             {
-                var res = storageApi.VectorSetLinks(key, element, withScores, ref idResult, ref distanceResult);
+                var res = storageApi.VectorSetLinks(key, element, ref idResult, ref distanceResult);
 
                 switch (res)
                 {
@@ -1720,9 +1720,42 @@ private bool NetworkVLINKS<TGarnetApi>(ref TGarnetApi storageApi)
 
                     case GarnetStatus.OK:
                         {
-                            // TODO: implement!
-                            while (!RespWriteUtils.TryWriteDirect(CmdStrings.RESP_OK, ref dcurr, dend))
-                                SendAndReset();
+                            var numLinks = distanceResult.Length / sizeof(float);
+
+                            WriteArrayLength(numLinks);
+
+                            var remainingIds = idResult.Span;
+                            var remainingScores = distanceResult.Span;
+                            for (var i = 0; i < numLinks; i++)
+                            {
+                                var idLen = BinaryPrimitives.ReadInt32LittleEndian(remainingIds);
+                                var id = remainingIds.Slice(sizeof(int), idLen);
+                                var score = BinaryPrimitives.ReadSingleLittleEndian(remainingScores);
+
+                                if (withScores)
+                                {
+                                    if (respProtocolVersion == 3)
+                                    {
+                                        WriteMapLength(2);
+                                    }
+                                    else
+                                    {
+                                        WriteArrayLength(2);
+                                    }
+
+                                    WriteArrayLength(2);
+                                    WriteBulkString(id);
+                                    WriteDoubleNumeric(score);
+                                }
+                                else
+                                {
+                                    WriteArrayLength(1);
+                                    WriteBulkString(id);
+                                }
+
+                                remainingIds = remainingIds[(sizeof(int) + idLen)..];
+                                remainingScores = remainingScores[sizeof(float)..];
+                            }
                         }
                         break;
                 }
@@ -1769,7 +1802,7 @@ private bool NetworkVRANDMEMBER<TGarnetApi>(ref TGarnetApi storageApi)
             try
             {
 
-                var res = storageApi.VectorSetRandomMembers(key, count, ref idResult);
+                var res = storageApi.VectorSetRandomMembers(key, count, ref idResult, out var actualCount);
 
                 switch (res)
                 {
@@ -1789,9 +1822,18 @@ private bool NetworkVRANDMEMBER<TGarnetApi>(ref TGarnetApi storageApi)
 
                     case GarnetStatus.OK:
                         {
-                            // TODO: implement!
-                            while (!RespWriteUtils.TryWriteDirect(CmdStrings.RESP_OK, ref dcurr, dend))
-                                SendAndReset();
+                            WriteArrayLength(actualCount);
+                            var remainingIds = idResult.ReadOnlySpan;
+
+                            while (!remainingIds.IsEmpty)
+                            {
+                                var idLen = BinaryPrimitives.ReadInt32LittleEndian(remainingIds);
+                                var id = remainingIds.Slice(sizeof(int), idLen);
+
+                                WriteBulkString(id);
+
+                                remainingIds = remainingIds[(sizeof(int) + idLen)..];
+                            }
                         }
                         break;
                 }