Skip to content

Commit ccfb815

Browse files
marcominervadluc
andauthored
Add support for SQL Server Vector Search (#722)
## Motivation and Context (Why the change? What's the scenario?) Azure SQL Database now provides native Vector Support, currently in EAP (https://devblogs.microsoft.com/azure-sql/announcing-eap-native-vector-support-in-azure-sql-database). This PR extends SQL Server Memory DB with a flag that allows to use this new feature. > [!IMPORTANT] > Remember that, at this time, Vector Support is available only on Azure SQL Database. On the other hand, the current SQL Server Memory DB requires a COLUMNSTORE INDEX that, on Azure, is available only on vCore databases and Standard databases in S3 and above pricing tiers (https://azure.microsoft.com/en-us/blog/columnstore-support-in-standard-tier-azure-sql-databases). --------- Co-authored-by: Devis Lucato <[email protected]>
1 parent 4c6c8c7 commit ccfb815

File tree

8 files changed

+766
-308
lines changed

8 files changed

+766
-308
lines changed

extensions/SQLServer/SQLServer/DependencyInjection.cs

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,10 @@ namespace Microsoft.KernelMemory;
1414
public static partial class KernelMemoryBuilderExtensions
1515
{
1616
/// <summary>
17-
/// Kernel Memory Builder extension method to add SqlServer memory connector.
17+
/// Kernel Memory Builder extension method to add SQL Server memory connector.
1818
/// </summary>
1919
/// <param name="builder">KM builder instance</param>
20-
/// <param name="config">SqlServer configuration</param>
20+
/// <param name="config">SQL Server configuration</param>
2121
public static IKernelMemoryBuilder WithSqlServerMemoryDb(
2222
this IKernelMemoryBuilder builder,
2323
SqlServerConfig config)
@@ -27,15 +27,17 @@ public static IKernelMemoryBuilder WithSqlServerMemoryDb(
2727
}
2828

2929
/// <summary>
30-
/// Kernel Memory Builder extension method to add SqlServer memory connector.
30+
/// Kernel Memory Builder extension method to add SQL Server memory connector.
3131
/// </summary>
3232
/// <param name="builder">KM builder instance</param>
33-
/// <param name="connString">SqlServer connection string</param>
33+
/// <param name="connString">SQL Server connection string</param>
34+
/// <param name="useNativeVectorSearch">Whether to use native vector search or not</param>
3435
public static IKernelMemoryBuilder WithSqlServerMemoryDb(
3536
this IKernelMemoryBuilder builder,
36-
string connString)
37+
string connString,
38+
bool useNativeVectorSearch = false)
3739
{
38-
builder.Services.AddSqlServerAsMemoryDb(connString);
40+
builder.Services.AddSqlServerAsMemoryDb(connString, useNativeVectorSearch);
3941
return builder;
4042
}
4143
}
@@ -46,10 +48,10 @@ public static IKernelMemoryBuilder WithSqlServerMemoryDb(
4648
public static partial class DependencyInjection
4749
{
4850
/// <summary>
49-
/// Inject SqlServer as the default implementation of IMemoryDb
51+
/// Inject SQL Server as the default implementation of IMemoryDb
5052
/// </summary>
5153
/// <param name="services">Service collection</param>
52-
/// <param name="config">Postgres configuration</param>
54+
/// <param name="config">SQL Server configuration</param>
5355
public static IServiceCollection AddSqlServerAsMemoryDb(
5456
this IServiceCollection services,
5557
SqlServerConfig config)
@@ -60,15 +62,17 @@ public static IServiceCollection AddSqlServerAsMemoryDb(
6062
}
6163

6264
/// <summary>
63-
/// Inject SqlServer as the default implementation of IMemoryDb
65+
/// Inject SQL Server as the default implementation of IMemoryDb
6466
/// </summary>
6567
/// <param name="services">Service collection</param>
6668
/// <param name="connString">SQL Server connection string</param>
69+
/// <param name="useNativeVectorSearch">Whether to use native vector search or not</param>
6770
public static IServiceCollection AddSqlServerAsMemoryDb(
6871
this IServiceCollection services,
69-
string connString)
72+
string connString,
73+
bool useNativeVectorSearch = false)
7074
{
71-
var config = new SqlServerConfig { ConnectionString = connString };
75+
var config = new SqlServerConfig { ConnectionString = connString, UseNativeVectorSearch = useNativeVectorSearch };
7276
return services.AddSqlServerAsMemoryDb(config);
7377
}
7478
}

extensions/SQLServer/SQLServer/QueryProviders/DefaultQueryProvider.cs

