.Net: Remove SqlDataReaderDictionary in SQL Server provider (#12241)

We currently write records to the database by first serializing them to
a Dictionary, and then writing that Dictionary out as SqlParameters on
the command. Similarly, we read records by allocating and populating a
Dictionary, only to then iterate over that Dictionary and populate the
.NET instance from it.

This PR does away with the intermediate Dictionary representation. This
optimizes serialization by removing the double-copy, reduces
allocations, opens the way for boxing-free serialization, and also just
simplifies the code (less conversions and stuff to follow).

Author: roji

dotnet/src/Connectors/Connectors.Memory.SqlServer/SqlDataReaderDictionary.cs

@@ -259,20 +259,18 @@ public override async Task DeleteAsync(IEnumerable<TKey> keys, CancellationToken
             key,
             includeVectors);
 
-        using SqlDataReader reader = await connection.ExecuteWithErrorHandlingAsync(
+        return await connection.ExecuteWithErrorHandlingAsync(
             this._collectionMetadata,
             operationName: "Get",
             async () =>
             {
-                SqlDataReader reader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);
+                using SqlDataReader reader = await command.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);
                 await reader.ReadAsync(cancellationToken).ConfigureAwait(false);
-                return reader;
+                return reader.HasRows
+                    ? this._mapper.MapFromStorageToDataModel(reader, includeVectors)
+                    : null;
             },
             cancellationToken).ConfigureAwait(false);
-
-        return reader.HasRows
-            ? this._mapper.MapFromStorageToDataModel(new SqlDataReaderDictionary(reader, this._model.VectorProperties), includeVectors)
-            : default;
     }
 
     /// <inheritdoc/>
@@ -320,12 +318,22 @@ public override async IAsyncEnumerable<TRecord> GetAsync(IEnumerable<TKey> keys,
                 () => command.ExecuteReaderAsync(cancellationToken),
                 cancellationToken).ConfigureAwait(false);
 
-            while (await reader.ReadWithErrorHandlingAsync(
-                this._collectionMetadata,
-                "GetBatch",
-                cancellationToken).ConfigureAwait(false))
+            while (true)
             {
-                yield return this._mapper.MapFromStorageToDataModel(new SqlDataReaderDictionary(reader, this._model.VectorProperties), includeVectors);
+                TRecord? record = await VectorStoreErrorHandler.RunOperationAsync<TRecord?, SqlException>(
+                    this._collectionMetadata,
+                    "GetBatch",
+                    async () => await reader.ReadAsync(cancellationToken).ConfigureAwait(false)
+                        ? this._mapper.MapFromStorageToDataModel(reader, includeVectors)
+                        : null)
+                        .ConfigureAwait(false);
+
+                if (record is null)
+                {
+                    break;
+                }
+
+                yield return record;
             }
         } while (command.Parameters.Count == SqlServerConstants.MaxParameterCount);
     }
