Implement latest SQL Server approximate vector search feature

roji · roji · commit 7244a28a6567 · 2026-04-10T11:32:59.000+02:00
Part of #36384
diff --git a/.gitignore b/.gitignore
@@ -342,3 +342,6 @@ ASALocalRun/
 # Local History for Visual Studio
 .localhistory/
 full.targets.txt
+
+# Language Server cache
+*.lscache
diff --git a/Directory.Packages.props b/Directory.Packages.props
@@ -34,6 +34,7 @@
     <PackageVersion Include="Microsoft.Azure.Cosmos" Version="3.58.0" />
     <!-- SQL Server dependencies -->
     <PackageVersion Include="Microsoft.Data.SqlClient" Version="7.0.0" />
+    <PackageVersion Include="Microsoft.Data.SqlClient.Extensions.Azure" Version="1.0.0" />
     <PackageVersion Include="Microsoft.SqlServer.Types" Version="170.1000.7" />
     <!-- external dependencies -->
     <PackageVersion Include="Castle.Core" Version="5.2.1" />
diff --git a/src/EFCore.SqlServer/Extensions/SqlServerQueryableExtensions.cs b/src/EFCore.SqlServer/Extensions/SqlServerQueryableExtensions.cs
@@ -28,7 +28,16 @@ public static class SqlServerQueryableExtensions
     ///     An ANN (Approximate Nearest Neighbor) index is used only if a matching ANN index, with the same metric and on the same column,
     ///     is found. If there are no compatible ANN indexes, a warning is raised and the KNN (k-Nearest Neighbor) algorithm is used.
     /// </param>
-    /// <param name="topN">The maximum number of similar vectors that must be returned. It must be a positive integer.</param>
+    /// <remarks>
+    ///     Compose the returned query with <c>OrderBy(r => r.Distance)</c> and <c>Take(...)</c> to limit the results as required
+    ///     for approximate vector search. For example:
+    ///     <code>
+    ///     var results = context.Set&lt;Blog&gt;()
+    ///         .VectorSearch(b => b.Embedding, embedding, "cosine")
+    ///         .OrderBy(r => r.Distance)
+    ///         .Take(10);
+    ///     </code>
+    /// </remarks>
     /// <seealso href="https://learn.microsoft.com/sql/t-sql/functions/vector-search-transact-sql">
     ///     SQL Server documentation for <c>VECTOR_SEARCH()</c>.
     /// </seealso>
@@ -38,8 +47,7 @@ public static IQueryable<VectorSearchResult<T>> VectorSearch<T, TVector>(
         this DbSet<T> source,
         Expression<Func<T, TVector>> vectorPropertySelector,
         TVector similarTo,
-        [NotParameterized] string metric,
-        int topN)
+        [NotParameterized] string metric)
         where T : class
         where TVector : unmanaged
     {
@@ -50,12 +58,11 @@ public static IQueryable<VectorSearchResult<T>> VectorSearch<T, TVector>(
             ? queryableSource.Provider.CreateQuery<VectorSearchResult<T>>(
                 Expression.Call(
                     // Note that the method used is the one below, accepting IQueryable<T>, not DbSet<T>
-                    method: new Func<IQueryable<T>, Expression<Func<T, TVector>>, TVector, string, int, IQueryable<VectorSearchResult<T>>>(VectorSearch).Method,
+                    method: new Func<IQueryable<T>, Expression<Func<T, TVector>>, TVector, string, IQueryable<VectorSearchResult<T>>>(VectorSearch).Method,
                     root,
                     Expression.Quote(vectorPropertySelector),
                     Expression.Constant(similarTo),
-                    Expression.Constant(metric),
-                    Expression.Constant(topN)))
+                    Expression.Constant(metric)))
             : throw new InvalidOperationException(CoreStrings.FunctionOnNonEfLinqProvider(nameof(VectorSearch)));
     }
 
@@ -67,8 +74,7 @@ private static IQueryable<VectorSearchResult<T>> VectorSearch<T, TVector>(
         this IQueryable<T> source,
         Expression<Func<T, TVector>> vectorPropertySelector,
         TVector similarTo,
-        [NotParameterized] string metric,
-        int topN)
+        [NotParameterized] string metric)
         where T : class
         where TVector : unmanaged
         => throw new UnreachableException();
