Skip to content

Commit d595e48

Browse files
rojiwestey-m
roji
and
westey-m
authored
.Net: [MEVD] Fix dynamic batching APIs for PostgreSQL (#13323)
Fixes #13303 --------- Co-authored-by: westey <[email protected]>
1 parent ad97c3a commit d595e48

File tree

3 files changed

+75
-55
lines changed

3 files changed

+75
-55
lines changed

‎dotnet/src/VectorData/PgVector/PostgresSqlBuilder.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,7 @@ internal static void BuildGetCommand<TKey>(NpgsqlCommand command, string schema,
269269
internal static void BuildGetBatchCommand<TKey>(NpgsqlCommand command, string schema, string tableName, CollectionModel model, List<TKey> keys, bool includeVectors = false)
270270
where TKey : notnull
271271
{
272-
NpgsqlDbType? keyType = PostgresPropertyMapping.GetNpgsqlDbType(typeof(TKey)) ?? throw new ArgumentException($"Unsupported key type {typeof(TKey).Name}");
272+
NpgsqlDbType? keyType = PostgresPropertyMapping.GetNpgsqlDbType(model.KeyProperty.Type) ?? throw new UnreachableException($"Unsupported key type {model.KeyProperty.Type.Name}");
273273

274274
// Generate the column names
275275
var columns = model.Properties

‎dotnet/test/VectorData/InMemory.ConformanceTests/ModelTests/InMemoryDynamicModelTests.cs‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,30 @@ public override async Task GetAsync_single_record(bool includeVectors)
2727
AssertEquivalent(expectedRecord, received, includeVectors: true, fixture.TestStore.VectorsComparable);
2828
}
2929

30+
public override async Task GetAsync_multiple_records(bool includeVectors)
31+
{
32+
if (includeVectors)
33+
{
34+
await base.GetAsync_multiple_records(includeVectors);
35+
return;
36+
}
37+
38+
// InMemory always returns the vectors (IncludeVectors = false isn't respected)
39+
var expectedRecords = fixture.TestData.Take(2);
40+
var ids = expectedRecords.Select(record => record[KeyPropertyName]!);
41+
42+
var received = await fixture.Collection.GetAsync(ids, new() { IncludeVectors = false }).ToArrayAsync();
43+
44+
foreach (var record in expectedRecords)
45+
{
46+
AssertEquivalent(
47+
record,
48+
received.Single(r => r[KeyPropertyName]!.Equals(record[KeyPropertyName])),
49+
includeVectors: true,
50+
fixture.TestStore.VectorsComparable);
51+
}
52+
}
53+
3054
public override async Task GetAsync_with_filter(bool includeVectors)
3155
{
3256
if (includeVectors)

‎dotnet/test/VectorData/VectorData.ConformanceTests/ModelTests/DynamicModelTests.cs‎

Lines changed: 50 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -24,24 +24,23 @@ public virtual async Task GetAsync_single_record(bool includeVectors)
2424
AssertEquivalent(expectedRecord, received, includeVectors, fixture.TestStore.VectorsComparable);
2525
}
2626

27-
// TODO: https://github.com/microsoft/semantic-kernel/issues/13303
28-
// [ConditionalTheory, MemberData(nameof(IncludeVectorsData))]
29-
// public virtual Task GetAsync_multiple_records(bool includeVectors)
30-
// {
31-
// var expectedRecords = fixture.TestData.Take(2);
32-
// var ids = expectedRecords.Select(record => record[KeyPropertyName]!);
33-
34-
// var received = await fixture.Collection.GetAsync(ids, new() { IncludeVectors = includeVectors }).ToArrayAsync();
35-
36-
// foreach (var record in expectedRecords)
37-
// {
38-
// AssertEquivalent(
39-
// record,
40-
// received.Single(r => r[KeyPropertyName]!.Equals(record[KeyPropertyName])),
41-
// includeVectors,
42-
// fixture.TestStore.VectorsComparable);
43-
// }
44-
// }
27+
[ConditionalTheory, MemberData(nameof(IncludeVectorsData))]
28+
public virtual async Task GetAsync_multiple_records(bool includeVectors)
29+
{
30+
var expectedRecords = fixture.TestData.Take(2);
31+
var ids = expectedRecords.Select(record => record[KeyPropertyName]!);
32+
33+
var received = await fixture.Collection.GetAsync(ids, new() { IncludeVectors = includeVectors }).ToArrayAsync();
34+
35+
foreach (var record in expectedRecords)
36+
{
37+
AssertEquivalent(
38+
record,
39+
received.Single(r => r[KeyPropertyName]!.Equals(record[KeyPropertyName])),
40+
includeVectors,
41+
fixture.TestStore.VectorsComparable);
42+
}
43+
}
4544

4645
[ConditionalFact]
4746
public virtual async Task GetAsync_throws_for_null_key()
@@ -241,42 +240,39 @@ public virtual async Task Update_single_record()
241240
AssertEquivalent(updated, received, includeVectors: true, fixture.TestStore.VectorsComparable);
242241
}
243242

244-
// TODO: https://github.com/microsoft/semantic-kernel/issues/13303
245-
// [ConditionalFact]
246-
// public virtual async Task Insert_multiple_records()
247-
// {
248-
// Dictionary<string, object?>[] newRecords =
249-
// [
250-
// new()
251-
// {
252-
// [KeyPropertyName] = fixture.GenerateNextKey<TKey>(),
253-
// [IntegerPropertyName] = 100,
254-
// [StringPropertyName] = "New record 1",
255-
// [EmbeddingPropertyName] = new ReadOnlyMemory<float>([10, 0, 1])
256-
// },
257-
// new()
258-
// {
259-
// [KeyPropertyName] = fixture.GenerateNextKey<TKey>(),
260-
// [IntegerPropertyName] = 101,
261-
// [StringPropertyName] = "New record 2",
262-
// [EmbeddingPropertyName] = new ReadOnlyMemory<float>([10, 0, 2])
263-
// },
264-
// ];
265-
266-
// var keys = newRecords.Select(record => record[KeyPropertyName]!).ToArray();
267-
// Assert.Empty(await this.Collection.GetAsync(keys).ToArrayAsync());
268-
269-
// await this.Collection.UpsertAsync(newRecords);
270-
271-
// var received = await this.Collection.GetAsync(keys, new() { IncludeVectors = true }).ToArrayAsync();
272-
273-
// Assert.Collection(
274-
// received.OrderBy(r => r[IntegerPropertyName]),
275-
// r => AssertEquivalent(newRecords[0], r, includeVectors: true, fixture.TestStore.VectorsComparable),
276-
// r => AssertEquivalent(newRecords[1], r, includeVectors: true, fixture.TestStore.VectorsComparable));
277-
278-
// Assert.Equal(fixture.TestData.Count + 2, await this.GetRecordCount());
279-
// }
243+
[ConditionalFact]
244+
public virtual async Task Insert_multiple_records()
245+
{
246+
Dictionary<string, object?>[] newRecords =
247+
[
248+
new()
249+
{
250+
[KeyPropertyName] = fixture.GenerateNextKey<TKey>(),
251+
[IntegerPropertyName] = 100,
252+
[StringPropertyName] = "New record 1",
253+
[EmbeddingPropertyName] = new ReadOnlyMemory<float>([10, 0, 1])
254+
},
255+
new()
256+
{
257+
[KeyPropertyName] = fixture.GenerateNextKey<TKey>(),
258+
[IntegerPropertyName] = 101,
259+
[StringPropertyName] = "New record 2",
260+
[EmbeddingPropertyName] = new ReadOnlyMemory<float>([10, 0, 2])
261+
},
262+
];
263+
264+
var keys = newRecords.Select(record => record[KeyPropertyName]!).ToArray();
265+
Assert.Empty(await this.Collection.GetAsync(keys).ToArrayAsync());
266+
267+
await this.Collection.UpsertAsync(newRecords);
268+
269+
var received = await this.Collection.GetAsync(keys, new() { IncludeVectors = true }).ToArrayAsync();
270+
271+
Assert.Collection(
272+
received.OrderBy(r => r[IntegerPropertyName]),
273+
r => AssertEquivalent(newRecords[0], r, includeVectors: true, fixture.TestStore.VectorsComparable),
274+
r => AssertEquivalent(newRecords[1], r, includeVectors: true, fixture.TestStore.VectorsComparable));
275+
}
280276

281277
#endregion Upsert
282278

0 commit comments

Comments
 (0)