AzureDocumentDB: fix CosineDistance scores and score threshold, support CosineSimilarity

MEVD/src/AzureDocumentDB/AzureDocumentDB.csproj

@@ -1,7 +1,7 @@
 <Project Sdk="Microsoft.NET.Sdk">
 
   <PropertyGroup>
-    <Version>1.0.2</Version>
+    <Version>1.1.0</Version>
     <AssemblyName>CommunityToolkit.VectorData.AzureDocumentDB</AssemblyName>
     <RootNamespace>$(AssemblyName)</RootNamespace>
     <TargetFrameworks>net10.0;net8.0;netstandard2.1;net472</TargetFrameworks>

MEVD/src/AzureDocumentDB/DocumentDBCollection.cs

@@ -415,7 +415,7 @@ _ when vectorProperty.EmbeddingGenerationDispatcher is not null
             () => _mongoCollection.AggregateAsync<BsonDocument>(pipeline, cancellationToken: cancellationToken)).ConfigureAwait(false);
         using var errorHandlingAsyncCursor = new ErrorHandlingAsyncCursor<BsonDocument>(cursor, _collectionMetadata, OperationName);
 
-        await foreach (var result in EnumerateAndMapSearchResultsAsync(errorHandlingAsyncCursor, options, cancellationToken).ConfigureAwait(false))
+        await foreach (var result in EnumerateAndMapSearchResultsAsync(errorHandlingAsyncCursor, options, vectorProperty.DistanceFunction, cancellationToken).ConfigureAwait(false))
         {
             yield return result;
         }