@@ -335,7 +343,7 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell
     {
         Verify.NotNull(record);
 
-        IReadOnlyList<Embedding>?[]? generatedEmbeddings = null;
+        Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>? generatedEmbeddings = null;
 
         var vectorPropertyCount = this._model.VectorProperties.Count;
         for (var i = 0; i < vectorPropertyCount; i++)
@@ -354,8 +362,8 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell
             // and generate embeddings for them in a single batch. That's some more complexity though.
             if (vectorProperty.TryGenerateEmbedding<TRecord, Embedding<float>>(record, cancellationToken, out var floatTask))
             {
-                generatedEmbeddings ??= new IReadOnlyList<Embedding>?[vectorPropertyCount];
-                generatedEmbeddings[i] = [await floatTask.ConfigureAwait(false)];
+                generatedEmbeddings ??= new Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>(vectorPropertyCount);
+                generatedEmbeddings[vectorProperty] = [await floatTask.ConfigureAwait(false)];
             }
             else
             {
@@ -370,7 +378,8 @@ public override async Task UpsertAsync(TRecord record, CancellationToken cancell
             this._schema,
             this.Name,
             this._model,
-            this._mapper.MapFromDataToStorageModel(record, recordIndex: 0, generatedEmbeddings));
+            record,
+            generatedEmbeddings);
 
         await connection.ExecuteWithErrorHandlingAsync(
            this._collectionMetadata,
@@ -393,7 +402,7 @@ public override async Task UpsertAsync(IEnumerable<TRecord> records, Cancellatio
         IReadOnlyList<TRecord>? recordsList = null;
 
         // If an embedding generator is defined, invoke it once per property for all records.
-        IReadOnlyList<Embedding>?[]? generatedEmbeddings = null;
+        Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>? generatedEmbeddings = null;
 
         var vectorPropertyCount = this._model.VectorProperties.Count;
         for (var i = 0; i < vectorPropertyCount; i++)
@@ -426,8 +435,8 @@ public override async Task UpsertAsync(IEnumerable<TRecord> records, Cancellatio
             // and generate embeddings for them in a single batch. That's some more complexity though.
             if (vectorProperty.TryGenerateEmbeddings<TRecord, Embedding<float>>(records, cancellationToken, out var floatTask))
             {
-                generatedEmbeddings ??= new IReadOnlyList<Embedding>?[vectorPropertyCount];
-                generatedEmbeddings[i] = (IReadOnlyList<Embedding<float>>)await floatTask.ConfigureAwait(false);
+                generatedEmbeddings ??= new Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>(vectorPropertyCount);
+                generatedEmbeddings[vectorProperty] = (IReadOnlyList<Embedding<float>>)await floatTask.ConfigureAwait(false);
             }
             else
             {
@@ -459,9 +468,9 @@ public override async Task UpsertAsync(IEnumerable<TRecord> records, Cancellatio
                         this._schema,
                         this.Name,
                         this._model,
-                        records.Skip(taken)
-                               .Take(SqlServerConstants.MaxParameterCount / parametersPerRecord)
-                               .Select((r, i) => this._mapper.MapFromDataToStorageModel(r, taken + i, generatedEmbeddings))))
+                        records.Skip(taken).Take(SqlServerConstants.MaxParameterCount / parametersPerRecord),
+                        firstRecordIndex: taken,
+                        generatedEmbeddings))
                     {
                         break; // records is empty
                     }
@@ -613,7 +622,7 @@ private async IAsyncEnumerable<VectorSearchResult<TRecord>> ReadVectorSearchResu
                 }
 
                 yield return new VectorSearchResult<TRecord>(
-                    this._mapper.MapFromStorageToDataModel(new SqlDataReaderDictionary(reader, vectorProperties), includeVectors),
+                    this._mapper.MapFromStorageToDataModel(reader, includeVectors),
                     reader.GetDouble(scoreIndex));
             }
         }
@@ -655,7 +664,7 @@ public override async IAsyncEnumerable<TRecord> GetAsync(Expression<Func<TRecord
                 operationName: "GetAsync",
                 cancellationToken).ConfigureAwait(false))
         {
-            yield return this._mapper.MapFromStorageToDataModel(new SqlDataReaderDictionary(reader, vectorProperties), options.IncludeVectors);
+            yield return this._mapper.MapFromStorageToDataModel(reader, options.IncludeVectors);
         }
     }
 }

dotnet/src/Connectors/Connectors.Memory.SqlServer/SqlServerCollection.cs

@@ -6,6 +6,7 @@
 using System.Text;
 using System.Text.Json;
 using Microsoft.Data.SqlClient;
+using Microsoft.Extensions.AI;
 using Microsoft.Extensions.VectorData;
 using Microsoft.Extensions.VectorData.ProviderServices;
 