diff --git a/src/EFCore.SqlServer/Extensions/VectorSearchResult.cs b/src/EFCore.SqlServer/Extensions/VectorSearchResult.cs
@@ -7,7 +7,7 @@ namespace Microsoft.EntityFrameworkCore;
 
 /// <summary>
 ///     Represents the results from a call to
-///     <see cref="SqlServerQueryableExtensions.VectorSearch{T, TVector}(DbSet{T}, Expression{Func{T, TVector}}, TVector, string, int)" />.
+///     <see cref="SqlServerQueryableExtensions.VectorSearch{T, TVector}(DbSet{T}, Expression{Func{T, TVector}}, TVector, string)" />.
 /// </summary>
 [Experimental(EFDiagnostics.SqlServerVectorSearch)]
 public readonly struct VectorSearchResult<T>(T value, double distance)
diff --git a/src/EFCore.SqlServer/Query/Internal/SqlServerQuerySqlGenerator.cs b/src/EFCore.SqlServer/Query/Internal/SqlServerQuerySqlGenerator.cs
@@ -155,16 +155,14 @@ protected override Expression VisitTableValuedFunction(TableValuedFunctionExpres
                     TableExpression table,
                     ColumnExpression column,
                     SqlExpression similarTo,
-                    SqlConstantExpression { Value: string } metric,
-                    SqlExpression topN
+                    SqlConstantExpression { Value: string } metric
                 ]
             }:
                 // VECTOR_SEARCH(
                 //     TABLE = [Articles] AS t,
                 //     COLUMN = [Vector],
                 //     SIMILAR_TO = @qv,
-                //     METRIC = 'Cosine',
-                //     TOP_N = 3
+                //     METRIC = 'Cosine'
                 // )
                 Sql.AppendLine("VECTOR_SEARCH(");
 
@@ -185,10 +183,6 @@ SqlExpression topN
 
                     Sql.Append("METRIC = ");
                     Visit(metric);
-                    Sql.AppendLine(",");
-
-                    Sql.Append("TOP_N = ");
-                    Visit(topN);
                     Sql.AppendLine();
                 }
 
@@ -558,6 +552,13 @@ protected override void GenerateTop(SelectExpression selectExpression)
             Visit(selectExpression.Limit);
 
             Sql.Append(") ");
+
+            // When performing approximate vector search with VECTOR_SEARCH(), SQL Server requires adding
+            // WITH APPROXIMATE: https://learn.microsoft.com/sql/t-sql/functions/vector-search-transact-sql
+            if (selectExpression.Tables.Any(t => t.UnwrapJoin() is TableValuedFunctionExpression { Name: "VECTOR_SEARCH" })
+            {
+                Sql.Append("WITH APPROXIMATE ");
+            }
         }
 
         _withinTable = parentWithinTable;
diff --git a/src/EFCore.SqlServer/Query/Internal/SqlServerQueryableMethodTranslatingExpressionVisitor.cs b/src/EFCore.SqlServer/Query/Internal/SqlServerQueryableMethodTranslatingExpressionVisitor.cs
@@ -91,8 +91,7 @@ when methodCallExpression.Arguments is
                         _, // source, translated above
                         UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression vectorPropertySelector },
                         var similarTo,
-                        var metric,
-                        var topN
+                        var metric
                     ]
                     && source is
                     {
@@ -113,8 +112,7 @@ var topN
                     }
 
                     if (TranslateExpression(similarTo) is not { } translatedSimilarTo
-                        || TranslateExpression(metric, applyDefaultTypeMapping: false) is not { } translatedMetric
-                        || TranslateExpression(topN) is not { } translatedTopN)
+                        || TranslateExpression(metric, applyDefaultTypeMapping: false) is not { } translatedMetric)
                     {
                         return QueryCompilationContext.NotTranslatedExpression;
                     }
@@ -135,8 +133,7 @@ var topN
                             // as required by SQL Server)
                             vectorColumn,
                             translatedSimilarTo,
