Skip to content

Commit 7a87f47

Browse files
rojiroji
committed
[MEVD] Fix Guid (and other) keys for dynamic collections
Closes #13340
1 parent a5810db commit 7a87f47

File tree

10 files changed

+111
-22
lines changed

10 files changed

+111
-22
lines changed

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

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,15 @@ public BsonDocument MapFromDataToStorageModel(Dictionary<string, object?> dataMo
3030
: keyValue switch
3131
{
3232
string s => s,
33+
Guid g => BsonValue.Create(g),
34+
ObjectId o => o,
35+
long i => i,
36+
int i => i,
37+
3338
null => throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."),
3439
_ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string.")
3540
};
3641

37-
document[MongoConstants.MongoReservedKeyPropertyName] = (string)(dataModel[model.KeyProperty.ModelName]
38-
?? throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."));
39-
4042
foreach (var property in model.DataProperties)
4143
{
4244
if (dataModel.TryGetValue(property.ModelName, out var dataValue))
@@ -88,9 +90,22 @@ Embedding<float> e
8890
switch (property)
8991
{
9092
case KeyPropertyModel keyProperty:
91-
result[keyProperty.ModelName] = storageModel.TryGetValue(MongoConstants.MongoReservedKeyPropertyName, out var keyValue)
92-
? keyValue.AsString
93-
: throw new InvalidOperationException("No key property was found in the record retrieved from storage.");
93+
if (!storageModel.TryGetValue(MongoConstants.MongoReservedKeyPropertyName, out var keyValue))
94+
{
95+
throw new InvalidOperationException("No key property was found in the record retrieved from storage.");
96+
}
97+
98+
result[keyProperty.ModelName] = keyProperty.Type switch
99+
{
100+
var t when t == typeof(string) => keyValue.AsString,
101+
var t when t == typeof(Guid) => keyValue.AsGuid,
102+
var t when t == typeof(ObjectId) => keyValue.AsObjectId,
103+
var t when t == typeof(long) => keyValue.AsInt64,
104+
var t when t == typeof(int) => keyValue.AsInt32,
105+
106+
_ => throw new UnreachableException()
107+
};
108+
94109
continue;
95110

96111
case DataPropertyModel dataProperty:

‎dotnet/src/VectorData/AzureAISearch/AzureAISearchDynamicMapper.cs‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,10 @@ public JsonObject MapFromDataToStorageModel(Dictionary<string, object?> dataMode
2929
: keyValue switch
3030
{
3131
string s => s,
32+
Guid g => g.ToString(),
33+
3234
null => throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."),
33-
_ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string.")
35+
_ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string or Guid.")
3436
};
3537

3638
foreach (var dataProperty in model.DataProperties)
@@ -97,9 +99,16 @@ public JsonObject MapFromDataToStorageModel(Dictionary<string, object?> dataMode
9799
switch (property)
98100
{
99101
case KeyPropertyModel keyProperty:
100-
result[keyProperty.ModelName] = (string?)storageModel[keyProperty.StorageName]
102+
var key = (string?)storageModel[keyProperty.StorageName]
101103
?? throw new InvalidOperationException($"The key property '{keyProperty.StorageName}' is missing from the record retrieved from storage.");
102104

105+
result[keyProperty.ModelName] = keyProperty.Type switch
106+
{
107+
var t when t == typeof(string) => key,
108+
var t when t == typeof(Guid) => Guid.Parse(key),
109+
_ => throw new UnreachableException()
110+
};
111+
103112
continue;
104113

105114
case DataPropertyModel dataProperty:

‎dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlCollection.cs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -860,7 +860,9 @@ private static IEnumerable<CosmosNoSqlCompositeKey> GetCompositeKeys(IEnumerable
860860
IEnumerable<object> k => k.Select(key => key switch
861861
{
862862
string s => new CosmosNoSqlCompositeKey(recordKey: s, partitionKey: s),
863+
Guid g when g.ToString() is var guidString => new CosmosNoSqlCompositeKey(recordKey: guidString, partitionKey: guidString),
863864
CosmosNoSqlCompositeKey ck => ck,
865+
864866
_ => throw new ArgumentException($"Invalid key type '{key.GetType().Name}'.")
865867
}),
866868

‎dotnet/src/VectorData/CosmosNoSql/CosmosNoSqlDynamicMapper.cs‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ public JsonObject MapFromDataToStorageModel(Dictionary<string, object?> dataMode
2929
: keyValue switch
3030
{
3131
string s => s,
32+
Guid g => g.ToString(),
33+
3234
null => throw new InvalidOperationException($"Key property '{model.KeyProperty.ModelName}' is null."),
3335
_ => throw new InvalidCastException($"Key property '{model.KeyProperty.ModelName}' must be a string.")
3436
};
@@ -121,9 +123,16 @@ static bool TryGetReadOnlyMemory<T>(object value, [NotNullWhen(true)] out ReadOn
121123
switch (property)
122124
{
123125
case KeyPropertyModel keyProperty:
124-
result[keyProperty.ModelName] = storageModel.TryGetPropertyValue(CosmosNoSqlConstants.ReservedKeyPropertyName, out var keyValue)
125-
? keyValue?.GetValue<string>()
126-
: throw new InvalidOperationException("No key property was found in the record retrieved from storage.");
126+
var key = (string?)storageModel[CosmosNoSqlConstants.ReservedKeyPropertyName]
127+
?? throw new InvalidOperationException($"The key property '{keyProperty.StorageName}' is missing from the record retrieved from storage.");
128+
129+
result[keyProperty.ModelName] = keyProperty.Type switch
130+
{
131+
var t when t == typeof(string) => key,
132+
var t when t == typeof(Guid) => Guid.Parse(key),
133+
_ => throw new UnreachableException()
134+
};
135+
127136
continue;
128137

129138
case DataPropertyModel dataProperty:

‎dotnet/src/VectorData/Redis/RedisJsonDynamicMapper.cs‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -86,22 +86,39 @@ internal class RedisJsonDynamicMapper(CollectionModel model, JsonSerializerOptio
8686
}
8787
}
8888

89-
return ((string)dataModel[model.KeyProperty.ModelName]!, jsonObject);
89+
var storageKey = dataModel[model.KeyProperty.ModelName] switch
90+
{
91+
string s => s,
92+
Guid g => g.ToString(),
93+
94+
_ => throw new UnreachableException()
95+
};
96+
97+
return (storageKey, jsonObject);
9098
}
9199

92100
/// <inheritdoc />
93101
public Dictionary<string, object?> MapFromStorageToDataModel((string Key, JsonNode Node) storageModel, bool includeVectors)
94102
{
95103
var dataModel = new Dictionary<string, object?>
96104
{
97-
[model.KeyProperty.ModelName] = storageModel.Key,
105+
[model.KeyProperty.ModelName] = model.KeyProperty.Type switch
106+
{
107+
Type t when t == typeof(string) => storageModel.Key,
108+
Type t when t == typeof(Guid) => Guid.Parse(storageModel.Key),
109+
110+
_ => throw new UnreachableException()
111+
},
98112
};
99113

100114
// 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.
115+
// If there's a single data property, we get a simple value (no object wrapper).
101116
var jsonObject = storageModel.Node switch
102117
{
103-
JsonObject topLevelJsonObject => topLevelJsonObject,
104-
JsonArray jsonArray and [JsonObject arrayEntryJsonObject] => arrayEntryJsonObject,
118+
JsonValue v when model.DataProperties is [var singleDataProperty] => new JsonObject([new(singleDataProperty.StorageName, v)]),
119+
JsonObject o => o,
120+
JsonArray a and [JsonObject arrayEntryJsonObject] => arrayEntryJsonObject,
121+
105122
_ => throw new InvalidOperationException($"Invalid data format for document with key '{storageModel.Key}'"),
106123
};
107124

‎dotnet/src/VectorData/Redis/RedisJsonDynamicModelBuilder.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ internal class RedisJsonDynamicModelBuilder(CollectionModelBuildingOptions optio
1919

2020
protected override bool IsKeyPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes)
2121
{
22-
supportedTypes = "string";
22+
supportedTypes = "string, Guid";
2323

24-
return type == typeof(string);
24+
return type == typeof(string) || type == typeof(Guid);
2525
}
2626

2727
protected override bool IsDataPropertyTypeValid(Type type, [NotNullWhen(false)] out string? supportedTypes)

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,9 @@ public class InMemoryKeyTypeTests(InMemoryKeyTypeTests.Fixture fixture)
2222
[ConditionalFact]
2323
public virtual Task String() => this.Test<string>("foo");
2424

25-
protected override async Task Test<TKey>(TKey mainValue)
25+
protected override async Task Test<TKey>(TKey keyValue)
2626
{
27-
await base.Test(mainValue);
27+
await base.Test(keyValue);
2828

2929
// For InMemory, delete the collection, otherwise the next test that runs will fail because the collection
3030
// already exists but with the previous key type.

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,5 +25,8 @@ public class RedisHashSetKeyTypeTests(RedisHashSetKeyTypeTests.Fixture fixture)
2525
// we seem to get key values from the previous collection despite having deleted and recreated it. So we uniquify the collection name instead.
2626
public override VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>()
2727
=> this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
28+
29+
public override VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>()
30+
=> this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
2831
}
2932
}

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,5 +25,8 @@ public class RedisJsonKeyTypeTests(RedisJsonKeyTypeTests.Fixture fixture)
2525
// we seem to get key values from the previous collection despite having deleted and recreated it. So we uniquify the collection name instead.
2626
public override VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>()
2727
=> this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
28+
29+
public override VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>()
30+
=> this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName + (++this._collectionCounter), this.CreateRecordDefinition<TKey>());
2831
}
2932
}

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

Lines changed: 34 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public abstract class KeyTypeTests(KeyTypeTests.Fixture fixture)
1414
[ConditionalFact]
1515
public virtual Task Guid() => this.Test<Guid>(new Guid("603840bf-cf91-4521-8b8e-8b6a2e75910a"));
1616

17-
protected virtual async Task Test<TKey>(TKey mainValue)
17+
protected virtual async Task Test<TKey>(TKey keyValue)
1818
where TKey : notnull
1919
{
2020
using var collection = fixture.CreateCollection<TKey>();
@@ -24,18 +24,45 @@ protected virtual async Task Test<TKey>(TKey mainValue)
2424

2525
var record = new Record<TKey>
2626
{
27-
Key = mainValue,
27+
Key = keyValue,
2828
Int = 8,
2929
Vector = new ReadOnlyMemory<float>([1, 2, 3])
3030
};
3131

3232
await collection.UpsertAsync(record);
3333
await fixture.TestStore.WaitForDataAsync(collection, recordCount: 1);
3434

35-
var result = await collection.GetAsync(mainValue);
35+
var result = await collection.GetAsync(keyValue);
3636

3737
Assert.NotNull(result);
38+
Assert.Equal(keyValue, result.Key);
3839
Assert.Equal(8, result.Int);
40+
41+
///////////////////////
42+
// Test dynamic mapping
43+
///////////////////////
44+
await collection.DeleteAsync(keyValue);
45+
await fixture.TestStore.WaitForDataAsync(collection, recordCount: 0);
46+
47+
using var dynamicCollection = fixture.CreateDynamicCollection<TKey>();
48+
await dynamicCollection.EnsureCollectionExistsAsync();
49+
50+
var dynamicRecord = new Dictionary<string, object?>
51+
{
52+
[nameof(Record<TKey>.Key)] = keyValue,
53+
[nameof(Record<TKey>.Int)] = 8,
54+
[nameof(Record<TKey>.Vector)] = new ReadOnlyMemory<float>([1, 2, 3])
55+
};
56+
57+
await dynamicCollection.UpsertAsync(dynamicRecord);
58+
await fixture.TestStore.WaitForDataAsync(dynamicCollection, recordCount: 1);
59+
60+
var dynamicResult = await dynamicCollection.GetAsync(keyValue);
61+
62+
Assert.NotNull(dynamicResult);
63+
Assert.IsType<TKey>(dynamicResult[nameof(Record<TKey>.Key)]);
64+
Assert.Equal(keyValue, (TKey)dynamicResult[nameof(Record<TKey>.Key)]!);
65+
Assert.Equal(8, dynamicResult[nameof(Record<TKey>.Int)]);
3966
}
4067

4168
public abstract class Fixture : VectorStoreFixture
@@ -47,6 +74,10 @@ public virtual VectorStoreCollection<TKey, Record<TKey>> CreateCollection<TKey>(
4774
where TKey : notnull
4875
=> this.TestStore.DefaultVectorStore.GetCollection<TKey, Record<TKey>>(this.CollectionName, this.CreateRecordDefinition<TKey>());
4976

77+
public virtual VectorStoreCollection<object, Dictionary<string, object?>> CreateDynamicCollection<TKey>()
78+
where TKey : notnull
79+
=> this.TestStore.DefaultVectorStore.GetDynamicCollection(this.CollectionName, this.CreateRecordDefinition<TKey>());
80+
5081
public virtual VectorStoreCollectionDefinition CreateRecordDefinition<TKey>()
5182
where TKey : notnull
5283
=> new()

0 commit comments

Comments
 (0)