.Net: [MEVD] Support BinaryEmbedding with PostgreSQL (#13322)

The main motivation here was to ensure that our infrastructure and tests
properly support binary embeddings (and more in general, non-float/half
ones).

Closes #13321

Author: roji

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<TRecord, BinaryEmbedding>(record, cancellationToken, out var binaryTask))
+            {
+                generatedEmbeddings ??= new Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>(vectorPropertyCount);
+                generatedEmbeddings[vectorProperty] = [await binaryTask.ConfigureAwait(false)];
+            }
             else
             {
                 throw new InvalidOperationException(
@@ -420,10 +425,9 @@ _ when vectorProperty.EmbeddingGenerator is IEmbeddingGenerator<TInput, Embeddin
 
             // Dense Binary
             BitArray b => b,
-            // TODO: Uncomment once we sync to the latest MEAI
-            // BinaryEmbedding e => e.Vector,
-            // _ when vectorProperty.EmbeddingGenerator is IEmbeddingGenerator<TVector, BinaryEmbedding> generator
-            //     => (await generator.GenerateEmbeddingAsync(value, cancellationToken: cancellationToken).ConfigureAwait(false)).Vector,
+            BinaryEmbedding e => e.Vector,
+            _ when vectorProperty.EmbeddingGenerator is IEmbeddingGenerator<TInput, BinaryEmbedding> generator
+                => await generator.GenerateAsync(searchValue, cancellationToken: cancellationToken).ConfigureAwait(false),
 
             // Sparse
             SparseVector sv => sv,

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;

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<float>, Embedding<float>, float[], ReadOnlyMemory<Half>, Embedding<Half>, Half[], BitArray, or SparseVector";
+    internal const string SupportedVectorTypes = "ReadOnlyMemory<float>, Embedding<float>, float[], ReadOnlyMemory<Half>, Embedding<Half>, 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<Half>) ||
             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<Embedding<Half>>(embeddingGenerator, userRequestedEmbeddingType)
 #endif
-        ;
+        ?? vectorProperty.ResolveEmbeddingType<BinaryEmbedding>(embeddingGenerator, userRequestedEmbeddingType);
 }

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.")
         };

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<Half>([(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<BitArray>(new BitArray(new bool[] { true, false, true }), distanceFunction: DistanceFunction.HammingDistance, embeddingGenerator: null);
+        => this.Test<BitArray>(
+            new BitArray([true, false, true]),
+            new BinaryEmbeddingGenerator(new BitArray([true, false, true])),
+            distanceFunction: DistanceFunction.HammingDistance);
+
+    [ConditionalFact]
+    public virtual Task BinaryEmbedding()
+        => this.Test<BinaryEmbedding>(
+            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<SparseVector>(new SparseVector(new ReadOnlyMemory<float>([1, 2, 3])), embeddingGenerator: null);
 
+    // TODO: Figure out the embedding generation story for sparsevec - need an Embedding wrapper
+
     public new class Fixture : EmbeddingTypeTests<int>.Fixture
     {
         public override TestStore TestStore => PostgresTestStore.Instance;

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<TVector>(
 
     protected sealed class ReadOnlyMemoryEmbeddingGenerator<T>(T[] data) : IEmbeddingGenerator<string, Embedding<T>>
     {
-        public Task<GeneratedEmbeddings<Embedding<T>>> GenerateAsync(IEnumerable<string> values, EmbeddingGenerationOptions? options = null, CancellationToken cancellationToken = default)
+        public Task<GeneratedEmbeddings<Embedding<T>>> GenerateAsync(
+            IEnumerable<string> values,
+            EmbeddingGenerationOptions? options = null,
+            CancellationToken cancellationToken = default)
             => Task.FromResult(new GeneratedEmbeddings<Embedding<T>>([new(data)]));
 
         public object? GetService(Type serviceType, object? serviceKey = null) => null;
         public void Dispose() { }
     }
 
+    protected sealed class BinaryEmbeddingGenerator(BitArray data) : IEmbeddingGenerator<string, BinaryEmbedding>
+    {
+        public Task<GeneratedEmbeddings<BinaryEmbedding>> GenerateAsync(
+            IEnumerable<string> values,
+            EmbeddingGenerationOptions? options = null,
+            CancellationToken cancellationToken = default)
+            => Task.FromResult(new GeneratedEmbeddings<BinaryEmbedding>([new(data)]));
+
+        public object? GetService(Type serviceType, object? serviceKey = null) => null;
+        public void Dispose() { }
+    }
+
     public class RecordWithString
     {
         public TKey Key { get; set; }