@@ -562,6 +562,7 @@ private async Task<IAsyncCursor<BsonDocument>> FindAsync(
     private async IAsyncEnumerable<VectorSearchResult<TRecord>> EnumerateAndMapSearchResultsAsync(
         ErrorHandlingAsyncCursor<BsonDocument> cursor,
         MEVD.VectorSearchOptions<TRecord> searchOptions,
+        string? distanceFunction,
         [EnumeratorCancellation] CancellationToken cancellationToken)
     {
         var skipCounter = 0;
@@ -572,7 +573,7 @@ private async IAsyncEnumerable<VectorSearchResult<TRecord>> EnumerateAndMapSearc
             {
                 if (skipCounter >= searchOptions.Skip)
                 {
-                    var score = response[ScorePropertyName].AsDouble;
+                    var score = DocumentDBCollectionSearchMapping.GetScore(response[ScorePropertyName].AsDouble, distanceFunction);
                     var record = _mapper.MapFromStorageToDataModel(response[DocumentPropertyName].AsBsonDocument, includeVectors: searchOptions.IncludeVectors);
 
                     yield return new VectorSearchResult<TRecord>(record, score);

MEVD/src/AzureDocumentDB/DocumentDBCollectionCreateMapping.cs

@@ -120,7 +120,7 @@ private static string GetIndexKind(string? indexKind, string vectorPropertyName)
     private static string GetDistanceFunction(string? distanceFunction, string vectorPropertyName)
         => DocumentDBCollectionSearchMapping.GetVectorPropertyDistanceFunction(distanceFunction) switch
         {
-            DistanceFunction.CosineDistance => "COS",
+            DistanceFunction.CosineDistance or DistanceFunction.CosineSimilarity => "COS",
             DistanceFunction.DotProductSimilarity => "IP",
             DistanceFunction.EuclideanDistance => "L2",
             _ => throw new NotSupportedException($"Distance function '{distanceFunction}' for {nameof(VectorStoreVectorProperty)} '{vectorPropertyName}' is not supported by the Azure DocumentDB VectorStore.")

MEVD/src/AzureDocumentDB/DocumentDBCollectionSearchMapping.cs

@@ -96,18 +96,27 @@ public static BsonDocument GetProjectionQuery(string scorePropertyName, string d
         };
     }
 
+    /// <summary>Returns the score of a result from the search score that Azure DocumentDB returns.</summary>
+    /// <remarks>
+    /// For COS Azure DocumentDB returns the cosine similarity: CosineSimilarity returns it as it is, and CosineDistance returns
+    /// 1 - similarity. IP and L2 return the dot product and the Euclidean distance as they are.
+    /// </remarks>
+    public static double GetScore(double searchScore, string? distanceFunction)
+        => GetVectorPropertyDistanceFunction(distanceFunction) == DistanceFunction.CosineDistance ? 1 - searchScore : searchScore;
+
     /// <summary>Returns a $match stage to filter results by score threshold.</summary>
     /// <remarks>
-    /// For COS and IP Azure DocumentDB returns a similarity score where higher values mean more similar,
-    /// so we filter with $gte. For L2 it returns a distance where lower values mean more similar,
-    /// so we filter with $lte.
+    /// The stage filters the search score that Azure DocumentDB returns. For COS and IP higher values mean more similar,
+    /// so we filter with $gte; for CosineDistance a distance of at most the threshold is a similarity of at least 1 - threshold.
+    /// For L2 it returns a distance where lower values mean more similar, so we filter with $lte.
     /// </remarks>
     public static BsonDocument GetScoreThresholdMatchQuery(string scorePropertyName, double scoreThreshold, string? distanceFunction)
     {
-        var comparisonOperator = GetVectorPropertyDistanceFunction(distanceFunction) switch
+        var (comparisonOperator, searchScore) = GetVectorPropertyDistanceFunction(distanceFunction) switch
         {
-            DistanceFunction.CosineDistance or DistanceFunction.DotProductSimilarity => "$gte",
-            DistanceFunction.EuclideanDistance => "$lte",
+            DistanceFunction.CosineDistance => ("$gte", 1 - scoreThreshold),
+            DistanceFunction.CosineSimilarity or DistanceFunction.DotProductSimilarity => ("$gte", scoreThreshold),
+            DistanceFunction.EuclideanDistance => ("$lte", scoreThreshold),
             _ => throw new NotSupportedException($"Score threshold is not supported for distance function '{distanceFunction}'.")
         };
 
@@ -116,7 +125,7 @@ public static BsonDocument GetScoreThresholdMatchQuery(string scorePropertyName,
             {
                 "$match", new BsonDocument
                 {
-                    { scorePropertyName, new BsonDocument { { comparisonOperator, scoreThreshold } } }
+                    { scorePropertyName, new BsonDocument { { comparisonOperator, searchScore } } }
                 }
             }
         };

MEVD/test/AzureDocumentDB.ConformanceTests/DocumentDBDistanceFunctionTests.cs

@@ -11,7 +11,6 @@ namespace AzureDocumentDB.ConformanceTests;
 public class DocumentDBDistanceFunctionTests(DocumentDBDistanceFunctionTests.Fixture fixture)
     : DistanceFunctionTests<int>(fixture), IClassFixture<DocumentDBDistanceFunctionTests.Fixture>
 {
-    public override Task CosineSimilarity() => Assert.ThrowsAsync<NotSupportedException>(base.CosineSimilarity);
     public override Task EuclideanSquaredDistance() => Assert.ThrowsAsync<NotSupportedException>(base.EuclideanSquaredDistance);
     public override Task NegativeDotProductSimilarity() => Assert.ThrowsAsync<NotSupportedException>(base.NegativeDotProductSimilarity);
     public override Task HammingDistance() => Assert.ThrowsAsync<NotSupportedException>(base.HammingDistance);
@@ -20,7 +19,5 @@ public class DocumentDBDistanceFunctionTests(DocumentDBDistanceFunctionTests.Fix
     public new class Fixture() : DistanceFunctionTests<int>.Fixture
     {
         public override TestStore TestStore => DocumentDBTestStore.Instance;
-
-        public override bool AssertScores { get; } = false;
     }
 }

MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionDistanceFunctionTests.cs

@@ -0,0 +1,55 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Generic;
+using CommunityToolkit.VectorData.AzureDocumentDB;
+using Microsoft.Extensions.VectorData;
+using Microsoft.Extensions.VectorData.ProviderServices;
+using MongoDB.Bson;
+using Xunit;
+
+namespace AzureDocumentDB.UnitTests;
+
+public sealed class DocumentDBCollectionDistanceFunctionTests
+{
+    [Theory]
+    [InlineData(DistanceFunction.CosineDistance, "COS")]
+    [InlineData(DistanceFunction.CosineSimilarity, "COS")]
+    [InlineData(DistanceFunction.DotProductSimilarity, "IP")]
+    [InlineData(DistanceFunction.EuclideanDistance, "L2")]
+    [InlineData(null, "COS")]
+    public void IndexUsesTheSimilarityOfTheDistanceFunction(string? distanceFunction, string expectedSimilarity)
+    {
+        var vectorProperty = new VectorPropertyModel("vector", typeof(ReadOnlyMemory<float>)) { Dimensions = 4, DistanceFunction = distanceFunction };
+
+        var indexes = DocumentDBCollectionCreateMapping.GetVectorIndexes([vectorProperty], new HashSet<string?>(), numLists: 1, efConstruction: 64);
+
+        Assert.Equal(expectedSimilarity, indexes[0]["cosmosSearchOptions"]["similarity"].AsString);
+    }
+
+    // COS returns the cosine similarity: CosineDistance gives 1 - similarity, and its threshold keeps the distances up to it.
+    [Theory]
+    [InlineData(DistanceFunction.CosineDistance, 0.25, 0.75)]
+    [InlineData(null, 0.25, 0.75)]
+    [InlineData(DistanceFunction.CosineSimilarity, 0.75, 0.75)]
+    [InlineData(DistanceFunction.DotProductSimilarity, 0.75, 0.75)]
+    [InlineData(DistanceFunction.EuclideanDistance, 0.75, 0.75)]
+    public void ScoreFollowsTheDistanceFunction(string? distanceFunction, double expectedScore, double searchScore)
+        => Assert.Equal(expectedScore, DocumentDBCollectionSearchMapping.GetScore(searchScore, distanceFunction), precision: 10);
+
+    [Theory]
+    [InlineData(DistanceFunction.CosineDistance, 0.25, "$gte", 0.75)]
+    [InlineData(null, 0.25, "$gte", 0.75)]
+    [InlineData(DistanceFunction.CosineSimilarity, 0.75, "$gte", 0.75)]
+    [InlineData(DistanceFunction.DotProductSimilarity, 0.75, "$gte", 0.75)]
+    [InlineData(DistanceFunction.EuclideanDistance, 0.75, "$lte", 0.75)]
+    public void ScoreThresholdFollowsTheDistanceFunction(string? distanceFunction, double scoreThreshold, string expectedOperator, double expectedSearchScore)
+    {
+        var match = DocumentDBCollectionSearchMapping.GetScoreThresholdMatchQuery("score", scoreThreshold, distanceFunction);
+
+        var condition = match["$match"]["score"].AsBsonDocument;
+        Assert.Equal(expectedOperator, condition.GetElement(0).Name);
+        Assert.Equal(expectedSearchScore, condition[0].AsDouble, precision: 10);
+    }
+}

MEVD/test/AzureDocumentDB.UnitTests/DocumentDBCollectionTests.cs

@@ -561,7 +561,9 @@ public async Task SearchReturnsRecordWithScoreAsync()
         Assert.NotNull(result);
         Assert.Equal("key", result.Record.HotelId);
         Assert.Equal("Test Name", result.Record.HotelName);
-        Assert.Equal(0.99f, result.Score);
+
+        // The model uses CosineDistance: the score is 1 - the cosine similarity that COS returns.
+        Assert.Equal(1 - 0.99f, result.Score!.Value, precision: 6);
     }
 
     public static TheoryData<List<string>, string, bool> CollectionExistsData => new()