-                            translatedMetric,
-                            translatedTopN
+                            translatedMetric
                         ]);
 
                     // We have the VECTOR_SEARCH() function call. Modify the SelectExpression and shaper to use it and project
diff --git a/test/EFCore.SqlServer.FunctionalTests/EFCore.SqlServer.FunctionalTests.csproj b/test/EFCore.SqlServer.FunctionalTests/EFCore.SqlServer.FunctionalTests.csproj
@@ -77,5 +77,6 @@
   <ItemGroup>
     <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" />
     <PackageReference Include="Microsoft.Extensions.Configuration.Json" />
+    <PackageReference Include="Microsoft.Data.SqlClient.Extensions.Azure" />
   </ItemGroup>
 </Project>
diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/Translations/VectorTranslationsSqlServerTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/Translations/VectorTranslationsSqlServerTest.cs
@@ -65,7 +65,9 @@ ORDER BY VECTOR_DISTANCE('cosine', [v].[Vector], CAST('[1,2,100]' AS VECTOR(3)))
 """);
     }
 
+    // The latest vector index version (required for VECTOR_SEARCH) is only available on Azure SQL (#36384).
     [ConditionalFact]
+    [SqlServerCondition(SqlServerCondition.IsAzureSql)]
     [Experimental("EF9105")]
     public async Task VectorSearch_project_entity_and_distance()
     {
@@ -74,28 +76,32 @@ public async Task VectorSearch_project_entity_and_distance()
         var vector = new SqlVector<float>(new float[] { 1, 2, 100 });
 
         var results = await ctx.VectorEntities
-            .VectorSearch(e => e.Vector, similarTo: vector, "cosine", topN: 1)
+            .VectorSearch(e => e.Vector, similarTo: vector, "cosine")
+            .OrderBy(e => e.Distance)
+            .Take(1)
             .ToListAsync();
 
         Assert.Equal(2, results.Single().Value.Id);
 
         AssertSql(
             """
-@p='Microsoft.Data.SqlTypes.SqlVector`1[System.Single]' (Size = 20) (DbType = Binary)
 @p1='1'
+@p='Microsoft.Data.SqlTypes.SqlVector`1[System.Single]' (Size = 20) (DbType = Binary)
 
-SELECT [v].[Id], [v0].[Distance]
+SELECT TOP(@p1) WITH APPROXIMATE [v].[Id], [v0].[Distance]
 FROM VECTOR_SEARCH(
     TABLE = [VectorEntities] AS [v],
     COLUMN = [Vector],
     SIMILAR_TO = @p,
-    METRIC = 'cosine',
-    TOP_N = @p1
+    METRIC = 'cosine'
 ) AS [v0]
+ORDER BY [v0].[Distance]
 """);
     }
 
+    // The latest vector index version (required for VECTOR_SEARCH) is only available on Azure SQL (#36384).
     [ConditionalFact]
+    [SqlServerCondition(SqlServerCondition.IsAzureSql)]
     [Experimental("EF9105")]
     public async Task VectorSearch_project_entity_only_with_distance_filter_and_ordering()
     {
@@ -104,10 +110,11 @@ public async Task VectorSearch_project_entity_only_with_distance_filter_and_orde
         var vector = new SqlVector<float>(new float[] { 1, 2, 100 });
 
         var results = await ctx.VectorEntities
-            .VectorSearch(e => e.Vector, similarTo: vector, "cosine", topN: 3)
+            .VectorSearch(e => e.Vector, similarTo: vector, "cosine")
             .Where(e => e.Distance < 0.01)
             .OrderBy(e => e.Distance)
             .Select(e => e.Value)
+            .Take(3)
             .ToListAsync();
 
         Assert.Collection(
@@ -117,16 +124,15 @@ public async Task VectorSearch_project_entity_only_with_distance_filter_and_orde
 
         AssertSql(
             """
-@p='Microsoft.Data.SqlTypes.SqlVector`1[System.Single]' (Size = 20) (DbType = Binary)
 @p1='3'
+@p='Microsoft.Data.SqlTypes.SqlVector`1[System.Single]' (Size = 20) (DbType = Binary)
 
