From c0d6b9f32b0864576661507894464c3346d34217 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Tue, 28 Jul 2026 15:45:24 -0400 Subject: [PATCH 01/12] add logCallback --- libs/server/Resp/Vector/DiskANNService.cs | 7 +- .../Resp/Vector/VectorManager.Callbacks.cs | 65 +++++++++++++++++++ .../Resp/Vector/VectorManager.Locking.cs | 6 +- .../Resp/Vector/VectorManager.Migration.cs | 2 +- .../DiskANN/DiskANNServiceTests.cs | 29 +++++++-- website/docs/dev/vector-sets.md | 13 +++- 6 files changed, 110 insertions(+), 12 deletions(-) diff --git a/libs/server/Resp/Vector/DiskANNService.cs b/libs/server/Resp/Vector/DiskANNService.cs index 92b8c74be1d..b6be8f81137 100644 --- a/libs/server/Resp/Vector/DiskANNService.cs +++ b/libs/server/Resp/Vector/DiskANNService.cs @@ -45,6 +45,7 @@ public nint CreateIndex( delegate* unmanaged[Cdecl] deleteCallback, delegate* unmanaged[Cdecl] readModifyWriteCallback, delegate* unmanaged[Cdecl] filterCallback, + delegate* unmanaged[Cdecl] 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"); @@ -74,9 +75,10 @@ public nint RecreateIndex( delegate* unmanaged[Cdecl] deleteCallback, delegate* unmanaged[Cdecl] readModifyWriteCallback, delegate* unmanaged[Cdecl] filterCallback, + delegate* unmanaged[Cdecl] 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) { @@ -345,6 +347,7 @@ public static partial nint create_index( nint deleteCallback, nint readModifyWriteCallback, nint filterCallback, + nint logCallback, [MarshalAs(UnmanagedType.U1)] out bool quantizationNeeded ); diff --git a/libs/server/Resp/Vector/VectorManager.Callbacks.cs b/libs/server/Resp/Vector/VectorManager.Callbacks.cs index 8cd0f82a020..8e6b32c84ec 100644 --- a/libs/server/Resp/Vector/VectorManager.Callbacks.cs +++ b/libs/server/Resp/Vector/VectorManager.Callbacks.cs @@ -6,7 +6,9 @@ using System.Diagnostics; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; +using System.Text; using Garnet.common; +using Microsoft.Extensions.Logging; using Tsavorite.core; namespace Garnet.server @@ -242,6 +244,7 @@ internal readonly void CompletePending(ref VectorBasicContext objectContext) private unsafe delegate* unmanaged[Cdecl] DeleteCallbackPtr { get; } = &DeleteCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] ReadModifyWriteCallbackPtr { get; } = &ReadModifyWriteCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] InlineFilterCallbackPtr { get; } = &FilterCallbackUnmanaged; + private unsafe delegate* unmanaged[Cdecl] LogCallbackPtr { get; } = &LogCallbackUnmanaged; /// /// Used to thread the active across p/invoke and reverse p/invoke boundaries into DiskANN. @@ -334,6 +337,68 @@ or VectorQuantType.XBin_I8 or VectorQuantType.XBin_U8 }; } + [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] + private static unsafe void LogCallbackUnmanaged(ulong context, nint logMessage, nuint logMessageLength) + { + const int MaxVectorSetLogNameLength = 64; + + if (ActiveThreadSession == null) + { + // Can't do anything here + return; + } + + var msgUtf8Raw = new ReadOnlySpan((byte*)logMessage, (int)logMessageLength); + var msg = Encoding.UTF8.GetString(msgUtf8Raw); + + var contextNoNs = context & ~(ContextStep - 1); + var nsBits = context & (ContextStep - 1); + + var ns = + nsBits switch + { + DiskANNService.Attributes => nameof(DiskANNService.Attributes), + DiskANNService.ExternalIdMap => nameof(DiskANNService.ExternalIdMap), + DiskANNService.FullVector => nameof(DiskANNService.FullVector), + DiskANNService.InternalIdMap => nameof(DiskANNService.InternalIdMap), + DiskANNService.NeighborList => nameof(DiskANNService.NeighborList), + DiskANNService.QuantizedVector => nameof(DiskANNService.QuantizedVector), + _ => $"!!UNKNOWN ({nsBits})!!", + }; + + string vectorSet; + int args; + if (ActiveThreadSession.parseState.Count > 0) + { + args = ActiveThreadSession.parseState.Count; + var probablyVectorSet = ActiveThreadSession.parseState.GetArgSliceByRef(0).ReadOnlySpan; + + if (probablyVectorSet.Length > MaxVectorSetLogNameLength) + { + vectorSet = $"(Escaped for length ({probablyVectorSet.Length}): {SpanByte.ToShortString(probablyVectorSet, MaxVectorSetLogNameLength)})"; + } + else + { + try + { + vectorSet = Encoding.UTF8.GetString(probablyVectorSet); + } + catch + { + vectorSet = $"(Escaped non-utf8: {SpanByte.ToShortString(probablyVectorSet, MaxVectorSetLogNameLength)})"; + } + } + } + else + { + vectorSet = ""; + args = 0; + } + + // TODO: It'd be nice to get the command in here as well + ActiveThreadSession.vectorManager.logger?.LogWarning("DiskANN Log Message={msg}, Context={contextNoNs}, Namespace={ns}, VectorSet={vectorSet}, CommandArgsCount={args}", msg, contextNoNs, ns, vectorSet, args); + } + [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static unsafe void ReadCallbackUnmanaged( ulong context, diff --git a/libs/server/Resp/Vector/VectorManager.Locking.cs b/libs/server/Resp/Vector/VectorManager.Locking.cs index 5f1b4a9de02..cd23650714c 100644 --- a/libs/server/Resp/Vector/VectorManager.Locking.cs +++ b/libs/server/Resp/Vector/VectorManager.Locking.cs @@ -175,7 +175,7 @@ internal VectorSetLock ReadVectorIndex(StorageSession storageSession, ReadOnlySp bool requestQuantization; unsafe { - newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); } input.header.cmd = RespCommand.VADD; @@ -371,7 +371,7 @@ out GarnetStatus status unsafe { - newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); } input.parseState.EnsureCapacity(12); @@ -403,7 +403,7 @@ out GarnetStatus status unsafe { - newlyAllocatedIndex = Service.CreateIndex(indexContext, dims, reduceDims, quantizer, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.CreateIndex(indexContext, dims, reduceDims, quantizer, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); } input.parseState.EnsureCapacity(12); diff --git a/libs/server/Resp/Vector/VectorManager.Migration.cs b/libs/server/Resp/Vector/VectorManager.Migration.cs index 782e9956293..1c11c54db03 100644 --- a/libs/server/Resp/Vector/VectorManager.Migration.cs +++ b/libs/server/Resp/Vector/VectorManager.Migration.cs @@ -192,7 +192,7 @@ public void HandleMigratedIndexKey( bool requestQuantization; unsafe { - newlyAllocatedIndex = Service.RecreateIndex(context, dimensions, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.RecreateIndex(context, dimensions, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); } var ctxArg = PinnedSpanByte.FromPinnedSpan(MemoryMarshal.Cast(MemoryMarshal.CreateSpan(ref context, 1))); diff --git a/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs b/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs index 3e84c1ec5b2..a767a10f271 100644 --- a/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs +++ b/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs @@ -9,6 +9,7 @@ using System.Linq; using System.Runtime.CompilerServices; using System.Runtime.InteropServices; +using System.Text; using Garnet.server; using NUnit.Framework; using NUnit.Framework.Legacy; @@ -24,6 +25,7 @@ public class DiskANNServiceTests : TestBase private delegate byte DeleteCallbackDelegate(ulong context, nint keyData, nuint keyLength); private delegate byte ReadModifyWriteCallbackDelegate(ulong context, nint keyData, nuint keyLength, nuint writeLength, nint dataCallback, nint dataCallbackContext); private delegate byte InlineFilterCallbackDelegate(ulong context, uint internalId); + private delegate void LogCallbackDelegate(ulong context, nint logMessage, nuint logMessageLength); private sealed class ContextAndKeyComparer : IEqualityComparer<(ulong Context, byte[] Data)> { @@ -159,24 +161,33 @@ unsafe byte ReadModifyWriteCallback(ulong context, nint keyData, nuint keyLength return 1; } - unsafe byte InlineFilterCallback(ulong context, uint internalId) + byte InlineFilterCallback(ulong context, uint internalId) { return 1; } + unsafe void LogCallback(ulong context, nint logMessage, nuint logMessageLength) + { + var text = Encoding.UTF8.GetString(new ReadOnlySpan((byte*)logMessage, (int)logMessageLength)); + + TestContext.Progress.WriteLine($"LogCallback({context}, \"{text}\")"); + } + ReadCallbackDelegate readDel = ReadCallback; WriteCallbackDelegate writeDel = WriteCallback; DeleteCallbackDelegate deleteDel = DeleteCallback; ReadModifyWriteCallbackDelegate rmwDel = ReadModifyWriteCallback; InlineFilterCallbackDelegate filterDel = InlineFilterCallback; + LogCallbackDelegate logDel = LogCallback; var readFuncPtr = Marshal.GetFunctionPointerForDelegate(readDel); var writeFuncPtr = Marshal.GetFunctionPointerForDelegate(writeDel); var deleteFuncPtr = Marshal.GetFunctionPointerForDelegate(deleteDel); var rmwFuncPtr = Marshal.GetFunctionPointerForDelegate(rmwDel); var filterFuncPtr = Marshal.GetFunctionPointerForDelegate(filterDel); + var logFuncPtr = Marshal.GetFunctionPointerForDelegate(logDel); - var rawIndex = NativeDiskANNMethods.create_index(Context, 75, 0, VectorQuantType.XNoQuant_U8, VectorDistanceMetricType.L2, 10, 10, readFuncPtr, writeFuncPtr, deleteFuncPtr, rmwFuncPtr, filterFuncPtr, out _); + var rawIndex = NativeDiskANNMethods.create_index(Context, 75, 0, VectorQuantType.XNoQuant_U8, VectorDistanceMetricType.L2, 10, 10, readFuncPtr, writeFuncPtr, deleteFuncPtr, rmwFuncPtr, filterFuncPtr, logFuncPtr, out _); Span id = [0, 1, 2, 3]; Span elem = Enumerable.Range(0, 75).Select(static x => (byte)x).ToArray(); @@ -361,25 +372,33 @@ unsafe byte ReadModifyWriteCallback(ulong context, nint keyData, nuint keyLength return 1; } - unsafe byte InlineFilterCallback(ulong context, uint internalId) + byte InlineFilterCallback(ulong context, uint internalId) { return 1; } + unsafe void LogCallback(ulong context, nint logMessage, nuint logMessageLength) + { + var text = Encoding.UTF8.GetString(new ReadOnlySpan((byte*)logMessage, (int)logMessageLength)); + + TestContext.Progress.WriteLine($"LogCallback({context}, \"{text}\")"); + } ReadCallbackDelegate readDel = ReadCallback; WriteCallbackDelegate writeDel = WriteCallback; DeleteCallbackDelegate deleteDel = DeleteCallback; ReadModifyWriteCallbackDelegate rmwDel = ReadModifyWriteCallback; InlineFilterCallbackDelegate filterDel = InlineFilterCallback; + LogCallbackDelegate logDel = LogCallback; var readFuncPtr = Marshal.GetFunctionPointerForDelegate(readDel); var writeFuncPtr = Marshal.GetFunctionPointerForDelegate(writeDel); var deleteFuncPtr = Marshal.GetFunctionPointerForDelegate(deleteDel); var rmwFuncPtr = Marshal.GetFunctionPointerForDelegate(rmwDel); var filterFuncPtr = Marshal.GetFunctionPointerForDelegate(filterDel); + var logFuncPtr = Marshal.GetFunctionPointerForDelegate(logDel); - var rawIndex = NativeDiskANNMethods.create_index(Context, 75, 0, VectorQuantType.XNoQuant_U8, VectorDistanceMetricType.L2, 10, 10, readFuncPtr, writeFuncPtr, deleteFuncPtr, rmwFuncPtr, filterFuncPtr, out _); + var rawIndex = NativeDiskANNMethods.create_index(Context, 75, 0, VectorQuantType.XNoQuant_U8, VectorDistanceMetricType.L2, 10, 10, readFuncPtr, writeFuncPtr, deleteFuncPtr, rmwFuncPtr, filterFuncPtr, logFuncPtr, out _); Span id = [0, 1, 2, 3]; Span elem = Enumerable.Range(0, 75).Select(static x => (byte)x).ToArray(); @@ -424,7 +443,7 @@ unsafe byte InlineFilterCallback(ulong context, uint internalId) { NativeDiskANNMethods.drop_index(Context, rawIndex); - rawIndex = NativeDiskANNMethods.create_index(Context, 75, 0, VectorQuantType.XNoQuant_U8, VectorDistanceMetricType.L2, 10, 10, readFuncPtr, writeFuncPtr, deleteFuncPtr, rmwFuncPtr, filterFuncPtr, out _); + rawIndex = NativeDiskANNMethods.create_index(Context, 75, 0, VectorQuantType.XNoQuant_U8, VectorDistanceMetricType.L2, 10, 10, readFuncPtr, writeFuncPtr, deleteFuncPtr, rmwFuncPtr, filterFuncPtr, logFuncPtr, out _); } // Search value diff --git a/website/docs/dev/vector-sets.md b/website/docs/dev/vector-sets.md index 0d0f933f114..ea833422252 100644 --- a/website/docs/dev/vector-sets.md +++ b/website/docs/dev/vector-sets.md @@ -384,11 +384,22 @@ Newly allocated values are guaranteed to be all zeros. The callback returns 1 if the key-value pair was found or created, and 0 if some error occurred. +### Log Callback + +A simple callback providing a way for DiskANN to log richer error messages into Garnet. Its signature is: +```csharp +void LogCallbackUnmanaged(ulong context, nint logMessage, nuint logMessageLength) +``` + +`context` identifies which Vector Set is being operated on AND the associated namespace, and `logMessage` and `logMessageLength` represent a `Span` of the log message. + +The log message is UTF8 encoded text. + ### DiskANN Functions Garnet calls into the following DiskANN functions: - - [x] `nint create_index(ulong context, uint dimensions, uint reduceDims, VectorQuantType quantType, VectorDistanceMetricType distanceMetric, uint buildExplorationFactor, uint numLinks, nint readCallback, nint writeCallback, nint deleteCallback, nint readModifyWriteCallback, nint filterCallback, out bool quantizationNeeded)` + - [x] `nint create_index(ulong context, uint dimensions, uint reduceDims, VectorQuantType quantType, VectorDistanceMetricType distanceMetric, uint buildExplorationFactor, uint numLinks, nint readCallback, nint writeCallback, nint deleteCallback, nint readModifyWriteCallback, nint filterCallback, nint logCallback, out bool quantizationNeeded)` - [x] `void drop_index(ulong context, nint index)` - [x] `DiskANNInsertResult insert(ulong context, nint index, nint id_data, nuint id_len, nint vector_data, nuint vector_len, nint attribute_data, nuint attribute_len)` - [x] `byte remove(ulong context, nint index, nint id_data, nuint id_len)` From 18391c89884fc4aa24f66952fe5f495ed8bc4192 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Tue, 28 Jul 2026 16:56:17 -0400 Subject: [PATCH 02/12] add valueLengthHint to ReadCallbackUnmanaged --- libs/server/Resp/Vector/DiskANNService.cs | 4 +- .../Resp/Vector/VectorManager.Callbacks.cs | 144 ++++-------------- .../Resp/Vector/VectorManager.Locking.cs | 4 - libs/server/Resp/Vector/VectorManager.cs | 11 +- .../RespVectorSetTests.cs | 12 +- website/docs/dev/vector-sets.md | 12 +- 6 files changed, 50 insertions(+), 137 deletions(-) diff --git a/libs/server/Resp/Vector/DiskANNService.cs b/libs/server/Resp/Vector/DiskANNService.cs index b6be8f81137..3fbc99bf897 100644 --- a/libs/server/Resp/Vector/DiskANNService.cs +++ b/libs/server/Resp/Vector/DiskANNService.cs @@ -40,7 +40,7 @@ public nint CreateIndex( uint buildExplorationFactor, uint numLinks, VectorDistanceMetricType distanceMetric, - delegate* unmanaged[Cdecl] readCallback, + delegate* unmanaged[Cdecl] readCallback, delegate* unmanaged[Cdecl] writeCallback, delegate* unmanaged[Cdecl] deleteCallback, delegate* unmanaged[Cdecl] readModifyWriteCallback, @@ -70,7 +70,7 @@ public nint RecreateIndex( uint buildExplorationFactor, uint numLinks, VectorDistanceMetricType distanceMetricType, - delegate* unmanaged[Cdecl] readCallback, + delegate* unmanaged[Cdecl] readCallback, delegate* unmanaged[Cdecl] writeCallback, delegate* unmanaged[Cdecl] deleteCallback, delegate* unmanaged[Cdecl] readModifyWriteCallback, diff --git a/libs/server/Resp/Vector/VectorManager.Callbacks.cs b/libs/server/Resp/Vector/VectorManager.Callbacks.cs index 8e6b32c84ec..e7938f1d5da 100644 --- a/libs/server/Resp/Vector/VectorManager.Callbacks.cs +++ b/libs/server/Resp/Vector/VectorManager.Callbacks.cs @@ -18,49 +18,36 @@ namespace Garnet.server /// public sealed partial class VectorManager { + /// + /// Per-record overhead (RecordInfo + key + length prefixes) added to the value size when computing the + /// initial disk-read size, so the whole record lands in one IO. Generous; the read is sector-aligned downstream. + /// + private const int VectorRecordReadOverheadBytes = 64; + public unsafe #if NET9_0_OR_GREATER ref #endif struct VectorReadBatch : IReadArgBatch { + /// + /// Total number of keys in batch. + /// public int Count { get; } public readonly ReadOnlySpan Parameters => default; - /// - /// Per-term initial disk read size. The big, fixed-size records (FullVector, and the adjacency - /// NeighborList) are sized to the active vector set's geometry () - /// so each lands in a single IO regardless of dimension / M — and different vector sets get different - /// optimal sizes. When the geometry is unset (paths that don't call SetActiveReadGeometry), FullVector - /// uses the configured store/session size (--initial-io-record-size) and the other terms use the - /// small default, avoiding over-reading a full-vector-sized block for a tiny record. - /// + /// public readonly int InitialIORecordSize { [MethodImpl(MethodImplOptions.AggressiveInlining)] - get - { - // Single thread-static read; the struct copy (3 ints) is cheaper than re-reading TLS per branch. - var geometry = ActiveReadGeometry; - switch (NamespaceBytes[0] & 7) - { - case DiskANNService.FullVector: - return geometry.FullVectorIOSize > 0 ? geometry.FullVectorIOSize : KVSettings.UseDefaultInitialIORecordSize; - case DiskANNService.NeighborList: - return geometry.NeighborListIOSize > 0 ? geometry.NeighborListIOSize : IStreamBuffer.DefaultInitialIORecordSize; - case DiskANNService.QuantizedVector: - return geometry.QuantizedVectorIOSize > 0 ? geometry.QuantizedVectorIOSize : IStreamBuffer.DefaultInitialIORecordSize; - default: - return IStreamBuffer.DefaultInitialIORecordSize; - } - } + get; } /// /// Per-term read-copy policy. The small per-element records (NeighborList adjacency, QuantizedVector, - /// internal/external id maps) are copied back into memory on disk read — to + /// internal/external id maps) are copied back into memory on disk read — to /// (the read cache when enabled, else the main-log tail) — so later hops and queries serve them from /// memory. The large raw FullVector and Attributes/Metadata are served from disk (CopyTo=None): quantized /// sets use the raw vector only for reranking, and for no-quant sets caching it yields no net gain once the @@ -69,19 +56,7 @@ public readonly int InitialIORecordSize public readonly ReadCopyOptions ReadCopyOptions { [MethodImpl(MethodImplOptions.AggressiveInlining)] - get - { - switch (NamespaceBytes[0] & 7) - { - case DiskANNService.NeighborList: - case DiskANNService.QuantizedVector: - case DiskANNService.InternalIdMap: - case DiskANNService.ExternalIdMap: - return new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ActiveThreadSession.vectorManager.StubReadCopyTo }; - default: - return new ReadCopyOptions { CopyFrom = ReadCopyFrom.None, CopyTo = ReadCopyTo.None }; - } - } + get; } private readonly ReadOnlySpan NamespaceBytes @@ -115,7 +90,7 @@ private readonly ReadOnlySpan NamespaceBytes private bool hasPending; - public VectorReadBatch(nint callback, nint callbackContext, uint keyCount, PinnedSpanByte lengthPrefixedKeys, ReadOnlySpan namespaceBytes) + public VectorReadBatch(nint callback, nint callbackContext, uint keyCount, PinnedSpanByte lengthPrefixedKeys, ReadOnlySpan namespaceBytes, ReadCopyOptions readOpts, int initialRecordSizeHint) { #if NET9_0_OR_GREATER this.namespaceBytes = namespaceBytes; @@ -133,6 +108,9 @@ public VectorReadBatch(nint callback, nint callbackContext, uint keyCount, Pinne currentPtr = this.lengthPrefixedKeys.ToPointer(); currentLen = *(int*)currentPtr; + + ReadCopyOptions = readOpts; + InitialIORecordSize = initialRecordSizeHint; } [MethodImpl(MethodImplOptions.AggressiveInlining)] @@ -239,7 +217,7 @@ internal readonly void CompletePending(ref VectorBasicContext objectContext) } } - private unsafe delegate* unmanaged[Cdecl] ReadCallbackPtr { get; } = &ReadCallbackUnmanaged; + private unsafe delegate* unmanaged[Cdecl] ReadCallbackPtr { get; } = &ReadCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] WriteCallbackPtr { get; } = &WriteCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] DeleteCallbackPtr { get; } = &DeleteCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] ReadModifyWriteCallbackPtr { get; } = &ReadModifyWriteCallbackUnmanaged; @@ -254,35 +232,6 @@ internal readonly void CompletePending(ref VectorBasicContext objectContext) [ThreadStatic] internal static StorageSession ActiveThreadSession; - /// - /// Per-term initial disk-read sizes (in bytes) for the vector set currently being operated on, so each - /// record is fetched in a single IO sized to its actual geometry. A field value of 0 means "not set" and - /// the read falls back to the normal default. - /// - internal struct VectorReadGeometry - { - /// Initial disk-read size for the FullVector record (term 0). - public int FullVectorIOSize; - - /// Initial disk-read size for the adjacency NeighborList record (term 1). - public int NeighborListIOSize; - - /// Initial disk-read size for the QuantizedVector record (term 2). - public int QuantizedVectorIOSize; - } - - /// - /// Per-term initial disk-read sizes for the vector set currently being operated on. Thread-static for the - /// same reason as (DiskANN runs single-threaded per operation): set on - /// entry to a search/add once the index's dimensions / links are known () - /// and reset to default when the index context is exited () so a - /// subsequent operation on a different set does not inherit stale sizes. Because the sizes are derived - /// per-index, two vector sets with different dimensions or M get different (each optimal) sizes within the - /// same Garnet instance. - /// - [ThreadStatic] - internal static VectorReadGeometry ActiveReadGeometry; - /// /// Destination for copying the small graph "stub" records (NeighborList adjacency, internal/external id /// maps, quantized vectors) back into memory when they are read from disk (see @@ -293,49 +242,7 @@ internal struct VectorReadGeometry /// settings in the same process do not clobber each other; reads reach it via /// .. /// - internal readonly ReadCopyTo StubReadCopyTo; - - /// - /// Per-record overhead (RecordInfo + key + length prefixes) added to the value size when computing the - /// initial disk-read size, so the whole record lands in one IO. Generous; the read is sector-aligned downstream. - /// - private const int VectorRecordReadOverheadBytes = 64; - - /// - /// Compute and stash the per-term initial disk-read sizes from the active vector set's geometry, so that - /// can size each read to the record it is fetching. - /// FullVector value = * (full bytes-per-element for ); - /// QuantizedVector value = effective-dims * (quantized bits-per-element / 8); NeighborList value = * sizeof(int). - /// - [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static void SetActiveReadGeometry(uint dimensions, uint numLinks, VectorQuantType quantType, uint reduceDims) - { - // The stored FullVector element size depends on the quantizer: the Redis quantizers (NoQuant/Bin/Q8) - // store F32 (4 bytes/dim), while the extended X* quantizers store 1 byte/dim. See the format mapping in - // VectorManager.TryGetEmbedding. Over-estimating only wastes bandwidth; under-estimating would force a - // second IO, so this must match the actual stored size. - var isBinaryQuant = quantType is VectorQuantType.Bin or VectorQuantType.XBin_I8 or VectorQuantType.XBin_U8; - var fullVectorElementBytes = quantType is VectorQuantType.XNoQuant_U8 or VectorQuantType.XNoQuant_I8 - or VectorQuantType.XBin_I8 or VectorQuantType.XBin_U8 - ? 1 - : sizeof(float); - - // QuantizedVector reads (term 2, used for the approximate-distance pass on quantized sets) are sized to - // the quantized width, which is much smaller than the full vector and differs by quantizer: the byte - // quantizers (Q8) store 1 byte/dim, while the binary quantizers (Bin) pack 1 bit/dim. Quantization is - // applied to the (optionally reduced) dimensions. Sizing per-quantizer avoids over-reading whole sectors - // for the tiny binary records; the overhead covers any per-vector scale, and an under-read just self- - // corrects with a second IO. - var quantizedDims = reduceDims != 0 ? reduceDims : dimensions; - var quantizedValueBytes = isBinaryQuant ? checked((int)quantizedDims + 7) / 8 : checked((int)quantizedDims); - - ActiveReadGeometry = new VectorReadGeometry - { - FullVectorIOSize = checked((int)dimensions * fullVectorElementBytes) + VectorRecordReadOverheadBytes, - NeighborListIOSize = checked((int)numLinks * sizeof(int)) + VectorRecordReadOverheadBytes, - QuantizedVectorIOSize = quantizedValueBytes + VectorRecordReadOverheadBytes, - }; - } + private readonly ReadCopyTo stubReadCopyTo; [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static unsafe void LogCallbackUnmanaged(ulong context, nint logMessage, nuint logMessageLength) @@ -403,6 +310,7 @@ private static unsafe void LogCallbackUnmanaged(ulong context, nint logMessage, private static unsafe void ReadCallbackUnmanaged( ulong context, uint numKeys, + uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, @@ -414,7 +322,17 @@ nint dataCallbackContext Span nsBytes = stackalloc byte[sizeof(uint)]; StoreContextInNamespace(context, ref nsBytes); - var enumerable = new VectorReadBatch(dataCallback, dataCallbackContext, numKeys, PinnedSpanByte.FromPinnedPointer((byte*)keysData, (int)keysLength), nsBytes); + // Calculate optimal read options for this batch + var readCopyOptions = + (context & (ContextStep - 1)) switch + { + DiskANNService.NeighborList or DiskANNService.QuantizedVector or DiskANNService.InternalIdMap or DiskANNService.ExternalIdMap => new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ActiveThreadSession.vectorManager.stubReadCopyTo }, + _ => new ReadCopyOptions { CopyFrom = ReadCopyFrom.None, CopyTo = ReadCopyTo.None }, + }; + + var valueLengthHintWithOverhead = valueLengthHint + VectorRecordReadOverheadBytes; + + var enumerable = new VectorReadBatch(dataCallback, dataCallbackContext, numKeys, PinnedSpanByte.FromPinnedPointer((byte*)keysData, (int)keysLength), nsBytes, readCopyOptions, (int)valueLengthHint); ref var ctx = ref ActiveThreadSession.vectorBasicContext; diff --git a/libs/server/Resp/Vector/VectorManager.Locking.cs b/libs/server/Resp/Vector/VectorManager.Locking.cs index cd23650714c..da8fbf712b3 100644 --- a/libs/server/Resp/Vector/VectorManager.Locking.cs +++ b/libs/server/Resp/Vector/VectorManager.Locking.cs @@ -39,10 +39,6 @@ public void Dispose() Debug.Assert(ActiveThreadSession != null, "Shouldn't exit context when not in one"); ActiveThreadSession = null; - // Clear the per-index read geometry so a subsequent operation on a different vector set - // (possibly with different dimensions / M) does not inherit stale sizes. - ActiveReadGeometry = default; - if (Unsafe.IsNullRef(in lockableCtx)) { return; diff --git a/libs/server/Resp/Vector/VectorManager.cs b/libs/server/Resp/Vector/VectorManager.cs index ae30a8f516b..2d9e89bed0a 100644 --- a/libs/server/Resp/Vector/VectorManager.cs +++ b/libs/server/Resp/Vector/VectorManager.cs @@ -183,7 +183,7 @@ public VectorManager(int dbId, GarnetServerOptions serverOptions, Func errorMsg ReadIndex(indexValue, out var context, out var dimensions, out var reduceDims, out var quantType, out _, out var numLinks, out var distanceMetric, out _, out var indexPtr); - // Size FullVector / NeighborList disk reads to this set's geometry (dimensions, M) for single-IO fetches. - SetActiveReadGeometry(dimensions, numLinks, quantType, reduceDims); - if (providedReduceDims != 0 && providedReduceDims != reduceDims) { errorMsg = "ERR Provided REDUCE does not match Vector Set definition"u8; @@ -742,9 +739,6 @@ ref SpanByteAndMemory filterBitmap ReadIndex(indexValue, out var context, out var dimensions, out var reduceDims, out var quantType, out _, out var numLinks, out _, out _, out var indexPtr); - // Size FullVector / NeighborList disk reads to this set's geometry (dimensions, M) for single-IO fetches. - SetActiveReadGeometry(dimensions, numLinks, quantType, reduceDims); - var effectiveEF = Math.Max(searchExplorationFactor, count); EnsureDistanceBufferSize(ref outputDistances, count); @@ -933,9 +927,6 @@ ref SpanByteAndMemory filterBitmap ReadIndex(indexValue, out var context, out var dimensions, out var reduceDims, out var quantType, out _, out var numLinks, out _, out _, out var indexPtr); - // Size FullVector / NeighborList disk reads to this set's geometry (dimensions, M) for single-IO fetches. - SetActiveReadGeometry(dimensions, numLinks, quantType, reduceDims); - var effectiveEF = Math.Max(searchExplorationFactor, count); EnsureDistanceBufferSize(ref outputDistances, count); diff --git a/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs b/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs index 80b76a6f48d..1695a7006ce 100644 --- a/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs +++ b/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs @@ -1736,7 +1736,7 @@ public unsafe void VectorReadBatchVariants() fixed (int* dataPtr = data) { var keyData = PinnedSpanByte.FromPinnedPointer((byte*)dataPtr, data.Length * sizeof(int)); - var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 1, keyData, namespaceBytes); + var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 1, keyData, namespaceBytes, new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ReadCopyTo.MainLog }, 0); var iters = 0; for (var i = 0; i < batch.Count; i++) @@ -1786,7 +1786,7 @@ public unsafe void VectorReadBatchVariants() fixed (int* dataPtr = data) { var keyData = PinnedSpanByte.FromPinnedPointer((byte*)dataPtr, data.Length * sizeof(int)); - var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 7, keyData, namespaceBytes); + var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 7, keyData, namespaceBytes, new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ReadCopyTo.MainLog }, 0); var iters = 0; for (var i = 0; i < batch.Count; i++) @@ -1840,7 +1840,7 @@ public unsafe void VectorReadBatchVariants() fixed (int* dataPtr = data) { var keyData = PinnedSpanByte.FromPinnedPointer((byte*)dataPtr, data.Length * sizeof(int)); - var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 7, keyData, namespaceBytes); + var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 7, keyData, namespaceBytes, new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ReadCopyTo.MainLog }, 0); var rand = new Random(2025_10_06_00); @@ -1898,7 +1898,7 @@ public unsafe void VectorReadBatchVariants() fixed (byte* dataPtr = data) { var keyData = PinnedSpanByte.FromPinnedPointer((byte*)dataPtr, data.Length); - var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 1, keyData, namespaceBytes); + var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 1, keyData, namespaceBytes, new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ReadCopyTo.MainLog }, 0); var iters = 0; for (var i = 0; i < batch.Count; i++) @@ -2017,7 +2017,7 @@ public unsafe void VectorReadBatchVariants() fixed (byte* dataPtr = data) { var keyData = PinnedSpanByte.FromPinnedPointer((byte*)dataPtr, data.Length); - var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 8, keyData, namespaceBytes); + var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 8, keyData, namespaceBytes, new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ReadCopyTo.MainLog }, 0); var iters = 0; for (var i = 0; i < batch.Count; i++) @@ -2150,7 +2150,7 @@ public unsafe void VectorReadBatchVariants() fixed (byte* dataPtr = data) { var keyData = PinnedSpanByte.FromPinnedPointer((byte*)dataPtr, data.Length); - var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 8, keyData, namespaceBytes); + var batch = new VectorManager.VectorReadBatch(input.Callback, input.CallbackContext, 8, keyData, namespaceBytes, new ReadCopyOptions { CopyFrom = ReadCopyFrom.AllImmutable, CopyTo = ReadCopyTo.MainLog }, 0); var rand = new Random(2025_10_06_01); diff --git a/website/docs/dev/vector-sets.md b/website/docs/dev/vector-sets.md index ea833422252..f925ee70035 100644 --- a/website/docs/dev/vector-sets.md +++ b/website/docs/dev/vector-sets.md @@ -320,14 +320,16 @@ All callbacks take a `ulong context` parameter which identifies the Vector Set i The most complicated of our callbacks, the signature is: ```csharp -void ReadCallbackUnmanaged(ulong context, uint numKeys, nint keysData, nuint keysLength, nint dataCallback, nint dataCallbackContext) +void ReadCallbackUnmanaged(ulong context, uint numKeys, uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, nint dataCallbackContext) ``` `context` identifies which Vector Set is being operated on AND the associated namespace, `numKeys` tells us how many keys have been encoded into `keysData`, `keysData` and `keysLength` define a `Span` of length prefixied keys, `dataCallback` is a `delegate* unmanaged[Cdecl, SuppressGCTransition]` used to push found keys back into DiskANN, and `dataCallbackContext` is passed back unaltered to `dataCallback`. +`valueLengthHint` is the number of bytes DiskANN expects _each_ record to be if found. This value is just a hint, it's OK if it's wrong. It is better to be too large than too small, within reason. + In the `Span` defined by `keysData` and `keysLength` the keys are length prefixed with a 4-byte little endian `int`. -As we find keys, we invoke `dataCallback(index, dataCallbackContext, keyPointer, keyLength)`. If a key is not found, its index is simply skipped. The benefits of this is that we don't copy data out of the Tsavorite log as part of reads, DiskANN is able to do distance calculations and traversal over in-place data. +As we find keys, we invoke `dataCallback(index, dataCallbackContext, dataPointer, dataLengthLength)`. Invocations may be out of order, so `index` must be used to correlate data with keys. If a key is not found, its index is simply skipped. The benefits of this is that we don't copy data out of the Tsavorite log as part of reads, DiskANN is able to do distance calculations and traversal over in-place data. > [!NOTE] > Each invocation of `dataCallback` is a managed -> native transition, which can add up very quickly. We've reduced that as much as possible with function points and `SuppressGCTransition`, but that comes with risks. @@ -395,6 +397,12 @@ void LogCallbackUnmanaged(ulong context, nint logMessage, nuint logMessageLength The log message is UTF8 encoded text. +This log message is enriched on the Garnet side with: + - The context without namespace bits + - A text version of the namespace + - Our best guess at the Vector Set currently operated on + - The number of arguments for the command being processed + ### DiskANN Functions Garnet calls into the following DiskANN functions: From c9d05a66536ac09e37901d934ef8e2b089070301 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Tue, 28 Jul 2026 17:08:10 -0400 Subject: [PATCH 03/12] update filterCallback to take attribute data directly --- libs/server/Resp/Vector/DiskANNService.cs | 4 +- .../Resp/Vector/VectorManager.Callbacks.cs | 6 +- .../Resp/Vector/VectorManager.Filter.cs | 57 +++++-------------- .../Resp/Vector/VectorManager.Locking.cs | 6 +- .../Resp/Vector/VectorManager.Migration.cs | 2 +- .../DiskANN/DiskANNServiceTests.cs | 10 ++-- website/docs/dev/vector-sets.md | 13 +++++ 7 files changed, 41 insertions(+), 57 deletions(-) diff --git a/libs/server/Resp/Vector/DiskANNService.cs b/libs/server/Resp/Vector/DiskANNService.cs index 3fbc99bf897..e9770e0af00 100644 --- a/libs/server/Resp/Vector/DiskANNService.cs +++ b/libs/server/Resp/Vector/DiskANNService.cs @@ -44,7 +44,7 @@ public nint CreateIndex( delegate* unmanaged[Cdecl] writeCallback, delegate* unmanaged[Cdecl] deleteCallback, delegate* unmanaged[Cdecl] readModifyWriteCallback, - delegate* unmanaged[Cdecl] filterCallback, + delegate* unmanaged[Cdecl] filterCallback, delegate* unmanaged[Cdecl] logCallback, out bool quantizationRequested ) @@ -74,7 +74,7 @@ public nint RecreateIndex( delegate* unmanaged[Cdecl] writeCallback, delegate* unmanaged[Cdecl] deleteCallback, delegate* unmanaged[Cdecl] readModifyWriteCallback, - delegate* unmanaged[Cdecl] filterCallback, + delegate* unmanaged[Cdecl] filterCallback, delegate* unmanaged[Cdecl] logCallback, out bool quantizationRequested ) diff --git a/libs/server/Resp/Vector/VectorManager.Callbacks.cs b/libs/server/Resp/Vector/VectorManager.Callbacks.cs index e7938f1d5da..57c3b38396a 100644 --- a/libs/server/Resp/Vector/VectorManager.Callbacks.cs +++ b/libs/server/Resp/Vector/VectorManager.Callbacks.cs @@ -221,7 +221,7 @@ internal readonly void CompletePending(ref VectorBasicContext objectContext) private unsafe delegate* unmanaged[Cdecl] WriteCallbackPtr { get; } = &WriteCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] DeleteCallbackPtr { get; } = &DeleteCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] ReadModifyWriteCallbackPtr { get; } = &ReadModifyWriteCallbackUnmanaged; - private unsafe delegate* unmanaged[Cdecl] InlineFilterCallbackPtr { get; } = &FilterCallbackUnmanaged; + private unsafe delegate* unmanaged[Cdecl] FilterCallbackPtr { get; } = &FilterCallbackUnmanaged; private unsafe delegate* unmanaged[Cdecl] LogCallbackPtr { get; } = &LogCallbackUnmanaged; /// @@ -398,9 +398,9 @@ private static byte ReadModifyWriteCallbackUnmanaged(ulong context, nint keyData } [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] - private static byte FilterCallbackUnmanaged(ulong context, uint internalId) + private static unsafe byte FilterCallbackUnmanaged(ulong context, nint valueData, nuint valueLength) { - return EvaluateCandidateFilter(context, internalId); + return EvaluateCandidateFilter(context, new ReadOnlySpan((byte*)valueData, (int)valueLength)); } private static unsafe bool ReadSizeUnknown(ulong context, bool forceAlignment, ReadOnlySpan key, ref SpanByteAndMemory value) diff --git a/libs/server/Resp/Vector/VectorManager.Filter.cs b/libs/server/Resp/Vector/VectorManager.Filter.cs index 562b94fd58a..927b7ec7c66 100644 --- a/libs/server/Resp/Vector/VectorManager.Filter.cs +++ b/libs/server/Resp/Vector/VectorManager.Filter.cs @@ -8,7 +8,6 @@ using System.Runtime.CompilerServices; #endif using System.Runtime.InteropServices; -using Tsavorite.core; namespace Garnet.server { @@ -263,7 +262,7 @@ internal ref struct InlineFilterState /// Shared filter evaluation logic for both single and batch callbacks. /// Reads the candidate's external ID and attributes, then evaluates the compiled filter. /// - private static unsafe byte EvaluateCandidateFilter(ulong context, uint internalId) + private static unsafe byte EvaluateCandidateFilter(ulong context, ReadOnlySpan attr) { Debug.Assert(InlineFilterStatePtr != null, "Shouldn't call without pinning a filter state"); ref var state @@ -273,53 +272,23 @@ ref var state = ref *InlineFilterStatePtr; #endif - // 1. Read external ID for this internal_id via ExtMap - Span iidKey = stackalloc byte[sizeof(uint)]; - BinaryPrimitives.WriteUInt32LittleEndian(iidKey, internalId); - - Span eidBuf = stackalloc byte[128]; - var eidMem = SpanByteAndMemory.FromPinnedSpan(eidBuf); - try + // 3. Rebuild ExprProgram from thread-static state pointers + var program = new ExprProgram { - if (!ReadSizeUnknown(context | DiskANNService.ExternalIdMap, true, iidKey, ref eidMem)) - return 0; // can't find external ID → exclude - - // 2. Read attributes by external ID - Span attrBuf = stackalloc byte[256]; - var attrMem = SpanByteAndMemory.FromPinnedSpan(attrBuf); - try - { - if (!ReadSizeUnknown(context | DiskANNService.Attributes, true, eidMem.ReadOnlySpan, ref attrMem)) - return 0; // no attributes → exclude + Instructions = state.InstrBuf, + TuplePool = state.TuplePoolBuf, + RuntimePool = state.RuntimePoolBuf, + RuntimePoolLength = 0, + }; - // 3. Rebuild ExprProgram from thread-static state pointers - var program = new ExprProgram - { - Instructions = state.InstrBuf, - TuplePool = state.TuplePoolBuf, - RuntimePool = state.RuntimePoolBuf, - RuntimePoolLength = 0, - }; - - program.ResetRuntimePool(); + program.ResetRuntimePool(); - AttributeExtractor.ExtractFields(attrMem.ReadOnlySpan, state.FilterBytes, state.SelectorRanges, state.ExtractedFields, ref program); + AttributeExtractor.ExtractFields(attr, state.FilterBytes, state.SelectorRanges, state.ExtractedFields, ref program); - var stack = new ExprStack(state.StackBuf); - var pass = ExprRunner.Run(ref program, attrMem.ReadOnlySpan, state.FilterBytes, state.SelectorRanges, state.ExtractedFields, ref stack); + var stack = new ExprStack(state.StackBuf); + var pass = ExprRunner.Run(ref program, attr, state.FilterBytes, state.SelectorRanges, state.ExtractedFields, ref stack); - return pass ? (byte)1 : (byte)0; - } - finally - { - attrMem.Memory?.Dispose(); - } - } - finally - { - eidMem.Memory?.Dispose(); - } + return pass ? (byte)1 : (byte)0; } - } } \ No newline at end of file diff --git a/libs/server/Resp/Vector/VectorManager.Locking.cs b/libs/server/Resp/Vector/VectorManager.Locking.cs index da8fbf712b3..6f01d26409e 100644 --- a/libs/server/Resp/Vector/VectorManager.Locking.cs +++ b/libs/server/Resp/Vector/VectorManager.Locking.cs @@ -171,7 +171,7 @@ internal VectorSetLock ReadVectorIndex(StorageSession storageSession, ReadOnlySp bool requestQuantization; unsafe { - newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, FilterCallbackPtr, LogCallbackPtr, out requestQuantization); } input.header.cmd = RespCommand.VADD; @@ -367,7 +367,7 @@ out GarnetStatus status unsafe { - newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.RecreateIndex(indexContext, dims, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, FilterCallbackPtr, LogCallbackPtr, out requestQuantization); } input.parseState.EnsureCapacity(12); @@ -399,7 +399,7 @@ out GarnetStatus status unsafe { - newlyAllocatedIndex = Service.CreateIndex(indexContext, dims, reduceDims, quantizer, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.CreateIndex(indexContext, dims, reduceDims, quantizer, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, FilterCallbackPtr, LogCallbackPtr, out requestQuantization); } input.parseState.EnsureCapacity(12); diff --git a/libs/server/Resp/Vector/VectorManager.Migration.cs b/libs/server/Resp/Vector/VectorManager.Migration.cs index 1c11c54db03..61418c1df03 100644 --- a/libs/server/Resp/Vector/VectorManager.Migration.cs +++ b/libs/server/Resp/Vector/VectorManager.Migration.cs @@ -192,7 +192,7 @@ public void HandleMigratedIndexKey( bool requestQuantization; unsafe { - newlyAllocatedIndex = Service.RecreateIndex(context, dimensions, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, InlineFilterCallbackPtr, LogCallbackPtr, out requestQuantization); + newlyAllocatedIndex = Service.RecreateIndex(context, dimensions, reduceDims, quantType, buildExplorationFactor, numLinks, distanceMetric, ReadCallbackPtr, WriteCallbackPtr, DeleteCallbackPtr, ReadModifyWriteCallbackPtr, FilterCallbackPtr, LogCallbackPtr, out requestQuantization); } var ctxArg = PinnedSpanByte.FromPinnedSpan(MemoryMarshal.Cast(MemoryMarshal.CreateSpan(ref context, 1))); diff --git a/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs b/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs index a767a10f271..2890c90aa8e 100644 --- a/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs +++ b/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs @@ -20,11 +20,11 @@ namespace Garnet.test [TestFixture] public class DiskANNServiceTests : TestBase { - private delegate void ReadCallbackDelegate(ulong context, uint numKeys, nint keysData, nuint keysLength, nint dataCallback, nint dataCallbackContext); + private delegate void ReadCallbackDelegate(ulong context, uint numKeys, uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, nint dataCallbackContext); private delegate byte WriteCallbackDelegate(ulong context, nint keyData, nuint keyLength, nint writeData, nuint writeLength); private delegate byte DeleteCallbackDelegate(ulong context, nint keyData, nuint keyLength); private delegate byte ReadModifyWriteCallbackDelegate(ulong context, nint keyData, nuint keyLength, nuint writeLength, nint dataCallback, nint dataCallbackContext); - private delegate byte InlineFilterCallbackDelegate(ulong context, uint internalId); + private delegate byte InlineFilterCallbackDelegate(ulong context, nint attrData, nuint attrDataLength); private delegate void LogCallbackDelegate(ulong context, nint logMessage, nuint logMessageLength); private sealed class ContextAndKeyComparer : IEqualityComparer<(ulong Context, byte[] Data)> @@ -68,6 +68,7 @@ public void CheckInternalId() unsafe void ReadCallback( ulong context, uint numKeys, + uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, @@ -161,7 +162,7 @@ unsafe byte ReadModifyWriteCallback(ulong context, nint keyData, nuint keyLength return 1; } - byte InlineFilterCallback(ulong context, uint internalId) + byte InlineFilterCallback(ulong context, nint attrData, nuint attrDataLength) { return 1; } @@ -279,6 +280,7 @@ public void Recreate() unsafe void ReadCallback( ulong context, uint numKeys, + uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, @@ -372,7 +374,7 @@ unsafe byte ReadModifyWriteCallback(ulong context, nint keyData, nuint keyLength return 1; } - byte InlineFilterCallback(ulong context, uint internalId) + byte InlineFilterCallback(ulong context, nint attrData, nuint attrDataLength) { return 1; } diff --git a/website/docs/dev/vector-sets.md b/website/docs/dev/vector-sets.md index f925ee70035..564c92ae95d 100644 --- a/website/docs/dev/vector-sets.md +++ b/website/docs/dev/vector-sets.md @@ -386,6 +386,19 @@ Newly allocated values are guaranteed to be all zeros. The callback returns 1 if the key-value pair was found or created, and 0 if some error occurred. +### Filter Callback + +A simple callback providing a way for _Garnet_ to tell DiskANN if an attribute matches the current filter. Its signature is: +```csharp +byte FilterCallbackUnmanaged(ulong context, nint valueData, nuint valueLength) +``` + +`context` identifies whcih Vector Set is being operated on (the associated namespace is ignored), and `valueData` and `valueLength` represent a `Span` of the attribute to check. + +The current filter is ambient state that Garnet has already parsed and validated. + +If the attribute matches the current filter, 1 is returned and otherwise 0 is returned. + ### Log Callback A simple callback providing a way for DiskANN to log richer error messages into Garnet. Its signature is: From 7f16d14a15b6b5919fc202652d13fd0c7aad0dff Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Tue, 28 Jul 2026 17:45:24 -0400 Subject: [PATCH 04/12] wire up continue_search --- libs/server/Resp/Vector/DiskANNService.cs | 21 ++- .../Resp/Vector/VectorManager.Quantization.cs | 7 +- libs/server/Resp/Vector/VectorManager.cs | 129 ++++++++++++++++-- 3 files changed, 135 insertions(+), 22 deletions(-) diff --git a/libs/server/Resp/Vector/DiskANNService.cs b/libs/server/Resp/Vector/DiskANNService.cs index e9770e0af00..0f1fcfad07e 100644 --- a/libs/server/Resp/Vector/DiskANNService.cs +++ b/libs/server/Resp/Vector/DiskANNService.cs @@ -115,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 id) { @@ -300,9 +298,18 @@ out nint continuation } } - public int ContinueSearch(ulong context, nint index, nint continuation, Span outputIds, Span outputDistances, out nint newContinuation) + public int ContinueSearch(ulong context, nint index, nint continuation, Span outputIds, Span outputDistances, out nint newContinuation) { - throw new NotImplementedException(); + 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 internalId) @@ -464,7 +471,7 @@ 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, diff --git a/libs/server/Resp/Vector/VectorManager.Quantization.cs b/libs/server/Resp/Vector/VectorManager.Quantization.cs index ffa54f4ad71..b158a586fbf 100644 --- a/libs/server/Resp/Vector/VectorManager.Quantization.cs +++ b/libs/server/Resp/Vector/VectorManager.Quantization.cs @@ -115,7 +115,12 @@ static async Task QuantizationTaskAsync(VectorManager self, ChannelReader= 0) + { + while (continuation != 0) + { + var additionalResults = ContinueSearch(context, indexPtr, continuation, found, ref outputIds, ref outputDistances, out continuation); + + if (additionalResults < 0) + { + found = additionalResults; + break; + } + + found += additionalResults; + } + } } finally { - ActiveThreadSession.scratchBufferBuilder.RewindScratchBuffer(bufferSlice); + _ = ActiveThreadSession.scratchBufferBuilder.RewindScratchBuffer(bufferSlice); unsafe { @@ -865,6 +881,22 @@ out continuation outputDistances, out continuation ); + + if (found >= 0) + { + while (continuation != 0) + { + var additionalResults = ContinueSearch(context, indexPtr, continuation, found, ref outputIds, ref outputDistances, out continuation); + + if (additionalResults < 0) + { + found = additionalResults; + break; + } + + found += additionalResults; + } + } } } @@ -889,12 +921,6 @@ out continuation _ = ApplyPostFilter(filter, found, outputAttributes.ReadOnlySpan, filterBitmap.Span, ActiveThreadSession.scratchBufferBuilder); } - if (continuation != 0) - { - // TODO: paged results! - throw new NotImplementedException(); - } - outputDistances.Length = sizeof(float) * found; // Default assumption is length prefixed @@ -1007,6 +1033,22 @@ ref SpanByteAndMemory filterBitmap out continuation ); + if (found >= 0) + { + while (continuation != 0) + { + var additionalResults = ContinueSearch(context, indexPtr, continuation, found, ref outputIds, ref outputDistances, out continuation); + + if (additionalResults < 0) + { + found = additionalResults; + break; + } + + found += additionalResults; + } + } + } finally { @@ -1032,7 +1074,23 @@ out continuation outputIds, outputDistances, out continuation - ); + ); + + if (found >= 0) + { + while (continuation != 0) + { + var additionalResults = ContinueSearch(context, indexPtr, continuation, found, ref outputIds, ref outputDistances, out continuation); + + if (additionalResults < 0) + { + found = additionalResults; + break; + } + + found += additionalResults; + } + } } if (found < 0) @@ -1055,12 +1113,6 @@ out continuation _ = ApplyPostFilter(filter, found, outputAttributes.ReadOnlySpan, filterBitmap.Span, ActiveThreadSession.scratchBufferBuilder); } - if (continuation != 0) - { - // TODO: paged results! - throw new NotImplementedException(); - } - outputDistances.Length = sizeof(float) * found; // Default assumption is length prefixed @@ -1069,6 +1121,55 @@ out continuation return VectorManagerResult.OK; } + /// + /// Continue a search that previously produced partial results. + /// + /// All search_xxx methods continue in the same way, so this method is held in common. + /// + /// Returns number of new results fetched, and sets to a non-0 value if additional calls are necessary. + /// + internal int ContinueSearch(ulong context, nint indexPtr, nint oldContinuation, int foundSoFar, ref SpanByteAndMemory outputIds, ref SpanByteAndMemory outputDistances, out nint continuation) + { + Debug.Assert(oldContinuation != 0, "Expected non-zero continuation"); + + // Only ids can grow, so double them each time + var newIdSpace = MemoryPool.Shared.Rent(outputIds.Span.Length * 2); + outputIds.ReadOnlySpan.CopyTo(newIdSpace.Memory.Span); + + // TODO: Could remember this offset? It's a relatively rare occurrence so maybe not worth optimizing + var writeIdsInto = newIdSpace.Memory.Span; + for (var i = 0; i < foundSoFar; i++) + { + var skip = BinaryPrimitives.ReadInt32LittleEndian(writeIdsInto); + writeIdsInto = writeIdsInto[(sizeof(int) + skip)..]; + } + + var writeDistancesInto = outputDistances.Span[(sizeof(float) * foundSoFar)..]; + Debug.Assert(!writeDistancesInto.IsEmpty, "Expected space for remaining distances"); + + int count; + unsafe + { + // Guarantee these are pinned + fixed (byte* idPtr = writeIdsInto) + fixed (byte* distancePtr = writeDistancesInto) + { + count = Service.ContinueSearch(context, indexPtr, oldContinuation, writeIdsInto, writeDistancesInto, out continuation); + } + } + + // Error case, terminate + if (count < 0) + { + Debug.Assert(continuation == 0, "Expected no additional continuations on error result"); + return count; + } + + // Update ids on success + outputIds = new(newIdSpace, newIdSpace.Memory.Length); + return count; + } + /// /// Fetch attributes for a single element id. /// From c2c6adbf2339001b714f23bcbb38337530851dc0 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 11:42:55 -0400 Subject: [PATCH 05/12] sketch out VRANDMEMBER and random_members --- libs/server/API/GarnetApi.cs | 5 +- libs/server/API/GarnetWatchApi.cs | 4 +- libs/server/API/IGarnetApi.cs | 2 +- libs/server/Resp/Vector/DiskANNService.cs | 28 +++ .../Resp/Vector/RespServerSessionVectors.cs | 17 +- libs/server/Resp/Vector/VectorManager.cs | 177 ++++++++++++++++++ .../Session/MainStore/VectorStoreOps.cs | 8 +- 7 files changed, 228 insertions(+), 13 deletions(-) diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index f8af395f63f..898397a705f 100644 --- a/libs/server/API/GarnetApi.cs +++ b/libs/server/API/GarnetApi.cs @@ -390,8 +390,9 @@ public GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, b => storageSession.VectorSetLinks(key, element, withScores, ref idResults, ref distanceResults); /// - 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); + /// 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 errorMsg) => storageSession.VectorSetAdd(key, reduceDims, valueType, values, element, quantizer, buildExplorationFactor, attributes, numLinks, distanceMetric, out result, out errorMsg); diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index e63c389f687..adb2e36a51c 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -645,10 +645,10 @@ public GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, b } /// - 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); } /// diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index 84541eedf52..9528f605c6a 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -2155,7 +2155,7 @@ public bool IterateStore(ref TScanFunctions scanFunctions, ref l /// /// On success, has length prefixed element names. /// - GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults); + GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults, out int actualCount); /// /// Perform a similarity search given a vector and these parameters. diff --git a/libs/server/Resp/Vector/DiskANNService.cs b/libs/server/Resp/Vector/DiskANNService.cs index 0f1fcfad07e..8664f14aca6 100644 --- a/libs/server/Resp/Vector/DiskANNService.cs +++ b/libs/server/Resp/Vector/DiskANNService.cs @@ -298,6 +298,25 @@ out nint continuation } } + public bool RandomMembers( + ulong context, + nint index, + int count, + Span 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 outputIds, Span outputDistances, out nint newContinuation) { var output_ids_data = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(outputIds)); @@ -477,5 +496,14 @@ public static partial byte backfill_quant_vectors( 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 + ); } } \ No newline at end of file diff --git a/libs/server/Resp/Vector/RespServerSessionVectors.cs b/libs/server/Resp/Vector/RespServerSessionVectors.cs index c3b9af0bf87..c08f49e25d8 100644 --- a/libs/server/Resp/Vector/RespServerSessionVectors.cs +++ b/libs/server/Resp/Vector/RespServerSessionVectors.cs @@ -1769,7 +1769,7 @@ private bool NetworkVRANDMEMBER(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 +1789,18 @@ private bool NetworkVRANDMEMBER(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; } diff --git a/libs/server/Resp/Vector/VectorManager.cs b/libs/server/Resp/Vector/VectorManager.cs index 2828bf2cc78..dd478776abc 100644 --- a/libs/server/Resp/Vector/VectorManager.cs +++ b/libs/server/Resp/Vector/VectorManager.cs @@ -5,6 +5,7 @@ using System.Buffers; using System.Buffers.Binary; using System.Collections.Concurrent; +using System.Collections.Generic; using System.Diagnostics; using System.Linq; using System.Runtime.InteropServices; @@ -1460,6 +1461,182 @@ internal bool IsMember(ReadOnlySpan indexSpan, ReadOnlySpan element) return false; } + /// + /// Get up to random elements from a vector set. + /// + internal VectorManagerResult RandomMembers(ReadOnlySpan indexSpan, int count, bool allowDuplicates, ref SpanByteAndMemory ids, out int finalCount) + { + // Limit the number of times we'll try to get new random members + const int MaximumAttempts = 5; + + ReadIndex(indexSpan, out var context, out _, out _, out _, out _, out _, out _, out _, out var indexPtr); + + GCHandle? idsPin = null; + + var remainingCount = count; + var remainingIds = ids.Span; + var attempts = 0; + + try + { + while (true) + { + // Guarantee we'll read negative values while processing results + remainingIds.Fill(255); + + if (!ids.IsSpanByte) + { + var getRes = MemoryMarshal.TryGetArray(ids.Memory.Memory, out var arrSeg); + Debug.Assert(getRes, "Should always be able to get array to pin"); + + idsPin = GCHandle.Alloc(arrSeg.Array, GCHandleType.Pinned); + } + else + { + idsPin = null; + } + + if (!Service.RandomMembers(context, indexPtr, remainingCount, remainingIds)) + { + logger?.LogError("RandomMembers failed for context {context}", context); + finalCount = 0; + return VectorManagerResult.BadParams; + } + + // Handle Redis-isms by deduplicating if and stopping if needed + ProcessResults(ids.Span, allowDuplicates, out var actualCount, out var validIdsLength); + + attempts++; + + if (actualCount == count) + { + // Got all the results we need, stop + + finalCount = actualCount; + ids.Length = validIdsLength; + break; + } + + var newRemainingCount = count - actualCount; + + if (newRemainingCount == remainingCount || attempts == MaximumAttempts) + { + // No progress was made, give up and return what we have + + finalCount = actualCount; + ids.Length = validIdsLength; + break; + } + + remainingCount = newRemainingCount; + + // Grow size of output buffer to hold more results + if (remainingIds.Length < (remainingCount * MinimumSpacePerId)) + { + idsPin?.Free(); + idsPin = null; + + var newIds = MemoryPool.Shared.Rent(ids.Length * 2); + ids.Span.CopyTo(newIds.Memory.Span); + + ids = new(newIds, newIds.Memory.Length); + + remainingIds = ids.Span; + for (var i = 0; i < (count - remainingCount); i++) + { + var idLen = BinaryPrimitives.ReadInt16LittleEndian(remainingIds); + if (idLen < 0) + { + break; + } + + remainingIds = remainingIds[(sizeof(int) + idLen)..]; + } + } + } + + return VectorManagerResult.OK; + } + finally + { + idsPin?.Free(); + } + + // Scan over ids and count them - deduplicating if required + static void ProcessResults(Span candidates, bool allowDuplicates, out int actualCount, out int validCandidateLength) + { + if (allowDuplicates) + { + var remaining = candidates; + + var count = 0; + while (!remaining.IsEmpty) + { + var idLen = BinaryPrimitives.ReadInt32LittleEndian(remaining); + if (idLen < 0) + { + break; + } + + count++; + remaining = remaining[(sizeof(int) + idLen)..]; + } + + actualCount = count; + validCandidateLength = candidates.Length - remaining.Length; + } + else + { + var dupeTracker = new HashSet(ByteArrayComparer.Instance); +#if NET9_0_OR_GREATER + var dupeTrackerLookup = dupeTracker.GetAlternateLookup>(); +#endif + + var remaining = candidates; + + var count = 0; + while (!remaining.IsEmpty) + { + var idLen = BinaryPrimitives.ReadInt32LittleEndian(remaining); + if (idLen < 0) + { + break; + } + + var id = remaining.Slice(sizeof(int), idLen); + byte[] idArr = null; + var isDupe = +#if NET9_0_OR_GREATER + dupeTrackerLookup.Contains(id) +#else + dupeTracker.Contains(idArr ??= id.ToArray()) +#endif + ; + + var afterId = remaining[(sizeof(int) + idLen)..]; + + if (isDupe) + { + afterId.CopyTo(remaining); + afterId[^(sizeof(int) + idLen)..].Fill(255); + } + else + { + idArr ??= id.ToArray(); + dupeTracker.Add(idArr); + + remaining = afterId; + + count++; + } + } + + actualCount = count; + validCandidateLength = candidates.Length - remaining.Length; + } + } + } + [Conditional("DEBUG")] private static void AssertHaveStorageSession() { diff --git a/libs/server/Storage/Session/MainStore/VectorStoreOps.cs b/libs/server/Storage/Session/MainStore/VectorStoreOps.cs index 1a014ae9f3a..3b39a323588 100644 --- a/libs/server/Storage/Session/MainStore/VectorStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/VectorStoreOps.cs @@ -526,7 +526,7 @@ internal GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, /// /// On success, has length prefixed element names. /// - internal GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults) + internal GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults, out int actualCount) { parseState.InitializeWithArgument(key); @@ -536,12 +536,12 @@ internal GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref { if (status != GarnetStatus.OK) { + actualCount = 0; return status; } - // TODO: Implement! - idResults.Length = 0; - return GarnetStatus.OK; + var result = vectorManager.RandomMembers(key, Math.Abs(count), allowDuplicates: count < 0, ref idResults, out actualCount); + return result == VectorManagerResult.OK ? GarnetStatus.OK : GarnetStatus.NOTFOUND; } } From 9f687a67de0fe386b796f782810b638dec3bf8b9 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 14:01:47 -0400 Subject: [PATCH 06/12] update vector-sets.md --- website/docs/dev/vector-sets.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/website/docs/dev/vector-sets.md b/website/docs/dev/vector-sets.md index 564c92ae95d..878fd677bd4 100644 --- a/website/docs/dev/vector-sets.md +++ b/website/docs/dev/vector-sets.md @@ -430,8 +430,9 @@ Garnet calls into the following DiskANN functions: - [ ] `int continue_search(ulong context, nint index, nint continuation, nint output_ids, nuint output_ids_len, nint output_distances, nuint output_distances_len, nint new_continuation)` - [ ] `ulong card(ulong context, nint index)` - [x] `byte check_internal_id_valid(ulong context, nint index, nint internal_id, nuint internal_id_len)` - - [x] `build_quant_table(ulong context, nint index)` - - [x] `backfill_quant_vectors(ulong context, nint index, nuint task_index, nuint task_count)` + - [x] `void build_quant_table(ulong context, nint index)` + - [x] `byte backfill_quant_vectors(ulong context, nint index, nuint task_index, nuint task_count)` + - [ ] `byte random_members(ulong context, nint index, uint count, nint output_ids, nuint output_ids_len)` Some non-obvious subtleties: - The number of results _requested_ from `search_vector` and `search_element` is indicated by `output_distances_len` From a6de8018a0d3f951c85b1b0628443d0176822bc8 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 14:26:16 -0400 Subject: [PATCH 07/12] sketch out VLINKS and search_neighbors --- libs/server/API/GarnetApi.cs | 4 +- libs/server/API/GarnetWatchApi.cs | 4 +- libs/server/API/IGarnetApi.cs | 4 +- libs/server/Resp/Vector/DiskANNService.cs | 87 +++++++++++++++++++ .../Resp/Vector/RespServerSessionVectors.cs | 41 ++++++++- libs/server/Resp/Vector/VectorManager.cs | 44 ++++++++++ .../Session/MainStore/VectorStoreOps.cs | 8 +- 7 files changed, 177 insertions(+), 15 deletions(-) diff --git a/libs/server/API/GarnetApi.cs b/libs/server/API/GarnetApi.cs index 898397a705f..e43d747ad85 100644 --- a/libs/server/API/GarnetApi.cs +++ b/libs/server/API/GarnetApi.cs @@ -386,8 +386,8 @@ public GarnetStatus VectorSetIsMember(PinnedSpanByte key, PinnedSpanByte element => storageSession.VectorSetIsMember(key, element); /// - 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); /// public GarnetStatus VectorSetRandomMembers(PinnedSpanByte key, int count, ref SpanByteAndMemory idResults, out int actualCount) diff --git a/libs/server/API/GarnetWatchApi.cs b/libs/server/API/GarnetWatchApi.cs index adb2e36a51c..afcff376681 100644 --- a/libs/server/API/GarnetWatchApi.cs +++ b/libs/server/API/GarnetWatchApi.cs @@ -638,10 +638,10 @@ public GarnetStatus VectorSetIsMember(PinnedSpanByte key, PinnedSpanByte element } /// - 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); } /// diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index 9528f605c6a..1142d52e3f9 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -2142,9 +2142,9 @@ public bool IterateStore(ref TScanFunctions scanFunctions, ref l /// /// For a given element, find all neighbors and (optionally) the distance to those neighbors. /// - /// On success, has length prefixed element names, and (if is true) has a float for each of those elements. + /// On success, has length prefixed element names, and has a float for each of those elements. /// - 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); /// /// Fetch random elements from the given Vector Set. diff --git a/libs/server/Resp/Vector/DiskANNService.cs b/libs/server/Resp/Vector/DiskANNService.cs index 8664f14aca6..1ced2811c1e 100644 --- a/libs/server/Resp/Vector/DiskANNService.cs +++ b/libs/server/Resp/Vector/DiskANNService.cs @@ -298,6 +298,80 @@ out nint continuation } } + public int SearchNeighbors( + ulong context, + nint index, + ReadOnlySpan id, + SpanByteAndMemory outputIds, + SpanByteAndMemory outputDistances, + out nint continuation + ) + { + 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(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(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, @@ -449,6 +523,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, diff --git a/libs/server/Resp/Vector/RespServerSessionVectors.cs b/libs/server/Resp/Vector/RespServerSessionVectors.cs index c08f49e25d8..a03da610581 100644 --- a/libs/server/Resp/Vector/RespServerSessionVectors.cs +++ b/libs/server/Resp/Vector/RespServerSessionVectors.cs @@ -1706,7 +1706,7 @@ private bool NetworkVLINKS(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(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; } diff --git a/libs/server/Resp/Vector/VectorManager.cs b/libs/server/Resp/Vector/VectorManager.cs index dd478776abc..ccea6544e81 100644 --- a/libs/server/Resp/Vector/VectorManager.cs +++ b/libs/server/Resp/Vector/VectorManager.cs @@ -1637,6 +1637,50 @@ static void ProcessResults(Span candidates, bool allowDuplicates, out int } } + /// + /// Get the neighbors of a in a Vector Set, along with distances to each neighbor. + /// + internal VectorManagerResult GetNeighbors(ReadOnlySpan indexSpan, ReadOnlySpan element, ref SpanByteAndMemory outputIds, ref SpanByteAndMemory outputDistances) + { + ReadIndex(indexSpan, out var context, out _, out _, out _, out _, out _, out _, out _, out var indexPtr); + + var found = Service.SearchNeighbors(context, indexPtr, element, outputIds, outputDistances, out var continuation); + + if (found < 0) + { + Debug.Assert(continuation == 0, "Shouldn't have more results after failure"); + + logger?.LogError("GetNeighbors failed with {res} for context {context}", found, context); + + return VectorManagerResult.BadParams; + } + + while (continuation != 0) + { + var additionalResults = ContinueSearch(context, indexPtr, continuation, found, ref outputIds, ref outputDistances, out continuation); + + if (additionalResults < 0) + { + Debug.Assert(continuation == 0, "Shouldn't have more results after failure"); + + logger?.LogError("GetNeighbors failed in ContinueSearch with {additionalResults} for context {context}", additionalResults, context); + + found = additionalResults; + break; + } + + found += additionalResults; + } + + if (found < 0) + { + return VectorManagerResult.BadParams; + } + + outputDistances.Length = sizeof(float) * found; + return VectorManagerResult.OK; + } + [Conditional("DEBUG")] private static void AssertHaveStorageSession() { diff --git a/libs/server/Storage/Session/MainStore/VectorStoreOps.cs b/libs/server/Storage/Session/MainStore/VectorStoreOps.cs index 3b39a323588..54b853f8fb4 100644 --- a/libs/server/Storage/Session/MainStore/VectorStoreOps.cs +++ b/libs/server/Storage/Session/MainStore/VectorStoreOps.cs @@ -497,7 +497,7 @@ internal GarnetStatus VectorSetIsMember(PinnedSpanByte key, PinnedSpanByte eleme /// /// Determine neighbors of a given element, and (optionally) the distance to each neighbor. /// - internal GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, bool withScores, ref SpanByteAndMemory idResults, ref SpanByteAndMemory memoryResults) + internal GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults) { parseState.InitializeWithArgument(key); @@ -510,10 +510,8 @@ internal GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, return status; } - // TODO: Implement! - idResults.Length = 0; - memoryResults.Length = 0; - return GarnetStatus.OK; + var res = vectorManager.GetNeighbors(indexSpan, element, ref idResults, ref distanceResults); + return res == VectorManagerResult.OK && distanceResults.Length > 0 ? GarnetStatus.OK : GarnetStatus.NOTFOUND; } } From b313b9f31df2e9338a455efd876159027261b772 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 14:29:45 -0400 Subject: [PATCH 08/12] add asserts that all keys coming FROM DiskANN are aligned (i.e. multiples of 4) --- libs/server/Resp/Vector/VectorManager.Callbacks.cs | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/libs/server/Resp/Vector/VectorManager.Callbacks.cs b/libs/server/Resp/Vector/VectorManager.Callbacks.cs index 57c3b38396a..e09c8b48d87 100644 --- a/libs/server/Resp/Vector/VectorManager.Callbacks.cs +++ b/libs/server/Resp/Vector/VectorManager.Callbacks.cs @@ -170,6 +170,7 @@ public void GetKey(int i, out VectorElementKey key) AdvanceTo(i); ReadOnlySpan keyBytes = new(currentPtr + 4, currentLen); + Debug.Assert((keyBytes.Length % 4) == 0, "Unaligned key provided by DiskANN"); key = new(NamespaceBytes, keyBytes); } @@ -364,6 +365,8 @@ private static unsafe byte WriteCallbackUnmanaged(ulong context, nint keyData, n [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static byte DeleteCallbackUnmanaged(ulong context, nint keyData, nuint keyLength) { + Debug.Assert((keyLength % 4) == 0, "Unaligned key provided by DiskANN"); + var keyWithNamespace = MakeVectorElementKey(context, keyData, keyLength); ref var ctx = ref ActiveThreadSession.vectorBasicContext; @@ -377,6 +380,8 @@ private static byte DeleteCallbackUnmanaged(ulong context, nint keyData, nuint k [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])] private static byte ReadModifyWriteCallbackUnmanaged(ulong context, nint keyData, nuint keyLength, nuint writeLength, nint dataCallback, nint dataCallbackContext) { + Debug.Assert((keyLength % 4) == 0, "Unaligned key provided by DiskANN"); + var keyWithNamespace = MakeVectorElementKey(context, keyData, keyLength); ref var ctx = ref ActiveThreadSession.vectorBasicContext; From 87408fdc8a47c6a59a6a291c32f0c4b13ca84a37 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 14:35:59 -0400 Subject: [PATCH 09/12] removing alignment hackery, assert that data passed back to DiskANN is aligned --- .../VectorStore/VectorSessionFunctions.cs | 398 +++++++----------- 1 file changed, 147 insertions(+), 251 deletions(-) diff --git a/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs b/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs index 1206c991c28..de542b3cdd2 100644 --- a/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs +++ b/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs @@ -15,8 +15,6 @@ namespace Garnet.server /// public readonly struct VectorSessionFunctions : ISessionFunctions { - private const int ValueAlignmentBytes = 4; - private readonly FunctionsState functionsState; private readonly ReadSessionState readSessionState; @@ -38,62 +36,57 @@ public readonly bool Reader(in TSourceLogRecord srcLogRecord, { Debug.Assert(srcLogRecord.HasNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); - var value = AlignOrPin(in srcLogRecord, ref input, out var pin); - try + var value = srcLogRecord.ValueSpan; + + if (input.IsMigrationRead) { - if (input.IsMigrationRead) - { - Debug.Assert(input.Callback == 0, "No callback expected"); + Debug.Assert(input.Callback == 0, "No callback expected"); - // We can't ship the log record over because of alignment shenanigans - // TODO: When alignment is handled at the Tsavorite level, we CAN start shipping the log over like everything else + // We can't ship the log record over because of alignment shenanigans + // TODO: When alignment is handled at the Tsavorite level, we CAN start shipping the log over like everything else - var neededSpace = VectorManager.GetMigratedElementKeySerializationSize(srcLogRecord.KeyBytes, value); + var neededSpace = VectorManager.GetMigratedElementKeySerializationSize(srcLogRecord.KeyBytes, value); - output.SpanByteAndMemory.EnsureHeapMemorySize(neededSpace); + output.SpanByteAndMemory.EnsureHeapMemorySize(neededSpace); - VectorManager.SerializeMigratedElementKey(output.SpanByteAndMemory.Span, srcLogRecord.NamespaceBytes, srcLogRecord.KeyBytes, value); + VectorManager.SerializeMigratedElementKey(output.SpanByteAndMemory.Span, srcLogRecord.NamespaceBytes, srcLogRecord.KeyBytes, value); - return true; - } + return true; + } - unsafe + unsafe + { + if (input.Callback != 0) { - if (input.Callback != 0) - { - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(value)); - var dataLen = (nuint)value.Length; + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(value)); + var dataLen = (nuint)value.Length; - callback(input.Index, input.CallbackContext, dataPtr, dataLen); - return true; - } + Debug.Assert((dataPtr % 4) == 0, "About to pass unaligned value to DiskANN"); + callback(input.Index, input.CallbackContext, dataPtr, dataLen); + return true; } + } - if (input.ReadDesiredSize > 0) - { - Debug.Assert(output.SpanByteAndMemory.Length >= value.Length, "Should always have space for vector point reads"); - - output.SpanByteAndMemory.Length = value.Length; - value.CopyTo(output.SpanByteAndMemory.Span); - } - else - { - input.ReadDesiredSize = value.Length; - if (output.SpanByteAndMemory.Length >= value.Length) - { - value.CopyTo(output.SpanByteAndMemory.Span); - output.SpanByteAndMemory.Length = value.Length; - } - } + if (input.ReadDesiredSize > 0) + { + Debug.Assert(output.SpanByteAndMemory.Length >= value.Length, "Should always have space for vector point reads"); - return true; + output.SpanByteAndMemory.Length = value.Length; + value.CopyTo(output.SpanByteAndMemory.Span); } - finally + else { - pin?.Free(); + input.ReadDesiredSize = value.Length; + if (output.SpanByteAndMemory.Length >= value.Length) + { + value.CopyTo(output.SpanByteAndMemory.Span); + output.SpanByteAndMemory.Length = value.Length; + } } + + return true; } /// @@ -108,17 +101,10 @@ public readonly bool InitialWriter(ref LogRecord logRecord, in RecordSizeInfo si { Debug.Assert(logRecord.HasNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); - var value = AlignOrPin(in logRecord, ref input, out var pin); - try - { - srcValue.CopyTo(value); + var value = logRecord.ValueSpan; + srcValue.CopyTo(value); - return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); - } - finally - { - pin?.Free(); - } + return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); } /// @@ -143,17 +129,11 @@ public readonly bool InPlaceWriter(ref LogRecord logRecord, ref VectorInput inpu if (!logRecord.TrySetContentLengths(sizeInfo.FieldInfo.ValueSize, in sizeInfo)) return false; - var value = AlignOrPin(in logRecord, ref input, out var pin); - try - { - newValue.CopyTo(value); + var value = logRecord.ValueSpan; - return true; - } - finally - { - pin?.Free(); - } + newValue.CopyTo(value); + + return true; } /// @@ -185,7 +165,7 @@ public readonly RecordFieldInfo GetRMWModifiedFieldInfo(in TSo // Constant size indicated if (needsAlignmentPadding) { - return new() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize + ValueAlignmentBytes, ExtendedNamespaceSize = GetExtendedNamespaceSize(in srcLogRecord) }; + return new() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in srcLogRecord) }; } else { @@ -215,7 +195,7 @@ public readonly RecordFieldInfo GetRMWInitialFieldInfo(TKey key, ref Vecto } else { - return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize + ValueAlignmentBytes, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; + return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; } } @@ -225,7 +205,7 @@ public readonly RecordFieldInfo GetUpsertFieldInfo(TKey key, ReadOnlySpan< #if NET9_0_OR_GREATER , allows ref struct #endif - => new() { KeySize = key.KeyBytes.Length, ValueSize = value.Length + ValueAlignmentBytes, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; + => new() { KeySize = key.KeyBytes.Length, ValueSize = value.Length, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; /// Length of value object, when populated by Upsert using given value and input public readonly RecordFieldInfo GetUpsertFieldInfo(TKey key, IHeapObject value, ref VectorInput input) @@ -265,52 +245,46 @@ public readonly bool InitialUpdater(ref LogRecord logRecord, in RecordSizeInfo s Debug.Assert(logRecord.HasNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); var key = logRecord.Key; - var alignedValue = AlignOrPin(in logRecord, ref input, out var pin); + var alignedValue = logRecord.ValueSpan; - try + if (input.Callback == 0) { + Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); + Debug.Assert(key.Length == sizeof(int), "Should have int sized key for ContextMetadata"); - if (input.Callback == 0) + // Operating on ContextMetadata + + PinnedSpanByte newMetadataValue; + unsafe { - Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should never write a non-namespaced value with VectorSessionFunctions"); - Debug.Assert(key.Length == sizeof(int), "Should have int sized key for ContextMetadata"); + newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); + } - // Operating on ContextMetadata + newMetadataValue.CopyTo(alignedValue); - PinnedSpanByte newMetadataValue; - unsafe - { - newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); - } + return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); + } + else + { + Debug.Assert(input.WriteDesiredSize <= alignedValue.Length, "Insufficient space for initial update, this should never happen"); - newMetadataValue.CopyTo(alignedValue); + // Must explicitly 0 before passing if we're doing an initial update + alignedValue.Clear(); - return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); - } - else + unsafe { - Debug.Assert(input.WriteDesiredSize <= alignedValue.Length, "Insufficient space for initial update, this should never happen"); + // Callback takes: dataCallbackContext, dataPtr, dataLength + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - // Must explicitly 0 before passing if we're doing an initial update - alignedValue.Clear(); + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(alignedValue)); + var dataLen = (nuint)input.WriteDesiredSize; - unsafe - { - // Callback takes: dataCallbackContext, dataPtr, dataLength - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + Debug.Assert((dataPtr % 4) == 0, "About to pass unaligned value to DiskANN"); + callback(input.CallbackContext, dataPtr, dataLen); - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(alignedValue)); - var dataLen = (nuint)input.WriteDesiredSize; - callback(input.CallbackContext, dataPtr, dataLen); - - return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); - } + return logRecord.TrySetContentLengths(logRecord.ValueSpan.Length, in sizeInfo); } } - finally - { - pin?.Free(); - } } #endregion InitialUpdater @@ -328,68 +302,61 @@ public readonly bool CopyUpdater(in TSourceLogRecord srcLogRec var key = srcLogRecord.Key; - var oldValueAligned = AlignOrPin(in srcLogRecord, ref input, out var srcPin); - var newValueAligned = AlignOrPin(in dstLogRecord, ref input, out var dstPin); + var oldValueAligned = srcLogRecord.ValueSpan; + var newValueAligned = dstLogRecord.ValueSpan; - try + if (input.Callback == 0) { - if (input.Callback == 0) - { - // We're doing a Metadata update - - Debug.Assert(srcLogRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); - Debug.Assert(key.Length == sizeof(int), "Should have int sized key for ContextMetadata"); + // We're doing a Metadata update - // Doing a Metadata update - Debug.Assert(srcLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); - Debug.Assert(dstLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); - Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); + Debug.Assert(srcLogRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); + Debug.Assert(key.Length == sizeof(int), "Should have int sized key for ContextMetadata"); - ref readonly var oldMetadata = ref MemoryMarshal.Cast(oldValueAligned)[0]; + // Doing a Metadata update + Debug.Assert(srcLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); + Debug.Assert(dstLogRecord.ValueSpan.Length == VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); + Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); - PinnedSpanByte newMetadataValue; - unsafe - { - newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); - } + ref readonly var oldMetadata = ref MemoryMarshal.Cast(oldValueAligned)[0]; - ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; + PinnedSpanByte newMetadataValue; + unsafe + { + newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); + } - if (newMetadata.Version < oldMetadata.Version) - { - rmwInfo.Action = RMWAction.CancelOperation; - return false; - } + ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; - newMetadataValue.CopyTo(newValueAligned); - return dstLogRecord.TrySetContentLengths(srcLogRecord.ValueSpan.Length, in sizeInfo); - } - else + if (newMetadata.Version < oldMetadata.Version) { - Debug.Assert(input.WriteDesiredSize <= newValueAligned.Length, "Insufficient space for copy update, this should never happen"); - Debug.Assert(input.WriteDesiredSize <= oldValueAligned.Length, "Insufficient space for copy update, this should never happen"); - - oldValueAligned.CopyTo(newValueAligned); + rmwInfo.Action = RMWAction.CancelOperation; + return false; + } - unsafe - { - // Callback takes: dataCallbackContext, dataPtr, dataLength - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + newMetadataValue.CopyTo(newValueAligned); + return dstLogRecord.TrySetContentLengths(srcLogRecord.ValueSpan.Length, in sizeInfo); + } + else + { + Debug.Assert(input.WriteDesiredSize <= newValueAligned.Length, "Insufficient space for copy update, this should never happen"); + Debug.Assert(input.WriteDesiredSize <= oldValueAligned.Length, "Insufficient space for copy update, this should never happen"); - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(newValueAligned)); - var dataLen = (nuint)input.WriteDesiredSize; + oldValueAligned.CopyTo(newValueAligned); - callback(input.CallbackContext, dataPtr, dataLen); - } + unsafe + { + // Callback takes: dataCallbackContext, dataPtr, dataLength + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - return true; + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(newValueAligned)); + var dataLen = (nuint)input.WriteDesiredSize; + Debug.Assert((dataPtr % 4) == 0, "About to pass unaligned value to DiskANN"); + callback(input.CallbackContext, dataPtr, dataLen); } - } - finally - { - srcPin?.Free(); - dstPin?.Free(); + + return true; + } } #endregion CopyUpdater @@ -402,60 +369,55 @@ public readonly bool InPlaceUpdater(ref LogRecord logRecord, ref VectorInput inp var key = logRecord.Key; - var alignedValue = AlignOrPin(in logRecord, ref input, out var pin); - try - { - if (input.Callback == 0) - { - // We're doing a Metadata update + var alignedValue = logRecord.ValueSpan; - Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); + if (input.Callback == 0) + { + // We're doing a Metadata update - // Doing a Metadata update - Debug.Assert(alignedValue.Length >= VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); - Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); - Debug.Assert(key.Length == sizeof(int), "Should have int sized key for ContextMetadata"); + Debug.Assert(logRecord.NamespaceBytes.Length == 1 && logRecord.NamespaceBytes[0] == VectorManager.MetadataNamespace, "Should be operating on special namespace"); - ref readonly var oldMetadata = ref MemoryMarshal.Cast(alignedValue)[0]; + // Doing a Metadata update + Debug.Assert(alignedValue.Length >= VectorManager.ContextMetadata.Size, "Should be ContextMetadata"); + Debug.Assert(input.CallbackContext != 0, "Should have data on VectorInput"); + Debug.Assert(key.Length == sizeof(int), "Should have int sized key for ContextMetadata"); - PinnedSpanByte newMetadataValue; - unsafe - { - newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); - } + ref readonly var oldMetadata = ref MemoryMarshal.Cast(alignedValue)[0]; - ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; + PinnedSpanByte newMetadataValue; + unsafe + { + newMetadataValue = PinnedSpanByte.FromPinnedPointer((byte*)input.CallbackContext, VectorManager.ContextMetadata.Size); + } - if (newMetadata.Version < oldMetadata.Version) - { - rmwInfo.Action = RMWAction.CancelOperation; - return false; - } + ref readonly var newMetadata = ref MemoryMarshal.Cast(newMetadataValue.ReadOnlySpan)[0]; - newMetadataValue.CopyTo(alignedValue); - return true; - } - else + if (newMetadata.Version < oldMetadata.Version) { - Debug.Assert(input.WriteDesiredSize <= alignedValue.Length, "Insufficient space for inplace update, this should never happen"); + rmwInfo.Action = RMWAction.CancelOperation; + return false; + } - unsafe - { - // Callback takes: dataCallbackContext, dataPtr, dataLength - var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; + newMetadataValue.CopyTo(alignedValue); + return true; + } + else + { + Debug.Assert(input.WriteDesiredSize <= alignedValue.Length, "Insufficient space for inplace update, this should never happen"); - var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(alignedValue)); - var dataLen = (nuint)input.WriteDesiredSize; + unsafe + { + // Callback takes: dataCallbackContext, dataPtr, dataLength + var callback = (delegate* unmanaged[Cdecl, SuppressGCTransition])input.Callback; - callback(input.CallbackContext, dataPtr, dataLen); - } + var dataPtr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(alignedValue)); + var dataLen = (nuint)input.WriteDesiredSize; - return true; + Debug.Assert((dataPtr % 4) == 0, "About to pass unaligned value to DiskANN"); + callback(input.CallbackContext, dataPtr, dataLen); } - } - finally - { - pin?.Free(); + + return true; } } #endregion InPlaceUpdater @@ -496,72 +458,6 @@ private static TReturn ObjectOperationsNotExpected([CallerMemberName] s private static TReturn LogRecordOperationsNotExpected([CallerMemberName] string callerName = null, [CallerLineNumber] int lineNum = -1) => throw new InvalidOperationException($"LogRecord related operations are not expected, was: {callerName} on {lineNum}"); - // TODO: Remove all this alignment hackery when Tsavorite can enforce it - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static unsafe Span AlignOrPin(in TSourceLogRecord logRecord, ref VectorInput input, out GCHandle? pin) - where TSourceLogRecord : ISourceLogRecord - { - var maybeUnaligned = logRecord.ValueSpan; - - // Alignment is expected if we're passing to DiskANN or Garnet code explicitly requested it - var inputRequiresAligment = input.AlignmentExpected || input.Callback != 0; - - if (inputRequiresAligment) - { - if (logRecord.IsPinnedValue) - { - // LogRecord itself is in POH, but value might not be aligned so we need to do some checking - - Span ret; - - var leading = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(maybeUnaligned)) % 4; - if (leading == 0) - { - ret = maybeUnaligned[..^ValueAlignmentBytes]; - } - else - { - var skip = (int)(ValueAlignmentBytes - leading); - var tail = ValueAlignmentBytes - skip; - ret = maybeUnaligned[skip..^tail]; - } - - AssertAlignment(ret); - - pin = null; - return ret; - } - else - { - // Value isn't in log record, it's on the (presumably unpinned) heap as a byte[] - // - // This guarantees it's aligned, but it might move during any callback so pin - - pin = logRecord.ValueOverflow.Pin(); - - // We over allocated (we don't know how Tsavorite is going to place the value in advance) so trim the extra allocation off the end. - var ret = maybeUnaligned[..^ValueAlignmentBytes]; - - AssertAlignment(ret); - - return ret; - } - } - else - { - pin = null; - return maybeUnaligned; - } - } - - [Conditional("DEBUG")] - private static unsafe void AssertAlignment(ReadOnlySpan aligned) - { - var ptr = (nint)Unsafe.AsPointer(ref MemoryMarshal.GetReference(aligned)); - Debug.Assert((ptr % ValueAlignmentBytes) == 0, "Must guarantee 4-byte alignment before invoking callback"); - } - #region Post operation callbacks /// public readonly void PostInitialWriter(ref LogRecord logRecord, in RecordSizeInfo sizeInfo, ref VectorInput input, ReadOnlySpan srcValue, ref VectorOutput output, ref UpsertInfo upsertInfo) From 7855ce331e3bb78a9761c4bff1582dd39640cbe6 Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 14:41:59 -0400 Subject: [PATCH 10/12] document aligment requirements --- website/docs/dev/vector-sets.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/website/docs/dev/vector-sets.md b/website/docs/dev/vector-sets.md index 878fd677bd4..0cb7d93c877 100644 --- a/website/docs/dev/vector-sets.md +++ b/website/docs/dev/vector-sets.md @@ -423,10 +423,13 @@ Garnet calls into the following DiskANN functions: - [x] `nint create_index(ulong context, uint dimensions, uint reduceDims, VectorQuantType quantType, VectorDistanceMetricType distanceMetric, uint buildExplorationFactor, uint numLinks, nint readCallback, nint writeCallback, nint deleteCallback, nint readModifyWriteCallback, nint filterCallback, nint logCallback, out bool quantizationNeeded)` - [x] `void drop_index(ulong context, nint index)` - [x] `DiskANNInsertResult insert(ulong context, nint index, nint id_data, nuint id_len, nint vector_data, nuint vector_len, nint attribute_data, nuint attribute_len)` + * `vector_data` must be aligned for the quantizers underlying type (i.e. 4-byte for NOQUANT, 1-byte for XBIN_U8, etc.) - [x] `byte remove(ulong context, nint index, nint id_data, nuint id_len)` - [x] `byte set_attribute(ulong context, nint index, nint id_data, nuint id_len, nint attribute_data, nuint attribute_len)` - [x] `int search_vector(ulong context, nint index, nint vector_data, nuint vector_len, float delta, int search_exploration_factor, nint filter_data, nuint filter_len, nuint max_filtering_effort, nint output_ids, nuint output_ids_len, nint output_distances, nuint output_distances_len, nint continuation)` + * `vector_data` must be aligned as with `insert(...)` - [x] `int search_element(ulong context, nint index, nint id_data, nuint id_len, float delta, int search_exploration_factor, nint filter_data, nuint filter_len, nuint max_filtering_effort, nint output_ids, nuint output_ids_len, nint output_distances, nuint output_distances_len, nint continuation)` + - [ ] `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)` - [ ] `int continue_search(ulong context, nint index, nint continuation, nint output_ids, nuint output_ids_len, nint output_distances, nuint output_distances_len, nint new_continuation)` - [ ] `ulong card(ulong context, nint index)` - [x] `byte check_internal_id_valid(ulong context, nint index, nint internal_id, nuint internal_id_len)` @@ -442,6 +445,9 @@ Garnet calls into the following DiskANN functions: - `index` is always a pointer created by DiskANN and returned from `create_index` - `context` is always the `Context` value created by Garnet and stored in [`Index`](#indexes) for a Vector Set, this implies it is always a non-0 multiple of 8 - `search_vector`, `search_element`, and `continue_search` all return the number of ids written into `output_ids`, and if there are more values to return they set the `nint` _pointed to by_ `continuation` or `new_continuation` + - DiskANN guarantees that any keys it provides are aligned in records, i.e. they are multiples of 4-bytes in length + - Garnet guarantees any values it provides to _`dataCallbacks`_ are 4-byte aligned + * Importantly Garnet does not guarantee id holding parameters to DiskANN functions are aligned unless otherwise noted ### Vector Filter Expressions (`VSIM ... FILTER`) From 40b938a97e1d134980bd89f33aa250a8a8c67c6e Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Wed, 29 Jul 2026 14:45:20 -0400 Subject: [PATCH 11/12] formatting --- libs/server/API/IGarnetApi.cs | 2 +- .../Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/server/API/IGarnetApi.cs b/libs/server/API/IGarnetApi.cs index 1142d52e3f9..01edfe2eab9 100644 --- a/libs/server/API/IGarnetApi.cs +++ b/libs/server/API/IGarnetApi.cs @@ -2144,7 +2144,7 @@ public bool IterateStore(ref TScanFunctions scanFunctions, ref l /// /// On success, has length prefixed element names, and has a float for each of those elements. /// - GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults); + GarnetStatus VectorSetLinks(PinnedSpanByte key, PinnedSpanByte element, ref SpanByteAndMemory idResults, ref SpanByteAndMemory distanceResults); /// /// Fetch random elements from the given Vector Set. diff --git a/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs b/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs index 2890c90aa8e..0d869707879 100644 --- a/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs +++ b/test/standalone/Garnet.test.extensions/DiskANN/DiskANNServiceTests.cs @@ -20,7 +20,7 @@ namespace Garnet.test [TestFixture] public class DiskANNServiceTests : TestBase { - private delegate void ReadCallbackDelegate(ulong context, uint numKeys, uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, nint dataCallbackContext); + private delegate void ReadCallbackDelegate(ulong context, uint numKeys, uint valueLengthHint, nint keysData, nuint keysLength, nint dataCallback, nint dataCallbackContext); private delegate byte WriteCallbackDelegate(ulong context, nint keyData, nuint keyLength, nint writeData, nuint writeLength); private delegate byte DeleteCallbackDelegate(ulong context, nint keyData, nuint keyLength); private delegate byte ReadModifyWriteCallbackDelegate(ulong context, nint keyData, nuint keyLength, nuint writeLength, nint dataCallback, nint dataCallbackContext); From 4dfec2117f8f16ee14ff73298e6f250a39c5b68a Mon Sep 17 00:00:00 2001 From: Kevin Montrose Date: Fri, 31 Jul 2026 13:59:23 -0400 Subject: [PATCH 12/12] remove more dead alignment code --- .../Server/Migration/MigrateOperation.cs | 1 - libs/server/InputHeader.cs | 2 -- .../Resp/Vector/VectorManager.Callbacks.cs | 6 +---- .../Resp/Vector/VectorManager.Migration.cs | 1 - libs/server/Resp/Vector/VectorManager.cs | 14 +++++------ .../VectorStore/VectorSessionFunctions.cs | 23 ++----------------- .../RespVectorSetTests.cs | 4 ++-- 7 files changed, 12 insertions(+), 39 deletions(-) diff --git a/libs/cluster/Server/Migration/MigrateOperation.cs b/libs/cluster/Server/Migration/MigrateOperation.cs index 1eda6351e39..9a0b2bd6f79 100644 --- a/libs/cluster/Server/Migration/MigrateOperation.cs +++ b/libs/cluster/Server/Migration/MigrateOperation.cs @@ -107,7 +107,6 @@ public async Task 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) diff --git a/libs/server/InputHeader.cs b/libs/server/InputHeader.cs index 0c7dd2202a6..6fd732fb6a8 100644 --- a/libs/server/InputHeader.cs +++ b/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; diff --git a/libs/server/Resp/Vector/VectorManager.Callbacks.cs b/libs/server/Resp/Vector/VectorManager.Callbacks.cs index e09c8b48d87..c9a4620902d 100644 --- a/libs/server/Resp/Vector/VectorManager.Callbacks.cs +++ b/libs/server/Resp/Vector/VectorManager.Callbacks.cs @@ -349,7 +349,6 @@ private static unsafe byte WriteCallbackUnmanaged(ulong context, nint keyData, n ref var ctx = ref ActiveThreadSession.vectorBasicContext; VectorInput input = new(); - input.AlignmentExpected = true; var valueSpan = SpanByte.FromPinnedPointer((byte*)writeData, (int)writeLength); VectorOutput outputSpan = new(); @@ -408,7 +407,7 @@ private static unsafe byte FilterCallbackUnmanaged(ulong context, nint valueData return EvaluateCandidateFilter(context, new ReadOnlySpan((byte*)valueData, (int)valueLength)); } - private static unsafe bool ReadSizeUnknown(ulong context, bool forceAlignment, ReadOnlySpan key, ref SpanByteAndMemory value) + private static unsafe bool ReadSizeUnknown(ulong context, ReadOnlySpan key, ref SpanByteAndMemory value) { Debug.Assert(context <= uint.MaxValue, "Contexts > 2^32-1 are not supported"); @@ -424,9 +423,6 @@ private static unsafe bool ReadSizeUnknown(ulong context, bool forceAlignment, R VectorInput input = new(); input.ReadDesiredSize = -1; - // Sometimes we read DiskANN written data from the .NET side - // If that's the case, we need to pad for alignment even though .NET doesn't require it - input.AlignmentExpected = forceAlignment; fixed (byte* ptr = value.Span) { VectorOutput asSpanByte = new(ptr, value.Length); diff --git a/libs/server/Resp/Vector/VectorManager.Migration.cs b/libs/server/Resp/Vector/VectorManager.Migration.cs index 61418c1df03..765dc3d7948 100644 --- a/libs/server/Resp/Vector/VectorManager.Migration.cs +++ b/libs/server/Resp/Vector/VectorManager.Migration.cs @@ -52,7 +52,6 @@ ReadOnlySpan value #endif VectorInput input = default; - input.AlignmentExpected = true; VectorOutput outputSpan = new(new SpanByteAndMemory()); // When we migrate a record we expand the namespace to always occupy 4-bytes diff --git a/libs/server/Resp/Vector/VectorManager.cs b/libs/server/Resp/Vector/VectorManager.cs index ccea6544e81..e39b9f284a4 100644 --- a/libs/server/Resp/Vector/VectorManager.cs +++ b/libs/server/Resp/Vector/VectorManager.cs @@ -1183,7 +1183,7 @@ internal VectorManagerResult FetchSingleVectorElementAttributes(ReadOnlySpan indexValue, ReadOnlySpan var internalIdBytes = SpanByteAndMemory.FromPinnedSpan(internalId); try { - if (!ReadSizeUnknown(context | DiskANNService.InternalIdMap, forceAlignment: true, element, ref internalIdBytes)) + if (!ReadSizeUnknown(context | DiskANNService.InternalIdMap, element, ref internalIdBytes)) { return false; } @@ -1326,7 +1326,7 @@ internal bool TryGetEmbedding(ReadOnlySpan indexValue, ReadOnlySpan var asBytes = SpanByteAndMemory.FromPinnedSpan(asBytesSpan); try { - if (!ReadSizeUnknown(context | DiskANNService.FullVector, forceAlignment: true, internalId, ref asBytes)) + if (!ReadSizeUnknown(context | DiskANNService.FullVector, internalId, ref asBytes)) { return false; } @@ -1390,7 +1390,7 @@ internal bool TryGetRawEmbedding(ReadOnlySpan indexValue, ReadOnlySpan indexValue, ReadOnlySpan indexSpan, ReadOnlySpan element) Span internalId = stackalloc byte[sizeof(int)]; var internalIdBytes = SpanByteAndMemory.FromPinnedSpan(internalId); - var foundInternalId = ReadSizeUnknown(context | DiskANNService.InternalIdMap, forceAlignment: true, element, ref internalIdBytes); + var foundInternalId = ReadSizeUnknown(context | DiskANNService.InternalIdMap, element, ref internalIdBytes); if (foundInternalId) { Debug.Assert(internalIdBytes.IsSpanByte, "Shouldn't have allocated for this op"); diff --git a/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs b/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs index de542b3cdd2..ebda1e593c4 100644 --- a/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs +++ b/libs/server/Storage/Functions/VectorStore/VectorSessionFunctions.cs @@ -160,17 +160,7 @@ public readonly RecordFieldInfo GetRMWModifiedFieldInfo(in TSo return new() { KeySize = srcLogRecord.Key.Length, ValueSize = value.Length + (-input.WriteDesiredSize), ExtendedNamespaceSize = GetExtendedNamespaceSize(in srcLogRecord) }; } - var needsAlignmentPadding = input.AlignmentExpected || input.Callback != 0; - - // Constant size indicated - if (needsAlignmentPadding) - { - return new() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in srcLogRecord) }; - } - else - { - return new() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in srcLogRecord) }; - } + return new() { KeySize = srcLogRecord.Key.Length, ValueSize = input.WriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in srcLogRecord) }; } /// Initial expected length of value object when populated by RMW using given input @@ -182,21 +172,12 @@ public readonly RecordFieldInfo GetRMWInitialFieldInfo(TKey key, ref Vecto { var effectiveWriteDesiredSize = input.WriteDesiredSize; - var needsAlignmentPadding = input.AlignmentExpected || input.Callback != 0; - if (effectiveWriteDesiredSize < 0) { effectiveWriteDesiredSize = -effectiveWriteDesiredSize; } - if (!needsAlignmentPadding) - { - return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; - } - else - { - return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; - } + return new() { KeySize = key.KeyBytes.Length, ValueSize = effectiveWriteDesiredSize, ExtendedNamespaceSize = GetExtendedNamespaceSize(in key) }; } /// Length of value object, when populated by Upsert using given value and input diff --git a/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs b/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs index 1695a7006ce..c1e96e14fec 100644 --- a/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs +++ b/test/standalone/Garnet.test.vectorset/RespVectorSetTests.cs @@ -3958,7 +3958,7 @@ void WriteAndVerify(string content) var elementKey = new VectorElementKey(new ReadOnlySpan(nsPtr, ns.Length), new ReadOnlySpan(keyPtr, key.Length)); { - var input = new VectorInput { AlignmentExpected = true }; + var input = new VectorInput(); var valueSpan = SpanByte.FromPinnedPointer(valuePtr, value.Length); var output = new VectorOutput(); @@ -3973,7 +3973,7 @@ void WriteAndVerify(string content) Span buffer = stackalloc byte[256]; fixed (byte* bufferPtr = buffer) { - var input = new VectorInput { AlignmentExpected = true, ReadDesiredSize = -1 }; + var input = new VectorInput { ReadDesiredSize = -1 }; var output = new VectorOutput(bufferPtr, buffer.Length); var status = context.Read(elementKey, ref input, ref output);