From ae52e9fa50624d14fe88b6a80f91074761899c2c Mon Sep 17 00:00:00 2001 From: Shay Rojansky Date: Mon, 3 Nov 2025 09:25:30 +0100 Subject: [PATCH] [MEVD] Support BinaryEmbedding with PostgreSQL Closes #13321 --- .../VectorData/PgVector/PostgresCollection.cs | 12 ++++++++---- .../src/VectorData/PgVector/PostgresMapper.cs | 4 ++++ .../PgVector/PostgresModelBuilder.cs | 5 +++-- .../PgVector/PostgresPropertyMapping.cs | 2 ++ .../TypeTests/PostgresEmbeddingTypeTests.cs | 19 ++++++++++++++----- .../TypeTests/EmbeddingTypeTests.cs | 18 +++++++++++++++++- 6 files changed, 48 insertions(+), 12 deletions(-) diff --git a/dotnet/src/VectorData/PgVector/PostgresCollection.cs b/dotnet/src/VectorData/PgVector/PostgresCollection.cs index 7cfcfc74ac5b..da650d150828 100644 --- a/dotnet/src/VectorData/PgVector/PostgresCollection.cs +++ b/dotnet/src/VectorData/PgVector/PostgresCollection.cs @@ -199,6 +199,11 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell generatedEmbeddings[vectorProperty] = [await halfTask.ConfigureAwait(false)]; } #endif + else if (vectorProperty.TryGenerateEmbedding(record, cancellationToken, out var binaryTask)) + { + generatedEmbeddings ??= new Dictionary>(vectorPropertyCount); + generatedEmbeddings[vectorProperty] = [await binaryTask.ConfigureAwait(false)]; + } else { throw new InvalidOperationException( @@ -420,10 +425,9 @@ _ when vectorProperty.EmbeddingGenerator is IEmbeddingGenerator b, - // TODO: Uncomment once we sync to the latest MEAI - // BinaryEmbedding e => e.Vector, - // _ when vectorProperty.EmbeddingGenerator is IEmbeddingGenerator generator - // => (await generator.GenerateEmbeddingAsync(value, cancellationToken: cancellationToken).ConfigureAwait(false)).Vector, + BinaryEmbedding e => e.Vector, + _ when vectorProperty.EmbeddingGenerator is IEmbeddingGenerator generator + => await generator.GenerateAsync(searchValue, cancellationToken: cancellationToken).ConfigureAwait(false), // Sparse SparseVector sv => sv, diff --git a/dotnet/src/VectorData/PgVector/PostgresMapper.cs b/dotnet/src/VectorData/PgVector/PostgresMapper.cs index 70aafb7e2ebb..f824f4f5894b 100644 --- a/dotnet/src/VectorData/PgVector/PostgresMapper.cs +++ b/dotnet/src/VectorData/PgVector/PostgresMapper.cs @@ -77,6 +77,10 @@ public TRecord MapFromStorageToDataModel(NpgsqlDataReader reader, bool includeVe } #endif + case BitArray bitArray when vectorProperty.Type == typeof(BinaryEmbedding): + vectorProperty.SetValueAsObject(record, new BinaryEmbedding(bitArray)); + continue; + case BitArray bitArray: vectorProperty.SetValueAsObject(record, bitArray); continue; diff --git a/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs b/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs index 4cadfd765bdb..5e3b18509e5f 100644 --- a/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs +++ b/dotnet/src/VectorData/PgVector/PostgresModelBuilder.cs @@ -12,7 +12,7 @@ namespace Microsoft.SemanticKernel.Connectors.PgVector; internal class PostgresModelBuilder() : CollectionModelBuilder(PostgresModelBuilder.ModelBuildingOptions) { - internal const string SupportedVectorTypes = "ReadOnlyMemory, Embedding, float[], ReadOnlyMemory, Embedding, Half[], BitArray, or SparseVector"; + internal const string SupportedVectorTypes = "ReadOnlyMemory, Embedding, float[], ReadOnlyMemory, Embedding, Half[], BinaryEmbedding, BitArray, or SparseVector"; public static readonly CollectionModelBuildingOptions ModelBuildingOptions = new() { @@ -80,6 +80,7 @@ internal static bool IsVectorPropertyTypeValidCore(Type type, [NotNullWhen(false type == typeof(Embedding) || type == typeof(Half[]) || #endif + type == typeof(BinaryEmbedding) || type == typeof(BitArray) || type == typeof(SparseVector); } @@ -93,5 +94,5 @@ internal static bool IsVectorPropertyTypeValidCore(Type type, [NotNullWhen(false #if NET8_0_OR_GREATER ?? vectorProperty.ResolveEmbeddingType>(embeddingGenerator, userRequestedEmbeddingType) #endif - ; + ?? vectorProperty.ResolveEmbeddingType(embeddingGenerator, userRequestedEmbeddingType); } diff --git a/dotnet/src/VectorData/PgVector/PostgresPropertyMapping.cs b/dotnet/src/VectorData/PgVector/PostgresPropertyMapping.cs index 1160cf543986..5130cbdfcce7 100644 --- a/dotnet/src/VectorData/PgVector/PostgresPropertyMapping.cs +++ b/dotnet/src/VectorData/PgVector/PostgresPropertyMapping.cs @@ -30,6 +30,7 @@ internal static class PostgresPropertyMapping #endif BitArray bitArray => bitArray, + BinaryEmbedding binaryEmbedding => binaryEmbedding.Vector, SparseVector sparseVector => sparseVector, null => null, @@ -136,6 +137,7 @@ public static (string PgType, bool IsNullable) GetPgVectorTypeName(VectorPropert Type t when t == typeof(SparseVector) => "SPARSEVEC", Type t when t == typeof(BitArray) => "BIT", + Type t when t == typeof(BinaryEmbedding) => "BIT", _ => throw new NotSupportedException($"Type {vectorProperty.EmbeddingType.Name} is not supported by this store.") }; diff --git a/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresEmbeddingTypeTests.cs b/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresEmbeddingTypeTests.cs index f5a03484e4fc..42c35efc2515 100644 --- a/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresEmbeddingTypeTests.cs +++ b/dotnet/test/VectorData/PgVector.ConformanceTests/TypeTests/PostgresEmbeddingTypeTests.cs @@ -1,9 +1,7 @@ // Copyright (c) Microsoft. All rights reserved. using System.Collections; -#if NET8_0_OR_GREATER using Microsoft.Extensions.AI; -#endif using Microsoft.Extensions.VectorData; using Pgvector; using PgVector.ConformanceTests.Support; @@ -41,16 +39,27 @@ public virtual Task Array_of_Half() new ReadOnlyMemoryEmbeddingGenerator([(byte)1, (byte)2, (byte)3])); #endif - // TODO: Figure out the embedding generation story for binaryvec/sparsevec - need an Embedding wrapper - [ConditionalFact] public virtual Task BitArray() - => this.Test(new BitArray(new bool[] { true, false, true }), distanceFunction: DistanceFunction.HammingDistance, embeddingGenerator: null); + => this.Test( + new BitArray([true, false, true]), + new BinaryEmbeddingGenerator(new BitArray([true, false, true])), + distanceFunction: DistanceFunction.HammingDistance); + + [ConditionalFact] + public virtual Task BinaryEmbedding() + => this.Test( + new BinaryEmbedding(new([true, false, true])), + new BinaryEmbeddingGenerator(new BitArray([true, false, true])), + distanceFunction: DistanceFunction.HammingDistance, + vectorEqualityAsserter: (e, a) => Assert.Equal(e.Vector, a.Vector)); [ConditionalFact] public virtual Task SparseVector() => this.Test(new SparseVector(new ReadOnlyMemory([1, 2, 3])), embeddingGenerator: null); + // TODO: Figure out the embedding generation story for sparsevec - need an Embedding wrapper + public new class Fixture : EmbeddingTypeTests.Fixture { public override TestStore TestStore => PostgresTestStore.Instance; diff --git a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs index f5521e831caf..2cfd5baafcdc 100644 --- a/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs +++ b/dotnet/test/VectorData/VectorData.ConformanceTests/TypeTests/EmbeddingTypeTests.cs @@ -1,5 +1,6 @@ // Copyright (c) Microsoft. All rights reserved. +using System.Collections; using Microsoft.Extensions.AI; using Microsoft.Extensions.VectorData; using VectorData.ConformanceTests.Support; @@ -178,13 +179,28 @@ protected virtual async Task Test( protected sealed class ReadOnlyMemoryEmbeddingGenerator(T[] data) : IEmbeddingGenerator> { - public Task>> GenerateAsync(IEnumerable values, EmbeddingGenerationOptions? options = null, CancellationToken cancellationToken = default) + public Task>> GenerateAsync( + IEnumerable values, + EmbeddingGenerationOptions? options = null, + CancellationToken cancellationToken = default) => Task.FromResult(new GeneratedEmbeddings>([new(data)])); public object? GetService(Type serviceType, object? serviceKey = null) => null; public void Dispose() { } } + protected sealed class BinaryEmbeddingGenerator(BitArray data) : IEmbeddingGenerator + { + public Task> GenerateAsync( + IEnumerable values, + EmbeddingGenerationOptions? options = null, + CancellationToken cancellationToken = default) + => Task.FromResult(new GeneratedEmbeddings([new(data)])); + + public object? GetService(Type serviceType, object? serviceKey = null) => null; + public void Dispose() { } + } + public class RecordWithString { public TKey Key { get; set; }