-SELECT [v].[Id]
+SELECT TOP(@p1) WITH APPROXIMATE [v].[Id]
 FROM VECTOR_SEARCH(
     TABLE = [VectorEntities] AS [v],
     COLUMN = [Vector],
     SIMILAR_TO = @p,
-    METRIC = 'cosine',
-    TOP_N = @p1
+    METRIC = 'cosine'
 ) AS [v0]
 WHERE [v0].[Distance] < 0.01E0
 ORDER BY [v0].[Distance]
@@ -167,23 +173,27 @@ public class VectorQueryContext(DbContextOptions options) : PoolableDbContext(op
 
         public static async Task SeedAsync(VectorQueryContext context)
         {
-            var vectorEntities = new VectorEntity[]
-            {
-                new() { Id = 1, Vector = new SqlVector<float>(new float[] { 1, 2, 3 }) },
-                new() { Id = 2, Vector = new SqlVector<float>(new float[] { 1, 2, 100 }) },
-                new() { Id = 3, Vector = new SqlVector<float>(new float[] { 1, 2, 1000 }) }
-            };
+            // SQL Server vector indexes require at least 100 rows.
+            var vectorEntities = Enumerable.Range(1, 100).Select(
+                i => new VectorEntity
+                {
+                    Id = i,
+                    Vector = new SqlVector<float>(new float[] { i * 0.01f, i * 0.02f, i * 0.03f })
+                }).ToList();
+
+            // Override specific rows we use in test assertions
+            vectorEntities[0] = new VectorEntity { Id = 1, Vector = new SqlVector<float>(new float[] { 1, 2, 3 }) };
+            vectorEntities[1] = new VectorEntity { Id = 2, Vector = new SqlVector<float>(new float[] { 1, 2, 100 }) };
+            vectorEntities[2] = new VectorEntity { Id = 3, Vector = new SqlVector<float>(new float[] { 1, 2, 1000 }) };
 
             context.VectorEntities.AddRange(vectorEntities);
             await context.SaveChangesAsync();
 
-            // TODO (#36384): Remove this once it's out of preview
             await context.Database.ExecuteSqlAsync($"ALTER DATABASE SCOPED CONFIGURATION SET PREVIEW_FEATURES = ON");
 
             await context.Database.ExecuteSqlAsync($"""
 CREATE VECTOR INDEX vec_idx ON VectorEntities(Vector)
-WITH (METRIC = 'Cosine', TYPE = 'DiskANN')
-ON [PRIMARY];
+WITH (METRIC = 'Cosine', TYPE = 'DiskANN');
 """);
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -91,8 +91,7 @@ when methodCallExpression.Arguments is`
`91`	`91`	`_, // source, translated above`
`92`	`92`	`UnaryExpression { NodeType: ExpressionType.Quote, Operand: LambdaExpression vectorPropertySelector },`
`93`	`93`	`var similarTo,`
`94`		`- var metric,`
`95`		`- var topN`
	`94`	`+ var metric`
`96`	`95`	`]`
`97`	`96`	`&& source is`
`98`	`97`	`{`
`@@ -113,8 +112,7 @@ var topN`
`113`	`112`	`}`
`114`	`113`
`115`	`114`	`if (TranslateExpression(similarTo) is not { } translatedSimilarTo`
`116`		`- \|\| TranslateExpression(metric, applyDefaultTypeMapping: false) is not { } translatedMetric`
`117`		`- \|\| TranslateExpression(topN) is not { } translatedTopN)`
	`115`	`+ \|\| TranslateExpression(metric, applyDefaultTypeMapping: false) is not { } translatedMetric)`
`118`	`116`	`{`
`119`	`117`	`return QueryCompilationContext.NotTranslatedExpression;`
`120`	`118`	`}`
`@@ -135,8 +133,7 @@ var topN`
`135`	`133`	`// as required by SQL Server)`
`136`	`134`	`vectorColumn,`
`137`	`135`	`translatedSimilarTo,`
`138`		`- translatedMetric,`
`139`		`- translatedTopN`
	`136`	`+ translatedMetric`
`140`	`137`	`]);`
`141`	`138`
`142`	`139`	`// We have the VECTOR_SEARCH() function call. Modify the SelectExpression and shaper to use it and project`