Lines changed: 319 additions & 0 deletions
Large diffs are not rendered by default.
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using Microsoft.Data.SqlClient;
4+
5+
namespace Microsoft.KernelMemory.MemoryDb.SQLServer.QueryProviders;
6+
7+
public interface ISqlServerQueryProvider
8+
{
9+
/// <summary>
10+
/// Return SQL used to create a new index
11+
/// </summary>
12+
public string PrepareCreateIndexQuery(
13+
int sqlServerVersion,
14+
string index,
15+
int vectorSize);
16+
17+
/// <summary>
18+
/// Return SQL used to get a list of indexes
19+
/// </summary>
20+
public string PrepareGetIndexesQuery();
21+
22+
/// <summary>
23+
/// Return SQL used to delete an index
24+
/// </summary>
25+
public string PrepareDeleteIndexQuery(string index);
26+
27+
/// <summary>
28+
/// Return SQL used to delete a memory record
29+
/// </summary>
30+
public string PrepareDeleteRecordQuery(string index);
31+
32+
/// <summary>
33+
/// Return SQL used to get a list of memory records
34+
/// </summary>
35+
public string PrepareGetRecordsListQuery(
36+
string index,
37+
ICollection<MemoryFilter>? filters,
38+
bool withEmbedding,
39+
SqlParameterCollection parameters);
40+
41+
/// <summary>
42+
/// Return SQL used to get a list of similar memory records
43+
/// </summary>
44+
public string PrepareGetSimilarRecordsListQuery(
45+
string index,
46+
ICollection<MemoryFilter>? filters,
47+
bool withEmbedding,
48+
SqlParameterCollection parameters);
49+
50+
/// <summary>
51+
/// Return SQL used to upsert a batch of memory records
52+
/// </summary>
53+
public string PrepareUpsertRecordsBatchQuery(string index);
54+
55+
/// <summary>
56+
/// Return SQL used to create all supporting tables
57+
/// </summary>
58+
public string PrepareCreateAllSupportingTablesQuery();
59+
}
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// Copyright (c) Microsoft. All rights reserved.
2+
3+
using System.Globalization;
4+
using System.Text;
5+
using Microsoft.Data.SqlClient;
6+
7+
namespace Microsoft.KernelMemory.MemoryDb.SQLServer.QueryProviders;
8+
9+
internal static class Utils
10+
{
11+
/// <summary>
12+
/// Gets the full table name with schema.
13+
/// </summary>
14+
/// <param name="config">Server settings</param>
15+
/// <param name="tableName">The table name.</param>
16+
internal static string GetFullTableName(SqlServerConfig config, string tableName)
17+
{
18+
return $"[{config.Schema}].[{tableName}]";
19+
}
20+
21+
/// <summary>
22+
/// Generates the filters as SQL commands and sets the SQL parameters
23+
/// </summary>
24+
/// <param name="config">Server settings</param>
25+
/// <param name="index">The index name.</param>
26+
/// <param name="parameters">The SQL parameters to populate.</param>
27+
/// <param name="filters">The filters to apply</param>
28+
internal static string GenerateFilters(
29+
SqlServerConfig config,
30+
string index,
31+
SqlParameterCollection parameters,
32+
ICollection<MemoryFilter>? filters)
33+
{
34+
var filterBuilder = new StringBuilder();
35+
36+
if (filters is null || filters.Count <= 0 || filters.All(f => f.Count <= 0))
37+
{
38+
return string.Empty;
39+
}
40+
41+
filterBuilder.Append("AND ( ");
42+
43+
for (int i = 0; i < filters.Count; i++)
44+
{
45+
var filter = filters.ElementAt(i);
46+
47+
if (i > 0)
48+
{
49+
filterBuilder.Append(" OR ");
50+
}
51+
52+
for (int j = 0; j < filter.Pairs.Count(); j++)
53+
{
54+
var value = filter.Pairs.ElementAt(j);
55+
56+
if (j > 0)
57+
{
58+
filterBuilder.Append(" AND ");
59+
}
60+
61+
filterBuilder.Append(" ( ");
62+
63+
filterBuilder.Append(CultureInfo.CurrentCulture, $@"EXISTS (
64+
SELECT
65+
1
66+
FROM {GetFullTableName(config, $"{config.TagsTableName}_{index}")} AS [tags]
67+
WHERE
68+
[tags].[memory_id] = {GetFullTableName(config, config.MemoryTableName)}.[id]
69+
AND [name] = @filter_{i}_{j}_name
70+
AND [value] = @filter_{i}_{j}_value
71+
)
72+
");
73+
74+
filterBuilder.Append(" ) ");
75+
76+
parameters.AddWithValue($"@filter_{i}_{j}_name", value.Key);
77+
parameters.AddWithValue($"@filter_{i}_{j}_value", value.Value);
78+
}
79+
}
80+
81+
filterBuilder.Append(" )");
82+
83+
return filterBuilder.ToString();
84+
}
85+
}

0 commit comments

Comments
 (0)