.Net: Add PostgresVectorStore Memory connector. #9324

lossyrob · 2024-10-18T14:10:10Z

This PR adds a PostgresVectorStore and related classes to Microsoft.SemanticKernel.Connectors.Postgres.

Motivation and Context

As part of the move to having memory connectors implement the new Microsoft.Extensions.VectorData.IVectorStore architecture (see https://github.com/microsoft/semantic-kernel/blob/main/docs/decisions/0050-updated-vector-store-design.md), each memory connector needs to be updated with the new architecture. This PR tackles updating the existing Microsoft.SemanticKernel.Connectors.Postgres package to include this implementation. This will supercede the PostgresMemoryStore implementation.

Some high level comments about design:

PostgresVectorStore and PostgresVectorStoreRecordCollection get injected with an IPostgresVectorStoreDbClient. This abstracts the database communication and allows for unit tests to mock database interactions.
The PostgresVectorStoreDbClient gets passed in a NpgsqlDataSource from the user, which is used to manage connections to the database. The responsibility of connection pool lifecycle management is on the user.
The IPostgresVectorStoreDbClient is designed to accept and produce the storage model, which in this case is a Dictionary<string, object?> . This is the intermediate type that is mapped to by the IVectorStoreRecordMapper.
The PostgresVectorStoreDbClient also takes a IPostgresVectorStoreCollectionSqlBuilder, which generates SQL command information for interacting with the database. This abstracts the SQL queries related to each task, and allows for future expansion. This is particularly targeted at creating a AzureDBForPostgre vector store that will enable alternate vector implementations like DiskANN, while leveraging the same database client as the Postgres connector.
The integration tests for the vector store utilize Docker.Net to bring up a pgvector/pgvector docker container, which test are run against.

Contribution Checklist

The code builds clean without any errors or warnings
The PR follows the SK Contribution Guidelines and the pre-submission formatting script raises no violations
All unit tests pass, and I have added new tests where possible
I didn't break anyone 😄

Work in progress, some methods are not implemented yet.

…ctor-store-dotnet

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresConstants.cs

…ctor-store-dotnet

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresVectorStoreRecordCollection.cs

…ctor-store-dotnet

e.g. Change type of property 'CollectionInts' from 'System.Collections.Generic.ICollection<int>?' to 'System.Collections.Generic.HashSet<int>?' for improved performance

lossyrob · 2024-10-28T23:48:42Z

@westey-m thanks for all the detailed feedback. I think I've addressed all the comments, and this is good for another review.

westey-m · 2024-10-30T18:02:24Z

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresConstants.cs

@@ -64,4 +74,15 @@ internal static class PostgresConstants
    /// <summary>The name of the column that returns distance value in the database.</summary>
    /// <remarks>It is used in the similarity search query. Must not conflict with model property.</remarks>
    public const string DistanceColumnName = "sk_pg_distance";
+
+    /// <summary>The default index kind.</summary>
+    public const string DefaultIndexKind = IndexKind.Hnsw;


I think it's also OK to use Flat here as the default. When I was suggesting using Hnsw earlier, I didn't realise that postgres can also work without, and that's a totally acceptable default.

westey-m · 2024-10-30T18:10:22Z

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresServiceCollectionExtensions.cs

    /// <param name="options">Optional options to further configure the <see cref="IVectorStore"/>.</param>
    /// <param name="serviceId">An optional service id to use as the service key.</param>
    /// <returns>The service collection.</returns>
-    public static IServiceCollection AddPostgresVectorStore(this IServiceCollection services, string connectionString, PostgresVectorStoreOptions? options = default, string? serviceId = default)
+    public static IServiceCollection AddPostgresVectorStore(this IServiceCollection services, NpgsqlDataSource dataSource, PostgresVectorStoreOptions? options = default, string? serviceId = default)


I'm not convinced that this overload is that useful. Where someone builds their DI container, they will not typically be able to create a NpgsqlDataSource in a way that will get it destroyed in the right timeframe. E.g. if they use a using clause it'll just get destroyed immediately before the app really starts. If they register it with DI to allow it to handle it's lifecycle, they will use a different overload.

FYI: We typically follow a philosophy of avoiding adding extra overloads or helpers unless we have a good justification, simply because it means more maintenance and more chances of breaking changes in future. So less is more 😄

westey-m · 2024-10-30T18:10:52Z

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresServiceCollectionExtensions.cs

+    /// <param name="options">Optional options to further configure the <see cref="IVectorStoreRecordCollection{TKey, TRecord}"/>.</param>
+    /// <param name="serviceId">An optional service id to use as the service key.</param>
+    /// <returns>Service collection.</returns>
+    public static IServiceCollection AddPostgresVectorStoreRecordCollection<TKey, TRecord>(


Same comment as above, not sure we need this one.

westey-m · 2024-10-30T18:12:19Z

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresVectorStore.cs

-        {
-            yield return collection;
-        }
+        return this._postgresClient.GetTablesAsync(cancellationToken);


Don't forget to also wrap any errors here in VectorStoreOperationException

westey-m · 2024-10-30T18:14:50Z

dotnet/src/Connectors/Connectors.Postgres.UnitTests/PostgresGenericDataModelMapperTests.cs

+    public void MapFromDataToStorageModelWithNumericKeyReturnsValidStorageModel()
+    {
+        // Arrange
+        var definition = GetRecordDefinition<string>();


Looks like this is also using a string key instead of numeric as the test name suggests.

westey-m · 2024-10-30T18:16:06Z

dotnet/src/Connectors/Connectors.Postgres.UnitTests/PostgresGenericDataModelMapperTests.cs

+        var result = mapper.MapFromStorageToDataModel(storageModel, new() { IncludeVectors = includeVectors });
+
+        // Assert
+        Assert.Equal("key", result.Key);


Same as above, looks like this should still be changed to be numeric?

westey-m · 2024-10-30T18:19:03Z

dotnet/src/Connectors/Connectors.Postgres.UnitTests/PostgresServiceCollectionExtensionsTests.cs

+        this._serviceCollection.AddSingleton<NpgsqlDataSource>(dataSource);
+
+        // Act
+        this._serviceCollection.AddPostgresVectorStoreRecordCollection<string, TestRecord>("testcollection");


Numeric Key?

westey-m · 2024-10-30T18:20:52Z

...src/Connectors/Connectors.Postgres.UnitTests/PostgresVectorStoreCollectionSqlBuilderTests.cs

+                new VectorStoreRecordDataProperty("code", typeof(int)),
+                new VectorStoreRecordDataProperty("rating", typeof(float?)),
+                new VectorStoreRecordDataProperty("description", typeof(string)),
+                new VectorStoreRecordDataProperty("parking_is_included", typeof(bool)),


Might be worth having a record definition where some of the properties have storage names that are different to the main property name, to make sure these are used correctly in the command. Same with the other tests here that have to build names into the command.

westey-m · 2024-10-30T18:26:54Z

dotnet/src/Connectors/Connectors.Postgres.UnitTests/PostgresVectorStoreRecordMapperTests.cs

+    {
+        return new VectorStoreRecordDefinition
+        {
+            Properties = new List<VectorStoreRecordProperty>


Might be good to add one or two enumerables to the property list as well to verify that they word as expected.

westey-m · 2024-10-30T18:32:16Z

dotnet/src/IntegrationTests/Connectors/Memory/Postgres/PostgresVectorStoreFixture.cs

+        {
+            Properties = new List<VectorStoreRecordProperty>
+            {
+                new VectorStoreRecordKeyProperty("HotelId", typeof(ulong)),


Are we using this anywhere? From the constants list I don't believe we support ulong, so I would expect using this to fail.

westey-m · 2024-10-30T18:40:26Z

.../src/IntegrationTests/Connectors/Memory/Postgres/PostgresVectorStoreRecordCollectionTests.cs

+    [InlineData(typeof(long), 7L)]
+    [InlineData(typeof(string), "key1")]
+    [InlineData(typeof(Guid), null)]
+    public async Task ItCanGetAndDeleteRecordAsync(Type idType, object? key)


It is also possible to make test methods generic, e.g. public async Task ItCanGetAndDeleteRecordAsync<TKey>(Type idType, TKey? key), which means you can use TKey in the method and you don't have to use dynamic.

westey-m · 2024-10-30T18:46:42Z

.../src/IntegrationTests/Connectors/Memory/Postgres/PostgresVectorStoreRecordCollectionTests.cs

+namespace SemanticKernel.IntegrationTests.Connectors.Memory.Postgres;
+
+[Collection("PostgresVectorStoreCollection")]
+public sealed class PostgresVectorStoreRecordCollectionTests(PostgresVectorStoreFixture fixture)


It would be nice to have a test that reads some data that wasn't upserted by the Collection. E.g. if you create a table and upsert some records using SQL, and then read them using the Collection. The reason being that we want to ensure that the data in the DB is actually what we want it to be. E.g. if we do two complimentary things wrong, on upsert and read, so it looks fine when using the Collection, but the data isn't stored in the way we need it to in the DB, a test where we didn't write the data via the Collection should catch this.

Add PostgresVectorStore Memory connector.

8778d5f

Work in progress, some methods are not implemented yet.

markwallace-microsoft added .NET Issue or Pull requests regarding .NET code kernel Issues or pull requests impacting the core kernel memory labels Oct 18, 2024

lossyrob had a problem deploying to integration October 18, 2024 14:10 — with GitHub Actions Failure

Rob Emanuele added 3 commits October 18, 2024 18:00

Add UpsertBatch, GetBatch, and DeleteBatch

ddad99a

Remove unused CreateMapping

5447815

Merge remote-tracking branch 'upstream/main' into feature/postgres-ve…

7533f8c

…ctor-store-dotnet

lossyrob had a problem deploying to integration October 18, 2024 22:04 — with GitHub Actions Failure

Rob Emanuele added 2 commits October 21, 2024 10:27

Merge remote-tracking branch 'upstream/main' into feature/postgres-ve…

9a4f836

…ctor-store-dotnet

Add vector search to PostgresVectorStore

68a000e

lossyrob had a problem deploying to integration October 22, 2024 20:30 — with GitHub Actions Failure

JadynWong reviewed Oct 23, 2024

View reviewed changes

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresConstants.cs Outdated Show resolved Hide resolved

Rob Emanuele added 2 commits October 23, 2024 11:24

create index on collection creation

317f6af

Support Guid, test distance functions

f4f5ba2

lossyrob temporarily deployed to integration October 23, 2024 15:27 — with GitHub Actions Inactive

Rob Emanuele added 2 commits October 23, 2024 11:44

Format tests

2acf118

Merge remote-tracking branch 'upstream/main' into feature/postgres-ve…

5db2c59

…ctor-store-dotnet

lossyrob temporarily deployed to integration October 23, 2024 15:46 — with GitHub Actions Inactive

JadynWong reviewed Oct 23, 2024

View reviewed changes

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresVectorStoreRecordCollection.cs Outdated Show resolved Hide resolved

dotnet/src/Connectors/Connectors.Memory.Postgres/PostgresVectorStoreRecordCollection.cs Outdated Show resolved Hide resolved

markwallace-microsoft assigned westey-m Oct 24, 2024

Rob Emanuele added 5 commits October 24, 2024 11:11

Add service and kernel extensions

f4b4dc5

Default to Euclidean distance if no distance function is specified

5c58400

Add Postgres sample to concepts

8ea21cd

Add docs for setting configuration in samples\Concepts

4dcd222

Merge remote-tracking branch 'upstream/main' into feature/postgres-ve…

74b3764

…ctor-store-dotnet

markwallace-microsoft added the documentation label Oct 24, 2024

lossyrob temporarily deployed to integration October 24, 2024 15:16 — with GitHub Actions Inactive

Rob Emanuele added 2 commits October 24, 2024 11:27

Enforce dimension size in index creation

5c3e63f

Create index for CreateTableIfNotExistsAsyc

6d9f1fd

lossyrob had a problem deploying to integration October 28, 2024 18:01 — with GitHub Actions Failure

Rob Emanuele added 2 commits October 28, 2024 14:55

Improve error handling

81e1805

Avoid CA1859 in test class

a587260

e.g. Change type of property 'CollectionInts' from 'System.Collections.Generic.ICollection<int>?' to 'System.Collections.Generic.HashSet<int>?' for improved performance

lossyrob had a problem deploying to integration October 28, 2024 19:17 — with GitHub Actions Failure

Account for ngpsql missing func in .net std 2.0

e8fe800

lossyrob had a problem deploying to integration October 28, 2024 19:50 — with GitHub Actions Failure

Rob Emanuele added 2 commits October 28, 2024 16:21

Fix servicecollection tests

96c088e

Logic for dimension max moved and tested elsewhere

0fc76f6

lossyrob temporarily deployed to integration October 28, 2024 20:23 — with GitHub Actions Inactive

Remove unused using statement

266310b

lossyrob temporarily deployed to integration October 28, 2024 20:34 — with GitHub Actions Inactive

Remove logger from PostgresVectorStoreRecordCollection

08f110c

lossyrob temporarily deployed to integration October 29, 2024 14:17 — with GitHub Actions Inactive

Merge branch 'main' into feature/postgres-vector-store-dotnet

26516c5

lossyrob temporarily deployed to integration October 30, 2024 16:34 — with GitHub Actions Inactive

Use Flat instead of None index kind

5b44a80

lossyrob temporarily deployed to integration October 30, 2024 17:39 — with GitHub Actions Inactive

westey-m reviewed Oct 30, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

.Net: Add PostgresVectorStore Memory connector. #9324

.Net: Add PostgresVectorStore Memory connector. #9324

lossyrob commented Oct 18, 2024 •

edited

Loading

lossyrob commented Oct 28, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024 •

edited

Loading

westey-m Oct 30, 2024 •

edited

Loading

westey-m Oct 30, 2024

westey-m Oct 30, 2024

westey-m Oct 30, 2024

.Net: Add PostgresVectorStore Memory connector. #9324

Are you sure you want to change the base?

.Net: Add PostgresVectorStore Memory connector. #9324

Conversation

lossyrob commented Oct 18, 2024 • edited Loading

Motivation and Context

Contribution Checklist

lossyrob commented Oct 28, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

westey-m Oct 30, 2024 • edited Loading

Choose a reason for hiding this comment

westey-m Oct 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

lossyrob commented Oct 18, 2024 •

edited

Loading

westey-m Oct 30, 2024 •

edited

Loading

westey-m Oct 30, 2024 •

edited

Loading