@@ -122,7 +123,8 @@ internal static SqlCommand MergeIntoSingle(
         string? schema,
         string tableName,
         CollectionModel model,
-        IDictionary<string, object?> record)
+        object record,
+        Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>? generatedEmbeddings)
     {
         SqlCommand command = connection.CreateCommand();
         StringBuilder sb = new(200);
@@ -134,8 +136,13 @@ internal static SqlCommand MergeIntoSingle(
 
         foreach (var property in model.Properties)
         {
-            sb.AppendParameterName(property, ref paramIndex, out string paramName).Append(',');
-            command.AddParameter(property, paramName, record[property.StorageName]);
+            sb.AppendParameterName(property, ref paramIndex, out var paramName).Append(',');
+
+            var value = property is VectorPropertyModel vectorProperty && generatedEmbeddings?.TryGetValue(vectorProperty, out var ge) == true
+                ? ge[0]
+                : property.GetValueAsObject(record);
+
+            command.AddParameter(property, paramName, value);
         }
 
         sb[sb.Length - 1] = ')'; // replace the last comma with a closing parenthesis
@@ -174,7 +181,9 @@ internal static bool MergeIntoMany(
         string? schema,
         string tableName,
         CollectionModel model,
-        IEnumerable<IDictionary<string, object?>> records)
+        IEnumerable<object> records,
+        int firstRecordIndex,
+        Dictionary<VectorPropertyModel, IReadOnlyList<Embedding>>? generatedEmbeddings)
     {
         StringBuilder sb = new(200);
         // The DECLARE statement creates a table variable to store the keys of the inserted rows.
@@ -189,11 +198,18 @@ internal static bool MergeIntoMany(
         foreach (var record in records)
         {
             sb.Append('(');
+
             foreach (var property in model.Properties)
             {
-                sb.AppendParameterName(property, ref paramIndex, out string paramName).Append(',');
-                command.AddParameter(property, paramName, record[property.StorageName]);
+                sb.AppendParameterName(property, ref paramIndex, out var paramName).Append(',');
+
+                var value = property is VectorPropertyModel vectorProperty && generatedEmbeddings?.TryGetValue(vectorProperty, out var ge) == true
+                    ? ge[firstRecordIndex + rowIndex]
+                    : property.GetValueAsObject(record);
+
+                command.AddParameter(property, paramName, value);
             }
+
             sb[sb.Length - 1] = ')'; // replace the last comma with a closing parenthesis
             sb.AppendLine(",");
             rowIndex++;
@@ -585,14 +601,19 @@ private static void AddParameter(this SqlCommand command, PropertyModel? propert
                 command.Parameters.Add(name, System.Data.SqlDbType.VarBinary).Value = DBNull.Value;
                 break;
             case null:
-            case ReadOnlyMemory<float> vector when vector.Length == 0:
                 command.Parameters.AddWithValue(name, DBNull.Value);
                 break;
             case byte[] buffer:
                 command.Parameters.Add(name, System.Data.SqlDbType.VarBinary).Value = buffer;
                 break;
             case ReadOnlyMemory<float> vector:
-                command.Parameters.AddWithValue(name, JsonSerializer.Serialize(vector));
+                command.Parameters.AddWithValue(name, JsonSerializer.Serialize(vector, SqlServerJsonSerializerContext.Default.ReadOnlyMemorySingle));
+                break;
+            case Embedding<float> { Vector: var vector }:
+                command.Parameters.AddWithValue(name, JsonSerializer.Serialize(vector, SqlServerJsonSerializerContext.Default.ReadOnlyMemorySingle));
+                break;
+            case float[] vectorArray:
+                command.Parameters.AddWithValue(name, JsonSerializer.Serialize(vectorArray, SqlServerJsonSerializerContext.Default.SingleArray));
                 break;
             default:
                 command.Parameters.AddWithValue(name, value);

dotnet/src/Connectors/Connectors.Memory.SqlServer/SqlServerCommandBuilder.cs

@@ -0,0 +1,11 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System;
+using System.Text.Json.Serialization;
+
+namespace Microsoft.SemanticKernel.Connectors.SqlServer;
+
+// Note: this is temporary - SQL Server will switch away from using JSON arrays to represent embeddings in the future.
+[JsonSerializable(typeof(float[]))]
+[JsonSerializable(typeof(ReadOnlyMemory<float>))]
+internal partial class SqlServerJsonSerializerContext : JsonSerializerContext;