.Net: [MEVD] Fix Guid (and other) keys for dynamic collections

dotnet/src/InternalUtilities/connectors/Memory/MongoDB/MongoDynamicMapper.cs

@@ -30,13 +30,15 @@ public BsonDocument MapFromDataToStorageModel(Dictionary<string, object?> dataMo
             : keyValue switch
             {
                 string s => s,
+                Guid g => BsonValue.Create(g),
+                ObjectId o => o,
+                long i => i,
+                int i => i,
+
                 null => throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."),
                 _ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string.")
             };
 
-        document[MongoConstants.MongoReservedKeyPropertyName] = (string)(dataModel[model.KeyProperty.ModelName]
-            ?? throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."));
-
         foreach (var property in model.DataProperties)
         {
             if (dataModel.TryGetValue(property.ModelName, out var dataValue))
@@ -88,9 +90,22 @@ Embedding<float> e
             switch (property)
             {
                 case KeyPropertyModel keyProperty:
-                    result[keyProperty.ModelName] = storageModel.TryGetValue(MongoConstants.MongoReservedKeyPropertyName, out var keyValue)
-                        ? keyValue.AsString
-                        : throw new InvalidOperationException("No key property was found in the record retrieved from storage.");
+                    if (!storageModel.TryGetValue(MongoConstants.MongoReservedKeyPropertyName, out var keyValue))
+                    {
+                        throw new InvalidOperationException("No key property was found in the record retrieved from storage.");
+                    }
+
+                    result[keyProperty.ModelName] = keyProperty.Type switch
+                    {
+                        var t when t == typeof(string) => keyValue.AsString,
+                        var t when t == typeof(Guid) => keyValue.AsGuid,
+                        var t when t == typeof(ObjectId) => keyValue.AsObjectId,
+                        var t when t == typeof(long) => keyValue.AsInt64,
+                        var t when t == typeof(int) => keyValue.AsInt32,
+
+                        _ => throw new UnreachableException()
+                    };
+
                     continue;
 
                 case DataPropertyModel dataProperty:

dotnet/src/VectorData/AzureAISearch/AzureAISearchDynamicMapper.cs

@@ -29,8 +29,10 @@ public JsonObject MapFromDataToStorageModel(Dictionary<string, object?> dataMode
             : keyValue switch
             {
                 string s => s,
+                Guid g => g.ToString(),
+
                 null => throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."),
-                _ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string.")
+                _ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string or Guid.")
             };
 
         foreach (var dataProperty in model.DataProperties)
@@ -97,9 +99,16 @@ public JsonObject MapFromDataToStorageModel(Dictionary<string, object?> dataMode
             switch (property)
             {
                 case KeyPropertyModel keyProperty:
-                    result[keyProperty.ModelName] = (string?)storageModel[keyProperty.StorageName]
+                    var key = (string?)storageModel[keyProperty.StorageName]
                         ?? throw new InvalidOperationException($"The key property '{keyProperty.StorageName}' is missing from the record retrieved from storage.");
 
+                    result[keyProperty.ModelName] = keyProperty.Type switch
+                    {
+                        var t when t == typeof(string) => key,
+                        var t when t == typeof(Guid) => Guid.Parse(key),
+                        _ => throw new UnreachableException()
+                    };
+
                     continue;
 
                 case DataPropertyModel dataProperty:

dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs

@@ -860,7 +860,9 @@ private static IEnumerable<CosmosNoSqlCompositeKey> GetCompositeKeys(IEnumerable
             IEnumerable<object> k => k.Select(key => key switch
             {
                 string s => new CosmosNoSqlCompositeKey(recordKey: s, partitionKey: s),
+                Guid g when g.ToString() is var guidString => new CosmosNoSqlCompositeKey(recordKey: guidString, partitionKey: guidString),
                 CosmosNoSqlCompositeKey ck => ck,
+
                 _ => throw new ArgumentException($"Invalid key type '{key.GetType().Name}'.")
             }),
 

dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlDynamicMapper.cs

@@ -29,6 +29,8 @@ public JsonObject MapFromDataToStorageModel(Dictionary<string, object?> dataMode
             : keyValue switch
             {
                 string s => s,
+                Guid g => g.ToString(),
+
                 null => throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."),
                 _ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string.")
             };
@@ -121,9 +123,16 @@ static bool TryGetReadOnlyMemory<T>(object value, [NotNullWhen(true)] out ReadOn
             switch (property)
             {
                 case KeyPropertyModel keyProperty:
-                    result[keyProperty.ModelName] = storageModel.TryGetPropertyValue(CosmosNoSqlConstants.ReservedKeyPropertyName, out var keyValue)
-                        ? keyValue?.GetValue<string>()
-                        : throw new InvalidOperationException("No key property was found in the record retrieved from storage.");
+                    var key = (string?)storageModel[CosmosNoSqlConstants.ReservedKeyPropertyName]
+                        ?? throw new InvalidOperationException($"The key property '{keyProperty.StorageName}' is missing from the record retrieved from storage.");
+
+                    result[keyProperty.ModelName] = keyProperty.Type switch
+                    {
+                        var t when t == typeof(string) => key,
+                        var t when t == typeof(Guid) => Guid.Parse(key),
+                        _ => throw new UnreachableException()
+                    };
+
                     continue;
 
                 case DataPropertyModel dataProperty:

dotnet/src/VectorData/Redis/RedisJsonDynamicMapper.cs

@@ -86,22 +86,39 @@ internal class RedisJsonDynamicMapper(CollectionModel model, JsonSerializerOptio
             }
         }
 
-        return ((string)dataModel[model.KeyProperty.ModelName]!, jsonObject);
+        var storageKey = dataModel[model.KeyProperty.ModelName] switch
+        {
+            string s => s,
+            Guid g => g.ToString(),
+
+            _ => throw new UnreachableException()
+        };
+
+        return (storageKey, jsonObject);
     }
 
     /// <inheritdoc />
     public Dictionary<string, object?> MapFromStorageToDataModel((string Key, JsonNode Node) storageModel, bool includeVectors)
     {
         var dataModel = new Dictionary<string, object?>
         {
-            [model.KeyProperty.ModelName] = storageModel.Key,
+            [model.KeyProperty.ModelName] = model.KeyProperty.Type switch
+            {
+                Type t when t == typeof(string) => storageModel.Key,
+                Type t when t == typeof(Guid) => Guid.Parse(storageModel.Key),
+
+                _ => throw new UnreachableException()
+            },
         };
 
         // The redis result can be either a single object or an array with a single object in the case where we are doing an MGET.
+        // If there's a single data property, we get a simple value (no object wrapper).
         var jsonObject = storageModel.Node switch
         {
-            JsonObject topLevelJsonObject => topLevelJsonObject,
-            JsonArray jsonArray and [JsonObject arrayEntryJsonObject] => arrayEntryJsonObject,
+            JsonValue v when model.DataProperties is [var singleDataProperty] => new JsonObject([new(singleDataProperty.StorageName, v)]),
+            JsonObject o => o,
+            JsonArray a and [JsonObject arrayEntryJsonObject] => arrayEntryJsonObject,
+
             _ => throw new InvalidOperationException($"Invalid data format for document with key '{storageModel.Key}'"),
         };
 

dotnet/src/VectorData/Redis/RedisJsonDynamicModelBuilder.cs

@@ -19,9 +19,9 @@ internal class RedisJsonDynamicModelBuilder(CollectionModelBuildingOptions optio
 
     protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes)
     {
-        supportedTypes = "string";
+        supportedTypes = "string, Guid";
 
-        return type == typeof(string);
+        return type == typeof(string) || type == typeof(Guid);
     }
 
     protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes)

dotnet/test/VectorData/InMemory.ConformanceTests/TypeTests/InMemoryKeyTypeTests.cs

@@ -22,9 +22,9 @@ public class InMemoryKeyTypeTests(InMemoryKeyTypeTests.Fixture fixture)
     [ConditionalFact]
     public virtual Task String() => this.Test<string>("foo");
 
-    protected override async Task Test<TKey>(TKey mainValue)
+    protected override async Task Test<TKey>(TKey keyValue)
     {
-        await base.Test(mainValue);
+        await base.Test(keyValue);
 
         // For InMemory, delete the collection, otherwise the next test that runs will fail because the collection
         // already exists but with the previous key type.

dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisHashSetKeyTypeTests.cs

@@ -25,5 +25,8 @@ public class RedisHashSetKeyTypeTests(RedisHashSetKeyTypeTests.Fixture fixture)
         // we seem to get key values from the previous collection despite having deleted and recreated it. So we uniquify the collection name instead.
         public override VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>()
             => this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
+
+        public override VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>()
+            => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
     }
 }

dotnet/test/VectorData/Redis.ConformanceTests/TypeTests/RedisJsonKeyTypeTests.cs

@@ -25,5 +25,8 @@ public class RedisJsonKeyTypeTests(RedisJsonKeyTypeTests.Fixture fixture)
         // we seem to get key values from the previous collection despite having deleted and recreated it. So we uniquify the collection name instead.
         public override VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>()
             => this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
+
+        public override VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>()
+            => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
     }
 }

dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/KeyTypeTests.cs

@@ -14,7 +14,7 @@ public abstract class KeyTypeTests(KeyTypeTests.Fixture fixture)
     [ConditionalFact]
     public virtual Task Guid() => this.Test<Guid>(new Guid("603840bf-cf91-4521-8b8e-8b6a2e75910a"));
 
-    protected virtual async Task Test<TKey>(TKey mainValue)
+    protected virtual async Task Test<TKey>(TKey keyValue)
         where TKey : notnull
     {
         using var collection = fixture.CreateCollection<TKey>();
@@ -24,18 +24,45 @@ protected virtual async Task Test<TKey>(TKey mainValue)
 
         var record = new Record<TKey>
         {
-            Key = mainValue,
+            Key = keyValue,
             Int = 8,
             Vector = new ReadOnlyMemory<float>([1, 2, 3])
         };
 
         await collection.UpsertAsync(record);
         await fixture.TestStore.WaitForDataAsync(collection, recordCount: 1);
 
-        var result = await collection.GetAsync(mainValue);
+        var result = await collection.GetAsync(keyValue);
 
         Assert.NotNull(result);
+        Assert.Equal(keyValue, result.Key);
         Assert.Equal(8, result.Int);
+
+        ///////////////////////
+        // Test dynamic mapping
+        ///////////////////////
+        await collection.DeleteAsync(keyValue);
+        await fixture.TestStore.WaitForDataAsync(collection, recordCount: 0);
+
+        using var dynamicCollection = fixture.CreateDynamicCollection<TKey>();
+        await dynamicCollection.EnsureCollectionExistsAsync();
+
+        var dynamicRecord = new Dictionary<string, object?>
+        {
+            [nameof(Record<TKey>.Key)] = keyValue,
+            [nameof(Record<TKey>.Int)] = 8,
+            [nameof(Record<TKey>.Vector)] = new ReadOnlyMemory<float>([1, 2, 3])
+        };
+
+        await dynamicCollection.UpsertAsync(dynamicRecord);
+        await fixture.TestStore.WaitForDataAsync(dynamicCollection, recordCount: 1);
+
+        var dynamicResult = await dynamicCollection.GetAsync(keyValue);
+
+        Assert.NotNull(dynamicResult);
+        Assert.IsType<TKey>(dynamicResult[nameof(Record<TKey>.Key)]);
+        Assert.Equal(keyValue, (TKey)dynamicResult[nameof(Record<TKey>.Key)]!);
+        Assert.Equal(8, dynamicResult[nameof(Record<TKey>.Int)]);
     }
 
     public abstract class Fixture : VectorStoreFixture
@@ -47,6 +74,10 @@ public virtual VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>(
             where TKey : notnull
             => this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName, this.CreateRecordDefinition<TKey>());
 
+        public virtual VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>()
+            where TKey : notnull
+            => this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition<TKey>());
+
         public virtual VectorStoreCollectionDefinition CreateRecordDefinition<TKey>()
             where TKey : notnull
             => new()