Skip to content

Commit fd6d4a1

Browse files
authored
.Net: Implemented generic data model support for Azure CosmosDB NoSQL connector (#8981)
### Motivation and Context <!-- Thank you for your contribution to the semantic-kernel repo! Please help reviewers and future users, providing the following information: 1. Why is this change required? 2. What problem does it solve? 3. What scenario does it contribute to? 4. If it fixes an open issue, please link to the issue here. --> Related: #6522 - Implemented `AzureCosmosDBNoSQLGenericDataModelMapper` class. - Added unit and integration tests. ### Contribution Checklist <!-- Before submitting this PR, please make sure: --> - [x] The code builds clean without any errors or warnings - [x] The PR follows the [SK Contribution Guidelines](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md) and the [pre-submission formatting script](https://github.com/microsoft/semantic-kernel/blob/main/CONTRIBUTING.md#development-scripts) raises no violations - [x] All unit tests pass, and I have added new tests where possible - [x] I didn't break anyone 😄
1 parent 1c72a26 commit fd6d4a1

File tree

6 files changed

+637
-16
lines changed

6 files changed

+637
-16
lines changed
Lines changed: 368 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,368 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System;
4+
using System.Collections.Generic;
5+
using System.Linq;
6+
using System.Text.Json;
7+
using System.Text.Json.Nodes;
8+
using Microsoft.SemanticKernel.Connectors.AzureCosmosDBNoSQL;
9+
using Microsoft.SemanticKernel.Data;
10+
using Xunit;
11+
12+
namespace SemanticKernel.Connectors.AzureCosmosDBNoSQL.UnitTests;
13+
14+
/// <summary>
15+
/// Unit tests for <see cref="AzureCosmosDBNoSQLGenericDataModelMapper"/> class.
16+
/// </summary>
17+
public sealed class AzureCosmosDBNoSQLGenericDataModelMapperTests
18+
{
19+
private static readonly JsonSerializerOptions s_jsonSerializerOptions = JsonSerializerOptions.Default;
20+
21+
private static readonly VectorStoreRecordDefinition s_vectorStoreRecordDefinition = new()
22+
{
23+
Properties = new List<VectorStoreRecordProperty>
24+
{
25+
new VectorStoreRecordKeyProperty("Key", typeof(string)),
26+
new VectorStoreRecordDataProperty("BoolDataProp", typeof(bool)),
27+
new VectorStoreRecordDataProperty("NullableBoolDataProp", typeof(bool?)),
28+
new VectorStoreRecordDataProperty("StringDataProp", typeof(string)),
29+
new VectorStoreRecordDataProperty("IntDataProp", typeof(int)),
30+
new VectorStoreRecordDataProperty("NullableIntDataProp", typeof(int?)),
31+
new VectorStoreRecordDataProperty("LongDataProp", typeof(long)),
32+
new VectorStoreRecordDataProperty("NullableLongDataProp", typeof(long?)),
33+
new VectorStoreRecordDataProperty("FloatDataProp", typeof(float)),
34+
new VectorStoreRecordDataProperty("NullableFloatDataProp", typeof(float?)),
35+
new VectorStoreRecordDataProperty("DoubleDataProp", typeof(double)),
36+
new VectorStoreRecordDataProperty("NullableDoubleDataProp", typeof(double?)),
37+
new VectorStoreRecordDataProperty("DateTimeOffsetDataProp", typeof(DateTimeOffset)),
38+
new VectorStoreRecordDataProperty("NullableDateTimeOffsetDataProp", typeof(DateTimeOffset?)),
39+
new VectorStoreRecordDataProperty("TagListDataProp", typeof(List<string>)),
40+
#if NET5_0_OR_GREATER
41+
new VectorStoreRecordVectorProperty("HalfVector", typeof(ReadOnlyMemory<Half>)),
42+
new VectorStoreRecordVectorProperty("NullableHalfVector", typeof(ReadOnlyMemory<Half>?)),
43+
#endif
44+
new VectorStoreRecordVectorProperty("FloatVector", typeof(ReadOnlyMemory<float>)),
45+
new VectorStoreRecordVectorProperty("NullableFloatVector", typeof(ReadOnlyMemory<float>?)),
46+
new VectorStoreRecordVectorProperty("ByteVector", typeof(ReadOnlyMemory<byte>)),
47+
new VectorStoreRecordVectorProperty("NullableByteVector", typeof(ReadOnlyMemory<byte>?)),
48+
new VectorStoreRecordVectorProperty("SByteVector", typeof(ReadOnlyMemory<sbyte>)),
49+
new VectorStoreRecordVectorProperty("NullableSByteVector", typeof(ReadOnlyMemory<sbyte>?)),
50+
},
51+
};
52+
53+
private static readonly Dictionary<string, string> s_storagePropertyNames =
54+
s_vectorStoreRecordDefinition.Properties.ToDictionary(
55+
k => k.DataModelPropertyName,
56+
v => v is VectorStoreRecordKeyProperty ? "id" : v.DataModelPropertyName);
57+
58+
#if NET5_0_OR_GREATER
59+
private static readonly Half[] s_halfVector = [(Half)1.0f, (Half)2.0f, (Half)3.0f];
60+
#endif
61+
private static readonly float[] s_floatVector = [1.0f, 2.0f, 3.0f];
62+
private static readonly byte[] s_byteVector = [1, 2, 3];
63+
private static readonly sbyte[] s_sbyteVector = [1, 2, 3];
64+
private static readonly List<string> s_taglist = ["tag1", "tag2"];
65+
66+
[Fact]
67+
public void MapFromDataToStorageModelMapsAllSupportedTypes()
68+
{
69+
// Arrange
70+
var test = JsonSerializer.Serialize(s_byteVector);
71+
72+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
73+
s_vectorStoreRecordDefinition,
74+
s_storagePropertyNames,
75+
s_jsonSerializerOptions);
76+
77+
var dataModel = new VectorStoreGenericDataModel<string>("key")
78+
{
79+
Data =
80+
{
81+
["BoolDataProp"] = true,
82+
["NullableBoolDataProp"] = false,
83+
["StringDataProp"] = "string",
84+
["IntDataProp"] = 1,
85+
["NullableIntDataProp"] = 2,
86+
["LongDataProp"] = 3L,
87+
["NullableLongDataProp"] = 4L,
88+
["FloatDataProp"] = 5.0f,
89+
["NullableFloatDataProp"] = 6.0f,
90+
["DoubleDataProp"] = 7.0,
91+
["NullableDoubleDataProp"] = 8.0,
92+
["DateTimeOffsetDataProp"] = new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero),
93+
["NullableDateTimeOffsetDataProp"] = new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero),
94+
["TagListDataProp"] = s_taglist,
95+
},
96+
Vectors =
97+
{
98+
#if NET5_0_OR_GREATER
99+
["HalfVector"] = new ReadOnlyMemory<Half>(s_halfVector),
100+
["NullableHalfVector"] = new ReadOnlyMemory<Half>(s_halfVector),
101+
#endif
102+
["FloatVector"] = new ReadOnlyMemory<float>(s_floatVector),
103+
["NullableFloatVector"] = new ReadOnlyMemory<float>(s_floatVector),
104+
["ByteVector"] = new ReadOnlyMemory<byte>(s_byteVector),
105+
["NullableByteVector"] = new ReadOnlyMemory<byte>(s_byteVector),
106+
["SByteVector"] = new ReadOnlyMemory<sbyte>(s_sbyteVector),
107+
["NullableSByteVector"] = new ReadOnlyMemory<sbyte>(s_sbyteVector)
108+
},
109+
};
110+
111+
// Act
112+
var storageModel = sut.MapFromDataToStorageModel(dataModel);
113+
114+
// Assert
115+
Assert.Equal("key", (string?)storageModel["id"]);
116+
Assert.Equal(true, (bool?)storageModel["BoolDataProp"]);
117+
Assert.Equal(false, (bool?)storageModel["NullableBoolDataProp"]);
118+
Assert.Equal("string", (string?)storageModel["StringDataProp"]);
119+
Assert.Equal(1, (int?)storageModel["IntDataProp"]);
120+
Assert.Equal(2, (int?)storageModel["NullableIntDataProp"]);
121+
Assert.Equal(3L, (long?)storageModel["LongDataProp"]);
122+
Assert.Equal(4L, (long?)storageModel["NullableLongDataProp"]);
123+
Assert.Equal(5.0f, (float?)storageModel["FloatDataProp"]);
124+
Assert.Equal(6.0f, (float?)storageModel["NullableFloatDataProp"]);
125+
Assert.Equal(7.0, (double?)storageModel["DoubleDataProp"]);
126+
Assert.Equal(8.0, (double?)storageModel["NullableDoubleDataProp"]);
127+
Assert.Equal(new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero), (DateTimeOffset?)storageModel["DateTimeOffsetDataProp"]);
128+
Assert.Equal(new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero), (DateTimeOffset?)storageModel["NullableDateTimeOffsetDataProp"]);
129+
Assert.Equal(s_taglist, storageModel["TagListDataProp"]!.AsArray().GetValues<string>().ToArray());
130+
#if NET5_0_OR_GREATER
131+
Assert.Equal(s_halfVector, storageModel["HalfVector"]!.AsArray().Select(l => (Half)(float)l!).ToArray());
132+
Assert.Equal(s_halfVector, storageModel["NullableHalfVector"]!.AsArray().Select(l => (Half)(float)l!).ToArray());
133+
#endif
134+
Assert.Equal(s_floatVector, storageModel["FloatVector"]!.AsArray().GetValues<float>().ToArray());
135+
Assert.Equal(s_floatVector, storageModel["NullableFloatVector"]!.AsArray().GetValues<float>().ToArray());
136+
Assert.Equal(s_byteVector, storageModel["ByteVector"]!.AsArray().GetValues<byte>().ToArray());
137+
Assert.Equal(s_byteVector, storageModel["NullableByteVector"]!.AsArray().GetValues<byte>().ToArray());
138+
Assert.Equal(s_sbyteVector, storageModel["SByteVector"]!.AsArray().GetValues<sbyte>().ToArray());
139+
Assert.Equal(s_sbyteVector, storageModel["NullableSByteVector"]!.AsArray().GetValues<sbyte>().ToArray());
140+
}
141+
142+
[Fact]
143+
public void MapFromDataToStorageModelMapsNullValues()
144+
{
145+
// Arrange
146+
VectorStoreRecordDefinition vectorStoreRecordDefinition = new()
147+
{
148+
Properties = new List<VectorStoreRecordProperty>
149+
{
150+
new VectorStoreRecordKeyProperty("Key", typeof(string)),
151+
new VectorStoreRecordDataProperty("StringDataProp", typeof(string)),
152+
new VectorStoreRecordDataProperty("NullableIntDataProp", typeof(int?)),
153+
new VectorStoreRecordVectorProperty("NullableFloatVector", typeof(ReadOnlyMemory<float>?)),
154+
},
155+
};
156+
157+
var dataModel = new VectorStoreGenericDataModel<string>("key")
158+
{
159+
Data =
160+
{
161+
["StringDataProp"] = null,
162+
["NullableIntDataProp"] = null,
163+
},
164+
Vectors =
165+
{
166+
["NullableFloatVector"] = null,
167+
},
168+
};
169+
170+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
171+
s_vectorStoreRecordDefinition,
172+
s_storagePropertyNames,
173+
s_jsonSerializerOptions);
174+
175+
// Act
176+
var storageModel = sut.MapFromDataToStorageModel(dataModel);
177+
178+
// Assert
179+
Assert.Null(storageModel["StringDataProp"]);
180+
Assert.Null(storageModel["NullableIntDataProp"]);
181+
Assert.Null(storageModel["NullableFloatVector"]);
182+
}
183+
184+
[Fact]
185+
public void MapFromStorageToDataModelMapsAllSupportedTypes()
186+
{
187+
// Arrange
188+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
189+
s_vectorStoreRecordDefinition,
190+
s_storagePropertyNames,
191+
s_jsonSerializerOptions);
192+
193+
var storageModel = new JsonObject
194+
{
195+
["id"] = "key",
196+
["BoolDataProp"] = true,
197+
["NullableBoolDataProp"] = false,
198+
["StringDataProp"] = "string",
199+
["IntDataProp"] = 1,
200+
["NullableIntDataProp"] = 2,
201+
["LongDataProp"] = 3L,
202+
["NullableLongDataProp"] = 4L,
203+
["FloatDataProp"] = 5.0f,
204+
["NullableFloatDataProp"] = 6.0f,
205+
["DoubleDataProp"] = 7.0,
206+
["NullableDoubleDataProp"] = 8.0,
207+
["DateTimeOffsetDataProp"] = new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero),
208+
["NullableDateTimeOffsetDataProp"] = new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero),
209+
["TagListDataProp"] = new JsonArray(s_taglist.Select(l => (JsonValue)l).ToArray()),
210+
#if NET5_0_OR_GREATER
211+
["HalfVector"] = new JsonArray(s_halfVector.Select(l => (JsonValue)(float)l).ToArray()),
212+
["NullableHalfVector"] = new JsonArray(s_halfVector.Select(l => (JsonValue)(float)l).ToArray()),
213+
#endif
214+
["FloatVector"] = new JsonArray(s_floatVector.Select(l => (JsonValue)l).ToArray()),
215+
["NullableFloatVector"] = new JsonArray(s_floatVector.Select(l => (JsonValue)l).ToArray()),
216+
["ByteVector"] = new JsonArray(s_byteVector.Select(l => (JsonValue)l).ToArray()),
217+
["NullableByteVector"] = new JsonArray(s_byteVector.Select(l => (JsonValue)l).ToArray()),
218+
["SByteVector"] = new JsonArray(s_sbyteVector.Select(l => (JsonValue)l).ToArray()),
219+
["NullableSByteVector"] = new JsonArray(s_sbyteVector.Select(l => (JsonValue)l).ToArray())
220+
};
221+
222+
// Act
223+
var dataModel = sut.MapFromStorageToDataModel(storageModel, new StorageToDataModelMapperOptions { IncludeVectors = true });
224+
225+
// Assert
226+
Assert.Equal("key", dataModel.Key);
227+
Assert.Equal(true, dataModel.Data["BoolDataProp"]);
228+
Assert.Equal(false, dataModel.Data["NullableBoolDataProp"]);
229+
Assert.Equal("string", dataModel.Data["StringDataProp"]);
230+
Assert.Equal(1, dataModel.Data["IntDataProp"]);
231+
Assert.Equal(2, dataModel.Data["NullableIntDataProp"]);
232+
Assert.Equal(3L, dataModel.Data["LongDataProp"]);
233+
Assert.Equal(4L, dataModel.Data["NullableLongDataProp"]);
234+
Assert.Equal(5.0f, dataModel.Data["FloatDataProp"]);
235+
Assert.Equal(6.0f, dataModel.Data["NullableFloatDataProp"]);
236+
Assert.Equal(7.0, dataModel.Data["DoubleDataProp"]);
237+
Assert.Equal(8.0, dataModel.Data["NullableDoubleDataProp"]);
238+
Assert.Equal(new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero), dataModel.Data["DateTimeOffsetDataProp"]);
239+
Assert.Equal(new DateTimeOffset(2021, 1, 1, 0, 0, 0, TimeSpan.Zero), dataModel.Data["NullableDateTimeOffsetDataProp"]);
240+
Assert.Equal(s_taglist, dataModel.Data["TagListDataProp"]);
241+
#if NET5_0_OR_GREATER
242+
Assert.Equal(s_halfVector, ((ReadOnlyMemory<Half>)dataModel.Vectors["HalfVector"]!).ToArray());
243+
Assert.Equal(s_halfVector, ((ReadOnlyMemory<Half>)dataModel.Vectors["NullableHalfVector"]!)!.ToArray());
244+
#endif
245+
Assert.Equal(s_floatVector, ((ReadOnlyMemory<float>)dataModel.Vectors["FloatVector"]!).ToArray());
246+
Assert.Equal(s_floatVector, ((ReadOnlyMemory<float>)dataModel.Vectors["NullableFloatVector"]!)!.ToArray());
247+
Assert.Equal(s_byteVector, ((ReadOnlyMemory<byte>)dataModel.Vectors["ByteVector"]!).ToArray());
248+
Assert.Equal(s_byteVector, ((ReadOnlyMemory<byte>)dataModel.Vectors["NullableByteVector"]!)!.ToArray());
249+
Assert.Equal(s_sbyteVector, ((ReadOnlyMemory<sbyte>)dataModel.Vectors["SByteVector"]!).ToArray());
250+
Assert.Equal(s_sbyteVector, ((ReadOnlyMemory<sbyte>)dataModel.Vectors["NullableSByteVector"]!)!.ToArray());
251+
}
252+
253+
[Fact]
254+
public void MapFromStorageToDataModelMapsNullValues()
255+
{
256+
// Arrange
257+
VectorStoreRecordDefinition vectorStoreRecordDefinition = new()
258+
{
259+
Properties = new List<VectorStoreRecordProperty>
260+
{
261+
new VectorStoreRecordKeyProperty("Key", typeof(string)),
262+
new VectorStoreRecordDataProperty("StringDataProp", typeof(string)),
263+
new VectorStoreRecordDataProperty("NullableIntDataProp", typeof(int?)),
264+
new VectorStoreRecordVectorProperty("NullableFloatVector", typeof(ReadOnlyMemory<float>?)),
265+
},
266+
};
267+
268+
var storageModel = new JsonObject
269+
{
270+
["id"] = "key",
271+
["StringDataProp"] = null,
272+
["NullableIntDataProp"] = null,
273+
["NullableFloatVector"] = null
274+
};
275+
276+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
277+
s_vectorStoreRecordDefinition,
278+
s_storagePropertyNames,
279+
s_jsonSerializerOptions);
280+
281+
// Act
282+
var dataModel = sut.MapFromStorageToDataModel(storageModel, new StorageToDataModelMapperOptions { IncludeVectors = true });
283+
284+
// Assert
285+
Assert.Equal("key", dataModel.Key);
286+
Assert.Null(dataModel.Data["StringDataProp"]);
287+
Assert.Null(dataModel.Data["NullableIntDataProp"]);
288+
Assert.Null(dataModel.Vectors["NullableFloatVector"]);
289+
}
290+
291+
[Fact]
292+
public void MapFromStorageToDataModelThrowsForMissingKey()
293+
{
294+
// Arrange
295+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
296+
s_vectorStoreRecordDefinition,
297+
s_storagePropertyNames,
298+
s_jsonSerializerOptions);
299+
300+
var storageModel = new JsonObject();
301+
302+
// Act & Assert
303+
var exception = Assert.Throws<VectorStoreRecordMappingException>(
304+
() => sut.MapFromStorageToDataModel(storageModel, new StorageToDataModelMapperOptions { IncludeVectors = true }));
305+
}
306+
307+
[Fact]
308+
public void MapFromDataToStorageModelSkipsMissingProperties()
309+
{
310+
// Arrange
311+
VectorStoreRecordDefinition vectorStoreRecordDefinition = new()
312+
{
313+
Properties = new List<VectorStoreRecordProperty>
314+
{
315+
new VectorStoreRecordKeyProperty("Key", typeof(string)),
316+
new VectorStoreRecordDataProperty("StringDataProp", typeof(string)),
317+
new VectorStoreRecordVectorProperty("FloatVector", typeof(ReadOnlyMemory<float>)),
318+
},
319+
};
320+
321+
var dataModel = new VectorStoreGenericDataModel<string>("key");
322+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
323+
s_vectorStoreRecordDefinition,
324+
s_storagePropertyNames,
325+
s_jsonSerializerOptions);
326+
327+
// Act
328+
var storageModel = sut.MapFromDataToStorageModel(dataModel);
329+
330+
// Assert
331+
Assert.Equal("key", (string?)storageModel["id"]);
332+
Assert.False(storageModel.ContainsKey("StringDataProp"));
333+
Assert.False(storageModel.ContainsKey("FloatVector"));
334+
}
335+
336+
[Fact]
337+
public void MapFromStorageToDataModelSkipsMissingProperties()
338+
{
339+
// Arrange
340+
VectorStoreRecordDefinition vectorStoreRecordDefinition = new()
341+
{
342+
Properties = new List<VectorStoreRecordProperty>
343+
{
344+
new VectorStoreRecordKeyProperty("Key", typeof(string)),
345+
new VectorStoreRecordDataProperty("StringDataProp", typeof(string)),
346+
new VectorStoreRecordVectorProperty("FloatVector", typeof(ReadOnlyMemory<float>)),
347+
},
348+
};
349+
350+
var storageModel = new JsonObject
351+
{
352+
["id"] = "key"
353+
};
354+
355+
var sut = new AzureCosmosDBNoSQLGenericDataModelMapper(
356+
s_vectorStoreRecordDefinition,
357+
s_storagePropertyNames,
358+
s_jsonSerializerOptions);
359+
360+
// Act
361+
var dataModel = sut.MapFromStorageToDataModel(storageModel, new StorageToDataModelMapperOptions { IncludeVectors = true });
362+
363+
// Assert
364+
Assert.Equal("key", dataModel.Key);
365+
Assert.False(dataModel.Data.ContainsKey("StringDataProp"));
366+
Assert.False(dataModel.Vectors.ContainsKey("FloatVector"));
367+
}
368+
}

0 commit comments

Comments
 (0)